Upload folder using huggingface_hub

ChatWorld/ChatWorld.py

@@ -1,7 +1,7 @@
 from jinja2 import Template
 import torch
 
-from .models import GLM
+from .models import GLM, GLM_api
 
 from .NaiveDB import NaiveDB
 from .utils import *
@@ -19,7 +19,7 @@ class ChatWorld:
 
         self.history = []
 
-        self.client = None
+        self.client = GLM_api()
         self.model = GLM()
         self.db = NaiveDB()
         self.prompt = Template(('Please be aware that your codename in this conversation is "{{model_role_name}}"'
@@ -81,17 +81,15 @@ class ChatWorld:
         return {"role": "system", "content": self.prompt.render(model_role_name=self.model_role_name, model_role_nickname=self.model_role_nickname, role_name=role_name, role_nickname=role_nick_name, RAG=rag)}
 
     def chat(self, text: str, user_role_name: str, user_role_nick_name: str = None, use_local_model=False):
+        self.history.append(
+            {"role": "user", "content": f"{user_role_name}:「{text}」"})
         message = [self.getSystemPrompt(text,
-                                        user_role_name, user_role_nick_name)] + self.history
-        print(message)
+                                        user_role_name, user_role_nick_name), {"role": "user", "content": f"{user_role_name}:「{text}」"}]
         if use_local_model:
             response = self.model.get_response(message)
         else:
-            response = self.client.chat(
-                user_role_name, text, user_role_nick_name)
+            response = self.client.chat(message)
 
-        self.history.append(
-            {"role": "user", "content": f"{user_role_name}:「{text}」"})
         self.history.append(
             {"role": "assistant", "content": f"{self.model_role_name}:「{response}」"})
         return response

ChatWorld/NaiveDB.py

@@ -81,7 +81,10 @@ class NaiveDB:
         similarities.sort(key=lambda x: x[0], reverse=True)
         self.last_search_ids = [x[1] for x in similarities[:n_results]]
 
-        
+        stories_length = len(self.stories)
+        search_id_range = [(max(0, i-3), min(i+4, stories_length))
+                           for i in self.last_search_ids]
 
-        top_stories = [self.stories[_id] for _id in self.last_search_ids]
+        top_stories = ["\n".join(self.stories[start:end+1])
+                       for start, end in search_id_range]
         return top_stories

ChatWorld/models.py

@@ -40,7 +40,7 @@ class GLM():
 
         self.client = client.eval()
 
-    def message2query(messages) -> str:
+    def message2query(self, messages) -> str:
         # [{'role': 'user', 'content': '老师: 同学请自我介绍一下'}]
         # <|system|>
         # You are ChatGLM3, a large language model trained by Zhipu.AI. Follow the user's instructions carefully. Respond using markdown.
@@ -53,7 +53,9 @@ class GLM():
         return "".join([template.substitute(message) for message in messages])
 
     def get_response(self, message):
-        response, history = self.client.chat(self.tokenizer, message)
+        response, history = self.client.chat(
+            self.tokenizer, self.message2query(message))
+        print(self.message2query(message))
         return response
 
 
@@ -62,7 +64,8 @@ class GLM_api:
         self.client = ZhipuAI(api_key=os.environ["ZHIPU_API_KEY"])
         self.model = model_name
 
-    def getResponse(self, message):
+    def chat(self, message):
+        print(message)
         response = self.client.chat.completions.create(
-            model=self.model, prompt=message)
-        return response.choices[0].message
+            model=self.model, messages=message)
+        return response.choices[0].message.content

app.py

@@ -11,6 +11,8 @@ logging.basicConfig(level=logging.INFO, filename="demo.log", filemode="w",
 
 chatWorld = ChatWorld()
 
+role_name_list_global = None
+
 
 def getContent(input_file):
     # 读取文件内容
@@ -31,33 +33,57 @@ def getContent(input_file):
     role_name_list = [i for i in role_name_set if i != ""]
     logging.info(f"role_name_list: {role_name_list}")
 
+    global role_name_list_global
+    role_name_list_global = role_name_list
+
     return gr.Radio(choices=role_name_list, interactive=True, value=role_name_list[0]), gr.Radio(choices=role_name_list, interactive=True, value=role_name_list[-1])
 
 
 def submit_message(message, history, model_role_name, role_name, model_role_nickname, role_nickname):
+    print(f"history: {history}")
     chatWorld.setRoleName(model_role_name, model_role_nickname)
     response = chatWorld.chat(message,
                               role_name, role_nickname, use_local_model=True)
     return response
 
 
+def submit_message_api(message, history, model_role_name, role_name, model_role_nickname, role_nickname):
+    print(f"history: {history}")
+    chatWorld.setRoleName(model_role_name, model_role_nickname)
+    response = chatWorld.chat(message,
+                              role_name, role_nickname, use_local_model=False)
+    return response
+
+
+def get_role_list():
+    global role_name_list_global
+    if role_name_list_global:
+        return role_name_list_global
+    else:
+        return []
+
+
 with gr.Blocks() as demo:
 
     upload_c = gr.File(label="上传文档文件")
 
     with gr.Row():
-        model_role_name = gr.Radio([], label="模型角色名")
+        model_role_name = gr.Radio(get_role_list(), label="模型角色名")
         model_role_nickname = gr.Textbox(label="模型角色昵称")
 
     with gr.Row():
-        role_name = gr.Radio([], label="角色名")
+        role_name = gr.Radio(get_role_list(), label="角色名")
         role_nickname = gr.Textbox(label="角色昵称")
 
     upload_c.upload(fn=getContent, inputs=upload_c,
                     outputs=[model_role_name, role_name])
 
-    chatBox = gr.ChatInterface(
-        submit_message, chatbot=gr.Chatbot(height=400, render=False), additional_inputs=[model_role_name, role_name, model_role_nickname, role_nickname])
+    with gr.Row():
+        chatBox_local = gr.ChatInterface(
+            submit_message, chatbot=gr.Chatbot(height=400, label="本地模型", render=False), additional_inputs=[model_role_name, role_name, model_role_nickname, role_nickname])
+
+        chatBox_api = gr.ChatInterface(
+            submit_message_api, chatbot=gr.Chatbot(height=400, label="API模型", render=False), additional_inputs=[model_role_name, role_name, model_role_nickname, role_nickname])
 
 
-demo.launch(share=True, debug=True, server_name="0.0.0.0")
+demo.launch(server_name="0.0.0.0")