Upload 2 files

app.py

@@ -74,6 +74,8 @@ with gr.Blocks(title=title) as demo:
         neighborhoods = mb.search(history[-1][0])
         messages += neighborhoods
         messages.append({"role": "user", "content": history[-1][0]})
+
+        print('messages:', messages)
         print('user:', {"role": "user", "content": history[-1][0]})
         print('neighborhoods:', neighborhoods)
 
@@ -101,6 +103,7 @@ with gr.Blocks(title=title) as demo:
     def reset_model():
         global messages
         messages = [{"role": "system", "content": prompt_default}]
+        mb.reset()
         print("已重置模型")
         return [[None, "已重置模型"]]
 
@@ -109,7 +112,7 @@ with gr.Blocks(title=title) as demo:
         user_content = text[0:text.find("|")]
         assistant_content = text[text.find("|") + 1:]
         mb.upsert([Dialogue(user_content, assistant_content)])
-        return [[None, f"已将样例存入记忆区块 ({user_content} -> {assistant_content})"]]
+        return [[None, f"已将对话样例存入记忆区块\n问: {user_content}\n答: {assistant_content}"]]
 
 
     msg.submit(user, [prompt, msg, chatbot], [msg, chatbot], queue=False).then(
@@ -121,3 +124,4 @@ with gr.Blocks(title=title) as demo:
 
 if __name__ == "__main__":
     demo.launch(height=520, show_api=False, auth=("luoshaoye", "19970812"))
+    # demo.launch(height=520, show_api=False)

memory.py

@@ -3,8 +3,7 @@
 # @Author  : BarryWang
 # @FileName: memory.py
 # @Github  : https://github.com/BarryWangQwQ
-
-# import time
+import uuid
 
 from txtai.embeddings import Embeddings
 
@@ -38,8 +37,8 @@ class MemoryBlocks:
     def upsert(self, dialogue_list):
         self.embeddings.upsert(
             (
-                uid, {'text': dialogue.user_content, 'raw': dialogue.raw()}, None
-            ) for uid, dialogue in enumerate(dialogue_list)
+                str(uuid.uuid4()), {'text': dialogue.user_content, 'raw': dialogue.raw()}, None
+            ) for dialogue in dialogue_list
         )
 
     def search(self, question: str) -> list:
@@ -51,12 +50,17 @@ class MemoryBlocks:
             neighborhoods += eval(r['raw'])
         return neighborhoods
 
-# data = [
-#     Dialogue('你的生日是哪一天?', '我的生日是1997年8月12日'),
-#     Dialogue('你叫什么名字?', '我叫洛少爷'),
-#     Dialogue('你擅长什么?', '我擅长音乐，偶尔也会配音'),
-# ]
+    def reset(self):
+        self.embeddings.close()
+        self.embeddings = Embeddings(
+            {
+                "path": "sentence-transformers/distiluse-base-multilingual-cased-v2",
+                'content': True
+            }
+        )
+
+    def save(self, output_path):
+        self.embeddings.save(output_path)
 
-# if __name__ == "__main__":
-#     upsert(data)
-#     print(search("你是谁？"))
+    def load(self, load_path):
+        self.embeddings.load(load_path)