优化Langchain RAG检索：提升PDF文档问答准确性

pip install langchain pypdf chromadb sentence-transformers transformers openai

from langchain.document_loaders import PyPDFLoader, DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter

# 假设您的PDF文件位于'/tmp/'目录下，或者指定单个文件路径
# loader = PyPDFLoader("doc.pdf") # 如果是单个文件
loader = DirectoryLoader('/tmp/', glob="./*.pdf", loader_cls=PyPDFLoader) # 适用于目录中的所有PDF
documents = loader.load()

# 使用RecursiveCharacterTextSplitter进行文本分割
# chunk_size和chunk_overlap的调整对检索效果有显著影响，需要根据文档特性进行试验
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
texts = text_splitter.split_documents(documents)

from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma

# 初始化HuggingFace嵌入模型
# 推荐使用在语义相似性任务上表现良好的模型
embeddings = HuggingFaceEmbeddings(
    model_name="bert-base-multilingual-cased" # 您可以尝试其他模型，如"sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
)

# 定义持久化目录，方便下次直接加载向量数据库
persist_directory = "/tmp/chromadb"

# 从分割后的文本块创建并持久化Chroma向量数据库
vectordb = Chroma.from_documents(documents=texts, embedding=embeddings, persist_directory=persist_directory)

# 确保向量数据库被持久化
vectordb.persist()
print(f"Vector database created and persisted at: {persist_directory}")

from langchain.chains import RetrievalQA
from langchain.llms import OpenAI # 示例使用OpenAI，您也可以替换为开源LLM
from langchain.prompts import PromptTemplate
from langchain.callbacks.manager import CallbackManager
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain_community.llms import HuggingFaceHub # 导入HuggingFaceHub for open-source LLMs

# 初始化大型语言模型 (LLM)
# 示例使用OpenAI，请确保已设置OPENAI_API_KEY环境变量
# llm = OpenAI(temperature=0, model_name="text-davinci-003")

# 或者使用HuggingFaceHub上的开源LLM (需要HUGGINGFACEHUB_API_TOKEN)
llm = HuggingFaceHub(repo_id="google/flan-t5-base",
                     model_kwargs={"temperature":0.6, "max_length": 500, "max_new_tokens": 200})

# 可选：自定义PromptTemplate以引导LLM的回答风格
# template = """使用以下上下文来回答最后的问题。如果你不知道答案，就说你不知道，不要试图编造答案。
# {context}
# 问题: {question}
# 有用的回答:"""
# QA_CHAIN_PROMPT = PromptTemplate.from_template(template)


# 构建RetrievalQA链
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    retriever=vectordb.as_retriever(), # 从向量数据库获取检索器
    chain_type="stuff", # 将所有检索到的文档填充到LLM的上下文
    # chain_type_kwargs={"prompt": QA_CHAIN_PROMPT}, # 如果使用了自定义PromptTemplate
    return_source_documents=True # 返回检索到的源文档
)

# 提出问题并获取答案
question = "请总结这本书的主要内容" # 替换为您的实际问题
response = qa_chain({"query": question})

print("\n--- 问答结果 ---")
print(f"问题: {question}")
print(f"答案: {response['result']}")
print("\n--- 检索到的源文档 ---")
for doc in response['source_documents']:
    print(f"文档内容: {doc.page_content[:200]}...") # 打印前200字
    print(f"文档来源: {doc.metadata}")

from langchain_community.llms import HuggingFaceHub

# 示例1: Google Flan-T5-base
# llm = HuggingFaceHub(repo_id="google/flan-t5-base",
#                      model_kwargs={"temperature":0.6,"max_length": 500, "max_new_tokens": 200})

# 示例2: EleutherAI/gpt-neo-2.7B
# llm = HuggingFaceHub(repo_id="EleutherAI/gpt-neo-2.7B",
#                      model_kwargs={"temperature":0.6,"max_length": 500, "max_new_tokens": 200})