From 4ec339777c20aa738df2dd4653d3fad6c9642ed6 Mon Sep 17 00:00:00 2001 From: imClumsyPanda Date: Tue, 11 Apr 2023 21:55:42 +0800 Subject: [PATCH] Add support for folder path as input --- README.md | 21 ++++++++++++++------ chatglm_llm.py | 40 ++++++++++++++++++++++++++------------ knowledge_based_chatglm.py | 33 +++++++++++++++++++++++++------ 3 files changed, 70 insertions(+), 24 deletions(-) diff --git a/README.md b/README.md index 4ee5537..41b5d2c 100644 --- a/README.md +++ b/README.md @@ -25,8 +25,10 @@ 4. 更正`README.md`中的代码错误（感谢 [@calcitem](https://github.com/calcitem)）。 **[2023/04/11]** -1. 加入 Web UI V0.1 版本（感谢 [@liangtongt](https://github.com/liangtongt)）。 - +1. 加入 Web UI V0.1 版本（感谢 [@liangtongt](https://github.com/liangtongt)）； +2. `README.md`中增加常见问题（感谢 [@calcitem](https://github.com/calcitem)）； +3. 增加 LLM 和 Embedding 模型运行设备是否可用`cuda`、`mps`、`cpu`的自动判断。 +4. 在`knowledge_based_chatglm.py`中增加对`filepath`的判断，在之前支持单个文件导入的基础上，现支持单个文件夹路径作为输入，输入后将会遍历文件夹中各个文件，并在命令行中显示每个文件是否成功加载。 ## 使用方式 @@ -71,11 +73,18 @@ python knowledge_based_chatglm.py ``` -### 已知问题 -- 目前已测试支持 txt、docx、md 格式文件，更多文件格式请参考 [langchain 文档](https://python.langchain.com/en/latest/modules/indexes/document_loaders/examples/unstructured_file.html)，目前已知文档中若含有特殊字符，可能存在文件无法加载的问题； -- 使用 macOS 运行本项目时，可能因为 macOS 版本为 13.3 及以上版本导致与 pytorch 不兼容，无法正常运行的情况。 - ### 常见问题 +Q: 本项目支持哪些文件格式？ + +A: 目前已测试支持 txt、docx、md 格式文件，更多文件格式请参考 [langchain 文档](https://python.langchain.com/en/latest/modules/indexes/document_loaders/examples/unstructured_file.html)。目前已知文档中若含有特殊字符，可能存在文件无法加载的问题。 + +Q: 读取特定格式文件时遇到缺少`detectron2`时如何解决？ + +A: 因该包安装过程中遇到问题较多，且仅部分格式文件需要，所以未加入`requirements.txt`。可以通过一下命令安装 + +```commandline +pip install "detectron2@git+https://github.com/facebookresearch/detectron2.git@v0.6#egg=detectron2" +``` Q: `Resource punkt not found.` 如何解决？ diff --git a/chatglm_llm.py b/chatglm_llm.py index 55c0177..c074c23 100644 --- a/chatglm_llm.py +++ b/chatglm_llm.py @@ -4,8 +4,8 @@ from langchain.llms.utils import enforce_stop_tokens from transformers import AutoTokenizer, AutoModel import torch -DEVICE = "cuda" -DEVICE_ID = "0" +DEVICE = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu" +DEVICE_ID = "0" if torch.cuda.is_available() else None CUDA_DEVICE = f"{DEVICE}:{DEVICE_ID}" if DEVICE_ID else DEVICE @@ -18,7 +18,7 @@ def torch_gc(): class ChatGLM(LLM): max_token: int = 10000 - temperature: float = 0.1 + temperature: float = 0.01 top_p = 0.9 history = [] tokenizer: object = None @@ -48,16 +48,32 @@ class ChatGLM(LLM): self.history = self.history+[[None, response]] return response - def load_model(self, - model_name_or_path: str = "THUDM/chatglm-6b"): + def load_model(self, model_name_or_path: str = "THUDM/chatglm-6b"): self.tokenizer = AutoTokenizer.from_pretrained( model_name_or_path, trust_remote_code=True ) - self.model = ( - AutoModel.from_pretrained( - model_name_or_path, - trust_remote_code=True) - .half() - .cuda() - ) + if torch.cuda.is_available(): + self.model = ( + AutoModel.from_pretrained( + model_name_or_path, + trust_remote_code=True) + .half() + .cuda() + ) + elif torch.backends.mps.is_available(): + self.model = ( + AutoModel.from_pretrained( + model_name_or_path, + trust_remote_code=True) + .float() + .to('mps') + ) + else: + self.model = ( + AutoModel.from_pretrained( + model_name_or_path, + trust_remote_code=True) + .float() + ) + self.model = self.model.eval() diff --git a/knowledge_based_chatglm.py b/knowledge_based_chatglm.py index fb34150..9ee3911 100644 --- a/knowledge_based_chatglm.py +++ b/knowledge_based_chatglm.py @@ -8,12 +8,17 @@ from langchain.embeddings.huggingface import HuggingFaceEmbeddings from langchain.vectorstores import FAISS from langchain.document_loaders import UnstructuredFileLoader from chatglm_llm import ChatGLM +import sentence_transformers +import torch +import os + # Global Parameters -EMBEDDING_MODEL = "text2vec" +EMBEDDING_MODEL = "local"#"text2vec" VECTOR_SEARCH_TOP_K = 6 -LLM_MODEL = "chatglm-6b" +LLM_MODEL = "local"#"chatglm-6b" LLM_HISTORY_LEN = 3 +DEVICE = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu" # Show reply with source text from input document REPLY_WITH_SOURCE = True @@ -22,12 +27,14 @@ embedding_model_dict = { "ernie-tiny": "nghuyong/ernie-3.0-nano-zh", "ernie-base": "nghuyong/ernie-3.0-base-zh", "text2vec": "GanymedeNil/text2vec-large-chinese", + "local": "/Users/liuqian/Downloads/ChatGLM-6B/text2vec-large-chinese" } llm_model_dict = { "chatglm-6b-int4-qe": "THUDM/chatglm-6b-int4-qe", "chatglm-6b-int4": "THUDM/chatglm-6b-int4", "chatglm-6b": "THUDM/chatglm-6b", + "local": "/Users/liuqian/Downloads/ChatGLM-6B/chatglm-6b" } @@ -39,12 +46,26 @@ def init_cfg(LLM_MODEL, EMBEDDING_MODEL, LLM_HISTORY_LEN, V_SEARCH_TOP_K=6): chatglm.load_model(model_name_or_path=llm_model_dict[LLM_MODEL]) chatglm.history_len = LLM_HISTORY_LEN - embeddings = HuggingFaceEmbeddings(model_name=embedding_model_dict[EMBEDDING_MODEL], ) + embeddings = HuggingFaceEmbeddings(model_name=embedding_model_dict[EMBEDDING_MODEL],) + embeddings.client = sentence_transformers.SentenceTransformer(embeddings.model_name, + device=DEVICE) -def init_knowledge_vector_store(filepath): - loader = UnstructuredFileLoader(filepath, mode="elements") - docs = loader.load() +def init_knowledge_vector_store(filepath:str): + if os.path.isfile(filepath): + loader = UnstructuredFileLoader(filepath, mode="elements") + docs = loader.load() + print(f"{os.path.split(filepath)[-1]} 已成功加载") + elif os.path.isdir(filepath): + docs = [] + for file in os.listdir(filepath): + fullfilepath = os.path.join(filepath, file) + try: + loader = UnstructuredFileLoader(fullfilepath, mode="elements") + docs += loader.load() + print(f"{file} 已成功加载") + except: + print(f"{file} 未能成功加载") vector_store = FAISS.from_documents(docs, embeddings) return vector_store