优化Langchain RAG检索：提升文档信息匹配准确性

from langchain.document_loaders import PyPDFLoader, DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.chains import RetrievalQA
from langchain.llms import OpenAI, HuggingFaceHub
import os

# 1. 文档加载
# 假设您的PDF文档位于 'data/' 目录下
# 可以将单个PDF文件复制到 /tmp/ 或指定目录
# loader = PyPDFLoader("doc.pdf") # 如果是单个文件
loader = DirectoryLoader('./data/', glob="./*.pdf", loader_cls=PyPDFLoader) # 从目录加载所有PDF
documents = loader.load()

# 2. 文档分块
# 调整分块大小和重叠，以更好地适应FAQ文档结构
# 尝试较小的chunk_size，确保一个问答对不会被过度分割
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
texts = text_splitter.split_documents(documents)

# 3. 嵌入模型选择与向量存储
# 使用HuggingFaceEmbeddings，并指定高性能模型
# 注意：首次运行时可能需要下载模型
embeddings = HuggingFaceEmbeddings(
    model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2" # 推荐的多语言模型
    # 或者尝试 "bert-base-multilingual-cased"
)

# 持久化向量数据库，避免每次运行时都重新创建
persist_directory = "./chromadb" # 建议指定一个明确的路径
vectordb = Chroma.from_documents(documents=texts, embedding=embeddings, persist_directory=persist_directory)
vectordb.persist() # 保存向量数据库到磁盘

# 4. LLM选择
# 选项一：使用OpenAI模型 (需要设置OPENAI_API_KEY环境变量)
# os.environ["OPENAI_API_KEY"] = "YOUR_OPENAI_API_KEY"
# llm = OpenAI(temperature=0, model_name="text-davinci-003")

# 选项二：使用HuggingFace Hub上的开源LLM (需要设置HUGGINGFACEHUB_API_TOKEN环境变量)
# os.environ["HUGGINGFACEHUB_API_TOKEN"] = "YOUR_HF_API_TOKEN"
llm = HuggingFaceHub(repo_id="google/flan-t5-base",
                     model_kwargs={"temperature": 0.6, "max_length": 500, "max_new_tokens": 200})
# 或者尝试其他模型，例如：
# llm = HuggingFaceHub(repo_id="EleutherAI/gpt-neo-2.7B",
#                      model_kwargs={"temperature": 0.6, "max_length": 500, "max_new_tokens": 200})

# 5. 构建RAG检索链
# chain_type="stuff" 是最简单的链类型，它将所有检索到的文档“塞入”LLM的上下文。
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    retriever=vectordb.as_retriever(),
    chain_type="stuff",
    return_source_documents=True # 返回源文档，便于调试和验证
)

# 6. 执行查询
question = "请总结这本书的主要内容" # 示例查询
response = qa_chain({"query": question})
print(response)

# 调试：查看检索到的源文档
# if 'source_documents' in response:
#     print("\n--- 检索到的源文档 ---")
#     for doc in response['source_documents']:
#         print(f"Content: {doc.page_content[:200]}...") # 打印前200字
#         print(f"Metadata: {doc.metadata}")
#         print("-" * 20)