基于LangChain+LLM的本地知识库问答：从企业单文档问答到批量文档问答_langchain 本地知识库

​初次接触的朋友一看这么多组件可能直接晕了(封装的东西非常多，感觉它想把LLM所需要用到的功能/工具都封装起来)，为方便理解，我们可以先从大的层面把整个langchain库划分为三个大层：基础层、能力层、应用层

import os
os.environ["OPENAI_API_KEY"] = '你的api key'
from langchain.llms import OpenAI
 
llm = OpenAI(model_name="text-davinci-003",max_tokens=1024)
llm("怎么评价人工智能")

        model_token_mapping = {
            "gpt-4": 8192,
            "gpt-4-0314": 8192,
            "gpt-4-0613": 8192,
            "gpt-4-32k": 32768,
            "gpt-4-32k-0314": 32768,
            "gpt-4-32k-0613": 32768,
            "gpt-3.5-turbo": 4096,
            "gpt-3.5-turbo-0301": 4096,
            "gpt-3.5-turbo-0613": 4096,
            "gpt-3.5-turbo-16k": 16385,
            "gpt-3.5-turbo-16k-0613": 16385,
            "text-ada-001": 2049,
            "ada": 2049,
            "text-babbage-001": 2040,
            "babbage": 2049,
            "text-curie-001": 2049,
            "curie": 2049,
            "davinci": 2049,
            "text-davinci-003": 4097,
            "text-davinci-002": 4097,
            "code-davinci-002": 8001,
            "code-davinci-001": 8001,
            "code-cushman-002": 2048,
            "code-cushman-001": 2048,
        }

# 默认的文本分割器函数
def _get_default_text_splitter() -> TextSplitter:
    return RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=0)

"""Graph Index Creator."""                     # 定义"图索引创建器"的描述
 
# 导入相关的模块和类型定义
from typing import Optional, Type              # 导入可选类型和类型的基础类型
from langchain import BasePromptTemplate       # 导入基础提示模板
from langchain.chains.llm import LLMChain      # 导入LLM链
from langchain.graphs.networkx_graph import NetworkxEntityGraph, parse_triples  # 导入Networkx实体图和解析三元组的功能
from langchain.indexes.prompts.knowledge_triplet_extraction import (  # 从知识三元组提取模块导入对应的提示
    KNOWLEDGE_TRIPLE_EXTRACTION_PROMPT,
)
from langchain.pydantic_v1 import BaseModel                      # 导入基础模型
from langchain.schema.language_model import BaseLanguageModel    # 导入基础语言模型的定义
 
class GraphIndexCreator(BaseModel):  # 定义图索引创建器类，继承自BaseModel
    """Functionality to create graph index."""   # 描述该类的功能为"创建图索引"
 
    llm: Optional[BaseLanguageModel] = None      # 定义可选的语言模型属性，默认为None
    graph_type: Type[NetworkxEntityGraph] = NetworkxEntityGraph  # 定义图的类型，默认为NetworkxEntityGraph
 
    def from_text(
        self, text: str, prompt: BasePromptTemplate = KNOWLEDGE_TRIPLE_EXTRACTION_PROMPT
    ) -> NetworkxEntityGraph:                  # 定义一个方法，从文本中创建图索引
        """Create graph index from text."""    # 描述该方法的功能
        if self.llm is None:                   # 如果语言模型为None，则抛出异常
            raise ValueError("llm should not be None")
        graph = self.graph_type()  # 创建一个新的图
        chain = LLMChain(llm=self.llm, prompt=prompt)  # 使用当前的语言模型和提示创建一个LLM链
        output = chain.predict(text=text)      # 使用LLM链对文本进行预测
        knowledge = parse_triples(output)      # 解析预测输出得到的三元组
        for triple in knowledge:               # 遍历所有的三元组
            graph.add_triple(triple)           # 将三元组添加到图中
        return graph  # 返回创建的图
 
    async def afrom_text(             # 定义一个异步版本的from_text方法
        self, text: str, prompt: BasePromptTemplate = KNOWLEDGE_TRIPLE_EXTRACTION_PROMPT
    ) -> NetworkxEntityGraph:
        """Create graph index from text asynchronously."""  # 描述该异步方法的功能
        if self.llm is None:          # 如果语言模型为None，则抛出异常
            raise ValueError("llm should not be None")
        graph = self.graph_type()     # 创建一个新的图
        chain = LLMChain(llm=self.llm, prompt=prompt)       # 使用当前的语言模型和提示创建一个LLM链
        output = await chain.apredict(text=text)   # 异步使用LLM链对文本进行预测
        knowledge = parse_triples(output)          # 解析预测输出得到的三元组
        for triple in knowledge:                   # 遍历所有的三元组
            graph.add_triple(triple)               # 将三元组添加到图中
        return graph                               # 返回创建的图

