添加保存对话为Markdown功能 (Akegarasu#9)

* Add button to save history as md file * Add function to save history as markdown
YunCheng66 · Mar 18, 2023 · 6b4268c · 6b4268c
1 parent 44061e6
commit 6b4268c
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 4 deletions.
diff --git a/modules/context.py b/modules/context.py
@@ -49,15 +49,27 @@ def limit_round(self):
             self.history = self.history[-self.max_rounds:]
             self.rh = self.rh[-self.max_rounds:]
 
+    # 保存为json文件，供webui读取使用
     def save_history(self):
-        if not os.path.exists("outputs"):
-            os.mkdir("outputs")
+        if not os.path.exists("chat_history"):
+            os.mkdir("chat_history")
 
         s = [{"q": i[0], "o": i[1]} for i in self.history]
-        filename = f"save-{int(time.time())}.json"
-        with open(os.path.join("outputs", filename), "w", encoding="utf-8") as f:
+        filename = f"history-{int(time.time())}.json"
+        with open(os.path.join("chat_history", filename), "w", encoding="utf-8") as f:
             f.write(json.dumps(s, ensure_ascii=False))
 
+    # 增加保存为markdown文档的功能，方便人类查看
+    def save_as_md(self):
+        if not os.path.exists("markdown"):
+            os.mkdir("markdown")
+
+        filename = f"history-{int(time.time())}.md"
+        with open(os.path.join("markdown", filename), "w", encoding="utf-8") as f:
+            for i in self.history:
+                f.write(f"# 我: {i[0]}\n\n")
+                f.write(f"ChatGLM: {i[1]}\n\n")
+
     def load_history(self, file):
         try:
             with open(file.name, "r", encoding='utf-8') as f:

diff --git a/modules/ui.py b/modules/ui.py
@@ -61,6 +61,7 @@ def create_ui():
 
                         with gr.Row():
                             save_his_btn = gr.Button("保存对话")
+                            save_md_btn = gr.Button("保存为MarkDown")
                             load_his_btn = gr.UploadButton("读取对话", file_types=['file'], file_count='single')
 
             with gr.Column(scale=7):
@@ -86,6 +87,7 @@ def create_ui():
         clear_input.click(lambda x: "", inputs=[input_message], outputs=[input_message])
 
         save_his_btn.click(ctx.save_history)
+        save_md_btn.click(ctx.save_as_md)
         load_his_btn.upload(ctx.load_history, inputs=[
             load_his_btn,
         ], outputs=[