Upload 254 files

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+img/langchain+chatglm.png filter=lfs diff=lfs merge=lfs -text
+img/official_wechat_mp_account.png filter=lfs diff=lfs merge=lfs -text
+knowledge_base/samples/content/llm/img/分布式训练技术原理-幕布图片-392521-261326.jpg filter=lfs diff=lfs merge=lfs -text

chains/llmchain_with_history.py

@@ -0,0 +1,22 @@
+from server.utils import get_ChatOpenAI
+from configs.model_config import LLM_MODELS, TEMPERATURE
+from langchain.chains import LLMChain
+from langchain.prompts.chat import (
+    ChatPromptTemplate,
+    HumanMessagePromptTemplate,
+)
+
+model = get_ChatOpenAI(model_name=LLM_MODELS[0], temperature=TEMPERATURE)
+
+
+human_prompt = "{input}"
+human_message_template = HumanMessagePromptTemplate.from_template(human_prompt)
+
+chat_prompt = ChatPromptTemplate.from_messages(
+    [("human", "我们来玩成语接龙，我先来，生龙活虎"),
+     ("ai", "虎头虎脑"),
+     ("human", "{input}")])
+
+
+chain = LLMChain(prompt=chat_prompt, llm=model, verbose=True)
+print(chain({"input": "恼羞成怒"}))

configs/__init__.py

@@ -0,0 +1,8 @@
+from .basic_config import *
+from .model_config import *
+from .kb_config import *
+from .server_config import *
+from .prompt_config import *
+
+
+VERSION = "v0.2.10"

configs/basic_config.py.example

@@ -0,0 +1,32 @@
+import logging
+import os
+import langchain
+import tempfile
+import shutil
+
+
+# 是否显示详细日志
+log_verbose = False
+langchain.verbose = False
+
+# 通常情况下不需要更改以下内容
+
+# 日志格式
+LOG_FORMAT = "%(asctime)s - %(filename)s[line:%(lineno)d] - %(levelname)s: %(message)s"
+logger = logging.getLogger()
+logger.setLevel(logging.INFO)
+logging.basicConfig(format=LOG_FORMAT)
+
+
+# 日志存储路径
+LOG_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), "logs")
+if not os.path.exists(LOG_PATH):
+    os.mkdir(LOG_PATH)
+
+# 临时文件目录，主要用于文件对话
+BASE_TEMP_DIR = os.path.join(tempfile.gettempdir(), "chatchat")
+try:
+    shutil.rmtree(BASE_TEMP_DIR)
+except Exception:
+    pass
+os.makedirs(BASE_TEMP_DIR, exist_ok=True)

configs/kb_config.py.example

@@ -0,0 +1,145 @@
+import os
+
+# 默认使用的知识库
+DEFAULT_KNOWLEDGE_BASE = "samples"
+
+# 默认向量库/全文检索引擎类型。可选：faiss, milvus(离线) & zilliz(在线), pgvector,全文检索引擎es
+DEFAULT_VS_TYPE = "faiss"
+
+# 缓存向量库数量（针对FAISS）
+CACHED_VS_NUM = 1
+
+# 缓存临时向量库数量（针对FAISS），用于文件对话
+CACHED_MEMO_VS_NUM = 10
+
+# 知识库中单段文本长度(不适用MarkdownHeaderTextSplitter)
+CHUNK_SIZE = 250
+
+# 知识库中相邻文本重合长度(不适用MarkdownHeaderTextSplitter)
+OVERLAP_SIZE = 50
+
+# 知识库匹配向量数量
+VECTOR_SEARCH_TOP_K = 3
+
+# 知识库匹配的距离阈值，一般取值范围在0-1之间，SCORE越小，距离越小从而相关度越高。
+# 但有用户报告遇到过匹配分值超过1的情况，为了兼容性默认设为1，在WEBUI中调整范围为0-2
+SCORE_THRESHOLD = 1.0
+
+# 默认搜索引擎。可选：bing, duckduckgo, metaphor
+DEFAULT_SEARCH_ENGINE = "duckduckgo"
+
+# 搜索引擎匹配结题数量
+SEARCH_ENGINE_TOP_K = 3
+
+
+# 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 = ""
+
+# metaphor搜索需要KEY
+METAPHOR_API_KEY = ""
+
+# 心知天气 API KEY，用于天气Agent。申请：https://www.seniverse.com/
+SENIVERSE_API_KEY = ""
+
+# 是否开启中文标题加强，以及标题增强的相关配置
+# 通过增加标题判断，判断哪些文本为标题，并在metadata中进行标记；
+# 然后将文本与往上一级的标题进行拼合，实现文本信息的增强。
+ZH_TITLE_ENHANCE = False
+
+# PDF OCR 控制：只对宽高超过页面一定比例（图片宽/页面宽，图片高/页面高）的图片进行 OCR。
+# 这样可以避免 PDF 中一些小图片的干扰，提高非扫描版 PDF 处理速度
+PDF_OCR_THRESHOLD = (0.6, 0.6)
+
+# 每个知识库的初始化介绍，用于在初始化知识库时显示和Agent调用，没写则没有介绍，不会被Agent调用。
+KB_INFO = {
+    "知识库名称": "知识库介绍",
+    "samples": "关于本项目issue的解答",
+}
+
+
+# 通常情况下不需要更改以下内容
+
+# 知识库默认存储路径
+KB_ROOT_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), "knowledge_base")
+if not os.path.exists(KB_ROOT_PATH):
+    os.mkdir(KB_ROOT_PATH)
+# 数据库默认存储路径。
+# 如果使用sqlite，可以直接修改DB_ROOT_PATH；如果使用其它数据库，请直接修改SQLALCHEMY_DATABASE_URI。
+DB_ROOT_PATH = os.path.join(KB_ROOT_PATH, "info.db")
+SQLALCHEMY_DATABASE_URI = f"sqlite:///{DB_ROOT_PATH}"
+
+# 可选向量库类型及对应配置
+kbs_config = {
+    "faiss": {
+    },
+    "milvus": {
+        "host": "127.0.0.1",
+        "port": "19530",
+        "user": "",
+        "password": "",
+        "secure": False,
+    },
+    "zilliz": {
+        "host": "in01-a7ce524e41e3935.ali-cn-hangzhou.vectordb.zilliz.com.cn",
+        "port": "19530",
+        "user": "",
+        "password": "",
+        "secure": True,
+        },
+    "pg": {
+        "connection_uri": "postgresql://postgres:postgres@127.0.0.1:5432/langchain_chatchat",
+    },
+
+    "es": {
+        "host": "127.0.0.1",
+        "port": "9200",
+        "index_name": "test_index",
+        "user": "",
+        "password": ""
+    },
+    "milvus_kwargs":{
+        "search_params":{"metric_type": "L2"}, #在此处增加search_params
+        "index_params":{"metric_type": "L2","index_type": "HNSW"} # 在此处增加index_params
+    }
+}
+
+# TextSplitter配置项，如果你不明白其中的含义，就不要修改。
+text_splitter_dict = {
+    "ChineseRecursiveTextSplitter": {
+        "source": "huggingface",   # 选择tiktoken则使用openai的方法
+        "tokenizer_name_or_path": "",
+    },
+    "SpacyTextSplitter": {
+        "source": "huggingface",
+        "tokenizer_name_or_path": "gpt2",
+    },
+    "RecursiveCharacterTextSplitter": {
+        "source": "tiktoken",
+        "tokenizer_name_or_path": "cl100k_base",
+    },
+    "MarkdownHeaderTextSplitter": {
+        "headers_to_split_on":
+            [
+                ("#", "head1"),
+                ("##", "head2"),
+                ("###", "head3"),
+                ("####", "head4"),
+            ]
+    },
+}
+
+# TEXT_SPLITTER 名称
+TEXT_SPLITTER_NAME = "ChineseRecursiveTextSplitter"
+
+# Embedding模型定制��语的词表文件
+EMBEDDING_KEYWORD_FILE = "embedding_keywords.txt"

