当前位置：首页 > news >正文

实现我的第一个本地文档问答机器人

news 2025/9/15 17:44:29

本地文档问答机器人

下面是一个完整的本地文档问答机器人实现，涵盖了阶段三的所有核心概念：文档加载、文本分割、向量存储和检索增强生成（RAG）。

完整代码实现

import os.path
from typing import Listfrom huggingface_hub import snapshot_download
from langchain.chains import RetrievalQA
from langchain_community.document_loaders import TextLoader, Docx2txtLoader, PyPDFLoader, UnstructuredFileLoader
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI
from langchain_text_splitters import RecursiveCharacterTextSplitterclass LocalDocumentQA:def __init__(self):# 初始化llmself.llm = ChatOpenAI(model_name="xop3qwen1b7",openai_api_base="https://域名/v1",openai_api_key="sk-kBa9GlpIxpWX281iA35a...",temperature=0.1,  # 降低温度值，提升问答精度max_tokens=1024)# # 初始化嵌入模型，这里使用HuggingFaceEmbedding# self.embeddings = HuggingFaceEmbeddings(#     model_name="BAAI/bge-small-zh-v1.5",#     model_kwargs={'device': 'cpu'},#     encode_kwargs={'normalize_embeddings': True}  # 归一化向量# )# 避免每次下载，这里采用预下载后，后续使用本地数据的方式# 提前下载模型# model_path = snapshot_download(#     repo_id="BAAI/bge-small-zh-v1.5",#     local_dir="./models/bge-small-zh-v1.5",  # 指定本地路径#     local_dir_use_symlinks=False# )# 然后使用本地路径self.embeddings = HuggingFaceEmbeddings(model_name="./models/bge-small-zh-v1.5",  # 使用本地路径model_kwargs={'device': 'cpu'},encode_kwargs={'normalize_embeddings': True})# 初始化文本分割器self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=500,  # 分片大小chunk_overlap=100,  # 块之间的重叠部分separators=["\n\n", "\n", "。", "！", "？", " ", ""]  # 分割符，优先按照分割符分割)self.vectorstore = Noneself.qa_chain = Nonedef load_documents(self, file_paths: List[str]):"""加载多种格式的文档:param file_paths::return:"""documents = []for file_path in file_paths:if not os.path.exists(file_path):print(f"文件不存在:{file_path}")continuetry:if file_path.endswith(".txt"):loader = TextLoader(file_path, encoding="utf-8")elif file_path.endswith(".docx"):loader = Docx2txtLoader(file_path)elif file_path.endswith('.pdf'):loader = PyPDFLoader(file_path)else:# 尝试使用UnstructuredLoader处理其他格式loader = UnstructuredFileLoader(file_path)loaded_docs = loader.load()documents.extend(loaded_docs)print(f"已加载: {file_path} ({len(loaded_docs)}个文档)")except Exception as e:print(f"加载文件时出错 {file_path}: {e}")return documentsdef process_documents(self, fil_paths: List[str], persist_directory: str = "./chroma_db"):"""处理文档并创建向量数据库:param fil_paths::param persist_directory: 持久化目录:return:"""# 1. 加载文档documents = self.load_documents(fil_paths)if not documents:print("未加载任何文档，请检查文件路径")return False# 2. 分割文本print("正在分割文本...")texts = self.text_splitter.split_documents(documents)print(f"文档已被分割成 {len(texts)} 个片段")# 3. 创建向量存储print("正在创建向量数据库...")self.vectorstore = Chroma.from_documents(documents=texts,embedding=self.embeddings,persist_directory=persist_directory)# 4. 创建检索QA链print("正在创建问答链...")self._create_qa_chain()def _create_qa_chain(self):"""创建检索QA链:return:"""# 自定义提示词模板qa_template = """你是一个专业的文档问答助手。请严格根据提供的上下文信息回答问题。如果上下文中有答案，请根据上下文用中文回答。如果上下文中没有答案，请如实告知"根据提供的资料，我无法找到相关答案。"上下文:{context}问题:{question}请提供准确、有帮助的回答:"""# 创建提示词qa_prompt = PromptTemplate(template=qa_template,input_variables=["context", "question"])# 创建向量检索器retriever = self.vectorstore.as_retriever(search_type="similarity",search_kwargs={"k": 4}  # 检索最相关的4个片段)# 创建QA链self.qa_chain = RetrievalQA.from_chain_type(llm=self.llm,chain_type="stuff",retriever=retriever,chain_type_kwargs={"prompt": qa_prompt},return_source_documents=True  # 开启返回关联的文章信息)def ask_question(self, question: str):"""向文档提问:param question::return:"""if not self.qa_chain:print("请先完成文档加载和处理")return Nonetry:# 了解这个是如何给赋值的？？？result = self.qa_chain.invoke({"query": question})return resultexcept Exception as e:print(f"提问时出错: {e}")return Nonedef load_existing_vectorstore(self, persist_directory: str = "./chroma_db"):"""加载已存在的向量数据库:param persist_directory::return:"""if not os.path.exists(persist_directory):print("向量数据库不存在，请先处理文档")return Falsetry:self.vectorstore = Chroma(persist_directory=persist_directory,embedding_function=self.embeddings)self._create_qa_chain()print("✅ 已加载现有向量数据库")return Trueexcept Exception as e:print(f"加载向量数据库时出错: {e}")return Falsedef create_sample_document():"""创建示例文档"""sample_content = """LangChain 学习指南LangChain是一个用于开发由语言模型驱动的应用程序的框架。核心概念:1. 模型I/O: 与语言模型交互的基础接口2. 检索: 从外部源获取相关信息3. 链: 将多个组件组合成工作流4. 代理: 使用LLM决定执行哪些动作5. 内存: 在对话中保持状态安装方法:使用pip安装: pip install langchain如果需要使用OpenAI模型: pip install langchain-openai常见应用场景:- 文档问答系统- 聊天机器人- 代码生成工具- 数据分析助手最佳实践:1. 使用合适的文本分割策略2. 选择合适的嵌入模型3. 优化提示词工程4. 实施适当的错误处理"""with open("sample_document.txt", "w", encoding="utf-8") as f:f.write(sample_content)return ["sample_document.txt"]def main():# 创建示例文档# print("创建示例文档...")# file_paths = create_sample_document()file_paths = ["sample_document.txt", "JMeterFile.docx"]# 初始化问答系统qa_system = LocalDocumentQA()# 检查是否已有向量数据库if not qa_system.load_existing_vectorstore():# 处理文档并创建向量数据库print("处理文档并创建向量数据库...")success = qa_system.process_documents(file_paths)if not success:returnprint("\n" + "=" * 50)print("📚 本地文档问答系统已就绪")print("💡 你可以向文档提问了")print("📝 输入'退出'结束程序")print("=" * 50)# 交互式问答while True:question = input("\n请输入你的问题: ").strip()if question.lower() in ['退出', 'quit', 'exit', 'q']:print("感谢使用，再见！")breakif not question:continue# 获取答案result = qa_system.ask_question(question)if result:print("\n✅ 答案:")print(result["result"])# 显示参考来源print("\n📚 参考来源:")for i, doc in enumerate(result["source_documents"], 1):source = doc.metadata.get('source', '未知')content_preview = doc.page_content[:100] + "..." if len(doc.page_content) > 100 else doc.page_contentprint(f"{i}. [{source}] {content_preview}")if __name__ == "__main__":main()

