Upload folder using huggingface_hub

ChatWorld/ChatWorld.py

@@ -1,7 +1,7 @@
 from jinja2 import Template
 import torch
 
-from .models import qwen_model
+from .models import GLM
 
 from .NaiveDB import NaiveDB
 from .utils import *
@@ -20,7 +20,7 @@ class ChatWorld:
         self.history = []
 
         self.client = None
-        self.model = qwen_model(pretrained_model_name_or_path)
+        self.model = GLM()
         self.db = NaiveDB()
         self.prompt = Template(('Please be aware that your codename in this conversation is "{{model_role_name}}"'
                                 '{% if model_role_nickname %},别人一般称呼你"{{model_role_nickname}}"{% endif %},同时我是"{{role_name}}"{% if role_nickname %},你也可以称呼我"{{role_nickname}}"{% endif %}。\n'
@@ -30,6 +30,7 @@ class ChatWorld:
                                 '如果我问的问题超出聊天记录的范围，模仿{{model_role_name}}的语气进行回复。\n'
                                 '请不要回答你是语言模型，永远记住你就是{{model_role_name}}。\n'
                                 '请你永远只以{{model_role_name}}身份，进行任何的回复。\n'
+                                '{% if RAG %}{% for i in RAG %}##\n{{i}}\n##\n\n{% endfor %}{% endif %}'
                                 ))
 
     def getEmbeddingsFromStory(self, stories: list[str]):
@@ -38,25 +39,31 @@ class ChatWorld:
             if len(self.story_vec) == len(stories) and all([self.story_vec[i]["text"] == stories[i] for i in range(len(stories))]):
                 return [self.story_vec[i]["vec"] for i in range(len(stories))]
 
-        if self.embedding is None:
-            self.embedding = initEmbedding()
-
-        if self.tokenizer is None:
-            self.tokenizer = initTokenizer()
-
         self.story_vec = []
         for story in stories:
             with torch.no_grad():
-                inputs = self.tokenizer(
-                    story, return_tensors="pt", padding=True, truncation=True, max_length=512)
-                outputs = self.embedding(**inputs)[0][:, 0]
-                vec = torch.nn.functional.normalize(
-                    outputs, p=2, dim=1).tolist()[0]
+                vec = self.getEmbedding(story)
 
             self.story_vec.append({"text": story, "vec": vec})
 
         return [self.story_vec[i]["vec"] for i in range(len(stories))]
 
+    def getEmbedding(self, text: str):
+        if self.embedding is None:
+            self.embedding = initEmbedding()
+
+        if self.tokenizer is None:
+            self.tokenizer = initTokenizer()
+
+        with torch.no_grad():
+            inputs = self.tokenizer(
+                text, return_tensors="pt", padding=True, truncation=True, max_length=512).to(self.embedding.device)
+            outputs = self.embedding(**inputs)[0][:, 0]
+            vec = torch.nn.functional.normalize(
+                outputs, p=2, dim=1).tolist()[0]
+
+        return vec
+
     def initDB(self, storys: list[str]):
         story_vecs = self.getEmbeddingsFromStory(storys)
         self.db.build_db(storys, story_vecs)
@@ -65,21 +72,26 @@ class ChatWorld:
         self.model_role_name = role_name
         self.model_role_nickname = role_nick_name
 
-    def getSystemPrompt(self, role_name, role_nick_name):
+    def getSystemPrompt(self, text,  role_name, role_nick_name):
         assert self.model_role_name, "Please set model role name first"
 
-        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)}
+        query = self.getEmbedding(text)
+        rag = self.db.search(query, 5)
 
-    def chat(self, user_role_name: str, text: str, user_role_nick_name: str = None, use_local_model=False):
-        message = [self.getSystemPrompt(
-            user_role_name, user_role_nick_name)] + self.history
+        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):
+        message = [self.getSystemPrompt(text,
+                                        user_role_name, user_role_nick_name)] + self.history
+        print(message)
         if use_local_model:
             response = self.model.get_response(message)
         else:
             response = self.client.chat(
                 user_role_name, text, user_role_nick_name)
 
-        self.history.append({"role": "user", "content": text})
-        self.history.append({"role": "model", "content": response})
+        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,5 +81,7 @@ class NaiveDB:
         similarities.sort(key=lambda x: x[0], reverse=True)
         self.last_search_ids = [x[1] for x in similarities[:n_results]]
 
-        top_indices = [x[1] for x in similarities[:n_results]]
-        return top_indices
+        
+
+        top_stories = [self.stories[_id] for _id in self.last_search_ids]
+        return top_stories

ChatWorld/models.py

@@ -1,4 +1,7 @@
+import os
+from string import Template
 from transformers import AutoTokenizer, AutoModelForCausalLM
+from zhipuai import ZhipuAI
 
 
 class qwen_model:
@@ -11,7 +14,9 @@ class qwen_model:
     def get_response(self, message):
         message = self.tokenizer.apply_chat_template(
             message, tokenize=False, add_generation_prompt=True)
-        model_inputs = self.tokenizer([message], return_tensors="pt")
+        print(message)
+        model_inputs = self.tokenizer(
+            [message], return_tensors="pt").to(self.model.device)
         generated_ids = self.model.generate(
             model_inputs.input_ids,
             max_new_tokens=512
@@ -22,4 +27,42 @@ class qwen_model:
 
         response = self.tokenizer.batch_decode(
             generated_ids, skip_special_tokens=True)[0]
+
         return response
+
+
+class GLM():
+    def __init__(self, model_name="silk-road/Haruhi-Zero-GLM3-6B-0_4"):
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_name, trust_remote_code=True)
+        client = AutoModelForCausalLM.from_pretrained(
+            model_name, trust_remote_code=True, device_map="auto")
+
+        client = client.eval()
+
+    def message2query(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.
+        # <|user|>
+        # Hello
+        # <|assistant|>
+        # Hello, I'm ChatGLM3. What can I assist you today?
+        template = Template("<|$role|>\n$content\n")
+
+        return "".join([template.substitute(message) for message in messages])
+
+    def get_response(self, message):
+        response, history = self.client.chat(self.tokenizer, message)
+        return response
+
+
+class GLM_api:
+    def __init__(self, model_name="glm-4"):
+        self.client = ZhipuAI(api_key=os.environ["ZHIPU_API_KEY"])
+        self.model = model_name
+
+    def getResponse(self, message):
+        response = self.client.chat.completions.create(
+            model=self.model, prompt=message)
+        return response.choices[0].message

app.py

@@ -38,8 +38,8 @@ def getContent(input_file):
 
 def submit_message(message, history, model_role_name, role_name, model_role_nickname, role_nickname):
     chatWorld.setRoleName(model_role_name, model_role_nickname)
-    response = chatWorld.chat(
-        role_name, message, role_nickname, use_local_model=True)
+    response = chatWorld.chat(message,
+                              role_name, role_nickname, use_local_model=True)
     return response
 
 

run_gradio.sh

@@ -1,5 +1,6 @@
 export CUDA_VISIBLE_DEVICES=0
 export HF_HOME="/workspace/jyh/.cache/huggingface"
+export HF_ENDPOINT="https://hf-mirror.com"
 
 # Start the gradio server
 /workspace/jyh/miniconda3/envs/ChatWorld/bin/python /workspace/jyh/Zero-Haruhi/app.py