Add support for folder path as input

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.` 如何解决？
 

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()

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], )
-
-
-def init_knowledge_vector_store(filepath):
-    loader = UnstructuredFileLoader(filepath, mode="elements")
-    docs = loader.load()
+    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: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