加入了运行时切换模型的功能

ChuanhuChatbot.py

@@ -10,7 +10,7 @@ from modules.config import *
 from modules.utils import *
 from modules.presets import *
 from modules.overwrites import *
-from modules.models import get_model
+from modules.models import ModelManager
 
 gr.Chatbot.postprocess = postprocess
 PromptHelper.compact_text_chunks = compact_text_chunks
@@ -22,7 +22,7 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
     user_name = gr.State("")
     promptTemplates = gr.State(load_template(get_template_names(plain=True)[0], mode=2))
     user_question = gr.State("")
-    current_model = gr.State(get_model(MODELS[DEFAULT_MODEL], my_api_key)[0])
+    current_model = gr.State(ModelManager(model_name = MODELS[DEFAULT_MODEL], access_key = my_api_key))
 
     topic = gr.State("未命名对话历史记录")
 
@@ -197,7 +197,7 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
                             interactive=True,
                             label="max context",
                         )
-                        max_tokens_slider = gr.Slider(
+                        max_generation_slider = gr.Slider(
                             minimum=1,
                             maximum=32768,
                             value=1000,
@@ -350,7 +350,7 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
     # LLM Models
     keyTxt.change(current_model.value.set_key, keyTxt, [status_display]).then(**get_usage_args)
     keyTxt.submit(**get_usage_args)
-    model_select_dropdown.change(get_model, [model_select_dropdown, keyTxt, temperature_slider, top_p_slider, systemPromptTxt], [current_model, status_display], show_progress=True)
+    model_select_dropdown.change(current_model.value.get_model, [model_select_dropdown, keyTxt, temperature_slider, top_p_slider, systemPromptTxt], [status_display], show_progress=True)
 
     # Template
     systemPromptTxt.change(current_model.value.set_system_prompt, [systemPromptTxt], None)
@@ -392,7 +392,7 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
     top_p_slider.change(current_model.value.set_top_p, [top_p_slider], None)
     n_choices_slider.change(current_model.value.set_n_choices, [n_choices_slider], None)
     stop_sequence_txt.change(current_model.value.set_stop_sequence, [stop_sequence_txt], None)
-    max_tokens_slider.change(current_model.value.set_max_tokens, [max_tokens_slider], None)
+    max_generation_slider.change(current_model.value.set_max_tokens, [max_generation_slider], None)
     presence_penalty_slider.change(current_model.value.set_presence_penalty, [presence_penalty_slider], None)
     frequency_penalty_slider.change(current_model.value.set_frequency_penalty, [frequency_penalty_slider], None)
     logit_bias_txt.change(current_model.value.set_logit_bias, [logit_bias_txt], None)

modules/base_model.py

@@ -126,13 +126,15 @@ class BaseLLMModel:
 
         stream_iter = self.get_answer_stream_iter()
 
-        self.history.append(construct_assistant(""))
         for partial_text in stream_iter:
-            self.history[-1] = construct_assistant(partial_text)
             chatbot[-1] = (chatbot[-1][0], partial_text + display_append)
             self.all_token_counts[-1] += 1
             status_text = self.token_message()
             yield get_return_value()
+            if self.interrupted:
+                    self.recover()
+                    break
+        self.history.append(construct_assistant(partial_text))
 
     def next_chatbot_at_once(self, inputs, chatbot, fake_input=None, display_append=""):
         if fake_input:
@@ -277,9 +279,6 @@ class BaseLLMModel:
             )
             for chatbot, status_text in iter:
                 yield chatbot, status_text
-                if self.interrupted:
-                    self.recover()
-                    break
         else:
             logging.debug("不使用流式传输")
             chatbot, status_text = self.next_chatbot_at_once(
@@ -326,13 +325,13 @@ class BaseLLMModel:
         files=None,
         reply_language="中文",
     ):
-        logging.info("重试中……")
+        logging.debug("重试中……")
         if len(self.history) == 0:
             yield chatbot, f"{STANDARD_ERROR_MSG}上下文是空的"
             return
 
+        inputs = self.history[-2]["content"]
         del self.history[-2:]
-        inputs = chatbot[-1][0]
         self.all_token_counts.pop()
         iter = self.predict(
             inputs,
@@ -344,7 +343,7 @@ class BaseLLMModel:
         )
         for x in iter:
             yield x
-        logging.info("重试完毕")
+        logging.debug("重试完毕")
 
     # def reduce_token_size(self, chatbot):
     #     logging.info("开始减少token数量……")

modules/models.py

@@ -247,7 +247,7 @@ class ChatGLM_Client(BaseLLMModel):
     def _get_glm_style_input(self):
         history = [x["content"] for x in self.history]
         query = history.pop()
-        logging.info(colorama.Fore.YELLOW + f"{history}" + colorama.Fore.RESET)
+        logging.debug(colorama.Fore.YELLOW + f"{history}" + colorama.Fore.RESET)
         assert (
             len(history) % 2 == 0
         ), f"History should be even length. current history is: {history}"
@@ -365,11 +365,12 @@ class LLaMA_Client(BaseLLMModel):
 
 class ModelManager:
     def __init__(self, **kwargs) -> None:
-        self.model, self.msg = self.get_model(**kwargs)
+        self.get_model(**kwargs)
 
     def get_model(
         self,
         model_name,
+        lora_model_path=None,
         access_key=None,
         temperature=None,
         top_p=None,
@@ -378,7 +379,6 @@ class ModelManager:
         msg = f"模型设置为了： {model_name}"
         logging.info(msg)
         model_type = ModelType.get_type(model_name)
-        print(model_type.name)
         if model_type == ModelType.OpenAI:
             model = OpenAIClient(
                 model_name=model_name,
@@ -389,7 +389,93 @@ class ModelManager:
             )
         elif model_type == ModelType.ChatGLM:
             model = ChatGLM_Client(model_name)
-        return model, msg
+        self.model = model
+        return msg
+
+    def predict(self, *args):
+        iter = self.model.predict(*args)
+        for i in iter:
+            yield i
+
+    def billing_info(self):
+        return self.model.billing_info()
+
+    def set_key(self, *args):
+        return self.model.set_key(*args)
+
+    def load_chat_history(self, *args):
+        return self.model.load_chat_history(*args)
+
+    def interrupt(self, *args):
+        return self.model.interrupt(*args)
+
+    def reset(self, *args):
+        return self.model.reset(*args)
+
+    def retry(self, *args):
+        iter = self.model.retry(*args)
+        for i in iter:
+            yield i
+
+    def delete_first_conversation(self, *args):
+        return self.model.delete_first_conversation(*args)
+
+    def delete_last_conversation(self, *args):
+        return self.model.delete_last_conversation(*args)
+
+    def set_system_prompt(self, *args):
+        return self.model.set_system_prompt(*args)
+
+    def save_chat_history(self, *args):
+        return self.model.save_chat_history(*args)
+
+    def export_markdown(self, *args):
+        return self.model.export_markdown(*args)
+
+    def load_chat_history(self, *args):
+        return self.model.load_chat_history(*args)
+
+    def set_token_upper_limit(self, *args):
+        return self.model.set_token_upper_limit(*args)
+
+    # temperature_slider.change(current_model.value.set_temperature, [temperature_slider], None)
+    # top_p_slider.change(current_model.value.set_top_p, [top_p_slider], None)
+    # n_choices_slider.change(current_model.value.set_n_choices, [n_choices_slider], None)
+    # stop_sequence_txt.change(current_model.value.set_stop_sequence, [stop_sequence_txt], None)
+    # max_tokens_slider.change(current_model.value.set_max_tokens, [max_tokens_slider], None)
+    # presence_penalty_slider.change(current_model.value.set_presence_penalty, [presence_penalty_slider], None)
+    # frequency_penalty_slider.change(current_model.value.set_frequency_penalty, [frequency_penalty_slider], None)
+    # logit_bias_txt.change(current_model.value.set_logit_bias, [logit_bias_txt], None)
+    # user_identifier_txt.change(current_model.value.set_user_identifier, [user_identifier_txt], None)
+
+    def set_temperature(self, *args):
+        self.model.set_temperature(*args)
+
+    def set_top_p(self, *args):
+        self.model.set_top_p(*args)
+
+    def set_n_choices(self, *args):
+        self.model.set_n_choices(*args)
+
+    def set_stop_sequence(self, *args):
+        self.model.set_stop_sequence(*args)
+
+    def set_max_tokens(self, *args):
+        self.model.set_max_tokens(*args)
+
+    def set_presence_penalty(self, *args):
+        self.model.set_presence_penalty(*args)
+
+    def set_frequency_penalty(self, *args):
+        self.model.set_frequency_penalty(*args)
+
+    def set_logit_bias(self, *args):
+        self.model.set_logit_bias(*args)
+
+    def set_user_identifier(self, *args):
+        self.model.set_user_identifier(*args)
+
+
 
 
 if __name__ == "__main__":
@@ -397,8 +483,8 @@ if __name__ == "__main__":
         openai_api_key = cjson.load(f)["openai_api_key"]
     # set logging level to debug
     logging.basicConfig(level=logging.DEBUG)
-    # client, _ = get_model("gpt-3.5-turbo", openai_api_key)
-    client, _ = get_model("chatglm-6b-int4")
+    # client = ModelManager(model_name="gpt-3.5-turbo", access_key=openai_api_key)
+    client = ModelManager(model_name="chatglm-6b-int4")
     chatbot = []
     stream = False
     # 测试账单功能