configs/model_config.py.example

@@ -0,0 +1,302 @@
+import os
+
+# 可以指定一个绝对路径，统一存放所有的Embedding和LLM模型。
+# 每个模型可以是一个单独的目录，也可以是某个目录下的二级子目录。
+# 如果模型目录名称和 MODEL_PATH 中的 key 或 value 相同，程序会自动检测加载，无需修改 MODEL_PATH 中的路径。
+MODEL_ROOT_PATH = ""
+
+# 选用的 Embedding 名称
+EMBEDDING_MODEL = "bge-large-zh-v1.5"
+
+# Embedding 模型运行设备。设为 "auto" 会自动检测(会有警告)，也可手动设定为 "cuda","mps","cpu","xpu" 其中之一。
+EMBEDDING_DEVICE = "auto"
+
+# 选用的reranker模型
+RERANKER_MODEL = "bge-reranker-large"
+# 是否启用reranker模型
+USE_RERANKER = False
+RERANKER_MAX_LENGTH = 1024
+
+# 如果需要在 EMBEDDING_MODEL 中增加自定义的关键字时配置
+EMBEDDING_KEYWORD_FILE = "keywords.txt"
+EMBEDDING_MODEL_OUTPUT_PATH = "output"
+
+# 要运行的 LLM 名称，可以包括本地模型和在线模型。列表中本地模型将在启动项目时全部加载。
+# 列表中第一个模型将作为 API 和 WEBUI 的默认模型。
+# 在这里，我们使用目前主流的两个离线模型，其中，chatglm3-6b 为默认加载模型。
+# 如果你的显存不足，可使用 Qwen-1_8B-Chat, 该模型 FP16 仅需 3.8G显存。
+
+LLM_MODELS = ["chatglm3-6b", "zhipu-api", "openai-api"]
+Agent_MODEL = None
+
+# LLM 模型运行设备。设为"auto"会自动检测(会有警告)，也可手动设定为 "cuda","mps","cpu","xpu" 其中之一。
+LLM_DEVICE = "auto"
+
+HISTORY_LEN = 3
+
+MAX_TOKENS = 2048
+
+TEMPERATURE = 0.7
+
+ONLINE_LLM_MODEL = {
+    "openai-api": {
+        "model_name": "gpt-4",
+        "api_base_url": "https://api.openai.com/v1",
+        "api_key": "",
+        "openai_proxy": "",
+    },
+
+    # 智谱AI API,具体注册及api key获取请前往 http://open.bigmodel.cn
+    "zhipu-api": {
+        "api_key": "",
+        "version": "glm-4",
+        "provider": "ChatGLMWorker",
+    },
+
+    # 具体注册及api key获取请前往 https://api.minimax.chat/
+    "minimax-api": {
+        "group_id": "",
+        "api_key": "",
+        "is_pro": False,
+        "provider": "MiniMaxWorker",
+    },
+
+    # 具体注册及api key获取请前往 https://xinghuo.xfyun.cn/
+    "xinghuo-api": {
+        "APPID": "",
+        "APISecret": "",
+        "api_key": "",
+        "version": "v3.0", # 你使用的讯飞星火大模型版本，可选包括 "v3.0", "v2.0", "v1.5"
+        "provider": "XingHuoWorker",
+    },
+
+    # 百度千帆 API，申请方式请参考 https://cloud.baidu.com/doc/WENXINWORKSHOP/s/4lilb2lpf
+    "qianfan-api": {
+        "version": "ERNIE-Bot",  # 注意大小写。当前支持 "ERNIE-Bot" 或 "ERNIE-Bot-turbo"， 更多的见官方文档。
+        "version_url": "",  # 也可以不填写version，直接填写在千帆申请模型发布的API地址
+        "api_key": "",
+        "secret_key": "",
+        "provider": "QianFanWorker",
+    },
+
+    # 火山方舟 API，文档参考 https://www.volcengine.com/docs/82379
+    "fangzhou-api": {
+        "version": "chatglm-6b-model",
+        "version_url": "",
+        "api_key": "",
+        "secret_key": "",
+        "provider": "FangZhouWorker",
+    },
+
+    # 阿里云通义千问 API，文档参考 https://help.aliyun.com/zh/dashscope/developer-reference/api-details
+    "qwen-api": {
+        "version": "qwen-max",
+        "api_key": "",
+        "provider": "QwenWorker",
+        "embed_model": "text-embedding-v1"  # embedding 模型名称
+    },
+
+    # 百川 API，申请方式请参考 https://www.baichuan-ai.com/home#api-enter
+    "baichuan-api": {
+        "version": "Baichuan2-53B",
+        "api_key": "",
+        "secret_key": "",
+        "provider": "BaiChuanWorker",
+    },
+
+    # Azure API
+    "azure-api": {
+        "deployment_name": "",  # 部署容器的名字
+        "resource_name": "",  # https://{resource_name}.openai.azure.com/openai/ 填写resource_name的部分，其他部分不要填写
+        "api_version": "",  # API的版本，不是模型版本
+        "api_key": "",
+        "provider": "AzureWorker",
+    },
+
+    # 昆仑万维天工 API https://model-platform.tiangong.cn/
+    "tiangong-api": {
+        "version": "SkyChat-MegaVerse",
+        "api_key": "",
+        "secret_key": "",
+        "provider": "TianGongWorker",
+    },
+    # Gemini API https://makersuite.google.com/app/apikey
+    "gemini-api": {
+        "api_key": "",
+        "provider": "GeminiWorker",
+    }
+
+}
+
+# 在以下字典中修改属性值，以指定本地embedding模型存储位置。支持3种设置方法：
+# 1、将对应的值修改为模型绝对路径
+# 2、不修改此处的值（以 text2vec 为例）：
+#       2.1 如果{MODEL_ROOT_PATH}下存在如下任一子目录：
+#           - text2vec
+#           - GanymedeNil/text2vec-large-chinese
+#           - text2vec-large-chinese
+#       2.2 如果以上本地路径不存在，则使用huggingface模型
+
+MODEL_PATH = {
+    "embed_model": {
+        "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",
+        "text2vec-paraphrase": "shibing624/text2vec-base-chinese-paraphrase",
+        "text2vec-sentence": "shibing624/text2vec-base-chinese-sentence",
+        "text2vec-multilingual": "shibing624/text2vec-base-multilingual",
+        "text2vec-bge-large-chinese": "shibing624/text2vec-bge-large-chinese",
+        "m3e-small": "moka-ai/m3e-small",
+        "m3e-base": "moka-ai/m3e-base",
+        "m3e-large": "moka-ai/m3e-large",
+        "bge-small-zh": "BAAI/bge-small-zh",
+        "bge-base-zh": "BAAI/bge-base-zh",
+        "bge-large-zh": "BAAI/bge-large-zh",
+        "bge-large-zh-noinstruct": "BAAI/bge-large-zh-noinstruct",
+        "bge-base-zh-v1.5": "BAAI/bge-base-zh-v1.5",
+        "bge-large-zh-v1.5": "BAAI/bge-large-zh-v1.5",
+        "piccolo-base-zh": "sensenova/piccolo-base-zh",
+        "piccolo-large-zh": "sensenova/piccolo-large-zh",
+        "nlp_gte_sentence-embedding_chinese-large": "damo/nlp_gte_sentence-embedding_chinese-large",
+        "text-embedding-ada-002": "your OPENAI_API_KEY",
+    },
+
+    "llm_model": {
+        "chatglm2-6b": "THUDM/chatglm2-6b",
+        "chatglm2-6b-32k": "THUDM/chatglm2-6b-32k",
+        "chatglm3-6b": "THUDM/chatglm3-6b",
+        "chatglm3-6b-32k": "THUDM/chatglm3-6b-32k",
+
+        "Orion-14B-Chat": "OrionStarAI/Orion-14B-Chat",
+        "Orion-14B-Chat-Plugin": "OrionStarAI/Orion-14B-Chat-Plugin",
+        "Orion-14B-LongChat": "OrionStarAI/Orion-14B-LongChat",
+
+        "Llama-2-7b-chat-hf": "meta-llama/Llama-2-7b-chat-hf",
+        "Llama-2-13b-chat-hf": "meta-llama/Llama-2-13b-chat-hf",
+        "Llama-2-70b-chat-hf": "meta-llama/Llama-2-70b-chat-hf",
+
+        "Qwen-1_8B-Chat": "Qwen/Qwen-1_8B-Chat",
+        "Qwen-7B-Chat": "Qwen/Qwen-7B-Chat",
+        "Qwen-14B-Chat": "Qwen/Qwen-14B-Chat",
+        "Qwen-72B-Chat": "Qwen/Qwen-72B-Chat",
+
+        "baichuan-7b-chat": "baichuan-inc/Baichuan-7B-Chat",
+        "baichuan-13b-chat": "baichuan-inc/Baichuan-13B-Chat",
+        "baichuan2-7b-chat": "baichuan-inc/Baichuan2-7B-Chat",
+        "baichuan2-13b-chat": "baichuan-inc/Baichuan2-13B-Chat",
+
+        "internlm-7b": "internlm/internlm-7b",
+        "internlm-chat-7b": "internlm/internlm-chat-7b",
+        "internlm2-chat-7b": "internlm/internlm2-chat-7b",
+        "internlm2-chat-20b": "internlm/internlm2-chat-20b",
+
+        "BlueLM-7B-Chat": "vivo-ai/BlueLM-7B-Chat",
+        "BlueLM-7B-Chat-32k": "vivo-ai/BlueLM-7B-Chat-32k",
+
+        "Yi-34B-Chat": "https://huggingface.co/01-ai/Yi-34B-Chat",
+
+        "agentlm-7b": "THUDM/agentlm-7b",
+        "agentlm-13b": "THUDM/agentlm-13b",
+        "agentlm-70b": "THUDM/agentlm-70b",
+
+        "falcon-7b": "tiiuae/falcon-7b",
+        "falcon-40b": "tiiuae/falcon-40b",
+        "falcon-rw-7b": "tiiuae/falcon-rw-7b",
+
+        "aquila-7b": "BAAI/Aquila-7B",
+        "aquilachat-7b": "BAAI/AquilaChat-7B",
+        "open_llama_13b": "openlm-research/open_llama_13b",
+        "vicuna-13b-v1.5": "lmsys/vicuna-13b-v1.5",
+        "koala": "young-geng/koala",
+        "mpt-7b": "mosaicml/mpt-7b",
+        "mpt-7b-storywriter": "mosaicml/mpt-7b-storywriter",
+        "mpt-30b": "mosaicml/mpt-30b",
+        "opt-66b": "facebook/opt-66b",
+        "opt-iml-max-30b": "facebook/opt-iml-max-30b",
+        "gpt2": "gpt2",
+        "gpt2-xl": "gpt2-xl",
+        "gpt-j-6b": "EleutherAI/gpt-j-6b",
+        "gpt4all-j": "nomic-ai/gpt4all-j",
+        "gpt-neox-20b": "EleutherAI/gpt-neox-20b",
+        "pythia-12b": "EleutherAI/pythia-12b",
+        "oasst-sft-4-pythia-12b-epoch-3.5": "OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5",
+        "dolly-v2-12b": "databricks/dolly-v2-12b",
+        "stablelm-tuned-alpha-7b": "stabilityai/stablelm-tuned-alpha-7b",
+    },
+
+    "reranker": {
+        "bge-reranker-large": "BAAI/bge-reranker-large",
+        "bge-reranker-base": "BAAI/bge-reranker-base",
+    }
+}
+
+# 通常情况下不需要更改以下内容
+
+# nltk 模型存储路径
+NLTK_DATA_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), "nltk_data")
+
+# 使用VLLM可能导致模型推理能力下降，无法完成Agent任务
+VLLM_MODEL_DICT = {
+    "chatglm2-6b": "THUDM/chatglm2-6b",
+    "chatglm2-6b-32k": "THUDM/chatglm2-6b-32k",
+    "chatglm3-6b": "THUDM/chatglm3-6b",
+    "chatglm3-6b-32k": "THUDM/chatglm3-6b-32k",
+
+    "Llama-2-7b-chat-hf": "meta-llama/Llama-2-7b-chat-hf",
+    "Llama-2-13b-chat-hf": "meta-llama/Llama-2-13b-chat-hf",
+    "Llama-2-70b-chat-hf": "meta-llama/Llama-2-70b-chat-hf",
+
+    "Qwen-1_8B-Chat": "Qwen/Qwen-1_8B-Chat",
+    "Qwen-7B-Chat": "Qwen/Qwen-7B-Chat",
+    "Qwen-14B-Chat": "Qwen/Qwen-14B-Chat",
+    "Qwen-72B-Chat": "Qwen/Qwen-72B-Chat",
+
+    "baichuan-7b-chat": "baichuan-inc/Baichuan-7B-Chat",
+    "baichuan-13b-chat": "baichuan-inc/Baichuan-13B-Chat",
+    "baichuan2-7b-chat": "baichuan-inc/Baichuan-7B-Chat",
+    "baichuan2-13b-chat": "baichuan-inc/Baichuan-13B-Chat",
+
+    "BlueLM-7B-Chat": "vivo-ai/BlueLM-7B-Chat",
+    "BlueLM-7B-Chat-32k": "vivo-ai/BlueLM-7B-Chat-32k",
+
+    "internlm-7b": "internlm/internlm-7b",
+    "internlm-chat-7b": "internlm/internlm-chat-7b",
+    "internlm2-chat-7b": "internlm/Models/internlm2-chat-7b",
+    "internlm2-chat-20b": "internlm/Models/internlm2-chat-20b",
+
+    "aquila-7b": "BAAI/Aquila-7B",
+    "aquilachat-7b": "BAAI/AquilaChat-7B",
+
+    "falcon-7b": "tiiuae/falcon-7b",
+    "falcon-40b": "tiiuae/falcon-40b",
+    "falcon-rw-7b": "tiiuae/falcon-rw-7b",
+    "gpt2": "gpt2",
+    "gpt2-xl": "gpt2-xl",
+    "gpt-j-6b": "EleutherAI/gpt-j-6b",
+    "gpt4all-j": "nomic-ai/gpt4all-j",
+    "gpt-neox-20b": "EleutherAI/gpt-neox-20b",
+    "pythia-12b": "EleutherAI/pythia-12b",
+    "oasst-sft-4-pythia-12b-epoch-3.5": "OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5",
+    "dolly-v2-12b": "databricks/dolly-v2-12b",
+    "stablelm-tuned-alpha-7b": "stabilityai/stablelm-tuned-alpha-7b",
+    "open_llama_13b": "openlm-research/open_llama_13b",
+    "vicuna-13b-v1.3": "lmsys/vicuna-13b-v1.3",
+    "koala": "young-geng/koala",
+    "mpt-7b": "mosaicml/mpt-7b",
+    "mpt-7b-storywriter": "mosaicml/mpt-7b-storywriter",
+    "mpt-30b": "mosaicml/mpt-30b",
+    "opt-66b": "facebook/opt-66b",
+    "opt-iml-max-30b": "facebook/opt-iml-max-30b",
+
+}
+
+SUPPORT_AGENT_MODEL = [
+    "openai-api",  # GPT4 模型
+    "qwen-api",  # Qwen Max模型
+    "zhipu-api",  # 智谱AI GLM4模型
+    "Qwen",  # 所有Qwen系列本地模型
+    "chatglm3-6b",
+    "internlm2-chat-20b",
+    "Orion-14B-Chat-Plugin",
+]