# 定义知识图谱对话记忆类
class ConversationKGMemory(BaseChatMemory):
    """知识图谱对话记忆类
    在对话中与外部知识图谱集成，存储和检索对话中的知识三元组信息。
    """
 
    k: int = 2  # 考虑的上下文对话数量
    human_prefix: str = "Human"  # 人类前缀
    ai_prefix: str = "AI"  # AI前缀
    kg: NetworkxEntityGraph = Field(default_factory=NetworkxEntityGraph)  # 知识图谱实例
    knowledge_extraction_prompt: BasePromptTemplate = KNOWLEDGE_TRIPLE_EXTRACTION_PROMPT          # 知识提取提示
    entity_extraction_prompt: BasePromptTemplate = ENTITY_EXTRACTION_PROMPT  # 实体提取提示
    llm: BaseLanguageModel                  # 基础语言模型
    summary_message_cls: Type[BaseMessage] = SystemMessage  # 总结消息类
    memory_key: str = "history"             # 历史记忆键
 
    def load_memory_variables(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        """返回历史缓冲区。"""
        entities = self._get_current_entities(inputs)  # 获取当前实体
 
        summary_strings = []
        for entity in entities:  # 对于每个实体
            knowledge = self.kg.get_entity_knowledge(entity)      # 获取与实体相关的知识
            if knowledge:
                summary = f"On {entity}: {'. '.join(knowledge)}."  # 构建总结字符串
                summary_strings.append(summary)
        context: Union[str, List]
        if not summary_strings:
            context = [] if self.return_messages else ""
        elif self.return_messages:
            context = [
                self.summary_message_cls(content=text) for text in summary_strings
            ]
        else:
            context = "\n".join(summary_strings)
 
        return {self.memory_key: context}
 
    @property
    def memory_variables(self) -> List[str]:
        """始终返回记忆变量列表。"""
        return [self.memory_key]
 
    def _get_prompt_input_key(self, inputs: Dict[str, Any]) -> str:
        """获取提示的输入键。"""
        if self.input_key is None:
            return get_prompt_input_key(inputs, self.memory_variables)
        return self.input_key
 
    def _get_prompt_output_key(self, outputs: Dict[str, Any]) -> str:
        """获取提示的输出键。"""
        if self.output_key is None:
            if len(outputs) != 1:
                raise ValueError(f"One output key expected, got {outputs.keys()}")
            return list(outputs.keys())[0]
        return self.output_key
 
    def get_current_entities(self, input_string: str) -> List[str]:
        """从输入字符串中获取当前实体。"""
        chain = LLMChain(llm=self.llm, prompt=self.entity_extraction_prompt)
        buffer_string = get_buffer_string(
            self.chat_memory.messages[-self.k * 2 :],
            human_prefix=self.human_prefix,
            ai_prefix=self.ai_prefix,
        )
        output = chain.predict(
            history=buffer_string,
            input=input_string,
        )
        return get_entities(output)
 
    def _get_current_entities(self, inputs: Dict[str, Any]) -> List[str]:
        """获取对话中的当前实体。"""
        prompt_input_key = self._get_prompt_input_key(inputs)
        return self.get_current_entities(inputs[prompt_input_key])
 
    def get_knowledge_triplets(self, input_string: str) -> List[KnowledgeTriple]:
        """从输入字符串中获取知识三元组。"""
        chain = LLMChain(llm=self.llm, prompt=self.knowledge_extraction_prompt)
        buffer_string = get_buffer_string(
            self.chat_memory.messages[-self.k * 2 :],
            human_prefix=self.human_prefix,
            ai_prefix=self.ai_prefix,
        )
        output = chain.predict(
            history=buffer_string,
            input=input_string,
            verbose=True,
        )
        knowledge = parse_triples(output)  # 解析三元组
        return knowledge
 
    def _get_and_update_kg(self, inputs: Dict[str, Any]) -> None:
        """从对话历史中获取并更新知识图谱。"""
        prompt_input_key = self._get_prompt_input_key(inputs)
        knowledge = self.get_knowledge_triplets(inputs[prompt_input_key])
        for triple in knowledge:
            self.kg.add_triple(triple)  # 向知识图谱中添加三元组
 
    def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str]) -> None:
        """将此对话的上下文保存到缓冲区。"""
        super().save_context(inputs, outputs)
        self._get_and_update_kg(inputs)
 
    def clear(self) -> None:
        """清除记忆内容。"""
        super().clear()
        self.kg.clear()  # 清除知识图谱内容

# pip install wikipedia
from langchain.agents import load_tools
from langchain.agents import initialize_agent
from langchain.agents import AgentType
 
tools = load_tools(["wikipedia", "llm-math"], llm=llm)
agent = initialize_agent(tools, 
                         llm, 
                         agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, 
                         verbose=True)
 
 
agent.run("奥巴马的生日是哪天? 到2023年他多少岁了?")

dim, measure = 64, faiss.METRIC_L2
param = 'Flat'
index = faiss.index_factory(dim, param, measure)

dim, measure = 64, faiss.METRIC_L2
param = 'Flat'
index = faiss.index_factory(dim, param, measure)
index.is_trained                                   # 输出为True
index.add(xb)                                      # 向index中添加向量

dim, measure = 64, faiss.METRIC_L2 
param = 'IVF100,Flat'                           # 代表k-means聚类中心为100,   
index = faiss.index_factory(dim, param, measure)
print(index.is_trained)                          # 此时输出为False，因为倒排索引需要训练k-means，
index.train(xb)                                  # 因此需要先训练index，再add向量
index.add(xb)                                     

