提交更

chains/local_doc_qa.py

@@ -117,6 +117,9 @@ class LocalDocQA:
     
     问题:
     {question}"""
+        if  vs_path is None or vs_path =="":# or (not os.path.exists(vs_path))
+            result = self.llm.chat(query)
+        else:
             prompt = PromptTemplate(
                 template=prompt_template,
                 input_variables=["context", "question"]

configs/model_config.py

@@ -19,10 +19,11 @@ llm_model_dict = {
     "chatglm-6b-int4-qe": "THUDM/chatglm-6b-int4-qe",
     "chatglm-6b-int4": "THUDM/chatglm-6b-int4",
     "chatglm-6b": "THUDM/chatglm-6b",
+    "chatyuan": "ClueAI/ChatYuan-large-v2",
 }
 
 # LLM model name
-LLM_MODEL = "chatglm-6b"
+LLM_MODEL = "chatyuan" #"chatglm-6b"
 
 # Use p-tuning-v2 PrefixEncoder
 USE_PTUNING_V2 = False

models/chatglm_llm.py

@@ -82,6 +82,19 @@ class ChatGLM(LLM):
         self.history = self.history+[[None, response]]
         return response
 
+    def chat(self,
+              prompt: str) -> str:
+        response, _ = self.model.chat(
+            self.tokenizer,
+            prompt,
+            history=[],#self.history[-self.history_len:] if self.history_len>0 else 
+            max_length=self.max_token,
+            temperature=self.temperature,
+        )
+        torch_gc()
+        self.history = self.history+[[None, response]]
+        return response
+        
     def load_model(self,
                    model_name_or_path: str = "THUDM/chatglm-6b",
                    llm_device=LLM_DEVICE,