企业本地知识库助手 大模型+本地知识库

首发地址：企业本地知识库助手 大模型+本地知识库 – 麒麒

打造企业本地知识库助手，结合了企业私有的知识库，又利用了大模型的生成能力。

整体流程：

1、使用 embeddings模型将企业知识文档向量化存入向量数据库Faiss中

2、接收用户问题，从本地知识库中查找相关性排名前N的知识块

3、让大模型依据本地知识回答用户问题

embeddings模型

embeddings文本向量化模型选择，为了减少下载与启动时长，快速跑起demo，我使用的是iic/nlp_gte_sentence-embedding_chinese-base模型，大概几百MB，下载比较快。可以对比通义千问系列其他模型，文件比较大，可能语文理解会更好。当量，也可以使用直接调用embeddings在线api，避免本地部署。

大模型

大模型对机器的配置要求更高，可直接调用大模型服务，或者运行小尺寸的大模型。上代码代码我在本地调试运行过，可直接copy过去尝试一下。可稍微完善一下，就是一个企业本地知识库的智能助手。需要本地部署的话，可以使得ollama部署大模型。想快速跑一下代码，就直接使用阿里百链或其他大模型服务。

import os
import sys
import time
import logging
import argparse
from typing import List, Tuple, Dict, Optional
from langchain_community.llms import Ollama
from langchain_openai import ChatOpenAI# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)# 检查必要的依赖库
try:from PyPDF2 import PdfReaderfrom langchain.chains.question_answering import load_qa_chainfrom langchain_openai import OpenAIfrom langchain_community.callbacks.manager import get_openai_callbackfrom langchain.text_splitter import RecursiveCharacterTextSplitterfrom langchain_community.embeddings import HuggingFaceEmbeddingsfrom langchain_community.vectorstores import FAISSfrom modelscope import snapshot_download, AutoModel, AutoTokenizerimport torchimport torch.nn.functional as Ffrom langchain.embeddings.base import Embeddings
except ImportError as e:print(f"错误: 缺少必要的依赖库: {str(e)}")print("请安装所需依赖: pip install PyPDF2 langchain langchain-openai langchain-community faiss-cpu modelscope torch")sys.exit(1)# 自定义 ModelScope 嵌入类
class ModelScopeEmbeddings(Embeddings):"""使用 ModelScope 加载的模型进行文本嵌入"""# 类变量用于缓存已加载的模型_models_cache = {}def __init__(self, model_id: str = "iic/nlp_gte_sentence-embedding_chinese-base", device: str = None, batch_size: int = 32, use_query_prefix: bool = True):"""初始化 ModelScope 嵌入模型参数:model_id: ModelScope 模型IDdevice: 计算设备，如 'cuda:0', 'cpu' 等，默认为自动选择batch_size: 批处理大小，用于处理大量文档use_query_prefix: 是否为查询添加前缀"""self.model_id = model_idself.batch_size = batch_sizeself.use_query_prefix = use_query_prefix# 设置设备if device is None:self.device = "cuda" if torch.cuda.is_available() else "cpu"else:self.device = device# 缓存键cache_key = f"{model_id}_{self.device}"try:# 检查模型是否已经加载if cache_key in self._models_cache:logger.info(f"使用缓存中的模型: {model_id}")self.tokenizer, self.model = self._models_cache[cache_key]else:# 检查模型是否已经下载try:# 尝试直接加载本地模型self.model_dir = snapshot_download(model_id, cache_dir=os.path.join(os.path.expanduser("~"), ".cache", "modelscope"))logger.info(f"模型已存在于本地缓存: {self.model_dir}")except Exception as e:logger.info(f"本地未找到模型，正在从 ModelScope 下载模型 {model_id}...")try:self.model_dir = snapshot_download(model_id)logger.info(f"模型已下载到: {self.model_dir}")except Exception as download_error:logger.error(f"下载模型时出错: {str(download_error)}")raise ValueError(f"无法下载模型 {model_id}: {str(download_error)}")# 加载模型和分词器logger.info(f"正在加载模型到 {self.device} 设备...")try:self.tokenizer = AutoTokenizer.from_pretrained(self.model_dir)self.model = AutoModel.from_pretrained(self.model_dir)self.model.to(self.device)# 缓存模型self._models_cache[cache_key] = (self.tokenizer, self.model)logger.info(f"模型已加载并缓存")except Exception as load_error:logger.error(f"加载模型时出错: {str(load_error)}")raise ValueError(f"无法加载模型: {str(load_error)}")except Exception as e:logger.error(f"初始化 ModelScope 嵌入模型时出错: {str(e)}")raise ValueError(f"初始化失败: {str(e)}")def _mean_pooling(self, model_output, attention_mask):"""平均池化获取句子嵌入"""token_embeddings = model_output[0]input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)def embed_documents(self, texts: List[str]) -> List[List[float]]:"""为文档生成嵌入向量参数:texts: 文档文本列表返回:embeddings: 嵌入向量列表"""if not texts:return []try:# 分批处理文档all_embeddings = []for i in range(0, len(texts), self.batch_size):batch_texts = texts[i:i+self.batch_size]# 对文本进行编码encoded_input = self.tokenizer(batch_texts, padding=True, truncation=True, max_length=512,  # 添加最大长度限制return_tensors='pt').to(self.device)# 生成嵌入with torch.no_grad():model_output = self.model(**encoded_input)# 池化并归一化sentence_embeddings = self._mean_pooling(model_output, encoded_input['attention_mask'])normalized_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)# 添加到结果列表all_embeddings.extend(normalized_embeddings.cpu().tolist())return all_embeddingsexcept Exception as e:logger.error(f"生成文档嵌入时出错: {str(e)}")# 在出错时返回零向量，避免程序崩溃if len(texts) > 0:try:# 尝试获取嵌入维度dummy_output = self.embed_documents([texts[0]])dim = len(dummy_output[0])return [[0.0] * dim] * len(texts)except:# 如果无法确定维度，使用默认维度return [[0.0] * 768] * len(texts)return []def embed_query(self, text: str) -> List[float]:"""为查询生成嵌入向量参数:text: 查询文本返回:embedding: 嵌入向量"""if not text:# 返回零向量return [0.0] * 768try:# 为查询添加提示词（如果启用）query_text = f"查询：{text}" if self.use_query_prefix else textreturn self.embed_documents([query_text])[0]except Exception as e:logger.error(f"生成查询嵌入时出错: {str(e)}")# 返回零向量return [0.0] * 768class PDFChatBot:"""PDF聊天机器人类，用于处理PDF文档并回答相关问题"""def __init__(self, pdf_path: str, api_key: str, base_url: str = "https://dashscope.aliyuncs.com/compatible-mode/v1", model_id: str = "iic/nlp_gte_sentence-embedding_chinese-base", device: str = None, use_query_prefix: bool = True,llm_model_name: str = "deepseek-r1:1.5b"):"""初始化PDF聊天机器人参数:pdf_path: PDF文件路径api_key: OpenAI API密钥base_url: API基础URLmodel_id: ModelScope模型ID，用于文本嵌入device: 计算设备，如 'cuda:0', 'cpu' 等，默认为自动选择use_query_prefix: 是否为查询添加前缀"""self.pdf_path = pdf_pathself.api_key = api_keyself.base_url = base_urlself.model_id = model_idself.device = deviceself.use_query_prefix = use_query_prefixself.llm_model_name = llm_model_nameself.knowledge_base = Noneself.page_count = 0# 检查PDF文件是否存在if not os.path.exists(pdf_path):raise FileNotFoundError(f"PDF文件不存在: {pdf_path}")logger.info(f"初始化PDF聊天机器人，使用文件: {pdf_path}")def load_pdf(self) -> None:"""加载PDF文件并创建知识库"""start_time = time.time()logger.info(f"开始加载PDF文件: {self.pdf_path}")try:# 读取PDF文件pdf_reader = PdfReader(self.pdf_path)self.page_count = len(pdf_reader.pages)logger.info(f"PDF文件共有 {self.page_count} 页")# 提取文本和页码信息text, page_numbers = self._extract_text_with_page_numbers(pdf_reader)logger.info(f"提取的文本长度: {len(text)} 个字符")if not text:raise ValueError("无法从PDF中提取文本，请检查PDF文件是否有效")# 处理文本并创建知识库self.knowledge_base = self._process_text_with_splitter(text, page_numbers)elapsed_time = time.time() - start_timelogger.info(f"PDF加载和知识库创建完成，耗时: {elapsed_time:.2f} 秒")except Exception as e:logger.error(f"加载PDF文件时出错: {str(e)}")raisedef _extract_text_with_page_numbers(self, pdf) -> Tuple[str, List[int]]:"""从PDF中提取文本并记录每页文本对应的页码参数:pdf: PDF文件对象返回:text: 提取的文本内容page_numbers: 每页文本对应的页码列表"""text = ""page_texts = []  # 存储每页的文本page_numbers = []  # 存储每页的页码for page_number, page in enumerate(pdf.pages, start=1):try:extracted_text = page.extract_text()if extracted_text:text += extracted_text + "\n\n"  # 添加页面分隔符page_texts.append(extracted_text)page_numbers.append(page_number)else:logger.warning(f"第 {page_number} 页未找到文本")except Exception as e:logger.error(f"提取第 {page_number} 页文本时出错: {str(e)}")return text, page_numbersdef _process_text_with_splitter(self, text: str, page_numbers: List[int]) -> FAISS:"""处理文本并创建向量存储参数:text: 提取的文本内容page_numbers: 每页文本对应的页码列表返回:knowledgeBase: 基于FAISS的向量存储对象"""try:# 创建文本分割器，用于将长文本分割成小块text_splitter = RecursiveCharacterTextSplitter(separators=["\n\n", "\n", ".", " ", ""],chunk_size=1000,chunk_overlap=200,length_function=len,)# 分割文本chunks = text_splitter.split_text(text)logger.info(f"文本被分割成 {len(chunks)} 个块")try:# 创建嵌入模型 - 使用ModelScope的中文嵌入模型logger.info(f"正在初始化 ModelScope 嵌入模型: {self.model_id}...")embeddings = ModelScopeEmbeddings(model_id=self.model_id,device=self.device,use_query_prefix=self.use_query_prefix)# 从文本块创建知识库knowledge_base = FAISS.from_texts(chunks, embeddings)logger.info("已从文本块创建知识库")# 为每个文本块分配页码knowledge_base.page_info = self._assign_page_numbers_to_chunks(text, chunks, page_numbers)return knowledge_baseexcept Exception as e:logger.error(f"创建嵌入模型或知识库时出错: {str(e)}")raise ValueError(f"创建嵌入模型或知识库失败: {str(e)}. 请确保已安装sentence-transformers和faiss-cpu库")except Exception as e:logger.error(f"处理文本时出错: {str(e)}")raisedef _assign_page_numbers_to_chunks(self, text: str, chunks: List[str], page_numbers: List[int]) -> Dict[str, int]:"""为每个文本块分配页码参数:text: 完整文本chunks: 分割后的文本块page_numbers: 页码列表返回:page_info: 文本块到页码的映射字典"""page_info = {}try:# 如果没有页码信息，则返回空字典if not page_numbers:logger.warning("没有页码信息，无法分配页码")return page_info# 简化页码分配算法，使用文本位置比例来估计页码text_length = len(text)for chunk in chunks:chunk_position = text.find(chunk)if chunk_position == -1:# 如果找不到，使用最后一页page_info[chunk] = page_numbers[-1]else:# 根据文本位置比例估计页码position_ratio = chunk_position / text_lengthpage_index = min(int(position_ratio * len(page_numbers)), len(page_numbers) - 1)page_info[chunk] = page_numbers[page_index]return page_infoexcept Exception as e:logger.error(f"分配页码时出错: {str(e)}")# 出错时返回空字典，但不中断程序return {}def answer_question(self, query: str, top_k: int = 4) -> Dict:"""回答关于PDF内容的问题参数:query: 问题top_k: 检索的相关文档数量返回:result: 包含回答和来源页码的字典"""if not self.knowledge_base:raise ValueError("知识库尚未初始化，请先调用load_pdf()")logger.info(f"处理问题: {query}")try:# 执行相似度搜索，找到与查询相关的文档docs = self.knowledge_base.similarity_search(query, k=top_k)llm = ChatOpenAI(openai_api_key= "sk-xxxxxx",  # 从环境变量获取 API Keyopenai_api_base="https://dashscope.aliyuncs.com/compatible-mode/v1",  # 百炼兼容端点model_name="qwen-plus",  # 指定阿里模型temperature=0.2  # 控制生成随机性)# llm = Ollama(#     model="deepseek-r1:1.5b",  # 直接指定模型名称#     base_url="http://localhost:11434",  # Ollama 服务地址#     temperature=0.2  # 控制随机性# )# 加载问答链chain = load_qa_chain(llm, chain_type="stuff")# 准备输入数据input_data = {"input_documents": docs, "question": query}# 使用回调函数跟踪API调用成本with get_openai_callback() as cost:# 执行问答链response = chain.invoke(input=input_data)logger.info(f"查询已处理。成本: {cost}")# 记录唯一的页码unique_pages = set()page_sources = []# 显示每个文档块的来源页码for i, doc in enumerate(docs):text_content = getattr(doc, "page_content", "")# 尝试直接匹配source_page = self.knowledge_base.page_info.get(text_content, None)# 如果直接匹配失败，尝试去除空白后匹配if source_page is None:for chunk, page in self.knowledge_base.page_info.items():if chunk.strip() == text_content.strip():source_page = pagebreak# 如果仍然找不到匹配，使用默认值if source_page is None:source_page = "未知"if source_page not in unique_pages:unique_pages.add(source_page)page_sources.append(source_page)# 按页码排序try:page_sources = sorted([p for p in page_sources if isinstance(p, int)])except Exception:# 如果排序失败，使用原始列表passreturn {"answer": response["output_text"],"sources": page_sources,"cost": str(cost)}except Exception as e:logger.error(f"回答问题时出错: {str(e)}")raisedef main():"""主函数，处理命令行参数并运行PDF聊天机器人"""parser = argparse.ArgumentParser(description="PDF聊天机器人 - 使用FAISS向量数据库回答PDF文档相关问题")parser.add_argument("--pdf", type=str, default="./xxx.pdf", help="PDF文件路径")parser.add_argument("--api_key", type=str, default="sk-xxxxxx", help="OpenAI API密钥")parser.add_argument("--base_url", type=str, default="http://localhost:11434/v1", help="API基础URL")parser.add_argument("--model_id", type=str, default="iic/nlp_gte_sentence-embedding_chinese-base", help="ModelScope模型ID")parser.add_argument("--device", type=str, help="计算设备，如 'cuda:0', 'cpu' 等")parser.add_argument("--no_query_prefix", action="store_true", help="不为查询添加前缀")parser.add_argument("--llm_model", type=str, default="deepseek-r1:1.5b", help="LLM模型名称")args = parser.parse_args()try:# 检查PDF文件是否存在if not os.path.exists(args.pdf):print(f"错误: PDF文件不存在: {args.pdf}")print("请提供有效的PDF文件路径")sys.exit(1)# 创建PDF聊天机器人chatbot = PDFChatBot(args.pdf, args.api_key, args.base_url, args.model_id,args.device,not args.no_query_prefix,args.llm_model)# 加载PDF文件print(f"正在加载PDF文件: {args.pdf}...")try:chatbot.load_pdf()except Exception as e:print(f"加载PDF文件时出错: {str(e)}")print("请确保PDF文件格式正确且可读取")sys.exit(1)print(f"\n成功加载PDF文件: {args.pdf}")print(f"PDF共有 {chatbot.page_count} 页")print("\n现在您可以开始提问了! 输入'退出'或'exit'结束对话。\n")# 交互式问答循环while True:try:query = input("\n请输入您的问题: ")if query.lower() in ["退出", "exit", "quit", "q"]:print("感谢使用PDF聊天机器人，再见!")breakif not query.strip():continuetry:print("正在思考中...")# 获取回答result = chatbot.answer_question(query)# 打印回答print("\n回答:")print(result["answer"])# 打印来源if result["sources"]:print("\n来源页码:", ", ".join(map(str, result["sources"])))# 打印成本print(f"\n查询成本: {result['cost']}")except Exception as e:print(f"处理问题时出错: {str(e)}")print("请尝试重新提问或检查API密钥是否有效")except KeyboardInterrupt:print("\n程序被用户中断。感谢使用PDF聊天机器人，再见!")breakexcept Exception as e:print(f"发生错误: {str(e)}")except Exception as e:print(f"程序运行出错: {str(e)}")if __name__ == "__main__":try:main()except Exception as e:print(f"程序启动失败: {str(e)}")print("请检查环境配置和依赖库安装情况")