dim, measure = 64, faiss.METRIC_L2 
param =  'PQ16' 
index = faiss.index_factory(dim, param, measure)
print(index.is_trained)                          # 此时输出为False，因为倒排索引需要训练k-means，
index.train(xb)                                  # 因此需要先训练index，再add向量
index.add(xb)          

dim, measure = 64, faiss.METRIC_L2  
param =  'IVF100,PQ16'
index = faiss.index_factory(dim, param, measure) 
print(index.is_trained)                          # 此时输出为False，因为倒排索引需要训练k-means， 
index.train(xb)                                  # 因此需要先训练index，再add向量 index.add(xb)       

dim, measure = 64, faiss.METRIC_L2  
param =  'LSH'
index = faiss.index_factory(dim, param, measure) 
print(index.is_trained)                          # 此时输出为True
index.add(xb)       

dim, measure = 64, faiss.METRIC_L2   
param =  'HNSW64' 
index = faiss.index_factory(dim, param, measure)  
print(index.is_trained)                          # 此时输出为True 
index.add(xb)

​占用显存情况：大约15个G

​显存占用约13G

​你会发现该项目主要由以下各大模块组成(注意，该项目的最新版已经变化很大，本第三部分可以认为是针对v0.1.14左右的版本，往上最多到v0.1.16，新版对很多功能做了更高的封装，而从原理理解的角度来说，看老版 更好理解些)

from langchain.agents import Tool          # 导入工具模块
from langchain.tools import BaseTool       # 导入基础工具类
from langchain import PromptTemplate, LLMChain      # 导入提示模板和语言模型链
from agent.custom_search import DeepSearch          # 导入自定义搜索模块
 
# 导入基础单动作代理，输出解析器，语言模型单动作代理和代理执行器
from langchain.agents import BaseSingleActionAgent, AgentOutputParser, LLMSingleActionAgent, AgentExecutor    
from typing import List, Tuple, Any, Union, Optional, Type      # 导入类型注释模块
from langchain.schema import AgentAction, AgentFinish           # 导入代理动作和代理完成模式
from langchain.prompts import StringPromptTemplate          # 导入字符串提示模板
from langchain.callbacks.manager import CallbackManagerForToolRun      # 导入工具运行回调管理器
from langchain.base_language import BaseLanguageModel       # 导入基础语言模型
import re                                                   # 导入正则表达式模块
 
# 定义一个代理模板字符串
agent_template = """
你现在是一个{role}。这里是一些已知信息：
{related_content}
{background_infomation}
{question_guide}：{input}
{answer_format}
"""
 
# 定义一个自定义提示模板类，继承自字符串提示模板
class CustomPromptTemplate(StringPromptTemplate):
    template: str          # 提示模板字符串
    tools: List[Tool]      # 工具列表
 
    # 定义一个格式化函数，根据提供的参数生成最终的提示模板
    def format(self, **kwargs) -> str:
        intermediate_steps = kwargs.pop("intermediate_steps")
        # 判断是否有互联网查询信息
        if len(intermediate_steps) == 0:
            # 如果没有，则给出默认的背景信息，角色，问题指导和回答格式
            background_infomation = "\n"
            role = "傻瓜机器人"
            question_guide = "我现在有一个问题"
            answer_format = "如果你知道答案，请直接给出你的回答！如果你不知道答案，请你只回答\"DeepSearch('搜索词')\"，并将'搜索词'替换为你认为需要搜索的关键词，除此之外不要回答其他任何内容。\n\n下面请回答我上面提出的问题！"
 
        else:
            # 否则，根据 intermediate_steps 中的 AgentAction 拼装 background_infomation
            background_infomation = "\n\n你还有这些已知信息作为参考：\n\n"
            action, observation = intermediate_steps[0]
            background_infomation += f"{observation}\n"
            role = "聪明的 AI 助手"
            question_guide = "请根据这些已知信息回答我的问题"
            answer_format = ""
 
        kwargs["background_infomation"] = background_infomation
        kwargs["role"] = role
        kwargs["question_guide"] = question_guide
        kwargs["answer_format"] = answer_format
        return self.template.format(**kwargs)  # 格式化模板并返回
 
# 定义一个自定义搜索工具类，继承自基础工具类
class CustomSearchTool(BaseTool):
    name: str = "DeepSearch"           # 工具名称
    description: str = ""              # 工具描述
 
    # 定义一个运行函数，接受一个查询字符串和一个可选的回调管理器作为参数，返回DeepSearch的搜索结果
    def _run(self, query: str, run_manager: Optional[CallbackManagerForToolRun] = None):
        return DeepSearch.search(query = query)
 
    # 定义一个异步运行函数，但由于DeepSearch不支持异步，所以直接抛出一个未实现错误
    async def _arun(self, query: str):
        raise NotImplementedError("DeepSearch does not support async")
 
# 定义一个自定义代理类，继承自基础单动作代理
class CustomAgent(BaseSingleActionAgent):
    # 定义一个输入键的属性
    @property
    def input_keys(self):
        return ["input"]
 
    # 定义一个计划函数，接受一组中间步骤和其他参数，返回一个代理动作或者代理完成
    def plan(self, intermedate_steps: List[Tuple[AgentAction, str]],
            **kwargs: Any) -> Union[AgentAction, AgentFinish]:
        return AgentAction(tool="DeepSearch", tool_input=kwargs["input"], log="")
 
