Update app.py

app.py

@@ -18,41 +18,46 @@ NAME_DICT = {'汤师爷': 'tangshiye', '慕容复': 'murongfu', '李云龙': 'li
              '乔峰': 'qiaofeng', '神里绫华': 'ayaka', '雷电将军': 'raidenShogun', '于谦': 'yuqian'}
 
 
-def get_response(user_name, user_text, role, chatbot):
-    # 下载role zip 并解压
-    role_en = NAME_DICT[role]
-    if not os.path.exists("characters_zip"):
-        os.makedirs("characters_zip")
-    if not os.path.exists("characters"):
-        os.makedirs("characters")
-    if NAME_DICT[role] not in os.listdir("./characters"):
-        file_url = f"https://github.com/LC1332/Haruhi-2-Dev/raw/main/data/character_in_zip/{role_en}.zip"
-        os.makedirs(f"characters/{role_en}")
-        destination_file = f"characters_zip/{role_en}.zip"
-        wget.download(file_url, destination_file)
-        destination_folder = f"characters/{role_en}"
-        with zipfile.ZipFile(destination_file, 'r') as zip_ref:
-            zip_ref.extractall(destination_folder)
-
-    db_folder = f"./characters/{role_en}/content/{role_en}"
-    system_prompt = f"./characters/{role_en}/content/system_prompt.txt"
-    haruhi = ChatHaruhi(system_prompt=system_prompt,
+try:
+  os.makedirs("characters_zip")
+except:
+  pass
+try:
+  os.makedirs("characters")
+except:
+  pass
+AI_ROLES_OBJ = {}
+for ai_role_en in NAME_DICT.values():
+  file_url = f"https://github.com/LC1332/Haruhi-2-Dev/raw/main/data/character_in_zip/{ai_role_en}.zip"
+  os.makedirs(f"characters/{ai_role_en}")
+  destination_file = f"characters_zip/{ai_role_en}.zip"
+  wget.download(file_url, destination_file)
+  destination_folder = f"characters/{ai_role_en}"
+  with zipfile.ZipFile(destination_file, 'r') as zip_ref:
+      zip_ref.extractall(destination_folder)
+  db_folder = f"./characters/{ai_role_en}/content/{ai_role_en}"
+  system_prompt = f"./characters/{ai_role_en}/content/system_prompt.txt"
+  AI_ROLES_OBJ[ai_role_en] = ChatHaruhi(system_prompt=system_prompt,
                         llm="openai",
                         story_db=db_folder,
                         verbose=True)
-    response = haruhi.chat(role=user_name, text=user_text)
-    user_msg = user_name + ':「' + user_text + '」'
+
+
+def get_response(user_role, user_text, ai_role, chatbot):
+    ai_role_en = NAME_DICT[ai_role]
+    response = AI_ROLES_OBJ[ai_role_en].chat(role=user_role, text=user_text)
+    user_msg = user_role + ':「' + user_text + '」'
     chatbot.append((user_msg, response))
     return chatbot, None
 
 
-def clear(user_name, user_text, chatbot):
+def clear(user_role, user_text, chatbot):
     return None, None, []
 
 
-def get_image(role):
-    role_en = NAME_DICT[role]
-    return Image.open(f'images/{role_en}.jpg'), None, None, []
+def get_image(ai_role):
+    ai_role_en = NAME_DICT[ai_role]
+    return Image.open(f'images/{ai_role_en}.jpg'), None, None, []
 
 
 with gr.Blocks() as demo:
@@ -74,12 +79,12 @@ with gr.Blocks() as demo:
         chatbot = gr.Chatbot()
         role_image = gr.Image(height=400, value="./images/haruhi.jpg")
     with gr.Row():
-        user_name = gr.Textbox(label="user_role")
+        user_role = gr.Textbox(label="user_role")
         user_text = gr.Textbox(label="user_text")
     with gr.Row():
         submit = gr.Button("Submit")
         clean = gr.ClearButton(value="Clear")
-    role = gr.Radio(['汤师爷', '慕容复', '李云龙',
+    ai_role = gr.Radio(['汤师爷', '慕容复', '李云龙',
                      'Luna', '王多鱼', 'Ron', '鸠摩智',
                      'Snape', '凉宫春日', 'Malfoy', '虚竹',
                      '萧峰', '段誉', 'Hermione', 'Dumbledore',
@@ -90,8 +95,8 @@ with gr.Blocks() as demo:
                      'Sheldon', 'Raj', 'Penny',
                      '韦小宝', '乔峰', '神里绫华',
                      '雷电将军', '于谦'], label="characters", value='凉宫春日')
-    role.change(get_image, role, [role_image, user_name, user_text, chatbot])
-    user_text.submit(fn=get_response, inputs=[user_name, user_text, role, chatbot], outputs=[chatbot, user_text])
-    submit.click(fn=get_response, inputs=[user_name, user_text, role, chatbot], outputs=[chatbot, user_text])
-    clean.click(clear, [user_name, user_text, chatbot], [user_name, user_text, chatbot])
+    role.change(get_image, ai_role, [role_image, user_role, user_text, chatbot])
+    user_text.submit(fn=get_response, inputs=[user_role, user_text, ai_role, chatbot], outputs=[chatbot, user_text])
+    submit.click(fn=get_response, inputs=[user_role, user_text, ai_role, chatbot], outputs=[chatbot, user_text])
+    clean.click(clear, [user_role, user_text, chatbot], [user_role, user_text, chatbot])
 demo.launch(debug=True)