configs/prompt_config.py.example

@@ -0,0 +1,127 @@
+# prompt模板使用Jinja2语法，简单点就是用双大括号代替f-string的单大括号
+# 本配置文件支持热加载，修改prompt模板后无需重启服务。
+
+# LLM对话支持的变量：
+#   - input: 用户输入内容
+
+# 知识库和搜索引擎对话支持的变量：
+#   - context: 从检索结果拼接的知识文本
+#   - question: 用户提出的问题
+
+# Agent对话支持的变量：
+
+#   - tools: 可用的工具列表
+#   - tool_names: 可用的工具名称列表
+#   - history: 用户和Agent的对话历史
+#   - input: 用户输入内容
+#   - agent_scratchpad: Agent的思维记录
+
+PROMPT_TEMPLATES = {
+    "llm_chat": {
+        "default":
+            '{{ input }}',
+
+        "with_history":
+            'The following is a friendly conversation between a human and an AI. '
+            'The AI is talkative and provides lots of specific details from its context. '
+            'If the AI does not know the answer to a question, it truthfully says it does not know.\n\n'
+            'Current conversation:\n'
+            '{history}\n'
+            'Human: {input}\n'
+            'AI:',
+
+        "py":
+            '你是一个聪明的代码助手，请你给我写出简单的py代码。 \n'
+            '{{ input }}',
+    },
+
+
+    "knowledge_base_chat": {
+        "default":
+            '<指令>根据已知信息，简洁和专业的来回答问题。如果无法从中得到答案，请说 “根据已知信息无法回答该问题”，'
+            '不允许在答案中添加编造成分，答案请使用中文。 </指令>\n'
+            '<已知信息>{{ context }}</已知信息>\n'
+            '<问题>{{ question }}</问题>\n',
+
+        "text":
+            '<指令>根据已知信息，简洁和专业的来回答问题。如果无法从中得到答案，请说 “根据已知信息无法回答该问题”，答案请使用中文。 </指令>\n'
+            '<已知信息>{{ context }}</已知信息>\n'
+            '<问题>{{ question }}</问题>\n',
+
+        "empty":  # 搜不到知识库的时候使用
+            '请你回答我的问题:\n'
+            '{{ question }}\n\n',
+    },
+
+
+    "search_engine_chat": {
+        "default":
+            '<指令>这是我搜索到的互联网信息，请你根据这些信息进行提取并有调理，简洁的回答问题。'
+            '如果无法从中得到答案，请说 “无法搜索到能回答问题的内容”。 </指令>\n'
+            '<已知信息>{{ context }}</已知信息>\n'
+            '<问题>{{ question }}</问题>\n',
+
+        "search":
+            '<指令>根据已知信息，简洁和专业的来回答问题。如果无法从中得到答案，请说 “根据已知信息无法回答该问题”，答案请使用中文。 </指令>\n'
+            '<已知信息>{{ context }}</已知信息>\n'
+            '<问题>{{ question }}</问题>\n',
+    },
+
+
+    "agent_chat": {
+        "default":
+            'Answer the following questions as best you can. If it is in order, you can use some tools appropriately. '
+            'You have access to the following tools:\n\n'
+            '{tools}\n\n'
+            'Use the following format:\n'
+            'Question: the input question you must answer1\n'
+            'Thought: you should always think about what to do and what tools to use.\n'
+            'Action: the action to take, should be one of [{tool_names}]\n'
+            'Action Input: the input to the action\n'
+            'Observation: the result of the action\n'
+            '... (this Thought/Action/Action Input/Observation can be repeated zero or more times)\n'
+            'Thought: I now know the final answer\n'
+            'Final Answer: the final answer to the original input question\n'
+            'Begin!\n\n'
+            'history: {history}\n\n'
+            'Question: {input}\n\n'
+            'Thought: {agent_scratchpad}\n',
+
+        "ChatGLM3":
+            'You can answer using the tools, or answer directly using your knowledge without using the tools. '
+            'Respond to the human as helpfully and accurately as possible.\n'
+            'You have access to the following tools:\n'
+            '{tools}\n'
+            'Use a json blob to specify a tool by providing an action key (tool name) '
+            'and an action_input key (tool input).\n'
+            'Valid "action" values: "Final Answer" or  [{tool_names}]'
+            'Provide only ONE action per $JSON_BLOB, as shown:\n\n'
+            '```\n'
+            '{{{{\n'
+            '  "action": $TOOL_NAME,\n'
+            '  "action_input": $INPUT\n'
+            '}}}}\n'
+            '```\n\n'
+            'Follow this format:\n\n'
+            'Question: input question to answer\n'
+            'Thought: consider previous and subsequent steps\n'
+            'Action:\n'
+            '```\n'
+            '$JSON_BLOB\n'
+            '```\n'
+            'Observation: action result\n'
+            '... (repeat Thought/Action/Observation N times)\n'
+            'Thought: I know what to respond\n'
+            'Action:\n'
+            '```\n'
+            '{{{{\n'
+            '  "action": "Final Answer",\n'
+            '  "action_input": "Final response to human"\n'
+            '}}}}\n'
+            'Begin! Reminder to ALWAYS respond with a valid json blob of a single action. Use tools if necessary. '
+            'Respond directly if appropriate. Format is Action:```$JSON_BLOB```then Observation:.\n'
+            'history: {history}\n\n'
+            'Question: {input}\n\n'
+            'Thought: {agent_scratchpad}',
+    }
+}