# 定义一个自定义输出解析器，继承自代理输出解析器
class CustomOutputParser(AgentOutputParser):
    # 定义一个解析函数，接受一个语言模型的输出字符串，返回一个代理动作或者代理完成
    def parse(self, llm_output: str) -> Union[AgentAction, AgentFinish]:
        # 使用正则表达式匹配输出字符串，group1是调用函数名字，group2是传入参数
        match = re.match(r'^[\s\w]*(DeepSearch)\(([^\)]+)\)', llm_output, re.DOTALL)
        print(match)
 
        # 如果语言模型没有返回 DeepSearch() 则认为直接结束指令
        if not match:
            return AgentFinish(
                return_values={"output": llm_output.strip()},
                log=llm_output,
            )
        # 否则的话都认为需要调用 Tool
        else:
            action = match.group(1).strip()
            action_input = match.group(2).strip()
            return AgentAction(tool=action, tool_input=action_input.strip(" ").strip('"'), log=llm_output)
 
 
# 定义一个深度代理类
class DeepAgent:
    tool_name: str = "DeepSearch"       # 工具名称
    agent_executor: any                 # 代理执行器
    tools: List[Tool]                   # 工具列表
    llm_chain: any                      # 语言模型链
 
    # 定义一个查询函数，接受一个相关内容字符串和一个查询字符串，返回执行器的运行结果
    def query(self, related_content: str = "", query: str = ""):
        tool_name =这段代码的主要目的是建立一个深度搜索的AI代理。AI代理首先通过接收一个问题输入，然后根据输入生成一个提示模板，然后通过该模板引导AI生成回答或进行更深入的搜索。现在，我将继续为剩余的代码添加中文注释
 
