import os
import sys
import time
from langchain_community.document_loaders import PDFPlumberLoader
from langchain.vectorstores import FAISS
from langchain.prompts import PromptTemplate
from langchain_ollama.llms import OllamaLLM
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.schema import Document
from langchain.retrievers import BM25Retriever, EnsembleRetriever
from sentence_transformers import CrossEncoder
import torch
import warnings
warnings.filterwarnings("ignore")

# 1. 디바이스 설정 (CUDA 사용 가능 시 사용)
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"임베딩 모델을 사용할 디바이스: {device}")

# 2. 문서 경로 설정
available_files = {
    'manual': r"C:Usersju980Desktoprag_testmanual.pdf"
}

print("n사용 가능한 문서 목록:")
for key in available_files.keys():
    print(f"- {key}")

selected_file_key = input("n열람할 문서의 이름을 입력하세요: ").strip()

if selected_file_key not in available_files:
    print("해당 이름의 문서가 없습니다. 프로그램을 종료합니다.")
    sys.exit()

file_path = available_files[selected_file_key]
print(f"n'{selected_file_key}' 문서에서 정보를 검색합니다.n")

# 3. 문서 로드 및 전체 내용 출력
loader = PDFPlumberLoader(file_path)
documents = loader.load()

# 4. 문서 분할
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=2500,        
    chunk_overlap=500,      
    separators=["nn", "n", " ", ""]
)

split_docs = text_splitter.split_documents(documents)
print(f"총 분할된 문서 조각 수: {len(split_docs)}n")

# 5. 한국어에 최적화된 임베딩 모델 사용
embedding_model_name = 'jhgan/ko-sroberta-multitask'  # 한국어에 특화된 임베딩 모델로 변경
embeddings = HuggingFaceEmbeddings(
    model_name=embedding_model_name,
    model_kwargs={"device": device},
)

# 6. 벡터스토어 생성
vectorstore = FAISS.from_documents(documents=split_docs, embedding=embeddings)

# 7. 리트리버 생성 (FAISS + BM25)
# FAISS 기반 리트리버
vector_retriever = vectorstore.as_retriever(search_kwargs={"k": 20})  # k 값을 20으로 늘림

# BM25 기반 리트리버
bm25_retriever = BM25Retriever.from_documents(split_docs)

# Ensemble Retriever 생성 (FAISS + BM25)
ensemble_retriever = EnsembleRetriever(
    retrievers=[vector_retriever, bm25_retriever],
    weights=[0.4, 0.6]  # BM25의 가중치를 높임
)

# 8. 리랭킹 모델 초기화
reranker_model_name = "cross-encoder/ms-marco-MiniLM-L-6-v2"
try:
    cross_encoder = CrossEncoder(reranker_model_name, device=device)
    print("리랭커 모델이 성공적으로 로드되었습니다.n")
except Exception as e:
    print(f"리랭커 모델 로드 중 오류 발생: {e}")
    sys.exit()

# 9. 프롬프트 템플릿 설정
prompt_template = """다음은 제공된 문서에서 발췌한 내용입니다. 이 정보를 바탕으로 주어진 질문에 대해     정확하고 간결하게 답변하세요. 답변은 제공된 문서의 내용만을 참고해야 하며, 외부 정보는 포함하지 마세요.

문서 내용:
{context}

질문:
{question}

답변:"""

prompt = PromptTemplate(
    template=prompt_template,
    input_variables=["context", "question"]
)

# 10. LLM 초기화
llm = OllamaLLM(
    base_url="http://localhost:11434",  # Ollama 서버 주소
    model="llama3.1",  
    temperature=0
)

print("질문을 입력해 주세요. 'exit'를 입력하면 대화를 종료합니다.n")

# 11. 질문 처리 루프
while True:
    question = input("질문: ")
    if question.strip().lower() == 'exit':
        print("대화를 종료합니다.")
        break

    # 11-1. 문서 검색
    start_time = time.time()
    retrieved_docs = ensemble_retriever.get_relevant_documents(question)
    retrieval_time = time.time() - start_time
    print(f"n문서 검색 시간: {retrieval_time:.2f}초n")

    # 11-2. 검색된 문서 조각 출력 (디버깅용)
    print("----- 검색된 문서 조각 -----n")
    for idx, doc in enumerate(retrieved_docs):
        page = doc.metadata.get('page', 'N/A')
        source = doc.metadata.get('source', 'Unknown')
        print(f"조각 {idx+1} (Source: {source}, Page: {page}):n{doc.page_content}n")
print("----- 검색된 문서 조각 종료 -----n")

# 11-3. 리랭킹 적용
try:
    pairs = [[question, doc.page_content] for doc in retrieved_docs]
    scores = cross_encoder.predict(pairs)
    scored_docs = sorted(zip(retrieved_docs, scores), key=lambda x: x[1], reverse=True)
    reranked_docs = [doc for doc, score in scored_docs]
    top_docs = reranked_docs[:5]  
    print("리랭킹을 통해 상위 5개 문서를 선택했습니다.n")
except Exception as e:
    print(f"리랭킹 중 오류 발생: {e}")
    top_docs = retrieved_docs[:5]

# 11-4. 리랭크된 문서 조각 출력 (디버깅용)
print("----- 리랭크 후 상위 문서 조각 -----n")
for idx, doc in enumerate(top_docs):
    page = doc.metadata.get('page', 'N/A')
    source = doc.metadata.get('source', 'Unknown')
    print(f"조각 {idx+1} (Source: {source}, Page: {page}):n{doc.page_content}n")
print("----- 리랭크 후 상위 문서 조각 종료 -----n")

# 11-5. 컨텍스트 생성
context = "nn".join([doc.page_content for doc in top_docs])

# 11-6. 프롬프트 생성
formatted_prompt = prompt.format(context=context, question=question)
print("----- LLM에 전달되는 프롬프트 -----n")
print(formatted_prompt)
print("----- 프롬프트 종료 -----n")

# 11-7. 답변 생성
try:
    answer = llm.invoke(formatted_prompt)
except Exception as e:
    print(f"Ollama 호출 중 오류 발생: {e}")
    answer = "죄송합니다. 현재 답변을 생성할 수 없습니다."

print(f"답변: {answer}n")

This is my code, and it doesn’t retrieve my PDF file properly. It retrieved something on my file, but that’s not what I asked for.

For example, I ask “How many vacation days can I take in a year?” It retrieves working days or something, and gives me back another answer like “Christmas is a holiday (also in my file)”. I want a answer like “15 days in a year”.

What’s wrong with my code? Please help me.
your text

I tried many embedding models, many retrievers and adjusted weights. But the result is the same.

I expected that my code would retrieve an exact answer.

Please let me know how to fix it.