实现我的第一个本地文档问答机器人
本地文档问答机器人
下面是一个完整的本地文档问答机器人实现，涵盖了阶段三的所有核心概念：文档加载、文本分割、向量存储和检索增强生成（RAG）。
完整代码实现
import os.path
from typing import List
from 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 RecursiveCharacterTextSplitter
class LocalDocumentQA:
def __init__(self):
# 初始化llm
self.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 = None
self.qa_chain = None
def 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}")
continue
try:
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 documents
def 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 None
try:
# 了解这个是如何给赋值的？？？
result = self.qa_chain.invoke({"query": question})
return result
except Exception as e:
print(f"提问时出错: {e}")
return None
def load_existing_vectorstore(self, persist_directory: str = "./chroma_db"):
"""
加载已存在的向量数据库
:param persist_directory:
:return:
"""
if not os.path.exists(persist_directory):
print("向量数据库不存在，请先处理文档")
return False
try:
self.vectorstore = Chroma(
persist_directory=persist_directory,
embedding_function=self.embeddings
)
self._create_qa_chain()
print("✅ 已加载现有向量数据库")
return True
except Exception as e:
print(f"加载向量数据库时出错: {e}")
return False
def 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:
return
print("\n" + "=" * 50)
print("📚 本地文档问答系统已就绪")
print("💡 你可以向文档提问了")
print("📝 输入'退出'结束程序")
print("=" * 50)
# 交互式问答
while True:
question = input("\n请输入你的问题: ").strip()
if question.lower() in ['退出', 'quit', 'exit', 'q']:
print("感谢使用，再见！")
break
if 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_content
print(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})
常见问题解决
- 内存不足：减小
chunk_size或使用GPU加速嵌入模型 - 处理速度慢：尝试使用更小的嵌入模型或增加
chunk_size - 答案不准确：调整
chunk_size和chunk_overlap，或增加检索的文档数量(k值) - 中文支持不佳：确保使用中文优化的嵌入模型（如text2vec-large-chinese）
浙公网安备 33010602011771号