configs/server_config.py.example

@@ -0,0 +1,137 @@
+import sys
+from configs.model_config import LLM_DEVICE
+
+# httpx 请求默认超时时间（秒）。如果加载模型或对话较慢，出现超时错误，可以适当加大该值。
+HTTPX_DEFAULT_TIMEOUT = 300.0
+
+# API 是否开启跨域，默认为False，如果需要开启，请设置为True
+# is open cross domain
+OPEN_CROSS_DOMAIN = False
+
+# 各服务器默认绑定host。如改为"0.0.0.0"需要修改下方所有XX_SERVER的host
+DEFAULT_BIND_HOST = "0.0.0.0" if sys.platform != "win32" else "127.0.0.1"
+
+# webui.py server
+WEBUI_SERVER = {
+    "host": DEFAULT_BIND_HOST,
+    "port": 8501,
+}
+
+# api.py server
+API_SERVER = {
+    "host": DEFAULT_BIND_HOST,
+    "port": 7861,
+}
+
+# fastchat openai_api server
+FSCHAT_OPENAI_API = {
+    "host": DEFAULT_BIND_HOST,
+    "port": 20000,
+}
+
+# fastchat model_worker server
+# 这些模型必须是在model_config.MODEL_PATH或ONLINE_MODEL中正确配置的。
+# 在启动startup.py时，可用通过`--model-name xxxx yyyy`指定模型，不指定则为LLM_MODELS
+FSCHAT_MODEL_WORKERS = {
+    # 所有模型共用的默认配置，可在模型专项配置中进行覆盖。
+    "default": {
+        "host": DEFAULT_BIND_HOST,
+        "port": 20002,
+        "device": LLM_DEVICE,
+        # False,'vllm',使用的推理加速框架,使用vllm如果出现HuggingFace通信问题，参见doc/FAQ
+        # vllm对一些模型支持还不成熟，暂时默认关闭
+        "infer_turbo": False,
+
+        # model_worker多卡加载需要配置的参数
+        # "gpus": None, # 使用的GPU，以str的格式指定，如"0,1"，如失效请使用CUDA_VISIBLE_DEVICES="0,1"等形式指定
+        # "num_gpus": 1, # 使用GPU的数量
+        # "max_gpu_memory": "20GiB", # 每个GPU占用的最大显存
+
+        # 以下为model_worker非常用参数，可根据需要配置
+        # "load_8bit": False, # 开启8bit量化
+        # "cpu_offloading": None,
+        # "gptq_ckpt": None,
+        # "gptq_wbits": 16,
+        # "gptq_groupsize": -1,
+        # "gptq_act_order": False,
+        # "awq_ckpt": None,
+        # "awq_wbits": 16,
+        # "awq_groupsize": -1,
+        # "model_names": LLM_MODELS,
+        # "conv_template": None,
+        # "limit_worker_concurrency": 5,
+        # "stream_interval": 2,
+        # "no_register": False,
+        # "embed_in_truncate": False,
+
+        # 以下为vllm_worker配置参数,注意使用vllm必须有gpu，仅在Linux测试通过
+
+        # tokenizer = model_path # 如果tokenizer与model_path不一致在此处添加
+        # 'tokenizer_mode':'auto',
+        # 'trust_remote_code':True,
+        # 'download_dir':None,
+        # 'load_format':'auto',
+        # 'dtype':'auto',
+        # 'seed':0,
+        # 'worker_use_ray':False,
+        # 'pipeline_parallel_size':1,
+        # 'tensor_parallel_size':1,
+        # 'block_size':16,
+        # 'swap_space':4 , # GiB
+        # 'gpu_memory_utilization':0.90,
+        # 'max_num_batched_tokens':2560,
+        # 'max_num_seqs':256,
+        # 'disable_log_stats':False,
+        # 'conv_template':None,
+        # 'limit_worker_concurrency':5,
+        # 'no_register':False,
+        # 'num_gpus': 1
+        # 'engine_use_ray': False,
+        # 'disable_log_requests': False
+
+    },
+    "Qwen-1_8B-Chat": {
+        "device": "cpu",
+    },
+    "chatglm3-6b": {
+        "device": "cuda",
+    },
+
+    # 以下配置可以不用修改，在model_config中设置启动的模型
+    "zhipu-api": {
+        "port": 21001,
+    },
+    "minimax-api": {
+        "port": 21002,
+    },
+    "xinghuo-api": {
+        "port": 21003,
+    },
+    "qianfan-api": {
+        "port": 21004,
+    },
+    "fangzhou-api": {
+        "port": 21005,
+    },
+    "qwen-api": {
+        "port": 21006,
+    },
+    "baichuan-api": {
+        "port": 21007,
+    },
+    "azure-api": {
+        "port": 21008,
+    },
+    "tiangong-api": {
+        "port": 21009,
+    },
+    "gemini-api": {
+        "port": 21010,
+    },
+}
+
+FSCHAT_CONTROLLER = {
+    "host": DEFAULT_BIND_HOST,
+    "port": 20001,
+    "dispatch_method": "shortest_queue",
+}