```python
        self.tool_name
        result = self.agent_executor.run(related_content=related_content, input=query ,tool_name=self.tool_name)
        return result       # 返回执行器的运行结果
 
    # 在初始化函数中，首先从DeepSearch工具创建一个工具实例，并添加到工具列表中
    def __init__(self, llm: BaseLanguageModel, **kwargs):
        tools = [
                    Tool.from_function(
                        func=DeepSearch.search,
                        name="DeepSearch",
                        description=""
                    )
                ]
        self.tools = tools      # 保存工具列表
        tool_names = [tool.name for tool in tools]    # 提取工具列表中的工具名称
        output_parser = CustomOutputParser()          # 创建一个自定义输出解析器实例
        # 创建一个自定义提示模板实例
        prompt = CustomPromptTemplate(template=agent_template,
                                      tools=tools,
                                      input_variables=["related_content","tool_name", "input", "intermediate_steps"])
        # 创建一个语言模型链实例
        llm_chain = LLMChain(llm=llm, prompt=prompt)
        self.llm_chain = llm_chain      # 保存语言模型链实例
 
        # 创建一个语言模型单动作代理实例
        agent = LLMSingleActionAgent(
            llm_chain=llm_chain,
            output_parser=output_parser,
            stop=["\nObservation:"],
            allowed_tools=tool_names
        )
 
        # 创建一个代理执行器实例
        agent_executor = AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, verbose=True)
        self.agent_executor = agent_executor         # 保存代理执行器实例

#coding=utf8
# 声明文件编码格式为 utf8
 
from langchain.utilities import BingSearchAPIWrapper
# 导入 BingSearchAPIWrapper 类，这个类用于与 Bing 搜索 API 进行交互
 
from configs.model_config import BING_SEARCH_URL, BING_SUBSCRIPTION_KEY
# 导入配置文件中的 Bing 搜索 URL 和 Bing 订阅密钥
 
def bing_search(text, result_len=3):
    # 定义一个名为 bing_search 的函数，该函数接收一个文本和结果长度的参数，默认结果长度为3
 
    if not (BING_SEARCH_URL and BING_SUBSCRIPTION_KEY):
        # 如果 Bing 搜索 URL 或 Bing 订阅密钥未设置，则返回一个错误信息的文档
        return [{"snippet": "please set BING_SUBSCRIPTION_KEY and BING_SEARCH_URL in os ENV",
                 "title": "env inof not fould",
                 "link": "https://python.langchain.com/en/latest/modules/agents/tools/examples/bing_search.html"}]
 
    search = BingSearchAPIWrapper(bing_subscription_key=BING_SUBSCRIPTION_KEY,
                                  bing_search_url=BING_SEARCH_URL)
    # 创建 BingSearchAPIWrapper 类的实例，该实例用于与 Bing 搜索 API 进行交互
 
    return search.results(text, result_len)
    # 返回搜索结果，结果的数量由 result_len 参数决定
 
if __name__ == "__main__":
    # 如果这个文件被直接运行，而不是被导入作为模块，那么就执行以下代码
 
    r = bing_search('python')
    # 使用 Bing 搜索 API 来搜索 "python" 这个词，并将结果保存在变量 r 中
 
    print(r)
    # 打印出搜索结果

from abc import ABC  # 导入抽象基类
from langchain.llms.base import LLM           # 导入语言学习模型基类
from typing import Optional, List             # 导入类型标注模块
from models.loader import LoaderCheckPoint    # 导入模型加载点
from models.base import (BaseAnswer,          # 导入基本回答模型
                         AnswerResult)        # 导入回答结果模型
 
 
class ChatGLM(BaseAnswer, LLM, ABC):  # 定义ChatGLM类，继承基础回答、语言学习模型和抽象基类
    max_token: int = 10000          # 最大的token数
    temperature: float = 0.01       # 温度参数，用于控制生成文本的随机性
    top_p = 0.9  # 排序前0.9的token会被保留
    checkPoint: LoaderCheckPoint = None  # 检查点模型
    # history = []          # 历史记录
    history_len: int = 10   # 历史记录长度
 
    def __init__(self, checkPoint: LoaderCheckPoint = None):  # 初始化方法
        super().__init__()  # 调用父类的初始化方法
        self.checkPoint = checkPoint  # 赋值检查点模型
 
    @property
    def _llm_type(self) -> str:  # 定义只读属性_llm_type，返回语言学习模型的类型
        return "ChatGLM"
 
    @property
    def _check_point(self) -> LoaderCheckPoint:  # 定义只读属性_check_point，返回检查点模型
        return self.checkPoint
 
    @property
    def _history_len(self) -> int:  # 定义只读属性_history_len，返回历史记录的长度
        return self.history_len
 
    def set_history_len(self, history_len: int = 10) -> None:  # 设置历史记录长度
        self.history_len = history_len
 
    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:  # 定义_call方法，实现模型的具体调用
        print(f"__call:{prompt}")  # 打印调用的提示信息
        response, _ = self.checkPoint.model.chat(  # 调用模型的chat方法，获取回答和其他信息
            self.checkPoint.tokenizer,  # 使用的分词器
            prompt,  # 提示信息
            history=[],  # 历史记录
            max_length=self.max_token,      # 最大长度
            temperature=self.temperature    # 温度参数
        )
        print(f"response:{response}")  # 打印回答信息
        print(f"+++++++++++++++++++++++++++++++++++")  # 打印分隔线
        return response  # 返回回答
 
    def generatorAnswer(self, prompt: str,
                         history: List[List[str]] = [],
                         streaming: bool = False):  # 定义生成回答的方法，可以处理流式输入
 
        if streaming:  # 如果是流式输入
            history += [[]]  # 在历史记录中添加新的空列表
            for inum, (stream_resp, _) in enumerate(self.checkPoint.model.stream_chat(  # 对模型的stream_chat方法返回的结果进行枚举
                    self.checkPoint.tokenizer,  # 使用的分词器
                    prompt,  # 提示信息
                    history=history[-self.history_len:-1] if self.history_len > 1 else [],  # 使用的历史记录
                    max_length=self.max_token,  # 最大长度
                    temperature=self.temperature  # 温度参数
            )):
                # self.checkPoint.clear_torch_cache()  # 清空缓存
                history[-1] = [prompt, stream_resp]  # 更新最后一个历史记录
                answer_result = AnswerResult()  # 创建回答结果对象
                answer_result.history = history  # 更新回答结果的历史记录
                answer_result.llm_output = {"answer": stream_resp}  # 更新回答结果的输出
                yield answer_result  # 生成回答结果
        else:  # 如果不是流式输入
            response, _ = self.checkPoint.model.chat(  # 调用模型的chat方法，获取回答和其他信息
                self.checkPoint.tokenizer,  # 使用的分词器
                prompt,  # 提示信息
                history=history[-self.history_len:] if self.history_len > 0 else [],  # 使用的历史记录
                max_length=self.max_token,  # 最大长度
                temperature=self.temperature  # 温度参数
            )
            self.checkPoint.clear_torch_cache()  # 清空缓存
            history += [[prompt, response]]  # 更新历史记录
            answer_result = AnswerResult()  # 创建回答结果对象
            answer_result.history = history  # 更新回答结果的历史记录
            answer_result.llm_output = {"answer": response}  # 更新回答结果的输出
            yield answer_result  # 生成回答结果

import sys      # 导入sys模块，通常用于与Python解释器进行交互
from typing import Any      # 从typing模块导入Any，用于表示任何类型
 
# 从models.loader.args模块导入parser，可能是解析命令行参数用
from models.loader.args import parser       
# 从models.loader模块导入LoaderCheckPoint，可能是模型加载点
from models.loader import LoaderCheckPoint  
 
# 从configs.model_config模块导入llm_model_dict和LLM_MODEL
from configs.model_config import (llm_model_dict, LLM_MODEL)  
# 从models.base模块导入BaseAnswer，即模型的基础类
from models.base import BaseAnswer  
 
# 定义一个名为loaderCheckPoint的变量，类型为LoaderCheckPoint，并初始化为None
loaderCheckPoint: LoaderCheckPoint = None  
 
 
def loaderLLM(llm_model: str = None, no_remote_model: bool = False, use_ptuning_v2: bool = False) -> Any:
    """
    初始化 llm_model_ins LLM
    :param llm_model: 模型名称
    :param no_remote_model: 是否使用远程模型，如果需要加载本地模型，则添加 `--no-remote-model
    :param use_ptuning_v2: 是否使用 p-tuning-v2 PrefixEncoder
    :return:
    """
    pre_model_name = loaderCheckPoint.model_name      # 获取loaderCheckPoint的模型名称
    llm_model_info = llm_model_dict[pre_model_name]   # 从模型字典中获取模型信息
 
    if no_remote_model:      # 如果不使用远程模型
        loaderCheckPoint.no_remote_model = no_remote_model  # 将loaderCheckPoint的no_remote_model设置为True
    if use_ptuning_v2:       # 如果使用p-tuning-v2
        loaderCheckPoint.use_ptuning_v2 = use_ptuning_v2    # 将loaderCheckPoint的use_ptuning_v2设置为True
 
    if llm_model:            # 如果指定了模型名称
        llm_model_info = llm_model_dict[llm_model]  # 从模型字典中获取指定的模型信息
 
    if loaderCheckPoint.no_remote_model:  # 如果不使用远程模型
        loaderCheckPoint.model_name = llm_model_info['name']  # 将loaderCheckPoint的模型名称设置为模型信息中的name
    else:  # 如果使用远程模型
        loaderCheckPoint.model_name = llm_model_info['pretrained_model_name']  # 将loaderCheckPoint的模型名称设置为模型信息中的pretrained_model_name
 
    loaderCheckPoint.model_path = llm_model_info["local_model_path"]  # 设置模型的本地路径
 
    if 'FastChatOpenAILLM' in llm_model_info["provides"]:  # 如果模型信息中的provides包含'FastChatOpenAILLM'
        loaderCheckPoint.unload_model()  # 卸载模型
    else:  # 如果不包含
        loaderCheckPoint.reload_model()  # 重新加载模型
 
    provides_class = getattr(sys.modules['models'], llm_model_info['provides'])  # 获取模型类
    modelInsLLM = provides_class(checkPoint=loaderCheckPoint)  # 创建模型实例
    if 'FastChatOpenAILLM' in llm_model_info["provides"]:      # 如果模型信息中的provides包含'FastChatOpenAILLM'
        modelInsLLM.set_api_base_url(llm_model_info['api_base_url'])  # 设置API基础URL
        modelInsLLM.call_model_name(llm_model_info['name'])    # 设置模型名称
    return modelInsLLM  # 返回模型实例

import torch.cuda
import torch.backends
import os
import logging
import uuid
 
LOG_FORMAT = "%(levelname) -5s %(asctime)s" "-1d: %(message)s"
logger = logging.getLogger()
logger.setLevel(logging.INFO)
logging.basicConfig(format=LOG_FORMAT)
 
# 在以下字典中修改属性值，以指定本地embedding模型存储位置
# 如将 "text2vec": "GanymedeNil/text2vec-large-chinese" 修改为 "text2vec": "User/Downloads/text2vec-large-chinese"
# 此处请写绝对路径
embedding_model_dict = {
    "ernie-tiny": "nghuyong/ernie-3.0-nano-zh",
    "ernie-base": "nghuyong/ernie-3.0-base-zh",
    "text2vec-base": "shibing624/text2vec-base-chinese",
    "text2vec": "GanymedeNil/text2vec-large-chinese",
    "m3e-small": "moka-ai/m3e-small",
    "m3e-base": "moka-ai/m3e-base",
}
 
# Embedding model name
EMBEDDING_MODEL = "text2vec"
 
# Embedding running device
EMBEDDING_DEVICE = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
 
 
# supported LLM models
# llm_model_dict 处理了loader的一些预设行为，如加载位置，模型名称，模型处理器实例
# 在以下字典中修改属性值，以指定本地 LLM 模型存储位置
# 如将 "chatglm-6b" 的 "local_model_path" 由 None 修改为 "User/Downloads/chatglm-6b"
# 此处请写绝对路径
llm_model_dict = {
    "chatglm-6b-int4-qe": {
        "name": "chatglm-6b-int4-qe",
        "pretrained_model_name": "THUDM/chatglm-6b-int4-qe",
        "local_model_path": None,
        "provides": "ChatGLM"
    },
    "chatglm-6b-int4": {
        "name": "chatglm-6b-int4",
        "pretrained_model_name": "THUDM/chatglm-6b-int4",
        "local_model_path": None,
        "provides": "ChatGLM"
    },
    "chatglm-6b-int8": {
        "name": "chatglm-6b-int8",
        "pretrained_model_name": "THUDM/chatglm-6b-int8",
        "local_model_path": None,
        "provides": "ChatGLM"
    },
    "chatglm-6b": {
        "name": "chatglm-6b",
        "pretrained_model_name": "THUDM/chatglm-6b",
        "local_model_path": None,
        "provides": "ChatGLM"
    },
    "chatglm2-6b": {
        "name": "chatglm2-6b",
        "pretrained_model_name": "THUDM/chatglm2-6b",
        "local_model_path": None,
        "provides": "ChatGLM"
    },
    "chatglm2-6b-int4": {
        "name": "chatglm2-6b-int4",
        "pretrained_model_name": "THUDM/chatglm2-6b-int4",
        "local_model_path": None,
        "provides": "ChatGLM"
    },
    "chatglm2-6b-int8": {
        "name": "chatglm2-6b-int8",
        "pretrained_model_name": "THUDM/chatglm2-6b-int8",
        "local_model_path": None,
        "provides": "ChatGLM"
    },
    "chatyuan": {
        "name": "chatyuan",
        "pretrained_model_name": "ClueAI/ChatYuan-large-v2",
        "local_model_path": None,
        "provides": None
    },
    "moss": {
        "name": "moss",
        "pretrained_model_name": "fnlp/moss-moon-003-sft",
        "local_model_path": None,
        "provides": "MOSSLLM"
    },
    "vicuna-13b-hf": {
        "name": "vicuna-13b-hf",
        "pretrained_model_name": "vicuna-13b-hf",
        "local_model_path": None,
        "provides": "LLamaLLM"
    },
 
    # 通过 fastchat 调用的模型请参考如下格式
    "fastchat-chatglm-6b": {
        "name": "chatglm-6b",             # "name"修改为fastchat服务中的"model_name"
        "pretrained_model_name": "chatglm-6b",
        "local_model_path": None,
        "provides": "FastChatOpenAILLM",  # 使用fastchat api时，需保证"provides"为"FastChatOpenAILLM"
        "api_base_url": "http://localhost:8000/v1"  # "name"修改为fastchat服务中的"api_base_url"
    },
    "fastchat-chatglm2-6b": {
        "name": "chatglm2-6b",              # "name"修改为fastchat服务中的"model_name"
        "pretrained_model_name": "chatglm2-6b",
        "local_model_path": None,
        "provides": "FastChatOpenAILLM",    # 使用fastchat api时，需保证"provides"为"FastChatOpenAILLM"
        "api_base_url": "http://localhost:8000/v1"  # "name"修改为fastchat服务中的"api_base_url"
    },
 
    # 通过 fastchat 调用的模型请参考如下格式
    "fastchat-vicuna-13b-hf": {
        "name": "vicuna-13b-hf",          # "name"修改为fastchat服务中的"model_name"
        "pretrained_model_name": "vicuna-13b-hf",
        "local_model_path": None,
        "provides": "FastChatOpenAILLM",  # 使用fastchat api时，需保证"provides"为"FastChatOpenAILLM"
        "api_base_url": "http://localhost:8000/v1"  # "name"修改为fastchat服务中的"api_base_url"
    },
}
 
# LLM 名称
LLM_MODEL = "chatglm-6b"
# 量化加载8bit 模型
LOAD_IN_8BIT = False
# Load the model with bfloat16 precision. Requires NVIDIA Ampere GPU.
BF16 = False
# 本地lora存放的位置
LORA_DIR = "loras/"
 
# LLM lora path，默认为空，如果有请直接指定文件夹路径
LLM_LORA_PATH = ""
USE_LORA = True if LLM_LORA_PATH else False
 
# LLM streaming reponse
STREAMING = True
 
# Use p-tuning-v2 PrefixEncoder
USE_PTUNING_V2 = False
 
# LLM running device
LLM_DEVICE = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
 
# 知识库默认存储路径
KB_ROOT_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), "knowledge_base")
 
# 基于上下文的prompt模版，请务必保留"{question}"和"{context}"
PROMPT_TEMPLATE = """已知信息：
{context} 
根据上述已知信息，简洁和专业的来回答用户的问题。如果无法从中得到答案，请说 “根据已知信息无法回答该问题” 或 “没有提供足够的相关信息”，不允许在答案中添加编造成分，答案请使用中文。 问题是：{question}"""
 
# 缓存知识库数量,如果是ChatGLM2,ChatGLM2-int4,ChatGLM2-int8模型若检索效果不好可以调成’10’
CACHED_VS_NUM = 1
 
# 文本分句长度
SENTENCE_SIZE = 100
 
# 匹配后单段上下文长度
CHUNK_SIZE = 250
 
# 传入LLM的历史记录长度
LLM_HISTORY_LEN = 3
 
# 知识库检索时返回的匹配内容条数
VECTOR_SEARCH_TOP_K = 5
 
# 知识检索内容相关度 Score, 数值范围约为0-1100，如果为0，则不生效，经测试设置为小于500时，匹配结果更精准
VECTOR_SEARCH_SCORE_THRESHOLD = 0
 
NLTK_DATA_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), "nltk_data")
 
FLAG_USER_NAME = uuid.uuid4().hex
 
logger.info(f"""
loading model config
llm device: {LLM_DEVICE}
embedding device: {EMBEDDING_DEVICE}
dir: {os.path.dirname(os.path.dirname(__file__))}
flagging username: {FLAG_USER_NAME}
""")
 
# 是否开启跨域，默认为False，如果需要开启，请设置为True
# is open cross domain
OPEN_CROSS_DOMAIN = False
 
# Bing 搜索必备变量
# 使用 Bing 搜索需要使用 Bing Subscription Key,需要在azure port中申请试用bing search
# 具体申请方式请见
# https://learn.microsoft.com/en-us/bing/search-apis/bing-web-search/create-bing-search-service-resource
# 使用python创建bing api 搜索实例详见:
# https://learn.microsoft.com/en-us/bing/search-apis/bing-web-search/quickstarts/rest/python
BING_SEARCH_URL = "https://api.bing.microsoft.com/v7.0/search"
# 注意不是bing Webmaster Tools的api key，
 
# 此外，如果是在服务器上，报Failed to establish a new connection: [Errno 110] Connection timed out
# 是因为服务器加了防火墙，需要联系管理员加白名单，如果公司的服务器的话，就别想了GG
BING_SUBSCRIPTION_KEY = ""
 
# 是否开启中文标题加强，以及标题增强的相关配置
# 通过增加标题判断，判断哪些文本为标题，并在metadata中进行标记；
# 然后将文本与往上一级的标题进行拼合，实现文本信息的增强。
ZH_TITLE_ENHANCE = False

# 导入类型提示模块，用于强化代码的可读性和健壮性
from typing import List
 
# 导入UnstructuredFileLoader，这是一个从非结构化文件中加载文档的类
from langchain.document_loaders.unstructured import UnstructuredFileLoader
 
# 导入PaddleOCR，这是一个开源的OCR工具，用于从图片中识别和读取文字
from paddleocr import PaddleOCR
 
# 导入os模块，用于处理文件和目录
import os
 
# 导入fitz模块，用于处理PDF文件
import fitz
 
# 导入nltk模块，用于处理文本数据
import nltk
 
# 导入模型配置文件中的NLTK_DATA_PATH，这是nltk数据的路径
from configs.model_config import NLTK_DATA_PATH
 
# 设置nltk数据的路径，将模型配置中的路径添加到nltk的数据路径中
nltk.data.path = [NLTK_DATA_PATH] + nltk.data.path
 
# 定义一个类，UnstructuredPaddlePDFLoader，该类继承自UnstructuredFileLoader
class UnstructuredPaddlePDFLoader(UnstructuredFileLoader):
 
    # 定义一个内部方法_get_elements，返回一个列表
    def _get_elements(self) -> List:
 
        # 定义一个内部函数pdf_ocr_txt，用于从pdf中进行OCR并输出文本文件
        def pdf_ocr_txt(filepath, dir_path="tmp_files"):
            # 将dir_path与filepath的目录部分合并成一个新的路径
            full_dir_path = os.path.join(os.path.dirname(filepath), dir_path)
 
            # 如果full_dir_path对应的目录不存在，则创建这个目录
            if not os.path.exists(full_dir_path):
                os.makedirs(full_dir_path)
            
            # 创建一个PaddleOCR实例，设置一些参数
            ocr = PaddleOCR(use_angle_cls=True, lang="ch", use_gpu=False, show_log=False)
 
            # 打开pdf文件
            doc = fitz.open(filepath)
 
            # 创建一个txt文件的路径
            txt_file_path = os.path.join(full_dir_path, f"{os.path.split(filepath)[-1]}.txt")
 
            # 创建一个临时的图片文件路径
            img_name = os.path.join(full_dir_path, 'tmp.png')
 
            # 打开txt_file_path对应的文件，并以写模式打开
            with open(txt_file_path, 'w', encoding='utf-8') as fout:
                # 遍历pdf的所有页面
                for i in range(doc.page_count):
                    # 获取当前页面
                    page = doc[i]
 
                    # 获取当前页面的文本内容，并写入txt文件
                    text = page.get_text("")
                    fout.write(text)
                    fout.write("\n")
 
                    # 获取当前页面的所有图片
                    img_list = page.get_images()
 
                    # 遍历所有图片
                    for img in img_list:
                        # 将图片转换为Pixmap对象
                        pix = fitz.Pixmap(doc, img[0])
 
                        # 如果图片有颜色信息，则将其转换为RGB格式
                        if pix.n - pix.alpha >= 4:
                            pix = fitz.Pixmap(fitz.csRGB, pix)
                        
                        # 保存图片
                        pix.save(img_name)
 
                        # 对图片进行OCR识别
                        result = ocr.ocr(img_name)
 
                        # 从OCR结果中提取文本，并写入txt文件
                        ocr_result = [i[1][0] for line in result for i in line]
                        fout.write("\n".join(ocr_result))
            
            # 如果图片文件存在，则删除它
            if os.path.exists(img_name):
                os.remove(img_name)
            
            # 返回txt文件的路径
            return txt_file_path
 
        # 调用上面定义的函数，获取txt文件的路径
        txt_file_path = pdf_ocr_txt(self.file_path)
 
        # 导入partition_text函数，该函数用于将文本文件分块
        from unstructured.partition.text import partition_text
 
        # 对txt文件进行分块，并返回分块结果
        return partition_text(filename=txt_file_path, **self.unstructured_kwargs)
 
# 运行入口
if __name__ == "__main__":
    # 导入sys模块，用于操作Python的运行环境
    import sys
 
    # 将当前文件的上一级目录添加到Python的搜索路径中
    sys.path.append(os.path.dirname(os.path.dirname(__file__)))
 
    # 定义一个pdf文件的路径
    filepath = os.path.join(os.path.dirname(os.path.dirname(__file__)), "knowledge_base", "samples", "content", "test.pdf")
 
    # 创建一个UnstructuredPaddlePDFLoader的实例
    loader = UnstructuredPaddlePDFLoader(filepath, mode="elements")
 
    # 加载文档
    docs = loader.load()
 
    # 遍历并打印所有文档
    for doc in docs:
        print(doc)