Added support for multi-modal Model: XMBot

ChuanhuChatbot.py

@@ -12,6 +12,7 @@ from modules.presets import *
 from modules.overwrites import *
 from modules.models import get_model
 
+gr.Chatbot._postprocess_chat_messages = postprocess_chat_messages
 gr.Chatbot.postprocess = postprocess
 PromptHelper.compact_text_chunks = compact_text_chunks
 
@@ -321,6 +322,8 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
     submitBtn.click(**transfer_input_args).then(**chatgpt_predict_args).then(**end_outputing_args)
     submitBtn.click(**get_usage_args)
 
+    index_files.change(handle_file_upload, [current_model, index_files, chatbot], [index_files, chatbot, status_display])
+
     emptyBtn.click(
         reset,
         inputs=[current_model],

modules/base_model.py

@@ -8,6 +8,7 @@ import os
 import sys
 import requests
 import urllib3
+import traceback
 
 from tqdm import tqdm
 import colorama
@@ -28,6 +29,7 @@ class ModelType(Enum):
     OpenAI = 0
     ChatGLM = 1
     LLaMA = 2
+    XMBot = 3
 
     @classmethod
     def get_type(cls, model_name: str):
@@ -39,6 +41,8 @@ class ModelType(Enum):
             model_type = ModelType.ChatGLM
         elif "llama" in model_name_lower or "alpaca" in model_name_lower:
             model_type = ModelType.LLaMA
+        elif "xmbot" in model_name_lower:
+            model_type = ModelType.XMBot
         else:
             model_type = ModelType.Unknown
         return model_type
@@ -164,10 +168,19 @@ class BaseLLMModel:
         status_text = self.token_message()
         return chatbot, status_text
 
-    def prepare_inputs(self, inputs, use_websearch, files, reply_language):
-        old_inputs = None
+    def handle_file_upload(self, files, chatbot):
+        """if the model accepts multi modal input, implement this function"""
+        status = gr.Markdown.update()
+        if files:
+            construct_index(self.api_key, file_src=files)
+            status = "索引构建完成"
+        return gr.Files.update(), chatbot, status
+
+    def prepare_inputs(self, real_inputs, use_websearch, files, reply_language, chatbot):
+        fake_inputs = None
         display_append = []
         limited_context = False
+        fake_inputs = real_inputs
         if files:
             from llama_index.indices.vector_store.base_query import GPTVectorStoreIndexQuery
             from llama_index.indices.query.schema import QueryBundle
@@ -180,12 +193,11 @@ class BaseLLMModel:
                 OpenAIEmbedding,
             )
             limited_context = True
-            old_inputs = inputs
-            msg = "加载索引中……（这可能需要几分钟）"
+            msg = "加载索引中……"
             logging.info(msg)
             # yield chatbot + [(inputs, "")], msg
             index = construct_index(self.api_key, file_src=files)
-            assert index is not None, "索引构建失败"
+            assert index is not None, "获取索引失败"
             msg = "索引获取成功，生成回答中……"
             logging.info(msg)
             if local_embedding or self.model_type != ModelType.OpenAI:
@@ -212,22 +224,21 @@ class BaseLLMModel:
                     vector_store=index._vector_store,
                     docstore=index._docstore,
                 )
-                query_bundle = QueryBundle(inputs)
+                query_bundle = QueryBundle(real_inputs)
                 nodes = query_object.retrieve(query_bundle)
             reference_results = [n.node.text for n in nodes]
             reference_results = add_source_numbers(reference_results, use_source=False)
             display_append = add_details(reference_results)
             display_append = "\n\n" + "".join(display_append)
-            inputs = (
+            real_inputs = (
                 replace_today(PROMPT_TEMPLATE)
-                .replace("{query_str}", inputs)
+                .replace("{query_str}", real_inputs)
                 .replace("{context_str}", "\n\n".join(reference_results))
                 .replace("{reply_language}", reply_language)
             )
         elif use_websearch:
             limited_context = True
-            search_results = ddg(inputs, max_results=5)
-            old_inputs = inputs
+            search_results = ddg(real_inputs, max_results=5)
             reference_results = []
             for idx, result in enumerate(search_results):
                 logging.debug(f"搜索结果{idx + 1}：{result}")
@@ -238,15 +249,15 @@ class BaseLLMModel:
                 )
             reference_results = add_source_numbers(reference_results)
             display_append = "\n\n" + "".join(display_append)