docs/ES部署指南.md

@@ -0,0 +1,29 @@
+
+# 实现基于ES的数据插入、检索、删除、更新
+```shell
+author: 唐国梁Tommy
+e-mail: flytang186@qq.com
+
+如果遇到任何问题，可以与我联系，我这边部署后服务是没有问题的。
+```
+
+## 第1步：ES docker部署
+```shell
+docker network create elastic
+docker run -id --name elasticsearch --net elastic -p 9200:9200 -p 9300:9300 -e "discovery.type=single-node" -e "xpack.security.enabled=false" -e "xpack.security.http.ssl.enabled=false" -t docker.elastic.co/elasticsearch/elasticsearch:8.8.2
+```
+
+### 第2步：Kibana docker部署
+**注意：Kibana版本与ES保持一致**
+```shell
+docker pull docker.elastic.co/kibana/kibana:{version} 
+docker run --name kibana --net elastic -p 5601:5601 docker.elastic.co/kibana/kibana:{version}
+```
+
+### 第3步：核心代码
+```shell
+1. 核心代码路径
+server/knowledge_base/kb_service/es_kb_service.py
+
+2. 需要在 configs/model_config.py 中 配置 ES参数（IP， PORT）等；
+```

document_loaders/FilteredCSVloader.py

@@ -0,0 +1,81 @@
+## 指定制定列的csv文件加载器
+
+from langchain.document_loaders import CSVLoader
+import csv
+from io import TextIOWrapper
+from typing import Dict, List, Optional
+from langchain.docstore.document import Document
+from langchain.document_loaders.helpers import detect_file_encodings
+
+
+class FilteredCSVLoader(CSVLoader):
+    def __init__(
+            self,
+            file_path: str,
+            columns_to_read: List[str],
+            source_column: Optional[str] = None,
+            metadata_columns: List[str] = [],
+            csv_args: Optional[Dict] = None,
+            encoding: Optional[str] = None,
+            autodetect_encoding: bool = False,
+    ):
+        super().__init__(
+            file_path=file_path,
+            source_column=source_column,
+            metadata_columns=metadata_columns,
+            csv_args=csv_args,
+            encoding=encoding,
+            autodetect_encoding=autodetect_encoding,
+        )
+        self.columns_to_read = columns_to_read
+
+    def load(self) -> List[Document]:
+        """Load data into document objects."""
+
+        docs = []
+        try:
+            with open(self.file_path, newline="", encoding=self.encoding) as csvfile:
+                docs = self.__read_file(csvfile)
+        except UnicodeDecodeError as e:
+            if self.autodetect_encoding:
+                detected_encodings = detect_file_encodings(self.file_path)
+                for encoding in detected_encodings:
+                    try:
+                        with open(
+                            self.file_path, newline="", encoding=encoding.encoding
+                        ) as csvfile:
+                            docs = self.__read_file(csvfile)
+                            break
+                    except UnicodeDecodeError:
+                        continue
+            else:
+                raise RuntimeError(f"Error loading {self.file_path}") from e
+        except Exception as e:
+            raise RuntimeError(f"Error loading {self.file_path}") from e
+
+        return docs
+
+    def __read_file(self, csvfile: TextIOWrapper) -> List[Document]:
+        docs = []
+        csv_reader = csv.DictReader(csvfile, **self.csv_args)  # type: ignore
+        for i, row in enumerate(csv_reader):
+            if self.columns_to_read[0] in row:
+                content = row[self.columns_to_read[0]]
+                # Extract the source if available
+                source = (
+                    row.get(self.source_column, None)
+                    if self.source_column is not None
+                    else self.file_path
+                )
+                metadata = {"source": source, "row": i}
+
+                for col in self.metadata_columns:
+                    if col in row:
+                        metadata[col] = row[col]
+
+                doc = Document(page_content=content, metadata=metadata)
+                docs.append(doc)
+            else:
+                raise ValueError(f"Column '{self.columns_to_read[0]}' not found in CSV file.")
+
+        return docs