安装依赖

在运行前，请安装所需依赖：

pip install langchain langchain-community langchain-openai chromadb sentence-transformers unstructured pypdf

功能特点

这个本地文档问答机器人具有以下特点：

多格式支持：可以处理TXT、PDF等多种文档格式
中文优化：使用中文优化的嵌入模型（text2vec-large-chinese）
持久化存储：向量数据库持久化，避免重复处理文档
来源追溯：显示答案的参考来源，提高可信度
错误处理：完善的异常处理机制

使用指南

1. 首次运行

首次运行时会：

创建一个示例文档
处理文档并创建向量数据库
进入交互式问答界面

2. 使用自己的文档

要使用自己的文档，修改main()函数中的file_paths：

def main():# 替换为你的文档路径file_paths = ["path/to/your/document1.pdf", "path/to/your/document2.txt"]# 其余代码保持不变...

3. 高级配置

你可以调整以下参数来优化性能：

# 文本分割参数
self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=500,      # 每个片段的大小chunk_overlap=100,   # 片段之间的重叠separators=["\n\n", "\n", "。", "！", "？", " ", ""]  # 分割符优先级
)# 检索参数
retriever = self.vectorstore.as_retriever(search_type="similarity",search_kwargs={"k": 4}  # 检索最相关的4个片段
)

扩展功能

你可以根据需要扩展这个基础系统：

1. 添加网页加载器

from langchain_community.document_loaders import WebBaseLoader# 添加网页加载功能
def load_webpage(self, url: str):loader = WebBaseLoader(url)return loader.load()

2. 添加混合搜索

# 使用MMR（最大边际相关性）搜索，平衡相关性和多样性
retriever = self.vectorstore.as_retriever(search_type="mmr",search_kwargs={"k": 6, "fetch_k": 10}
)

3. 添加对话历史

from langchain.memory import ConversationBufferMemory# 添加对话记忆
memory = ConversationBufferMemory(memory_key="chat_history",return_messages=True
)# 修改QA链以支持对话
self.qa_chain = ConversationalRetrievalChain.from_llm(llm=self.llm,retriever=retriever,memory=memory,combine_docs_chain_kwargs={"prompt": QA_PROMPT}
)

4. 添加评估功能

def evaluate_answer(self, question: str, answer: str):"""评估答案质量"""evaluation_prompt = """请评估以下问答对的质量：问题: {question}答案: {answer}请从准确性、完整性和相关性三个方面评分（1-5分），并给出简要评价："""prompt = PromptTemplate(template=evaluation_prompt,input_variables=["question", "answer"])evaluation_chain = LLMChain(llm=self.llm, prompt=prompt)return evaluation_chain.run({"question": question, "answer": answer})