-            inputs = (
+            real_inputs = (
                 replace_today(WEBSEARCH_PTOMPT_TEMPLATE)
-                .replace("{query}", inputs)
+                .replace("{query}", real_inputs)
                 .replace("{web_results}", "\n\n".join(reference_results))
                 .replace("{reply_language}", reply_language)
             )
         else:
             display_append = ""
-        return limited_context, old_inputs, display_append, inputs
+        return limited_context, fake_inputs, display_append, real_inputs, chatbot
 
     def predict(
         self,
@@ -259,16 +270,17 @@ class BaseLLMModel:
         should_check_token_count=True,
     ):  # repetition_penalty, top_k
 
-
+        status_text = "开始生成回答……"
         logging.info(
             "输入为：" + colorama.Fore.BLUE + f"{inputs}" + colorama.Style.RESET_ALL
         )
         if should_check_token_count:
-            yield chatbot + [(inputs, "")], "开始生成回答……"
+            yield chatbot + [(inputs, "")], status_text
         if reply_language == "跟随问题语言（不稳定）":
             reply_language = "the same language as the question, such as English, 中文, 日本語, Español, Français, or Deutsch."
 
-        limited_context, old_inputs, display_append, inputs = self.prepare_inputs(inputs=inputs, use_websearch=use_websearch, files=files, reply_language=reply_language)
+        limited_context, fake_inputs, display_append, inputs, chatbot = self.prepare_inputs(real_inputs=inputs, use_websearch=use_websearch, files=files, reply_language=reply_language, chatbot=chatbot)
+        yield chatbot + [(fake_inputs, "")], status_text
 
         if (
             self.need_api_key and
@@ -303,7 +315,7 @@ class BaseLLMModel:
                 iter = self.stream_next_chatbot(
                     inputs,
                     chatbot,
-                    fake_input=old_inputs,
+                    fake_input=fake_inputs,
                     display_append=display_append,
                 )
                 for chatbot, status_text in iter:
@@ -313,11 +325,12 @@ class BaseLLMModel:
                 chatbot, status_text = self.next_chatbot_at_once(
                     inputs,
                     chatbot,
-                    fake_input=old_inputs,
+                    fake_input=fake_inputs,
                     display_append=display_append,
                 )
                 yield chatbot, status_text
         except Exception as e:
+            traceback.print_exc()
             status_text = STANDARD_ERROR_MSG + str(e)
             yield chatbot, status_text
 

modules/models.py

@@ -16,6 +16,7 @@ from duckduckgo_search import ddg
 import asyncio
 import aiohttp
 from enum import Enum
+import uuid
 
 from .presets import *
 from .llama_func import *
@@ -75,7 +76,8 @@ class OpenAIClient(BaseLLMModel):
     def billing_info(self):
         try:
             curr_time = datetime.datetime.now()
-            last_day_of_month = get_last_day_of_month(curr_time).strftime("%Y-%m-%d")
+            last_day_of_month = get_last_day_of_month(
+                curr_time).strftime("%Y-%m-%d")
             first_day_of_month = curr_time.replace(day=1).strftime("%Y-%m-%d")
             usage_url = f"{shared.state.usage_api_url}?start_date={first_day_of_month}&end_date={last_day_of_month}"
             try:
@@ -112,7 +114,8 @@ class OpenAIClient(BaseLLMModel):
         openai_api_key = self.api_key
         system_prompt = self.system_prompt
         history = self.history
-        logging.debug(colorama.Fore.YELLOW + f"{history}" + colorama.Fore.RESET)
+        logging.debug(colorama.Fore.YELLOW +
+                      f"{history}" + colorama.Fore.RESET)
         headers = {
             "Content-Type": "application/json",
             "Authorization": f"Bearer {openai_api_key}",
@@ -217,7 +220,7 @@ class ChatGLM_Client(BaseLLMModel):
         global CHATGLM_TOKENIZER, CHATGLM_MODEL
         if CHATGLM_TOKENIZER is None or CHATGLM_MODEL is None:
             system_name = platform.system()
-            model_path=None
+            model_path = None
             if os.path.exists("models"):
                 model_dirs = os.listdir("models")
                 if model_name in model_dirs:
@@ -257,16 +260,19 @@ class ChatGLM_Client(BaseLLMModel):
     def _get_glm_style_input(self):
         history = [x["content"] for x in self.history]
         query = history.pop()
-        logging.debug(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}"
-        history = [[history[i], history[i + 1]] for i in range(0, len(history), 2)]
+        history = [[history[i], history[i + 1]]
+                   for i in range(0, len(history), 2)]
         return history, query
 
     def get_answer_at_once(self):
         history, query = self._get_glm_style_input()
-        response, _ = CHATGLM_MODEL.chat(CHATGLM_TOKENIZER, query, history=history)
+        response, _ = CHATGLM_MODEL.chat(
+            CHATGLM_TOKENIZER, query, history=history)
         return response, len(response)
 
     def get_answer_stream_iter(self):
@@ -315,8 +321,10 @@ class LLaMA_Client(BaseLLMModel):
                 # raise Exception(f"models目录下没有这个模型: {model_name}")
             if lora_path is not None:
                 lora_path = f"lora/{lora_path}"
-            model_args = ModelArguments(model_name_or_path=model_source, lora_model_path=lora_path, model_type=None, config_overrides=None, config_name=None, tokenizer_name=None, cache_dir=None, use_fast_tokenizer=True, model_revision='main', use_auth_token=False, torch_dtype=None, use_lora=False, lora_r=8, lora_alpha=32, lora_dropout=0.1, use_ram_optimized_load=True)
-            pipeline_args = InferencerArguments(local_rank=0, random_seed=1, deepspeed='configs/ds_config_chatbot.json', mixed_precision='bf16')
+            model_args = ModelArguments(model_name_or_path=model_source, lora_model_path=lora_path, model_type=None, config_overrides=None, config_name=None, tokenizer_name=None, cache_dir=None,
+                                        use_fast_tokenizer=True, model_revision='main', use_auth_token=False, torch_dtype=None, use_lora=False, lora_r=8, lora_alpha=32, lora_dropout=0.1, use_ram_optimized_load=True)
+            pipeline_args = InferencerArguments(
+                local_rank=0, random_seed=1, deepspeed='configs/ds_config_chatbot.json', mixed_precision='bf16')
 
             with open(pipeline_args.deepspeed, "r") as f:
                 ds_config = json.load(f)
@@ -341,7 +349,6 @@ class LLaMA_Client(BaseLLMModel):
         #     " unconditionally."
         # )
 
-
     def _get_llama_style_input(self):
         history = []
         instruction = ""
@@ -379,7 +386,8 @@ class LLaMA_Client(BaseLLMModel):
         step = 1
         for _ in range(0, self.max_generation_token, step):
             input_dataset = self.dataset.from_dict(
-                {"type": "text_only", "instances": [{"text": context + partial_text}]}
+                {"type": "text_only", "instances": [
+                    {"text": context + partial_text}]}
             )
             output_dataset = LLAMA_INFERENCER.inference(
                 model=LLAMA_MODEL,
@@ -394,6 +402,94 @@ class LLaMA_Client(BaseLLMModel):
             yield partial_text
 
 
+class XMBot_Client(BaseLLMModel):
+    def __init__(self, api_key):
+        super().__init__(model_name="xmbot")
+        self.api_key = api_key
+        self.session_id = None
+        self.reset()
+        self.image_bytes = None
+        self.image_path = None
+        self.xm_history = []
+        self.url = "https://xmbot.net/web"
+
+    def reset(self):
+        self.session_id = str(uuid.uuid4())
+        return [], "已重置"
+
+    def try_read_image(self, filepath):
+        import base64
+
+        def is_image_file(filepath):
+            # 判断文件是否为图片
+            valid_image_extensions = [".jpg", ".jpeg", ".png", ".bmp", ".gif", ".tiff"]
+            file_extension = os.path.splitext(filepath)[1].lower()
+            return file_extension in valid_image_extensions
+
+        def read_image_as_bytes(filepath):
+            # 读取图片文件并返回比特流
+            with open(filepath, "rb") as f:
+                image_bytes = f.read()
+            return image_bytes
+
+        if is_image_file(filepath):
+            logging.info(f"读取图片文件: {filepath}")
+            image_bytes = read_image_as_bytes(filepath)
+            base64_encoded_image = base64.b64encode(image_bytes).decode()
+            self.image_bytes = base64_encoded_image
+            self.image_path = filepath
+        else:
+            self.image_bytes = None
+            self.image_path = None
+
+    def prepare_inputs(self, real_inputs, use_websearch, files, reply_language, chatbot):
+        fake_inputs = real_inputs
+        display_append = ""
+        limited_context = False
+        return limited_context, fake_inputs, display_append, real_inputs, chatbot
+
+    def handle_file_upload(self, files, chatbot):
+        """if the model accepts multi modal input, implement this function"""
+        if files:
+            for file in files:
+                if file.name:
+                    logging.info(f"尝试读取图像: {file.name}")
+                    self.try_read_image(file.name)
+            if self.image_path is not None:
+                chatbot = chatbot + [((self.image_path,), None)]
+            if self.image_bytes is not None:
+                logging.info("使用图片作为输入")
+                conv_id = str(uuid.uuid4())
+                data = {
+                    "user_id": self.api_key,
+                    "session_id": self.session_id,
+                    "uuid": conv_id,
+                    "data_type": "imgbase64",
+                    "data": self.image_bytes
+                }
+                # response = requests.post(self.url, json=data)
+                # response = json.loads(response.text)
+                # logging.info(f"图片回复: {response['data']}")
+                logging.info("发送了图片")
+        return None, chatbot, None
+
+    def get_answer_at_once(self):
+        question = self.history[-1]["content"]
+        conv_id = str(uuid.uuid4())
+        data = {
+            "user_id": self.api_key,
+            "session_id": self.session_id,
+            "uuid": conv_id,
+            "data_type": "text",
+            "data": question
+        }
+        response = requests.post(self.url, json=data)
+        response = json.loads(response.text)
+        return response["data"], len(response["data"])
+
+
+
+
 def get_model(
     model_name,
     lora_model_path=None,
@@ -429,7 +525,8 @@ def get_model(
             logging.info(msg)
             lora_selector_visibility = True
             if os.path.isdir("lora"):
-                lora_choices = get_file_names("lora", plain=True, filetypes=[""])
+                lora_choices = get_file_names(
+                    "lora", plain=True, filetypes=[""])
             lora_choices = ["No LoRA"] + lora_choices
         elif model_type == ModelType.LLaMA and lora_model_path != "":
             logging.info(f"正在加载LLaMA模型: {model_name} + {lora_model_path}")
@@ -440,6 +537,8 @@ def get_model(
             else:
                 msg += f" + {lora_model_path}"
             model = LLaMA_Client(model_name, lora_model_path)
+        elif model_type == ModelType.XMBot:
+            model = XMBot_Client(api_key=access_key)
         elif model_type == ModelType.Unknown:
             raise ValueError(f"未知模型: {model_name}")
         logging.info(msg)

modules/overwrites.py

@@ -4,6 +4,7 @@ import logging
 from llama_index import Prompt
 from typing import List, Tuple
 import mdtex2html
+from gradio_client import utils as client_utils
 
 from modules.presets import *
 from modules.llama_func import *
@@ -20,23 +21,60 @@ def compact_text_chunks(self, prompt: Prompt, text_chunks: List[str]) -> List[st
 
 
 def postprocess(
-    self, y: List[Tuple[str | None, str | None]]
-) -> List[Tuple[str | None, str | None]]:
-    """
-    Parameters:
-        y: List of tuples representing the message and response pairs. Each message and response should be a string, which may be in Markdown format.
-    Returns:
-        List of tuples representing the message and response. Each message and response will be a string of HTML.
-    """
-    if y is None or y == []:
-        return []
-    user, bot = y[-1]
-    if not detect_converted_mark(user):
-        user = convert_asis(user)
-    if not detect_converted_mark(bot):
-        bot = convert_mdtext(bot)
-    y[-1] = (user, bot)
-    return y
+        self,
+        y: List[List[str | Tuple[str] | Tuple[str, str] | None] | Tuple],
+    ) -> List[List[str | Dict | None]]:
+        """
+        Parameters:
+            y: List of lists representing the message and response pairs. Each message and response should be a string, which may be in Markdown format.  It can also be a tuple whose first element is a string filepath or URL to an image/video/audio, and second (optional) element is the alt text, in which case the media file is displayed. It can also be None, in which case that message is not displayed.
+        Returns:
+            List of lists representing the message and response. Each message and response will be a string of HTML, or a dictionary with media information. Or None if the message is not to be displayed.
+        """
+        if y is None:
+            return []
+        processed_messages = []
+        for message_pair in y:
+            assert isinstance(
+                message_pair, (tuple, list)
+            ), f"Expected a list of lists or list of tuples. Received: {message_pair}"
+            assert (
+                len(message_pair) == 2
+            ), f"Expected a list of lists of length 2 or list of tuples of length 2. Received: {message_pair}"
+
+            processed_messages.append(
+                [
+                    self._postprocess_chat_messages(message_pair[0], "user"),
+                    self._postprocess_chat_messages(message_pair[1], "bot"),
+                ]
+            )
+        return processed_messages
+
+def postprocess_chat_messages(
+        self, chat_message: str | Tuple | List | None, message_type: str
+    ) -> str | Dict | None:
+        if chat_message is None:
+            return None
+        elif isinstance(chat_message, (tuple, list)):
+            filepath = chat_message[0]
+            mime_type = client_utils.get_mimetype(filepath)
+            filepath = self.make_temp_copy_if_needed(filepath)
+            return {
+                "name": filepath,
+                "mime_type": mime_type,
+                "alt_text": chat_message[1] if len(chat_message) > 1 else None,
+                "data": None,  # These last two fields are filled in by the frontend
+                "is_file": True,
+            }
+        elif isinstance(chat_message, str):
+            if message_type == "bot":
+                if not detect_converted_mark(chat_message):
+                    chat_message = convert_mdtext(chat_message)
+            elif message_type == "user":
+                if not detect_converted_mark(chat_message):
+                    chat_message = convert_asis(chat_message)
+            return chat_message
+        else:
+            raise ValueError(f"Invalid message for Chatbot component: {chat_message}")
 
 with open("./assets/custom.js", "r", encoding="utf-8") as f, open("./assets/Kelpy-Codos.js", "r", encoding="utf-8") as f2:
     customJS = f.read()

modules/presets.py

@@ -29,7 +29,7 @@ PROXY_ERROR_MSG = "代理错误，无法获取对话。"  # 代理错误
 SSL_ERROR_PROMPT = "SSL错误，无法获取对话。"  # SSL 错误
 NO_APIKEY_MSG = "API key为空，请检查是否输入正确。"  # API key 长度不足 51 位
 NO_INPUT_MSG = "请输入对话内容。"  # 未输入对话内容
-BILLING_NOT_APPLICABLE_MSG = "模型本地运行中" # 本地运行的模型返回的账单信息
+BILLING_NOT_APPLICABLE_MSG = "账单信息不适用" # 本地运行的模型返回的账单信息
 
 TIMEOUT_STREAMING = 60  # 流式对话时的超时时间
 TIMEOUT_ALL = 200  # 非流式对话时的超时时间
@@ -72,6 +72,7 @@ MODELS = [
     "gpt-4-0314",
     "gpt-4-32k",
     "gpt-4-32k-0314",
+    "xmbot",
     "chatglm-6b",
     "chatglm-6b-int4",
     "chatglm-6b-int4-qe",
@@ -85,6 +86,8 @@ MODELS = [
     "llama-65b-hf",
 ]  # 可选的模型
 
+DEFAULT_MODEL = 0  # 默认的模型在MODELS中的序号，从0开始数
+
 os.makedirs("models", exist_ok=True)
 os.makedirs("lora", exist_ok=True)
 os.makedirs("history", exist_ok=True)
@@ -93,8 +96,6 @@ for dir_name in os.listdir("models"):
         if dir_name not in MODELS:
             MODELS.append(dir_name)
 
-DEFAULT_MODEL = 0  # 默认的模型在MODELS中的序号，从0开始数
-
 MODEL_TOKEN_LIMIT = {
     "gpt-3.5-turbo": 4096,
     "gpt-3.5-turbo-0301": 4096,

modules/utils.py

@@ -33,7 +33,7 @@ if TYPE_CHECKING:
     class DataframeData(TypedDict):
         headers: List[str]
         data: List[List[str | int | bool]]
-        
+
 def predict(current_model, *args):
     iter = current_model.predict(*args)
     for i in iter:
@@ -110,6 +110,9 @@ def set_user_identifier(current_model, *args):
 def set_single_turn(current_model, *args):
     current_model.set_single_turn(*args)
 
+def handle_file_upload(current_model, *args):
+    return current_model.handle_file_upload(*args)
+
 
 def count_token(message):
     encoding = tiktoken.get_encoding("cl100k_base")
@@ -197,10 +200,13 @@ def convert_asis(userinput):
 
 
 def detect_converted_mark(userinput):
-    if userinput.endswith(ALREADY_CONVERTED_MARK):
+    try:
+        if userinput.endswith(ALREADY_CONVERTED_MARK):
+            return True
+        else:
+            return False
+    except:
         return True
-    else:
-        return False
 
 
 def detect_language(code):