document_loaders/__init__.py

@@ -0,0 +1,4 @@
+from .mypdfloader import RapidOCRPDFLoader
+from .myimgloader import RapidOCRLoader
+from .mydocloader import RapidOCRDocLoader
+from .mypptloader import RapidOCRPPTLoader

document_loaders/mydocloader.py

@@ -0,0 +1,71 @@
+from langchain.document_loaders.unstructured import UnstructuredFileLoader
+from typing import List
+import tqdm
+
+
+class RapidOCRDocLoader(UnstructuredFileLoader):
+    def _get_elements(self) -> List:
+        def doc2text(filepath):
+            from docx.table import _Cell, Table
+            from docx.oxml.table import CT_Tbl
+            from docx.oxml.text.paragraph import CT_P
+            from docx.text.paragraph import Paragraph
+            from docx import Document, ImagePart
+            from PIL import Image
+            from io import BytesIO
+            import numpy as np
+            from rapidocr_onnxruntime import RapidOCR
+            ocr = RapidOCR()
+            doc = Document(filepath)
+            resp = ""
+
+            def iter_block_items(parent):
+                from docx.document import Document
+                if isinstance(parent, Document):
+                    parent_elm = parent.element.body
+                elif isinstance(parent, _Cell):
+                    parent_elm = parent._tc
+                else:
+                    raise ValueError("RapidOCRDocLoader parse fail")
+
+                for child in parent_elm.iterchildren():
+                    if isinstance(child, CT_P):
+                        yield Paragraph(child, parent)
+                    elif isinstance(child, CT_Tbl):
+                        yield Table(child, parent)
+
+            b_unit = tqdm.tqdm(total=len(doc.paragraphs)+len(doc.tables),
+                               desc="RapidOCRDocLoader block index: 0")
+            for i, block in enumerate(iter_block_items(doc)):
+                b_unit.set_description(
+                    "RapidOCRDocLoader  block index: {}".format(i))
+                b_unit.refresh()
+                if isinstance(block, Paragraph):
+                    resp += block.text.strip() + "\n"
+                    images = block._element.xpath('.//pic:pic')  # 获取所有图片
+                    for image in images:
+                        for img_id in image.xpath('.//a:blip/@r:embed'):  # 获取图片id
+                            part = doc.part.related_parts[img_id]  # 根据图片id获取对应的图片
+                            if isinstance(part, ImagePart):
+                                image = Image.open(BytesIO(part._blob))
+                                result, _ = ocr(np.array(image))
+                                if result:
+                                    ocr_result = [line[1] for line in result]
+                                    resp += "\n".join(ocr_result)
+                elif isinstance(block, Table):
+                    for row in block.rows:
+                        for cell in row.cells:
+                            for paragraph in cell.paragraphs:
+                                resp += paragraph.text.strip() + "\n"
+                b_unit.update(1)
+            return resp
+
+        text = doc2text(self.file_path)
+        from unstructured.partition.text import partition_text
+        return partition_text(text=text, **self.unstructured_kwargs)
+
+
+if __name__ == '__main__':
+    loader = RapidOCRDocLoader(file_path="../tests/samples/ocr_test.docx")
+    docs = loader.load()
+    print(docs)

document_loaders/myimgloader.py

@@ -0,0 +1,25 @@
+from typing import List
+from langchain.document_loaders.unstructured import UnstructuredFileLoader
+from document_loaders.ocr import get_ocr
+
+
+class RapidOCRLoader(UnstructuredFileLoader):
+    def _get_elements(self) -> List:
+        def img2text(filepath):
+            resp = ""
+            ocr = get_ocr()
+            result, _ = ocr(filepath)
+            if result:
+                ocr_result = [line[1] for line in result]
+                resp += "\n".join(ocr_result)
+            return resp
+
+        text = img2text(self.file_path)
+        from unstructured.partition.text import partition_text
+        return partition_text(text=text, **self.unstructured_kwargs)
+
+
+if __name__ == "__main__":
+    loader = RapidOCRLoader(file_path="../tests/samples/ocr_test.jpg")
+    docs = loader.load()
+    print(docs)

document_loaders/mypdfloader.py

@@ -0,0 +1,51 @@
+from typing import List
+from langchain.document_loaders.unstructured import UnstructuredFileLoader
+from configs import PDF_OCR_THRESHOLD
+from document_loaders.ocr import get_ocr
+import tqdm
+
+
+class RapidOCRPDFLoader(UnstructuredFileLoader):
+    def _get_elements(self) -> List:
+        def pdf2text(filepath):
+            import fitz # pyMuPDF里面的fitz包，不要与pip install fitz混淆
+            import numpy as np
+            ocr = get_ocr()
+            doc = fitz.open(filepath)
+            resp = ""
+
+            b_unit = tqdm.tqdm(total=doc.page_count, desc="RapidOCRPDFLoader context page index: 0")
+            for i, page in enumerate(doc):
+                b_unit.set_description("RapidOCRPDFLoader context page index: {}".format(i))
+                b_unit.refresh()
+                text = page.get_text("")
+                resp += text + "\n"
+
+                img_list = page.get_image_info(xrefs=True)
+                for img in img_list:
+                    if xref := img.get("xref"):
+                        bbox = img["bbox"]
+                        # 检查图片尺寸是否超过设定的阈值
+                        if ((bbox[2] - bbox[0]) / (page.rect.width) < PDF_OCR_THRESHOLD[0]
+                            or (bbox[3] - bbox[1]) / (page.rect.height) < PDF_OCR_THRESHOLD[1]):
+                            continue
+                        pix = fitz.Pixmap(doc, xref)
+                        img_array = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, -1)
+                        result, _ = ocr(img_array)
+                        if result:
+                            ocr_result = [line[1] for line in result]
+                            resp += "\n".join(ocr_result)
+
+                # 更新进度
+                b_unit.update(1)
+            return resp
+
+        text = pdf2text(self.file_path)
+        from unstructured.partition.text import partition_text
+        return partition_text(text=text, **self.unstructured_kwargs)
+
+
+if __name__ == "__main__":
+    loader = RapidOCRPDFLoader(file_path="../tests/samples/ocr_test.pdf")
+    docs = loader.load()
+    print(docs)

document_loaders/mypptloader.py

@@ -0,0 +1,59 @@
+from langchain.document_loaders.unstructured import UnstructuredFileLoader
+from typing import List
+import tqdm
+
+
+class RapidOCRPPTLoader(UnstructuredFileLoader):
+    def _get_elements(self) -> List:
+        def ppt2text(filepath):
+            from pptx import Presentation
+            from PIL import Image
+            import numpy as np
+            from io import BytesIO
+            from rapidocr_onnxruntime import RapidOCR
+            ocr = RapidOCR()
+            prs = Presentation(filepath)
+            resp = ""
+
+            def extract_text(shape):
+                nonlocal resp
+                if shape.has_text_frame:
+                    resp += shape.text.strip() + "\n"
+                if shape.has_table:
+                    for row in shape.table.rows:
+                        for cell in row.cells:
+                            for paragraph in cell.text_frame.paragraphs:
+                                resp += paragraph.text.strip() + "\n"
+                if shape.shape_type == 13:  # 13 表示图片
+                    image = Image.open(BytesIO(shape.image.blob))
+                    result, _ = ocr(np.array(image))
+                    if result:
+                        ocr_result = [line[1] for line in result]
+                        resp += "\n".join(ocr_result)
+                elif shape.shape_type == 6:  # 6 表示组合
+                    for child_shape in shape.shapes:
+                        extract_text(child_shape)
+
+            b_unit = tqdm.tqdm(total=len(prs.slides),
+                               desc="RapidOCRPPTLoader slide index: 1")
+            # 遍历所有幻灯片
+            for slide_number, slide in enumerate(prs.slides, start=1):
+                b_unit.set_description(
+                    "RapidOCRPPTLoader slide index: {}".format(slide_number))
+                b_unit.refresh()
+                sorted_shapes = sorted(slide.shapes,
+                                       key=lambda x: (x.top, x.left))  # 从上到下、从左到右遍历
+                for shape in sorted_shapes:
+                    extract_text(shape)
+                b_unit.update(1)
+            return resp
+
+        text = ppt2text(self.file_path)
+        from unstructured.partition.text import partition_text
+        return partition_text(text=text, **self.unstructured_kwargs)
+
+
+if __name__ == '__main__':
+    loader = RapidOCRPPTLoader(file_path="../tests/samples/ocr_test.pptx")
+    docs = loader.load()
+    print(docs)

document_loaders/ocr.py

@@ -0,0 +1,18 @@
+from typing import TYPE_CHECKING
+
+
+if TYPE_CHECKING:
+    try:
+        from rapidocr_paddle import RapidOCR
+    except ImportError:
+        from rapidocr_onnxruntime import RapidOCR
+
+
+def get_ocr(use_cuda: bool = True) -> "RapidOCR":
+    try:
+        from rapidocr_paddle import RapidOCR
+        ocr = RapidOCR(det_use_cuda=use_cuda, cls_use_cuda=use_cuda, rec_use_cuda=use_cuda)
+    except ImportError:
+        from rapidocr_onnxruntime import RapidOCR
+        ocr = RapidOCR()
+    return ocr

embeddings/add_embedding_keywords.py

@@ -0,0 +1,79 @@
+'''
+该功能是为了将关键词加入到embedding模型中，以便于在embedding模型中进行关键词的embedding
+该功能的实现是通过修改embedding模型的tokenizer来实现的
+该功能仅仅对EMBEDDING_MODEL参数对应的的模型有效，输出后的模型保存在原本模型
+感谢@CharlesJu1和@charlesyju的贡献提出了想法和最基础的PR
+
+保存的模型的位置位于原本嵌入模型的目录下，模型的名称为原模型名称+Merge_Keywords_时间戳
+'''
+import sys
+
+sys.path.append("..")
+import os
+import torch
+
+from datetime import datetime
+from configs import (
+    MODEL_PATH,
+    EMBEDDING_MODEL,
+    EMBEDDING_KEYWORD_FILE,
+)
+
+from safetensors.torch import save_model
+from sentence_transformers import SentenceTransformer
+from langchain_core._api import deprecated
+
+
+@deprecated(
+        since="0.3.0",
+        message="自定义关键词 Langchain-Chatchat 0.3.x 重写, 0.2.x中相关功能将废弃",
+        removal="0.3.0"
+    )
+def get_keyword_embedding(bert_model, tokenizer, key_words):
+    tokenizer_output = tokenizer(key_words, return_tensors="pt", padding=True, truncation=True)
+    input_ids = tokenizer_output['input_ids']
+    input_ids = input_ids[:, 1:-1]
+
+    keyword_embedding = bert_model.embeddings.word_embeddings(input_ids)
+    keyword_embedding = torch.mean(keyword_embedding, 1)
+    return keyword_embedding
+
+
+def add_keyword_to_model(model_name=EMBEDDING_MODEL, keyword_file: str = "", output_model_path: str = None):
+    key_words = []
+    with open(keyword_file, "r") as f:
+        for line in f:
+            key_words.append(line.strip())
+
+    st_model = SentenceTransformer(model_name)
+    key_words_len = len(key_words)
+    word_embedding_model = st_model._first_module()
+    bert_model = word_embedding_model.auto_model
+    tokenizer = word_embedding_model.tokenizer
+    key_words_embedding = get_keyword_embedding(bert_model, tokenizer, key_words)
+
+    embedding_weight = bert_model.embeddings.word_embeddings.weight
+    embedding_weight_len = len(embedding_weight)
+    tokenizer.add_tokens(key_words)
+    bert_model.resize_token_embeddings(len(tokenizer), pad_to_multiple_of=32)
+    embedding_weight = bert_model.embeddings.word_embeddings.weight
+    with torch.no_grad():
+        embedding_weight[embedding_weight_len:embedding_weight_len + key_words_len, :] = key_words_embedding
+
+    if output_model_path:
+        os.makedirs(output_model_path, exist_ok=True)
+        word_embedding_model.save(output_model_path)
+        safetensors_file = os.path.join(output_model_path, "model.safetensors")
+        metadata = {'format': 'pt'}
+        save_model(bert_model, safetensors_file, metadata)
+        print("save model to {}".format(output_model_path))
+
+
+def add_keyword_to_embedding_model(path: str = EMBEDDING_KEYWORD_FILE):
+    keyword_file = os.path.join(path)
+    model_name = MODEL_PATH["embed_model"][EMBEDDING_MODEL]
+    model_parent_directory = os.path.dirname(model_name)
+    current_time = datetime.now().strftime('%Y%m%d_%H%M%S')
+    output_model_name = "{}_Merge_Keywords_{}".format(EMBEDDING_MODEL, current_time)
+    output_model_path = os.path.join(model_parent_directory, output_model_name)
+    add_keyword_to_model(model_name, keyword_file, output_model_path)

embeddings/embedding_keywords.txt

@@ -0,0 +1,3 @@
+Langchain-Chatchat
+数据科学与大数据技术
+人工智能与先进计算