支持StableLM

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.models import get_model
 
 
 gr.Chatbot._postprocess_chat_messages = postprocess_chat_messages

modules/base_model.py

@@ -1,561 +0,0 @@
-from __future__ import annotations
-from typing import TYPE_CHECKING, List
-
-import logging
-import json
-import commentjson as cjson
-import os
-import sys
-import requests
-import urllib3
-import traceback
-
-from tqdm import tqdm
-import colorama
-from duckduckgo_search import ddg
-import asyncio
-import aiohttp
-from enum import Enum
-
-from .presets import *
-from .llama_func import *
-from .utils import *
-from . import shared
-from .config import retrieve_proxy
-
-
-class ModelType(Enum):
-    Unknown = -1
-    OpenAI = 0
-    ChatGLM = 1
-    LLaMA = 2
-    XMChat = 3
-
-    @classmethod
-    def get_type(cls, model_name: str):
-        model_type = None
-        model_name_lower = model_name.lower()
-        if "gpt" in model_name_lower:
-            model_type = ModelType.OpenAI
-        elif "chatglm" in model_name_lower:
-            model_type = ModelType.ChatGLM
-        elif "llama" in model_name_lower or "alpaca" in model_name_lower:
-            model_type = ModelType.LLaMA
-        elif "xmchat" in model_name_lower:
-            model_type = ModelType.XMChat
-        else:
-            model_type = ModelType.Unknown
-        return model_type
-
-
-class BaseLLMModel:
-    def __init__(
-        self,
-        model_name,
-        system_prompt="",
-        temperature=1.0,
-        top_p=1.0,
-        n_choices=1,
-        stop=None,
-        max_generation_token=None,
-        presence_penalty=0,
-        frequency_penalty=0,
-        logit_bias=None,
-        user="",
-    ) -> None:
-        self.history = []
-        self.all_token_counts = []
-        self.model_name = model_name
-        self.model_type = ModelType.get_type(model_name)
-        try:
-            self.token_upper_limit = MODEL_TOKEN_LIMIT[model_name]
-        except KeyError:
-            self.token_upper_limit = DEFAULT_TOKEN_LIMIT
-        self.interrupted = False
-        self.system_prompt = system_prompt
-        self.api_key = None
-        self.need_api_key = False
-        self.single_turn = False
-
-        self.temperature = temperature
-        self.top_p = top_p
-        self.n_choices = n_choices
-        self.stop_sequence = stop
-        self.max_generation_token = None
-        self.presence_penalty = presence_penalty
-        self.frequency_penalty = frequency_penalty
-        self.logit_bias = logit_bias
-        self.user_identifier = user
-
-    def get_answer_stream_iter(self):
-        """stream predict, need to be implemented
-        conversations are stored in self.history, with the most recent question, in OpenAI format
-        should return a generator, each time give the next word (str) in the answer
-        """
-        logging.warning("stream predict not implemented, using at once predict instead")
-        response, _ = self.get_answer_at_once()
-        yield response
-
-    def get_answer_at_once(self):
-        """predict at once, need to be implemented
-        conversations are stored in self.history, with the most recent question, in OpenAI format
-        Should return:
-        the answer (str)
-        total token count (int)
-        """
-        logging.warning("at once predict not implemented, using stream predict instead")
-        response_iter = self.get_answer_stream_iter()
-        count = 0
-        for response in response_iter:
-            count += 1
-        return response, sum(self.all_token_counts) + count
-
-    def billing_info(self):
-        """get billing infomation, inplement if needed"""
-        logging.warning("billing info not implemented, using default")
-        return BILLING_NOT_APPLICABLE_MSG
-
-    def count_token(self, user_input):
-        """get token count from input, implement if needed"""
-        logging.warning("token count not implemented, using default")
-        return len(user_input)
-
-    def stream_next_chatbot(self, inputs, chatbot, fake_input=None, display_append=""):
-        def get_return_value():
-            return chatbot, status_text
-
-        status_text = i18n("开始实时传输回答……")
-        if fake_input:
-            chatbot.append((fake_input, ""))
-        else:
-            chatbot.append((inputs, ""))
-
-        user_token_count = self.count_token(inputs)
-        self.all_token_counts.append(user_token_count)
-        logging.debug(f"输入token计数: {user_token_count}")
-
-        stream_iter = self.get_answer_stream_iter()
-
-        for partial_text in stream_iter:
-            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:
-            chatbot.append((fake_input, ""))
-        else:
-            chatbot.append((inputs, ""))
-        if fake_input is not None:
-            user_token_count = self.count_token(fake_input)
-        else:
-            user_token_count = self.count_token(inputs)
-        self.all_token_counts.append(user_token_count)
-        ai_reply, total_token_count = self.get_answer_at_once()
-        self.history.append(construct_assistant(ai_reply))
-        if fake_input is not None:
-            self.history[-2] = construct_user(fake_input)
-        chatbot[-1] = (chatbot[-1][0], ai_reply + display_append)
-        if fake_input is not None:
-            self.all_token_counts[-1] += count_token(construct_assistant(ai_reply))
-        else:
-            self.all_token_counts[-1] = total_token_count - sum(self.all_token_counts)
-        status_text = self.token_message()
-        return chatbot, status_text
-
-    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
-            from langchain.embeddings.huggingface import HuggingFaceEmbeddings
-            from langchain.chat_models import ChatOpenAI
-            from llama_index import (
-                GPTSimpleVectorIndex,
-                ServiceContext,
-                LangchainEmbedding,
-                OpenAIEmbedding,
-            )
-            limited_context = True
-            msg = "加载索引中……"
-            logging.info(msg)
-            # yield chatbot + [(inputs, "")], msg
-            index = construct_index(self.api_key, file_src=files)
-            assert index is not None, "获取索引失败"
-            msg = "索引获取成功，生成回答中……"
-            logging.info(msg)
-            if local_embedding or self.model_type != ModelType.OpenAI:
-                embed_model = LangchainEmbedding(HuggingFaceEmbeddings(model_name = "sentence-transformers/distiluse-base-multilingual-cased-v2"))
-            else:
-                embed_model = OpenAIEmbedding()
-            # yield chatbot + [(inputs, "")], msg
-            with retrieve_proxy():
-                prompt_helper = PromptHelper(
-                    max_input_size=4096,
-                    num_output=5,
-                    max_chunk_overlap=20,
-                    chunk_size_limit=600,
-                )
-                from llama_index import ServiceContext
-
-                service_context = ServiceContext.from_defaults(
-                    prompt_helper=prompt_helper, embed_model=embed_model
-                )
-                query_object = GPTVectorStoreIndexQuery(
-                    index.index_struct,
-                    service_context=service_context,
-                    similarity_top_k=5,
-                    vector_store=index._vector_store,
-                    docstore=index._docstore,
-                )
-                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)
-            real_inputs = (
-                replace_today(PROMPT_TEMPLATE)
-                .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(real_inputs, max_results=5)
-            reference_results = []
-            for idx, result in enumerate(search_results):
-                logging.debug(f"搜索结果{idx + 1}：{result}")
-                domain_name = urllib3.util.parse_url(result["href"]).host
-                reference_results.append([result["body"], result["href"]])
-                display_append.append(
-                    # f"{idx+1}. [{domain_name}]({result['href']})\n"
-                    f"<li><a href=\"{result['href']}\" target=\"_blank\">{domain_name}</a></li>\n"
-                )
-            reference_results = add_source_numbers(reference_results)
-            display_append = "<ol>\n\n" + "".join(display_append) + "</ol>"
-            real_inputs = (
-                replace_today(WEBSEARCH_PTOMPT_TEMPLATE)
-                .replace("{query}", real_inputs)
-                .replace("{web_results}", "\n\n".join(reference_results))
-                .replace("{reply_language}", reply_language)
-            )
-        else:
-            display_append = ""
-        return limited_context, fake_inputs, display_append, real_inputs, chatbot
-
-    def predict(
-        self,
-        inputs,
-        chatbot,
-        stream=False,
-        use_websearch=False,
-        files=None,
-        reply_language="中文",
-        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, "")], status_text
-        if reply_language == "跟随问题语言（不稳定）":
-            reply_language = "the same language as the question, such as English, 中文, 日本語, Español, Français, or Deutsch."
-
-        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
-            self.api_key is None
-            and not shared.state.multi_api_key
-        ):
-            status_text = STANDARD_ERROR_MSG + NO_APIKEY_MSG
-            logging.info(status_text)
-            chatbot.append((inputs, ""))
-            if len(self.history) == 0:
-                self.history.append(construct_user(inputs))
-                self.history.append("")
-                self.all_token_counts.append(0)
-            else:
-                self.history[-2] = construct_user(inputs)
-            yield chatbot + [(inputs, "")], status_text
-            return
-        elif len(inputs.strip()) == 0:
-            status_text = STANDARD_ERROR_MSG + NO_INPUT_MSG
-            logging.info(status_text)
-            yield chatbot + [(inputs, "")], status_text
-            return
-
-        if self.single_turn:
-            self.history = []
-            self.all_token_counts = []
-        self.history.append(construct_user(inputs))
-
-        try:
-            if stream:
-                logging.debug("使用流式传输")
-                iter = self.stream_next_chatbot(
-                    inputs,
-                    chatbot,
-                    fake_input=fake_inputs,
-                    display_append=display_append,
-                )
-                for chatbot, status_text in iter:
-                    yield chatbot, status_text
-            else:
-                logging.debug("不使用流式传输")
-                chatbot, status_text = self.next_chatbot_at_once(
-                    inputs,
-                    chatbot,
-                    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
-
-        if len(self.history) > 1 and self.history[-1]["content"] != inputs:
-            logging.info(
-                "回答为："
-                + colorama.Fore.BLUE
-                + f"{self.history[-1]['content']}"
-                + colorama.Style.RESET_ALL
-            )
-
-        if limited_context:
-            # self.history = self.history[-4:]
-            # self.all_token_counts = self.all_token_counts[-2:]
-            self.history = []
-            self.all_token_counts = []
-
-        max_token = self.token_upper_limit - TOKEN_OFFSET
-
-        if sum(self.all_token_counts) > max_token and should_check_token_count:
-            count = 0
-            while (
-                sum(self.all_token_counts)
-                > self.token_upper_limit * REDUCE_TOKEN_FACTOR
-                and sum(self.all_token_counts) > 0
-            ):
-                count += 1
-                del self.all_token_counts[0]
-                del self.history[:2]
-            logging.info(status_text)
-            status_text = f"为了防止token超限，模型忘记了早期的 {count} 轮对话"
-            yield chatbot, status_text
-
-    def retry(
-        self,
-        chatbot,
-        stream=False,
-        use_websearch=False,
-        files=None,
-        reply_language="中文",
-    ):
-        logging.debug("重试中……")
-        if len(self.history) > 0:
-            inputs = self.history[-2]["content"]
-            del self.history[-2:]
-            self.all_token_counts.pop()
-        elif len(chatbot) > 0:
-            inputs = chatbot[-1][0]
-        else:
-            yield chatbot, f"{STANDARD_ERROR_MSG}上下文是空的"
-            return
-
-        iter = self.predict(
-            inputs,
-            chatbot,
-            stream=stream,
-            use_websearch=use_websearch,
-            files=files,
-            reply_language=reply_language,
-        )
-        for x in iter:
-            yield x
-        logging.debug("重试完毕")
-
-    # def reduce_token_size(self, chatbot):
-    #     logging.info("开始减少token数量……")
-    #     chatbot, status_text = self.next_chatbot_at_once(
-    #         summarize_prompt,
-    #         chatbot
-    #     )
-    #     max_token_count = self.token_upper_limit * REDUCE_TOKEN_FACTOR
-    #     num_chat = find_n(self.all_token_counts, max_token_count)
-    #     logging.info(f"previous_token_count: {self.all_token_counts}, keeping {num_chat} chats")
-    #     chatbot = chatbot[:-1]
-    #     self.history = self.history[-2*num_chat:] if num_chat > 0 else []
-    #     self.all_token_counts = self.all_token_counts[-num_chat:] if num_chat > 0 else []
-    #     msg = f"保留了最近{num_chat}轮对话"
-    #     logging.info(msg)
-    #     logging.info("减少token数量完毕")
-    #     return chatbot, msg + "，" + self.token_message(self.all_token_counts if len(self.all_token_counts) > 0 else [0])
-
-    def interrupt(self):
-        self.interrupted = True
-
-    def recover(self):
-        self.interrupted = False
-
-    def set_token_upper_limit(self, new_upper_limit):
-        self.token_upper_limit = new_upper_limit
-        print(f"token上限设置为{new_upper_limit}")
-
-    def set_temperature(self, new_temperature):
-        self.temperature = new_temperature
-
-    def set_top_p(self, new_top_p):
-        self.top_p = new_top_p
-
-    def set_n_choices(self, new_n_choices):
-        self.n_choices = new_n_choices
-
-    def set_stop_sequence(self, new_stop_sequence: str):
-        new_stop_sequence = new_stop_sequence.split(",")
-        self.stop_sequence = new_stop_sequence
-
-    def set_max_tokens(self, new_max_tokens):
-        self.max_generation_token = new_max_tokens
-
-    def set_presence_penalty(self, new_presence_penalty):
-        self.presence_penalty = new_presence_penalty
-
-    def set_frequency_penalty(self, new_frequency_penalty):
-        self.frequency_penalty = new_frequency_penalty
-
-    def set_logit_bias(self, logit_bias):
-        logit_bias = logit_bias.split()
-        bias_map = {}
-        encoding = tiktoken.get_encoding("cl100k_base")
-        for line in logit_bias:
-            word, bias_amount = line.split(":")
-            if word:
-                for token in encoding.encode(word):
-                    bias_map[token] = float(bias_amount)
-        self.logit_bias = bias_map
-
-    def set_user_identifier(self, new_user_identifier):
-        self.user_identifier = new_user_identifier
-
-    def set_system_prompt(self, new_system_prompt):
-        self.system_prompt = new_system_prompt
-
-    def set_key(self, new_access_key):
-        self.api_key = new_access_key.strip()
-        msg = i18n("API密钥更改为了") + hide_middle_chars(self.api_key)
-        logging.info(msg)
-        return self.api_key, msg
-
-    def set_single_turn(self, new_single_turn):
-        self.single_turn = new_single_turn
-
-    def reset(self):
-        self.history = []
-        self.all_token_counts = []
-        self.interrupted = False
-        return [], self.token_message([0])
-
-    def delete_first_conversation(self):
-        if self.history:
-            del self.history[:2]
-            del self.all_token_counts[0]
-        return self.token_message()
-
-    def delete_last_conversation(self, chatbot):
-        if len(chatbot) > 0 and STANDARD_ERROR_MSG in chatbot[-1][1]:
-            msg = "由于包含报错信息，只删除chatbot记录"
-            chatbot.pop()
-            return chatbot, self.history
-        if len(self.history) > 0:
-            self.history.pop()
-            self.history.pop()
-        if len(chatbot) > 0:
-            msg = "删除了一组chatbot对话"
-            chatbot.pop()
-        if len(self.all_token_counts) > 0:
-            msg = "删除了一组对话的token计数记录"
-            self.all_token_counts.pop()
-        msg = "删除了一组对话"
-        return chatbot, msg
-
-    def token_message(self, token_lst=None):
-        if token_lst is None:
-            token_lst = self.all_token_counts
-        token_sum = 0
-        for i in range(len(token_lst)):
-            token_sum += sum(token_lst[: i + 1])
-        return i18n("Token 计数: ") + f"{sum(token_lst)}" + i18n("，本次对话累计消耗了 ") + f"{token_sum} tokens"
-
-    def save_chat_history(self, filename, chatbot, user_name):
-        if filename == "":
-            return
-        if not filename.endswith(".json"):
-            filename += ".json"
-        return save_file(filename, self.system_prompt, self.history, chatbot, user_name)
-
-    def export_markdown(self, filename, chatbot, user_name):
-        if filename == "":
-            return
-        if not filename.endswith(".md"):
-            filename += ".md"
-        return save_file(filename, self.system_prompt, self.history, chatbot, user_name)
-
-    def load_chat_history(self, filename, chatbot, user_name):
-        logging.debug(f"{user_name} 加载对话历史中……")
-        if type(filename) != str:
-            filename = filename.name
-        try:
-            with open(os.path.join(HISTORY_DIR, user_name, filename), "r") as f:
-                json_s = json.load(f)
-            try:
-                if type(json_s["history"][0]) == str:
-                    logging.info("历史记录格式为旧版，正在转换……")
-                    new_history = []
-                    for index, item in enumerate(json_s["history"]):
-                        if index % 2 == 0:
-                            new_history.append(construct_user(item))
-                        else:
-                            new_history.append(construct_assistant(item))
-                    json_s["history"] = new_history
-                    logging.info(new_history)
-            except:
-                # 没有对话历史
-                pass
-            logging.debug(f"{user_name} 加载对话历史完毕")
-            self.history = json_s["history"]
-            return filename, json_s["system"], json_s["chatbot"]
-        except FileNotFoundError:
-            logging.warning(f"{user_name} 没有找到对话历史文件，不执行任何操作")
-            return filename, self.system_prompt, chatbot
-
-    def like(self):
-        """like the last response, implement if needed
-        """
-        return gr.update()
-
-    def dislike(self):
-        """dislike the last response, implement if needed
-        """
-        return gr.update()

modules/models.py

@@ -1,625 +0,0 @@
-from __future__ import annotations
-from typing import TYPE_CHECKING, List
-
-import logging
-import json
-import commentjson as cjson
-import os
-import sys
-import requests
-import urllib3
-import platform
-import base64
-from io import BytesIO
-from PIL import Image
-
-from tqdm import tqdm
-import colorama
-from duckduckgo_search import ddg
-import asyncio
-import aiohttp
-from enum import Enum
-import uuid
-
-from .presets import *
-from .llama_func import *
-from .utils import *
-from . import shared
-from .config import retrieve_proxy
-from modules import config
-from .base_model import BaseLLMModel, ModelType
-
-
-class OpenAIClient(BaseLLMModel):
-    def __init__(
-        self,
-        model_name,
-        api_key,
-        system_prompt=INITIAL_SYSTEM_PROMPT,
-        temperature=1.0,
-        top_p=1.0,
-    ) -> None:
-        super().__init__(
-            model_name=model_name,
-            temperature=temperature,
-            top_p=top_p,
-            system_prompt=system_prompt,
-        )
-        self.api_key = api_key
-        self.need_api_key = True
-        self._refresh_header()
-
-    def get_answer_stream_iter(self):
-        response = self._get_response(stream=True)
-        if response is not None:
-            iter = self._decode_chat_response(response)
-            partial_text = ""
-            for i in iter:
-                partial_text += i
-                yield partial_text
-        else:
-            yield STANDARD_ERROR_MSG + GENERAL_ERROR_MSG
-
-    def get_answer_at_once(self):
-        response = self._get_response()
-        response = json.loads(response.text)
-        content = response["choices"][0]["message"]["content"]
-        total_token_count = response["usage"]["total_tokens"]
-        return content, total_token_count
-
-    def count_token(self, user_input):
-        input_token_count = count_token(construct_user(user_input))
-        if self.system_prompt is not None and len(self.all_token_counts) == 0:
-            system_prompt_token_count = count_token(
-                construct_system(self.system_prompt)
-            )
-            return input_token_count + system_prompt_token_count
-        return input_token_count
-
-    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")
-            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:
-                usage_data = self._get_billing_data(usage_url)
-            except Exception as e:
-                logging.error(f"获取API使用情况失败:" + str(e))
-                return i18n("**获取API使用情况失败**")
-            rounded_usage = "{:.5f}".format(usage_data["total_usage"] / 100)
-            return i18n("**本月使用金额** ") + f"\u3000 ${rounded_usage}"
-        except requests.exceptions.ConnectTimeout:
-            status_text = (
-                STANDARD_ERROR_MSG + CONNECTION_TIMEOUT_MSG + ERROR_RETRIEVE_MSG
-            )
-            return status_text
-        except requests.exceptions.ReadTimeout:
-            status_text = STANDARD_ERROR_MSG + READ_TIMEOUT_MSG + ERROR_RETRIEVE_MSG
-            return status_text
-        except Exception as e:
-            import traceback
-            traceback.print_exc()
-            logging.error(i18n("获取API使用情况失败:") + str(e))
-            return STANDARD_ERROR_MSG + ERROR_RETRIEVE_MSG
-
-    def set_token_upper_limit(self, new_upper_limit):
-        pass
-
-    @shared.state.switching_api_key  # 在不开启多账号模式的时候，这个装饰器不会起作用
-    def _get_response(self, stream=False):
-        openai_api_key = self.api_key
-        system_prompt = self.system_prompt
-        history = self.history
-        logging.debug(colorama.Fore.YELLOW +
-                      f"{history}" + colorama.Fore.RESET)
-        headers = {
-            "Content-Type": "application/json",
-            "Authorization": f"Bearer {openai_api_key}",
-        }
-
-        if system_prompt is not None:
-            history = [construct_system(system_prompt), *history]
-
-        payload = {
-            "model": self.model_name,
-            "messages": history,
-            "temperature": self.temperature,
-            "top_p": self.top_p,
-            "n": self.n_choices,
-            "stream": stream,
-            "presence_penalty": self.presence_penalty,
-            "frequency_penalty": self.frequency_penalty,
-        }
-
-        if self.max_generation_token is not None:
-            payload["max_tokens"] = self.max_generation_token
-        if self.stop_sequence is not None:
-            payload["stop"] = self.stop_sequence
-        if self.logit_bias is not None:
-            payload["logit_bias"] = self.logit_bias
-        if self.user_identifier is not None:
-            payload["user"] = self.user_identifier
-
-        if stream:
-            timeout = TIMEOUT_STREAMING
-        else:
-            timeout = TIMEOUT_ALL
-
-        # 如果有自定义的api-host，使用自定义host发送请求，否则使用默认设置发送请求
-        if shared.state.completion_url != COMPLETION_URL:
-            logging.info(f"使用自定义API URL: {shared.state.completion_url}")
-
-        with retrieve_proxy():
-            try:
-                response = requests.post(
-                    shared.state.completion_url,
-                    headers=headers,
-                    json=payload,
-                    stream=stream,
-                    timeout=timeout,
-                )
-            except:
-                return None
-        return response
-
-    def _refresh_header(self):
-        self.headers = {
-            "Content-Type": "application/json",
-            "Authorization": f"Bearer {self.api_key}",
-        }
-
-    def _get_billing_data(self, billing_url):
-        with retrieve_proxy():
-            response = requests.get(
-                billing_url,
-                headers=self.headers,
-                timeout=TIMEOUT_ALL,
-            )
-
-        if response.status_code == 200:
-            data = response.json()
-            return data
-        else:
-            raise Exception(
-                f"API request failed with status code {response.status_code}: {response.text}"
-            )
-
-    def _decode_chat_response(self, response):
-        error_msg = ""
-        for chunk in response.iter_lines():
-            if chunk:
-                chunk = chunk.decode()
-                chunk_length = len(chunk)
-                try:
-                    chunk = json.loads(chunk[6:])
-                except json.JSONDecodeError:
-                    print(i18n("JSON解析错误,收到的内容: ") + f"{chunk}")
-                    error_msg += chunk
-                    continue
-                if chunk_length > 6 and "delta" in chunk["choices"][0]:
-                    if chunk["choices"][0]["finish_reason"] == "stop":
-                        break
-                    try:
-                        yield chunk["choices"][0]["delta"]["content"]
-                    except Exception as e:
-                        # logging.error(f"Error: {e}")
-                        continue
-        if error_msg:
-            raise Exception(error_msg)
-
-    def set_key(self, new_access_key):
-        ret = super().set_key(new_access_key)
-        self._refresh_header()
-        return ret
-
-
-class ChatGLM_Client(BaseLLMModel):
-    def __init__(self, model_name) -> None:
-        super().__init__(model_name=model_name)
-        from transformers import AutoTokenizer, AutoModel
-        import torch
-        global CHATGLM_TOKENIZER, CHATGLM_MODEL
-        if CHATGLM_TOKENIZER is None or CHATGLM_MODEL is None:
-            system_name = platform.system()
-            model_path = None
-            if os.path.exists("models"):
-                model_dirs = os.listdir("models")
-                if model_name in model_dirs:
-                    model_path = f"models/{model_name}"
-            if model_path is not None:
-                model_source = model_path
-            else:
-                model_source = f"THUDM/{model_name}"
-            CHATGLM_TOKENIZER = AutoTokenizer.from_pretrained(
-                model_source, trust_remote_code=True
-            )
-            quantified = False
-            if "int4" in model_name:
-                quantified = True
-            model = AutoModel.from_pretrained(
-                    model_source, trust_remote_code=True
-                )
-            if torch.cuda.is_available():
-                # run on CUDA
-                logging.info("CUDA is available, using CUDA")
-                model = model.half().cuda()
-            # mps加速还存在一些问题，暂时不使用
-            elif system_name == "Darwin" and model_path is not None and not quantified:
-                logging.info("Running on macOS, using MPS")
-                # running on macOS and model already downloaded
-                model = model.half().to("mps")
-            else:
-                logging.info("GPU is not available, using CPU")
-                model = model.float()
-            model = model.eval()
-            CHATGLM_MODEL = model
-
-    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)
-        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)]
-        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)
-        return response, len(response)
-
-    def get_answer_stream_iter(self):
-        history, query = self._get_glm_style_input()
-        for response, history in CHATGLM_MODEL.stream_chat(
-            CHATGLM_TOKENIZER,
-            query,
-            history,
-            max_length=self.token_upper_limit,
-            top_p=self.top_p,
-            temperature=self.temperature,
-        ):
-            yield response
-
-
-class LLaMA_Client(BaseLLMModel):
-    def __init__(
-        self,
-        model_name,
-        lora_path=None,
-    ) -> None:
-        super().__init__(model_name=model_name)
-        from lmflow.datasets.dataset import Dataset
-        from lmflow.pipeline.auto_pipeline import AutoPipeline
-        from lmflow.models.auto_model import AutoModel
-        from lmflow.args import ModelArguments, DatasetArguments, InferencerArguments
-
-        self.max_generation_token = 1000
-        self.end_string = "\n\n"
-        # We don't need input data
-        data_args = DatasetArguments(dataset_path=None)
-        self.dataset = Dataset(data_args)
-        self.system_prompt = ""
-
-        global LLAMA_MODEL, LLAMA_INFERENCER
-        if LLAMA_MODEL is None or LLAMA_INFERENCER is None:
-            model_path = None
-            if os.path.exists("models"):
-                model_dirs = os.listdir("models")
-                if model_name in model_dirs:
-                    model_path = f"models/{model_name}"
-            if model_path is not None:
-                model_source = model_path
-            else:
-                model_source = f"decapoda-research/{model_name}"
-                # 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')
-
-            with open(pipeline_args.deepspeed, "r") as f:
-                ds_config = json.load(f)
-            LLAMA_MODEL = AutoModel.get_model(
-                model_args,
-                tune_strategy="none",
-                ds_config=ds_config,
-            )
-            LLAMA_INFERENCER = AutoPipeline.get_pipeline(
-                pipeline_name="inferencer",
-                model_args=model_args,
-                data_args=data_args,
-                pipeline_args=pipeline_args,
-            )
-
-    def _get_llama_style_input(self):
-        history = []
-        instruction = ""
-        if self.system_prompt:
-            instruction = (f"Instruction: {self.system_prompt}\n")
-        for x in self.history:
-            if x["role"] == "user":
-                history.append(f"{instruction}Input: {x['content']}")
-            else:
-                history.append(f"Output: {x['content']}")
-        context = "\n\n".join(history)
-        context += "\n\nOutput: "
-        return context
-
-    def get_answer_at_once(self):
-        context = self._get_llama_style_input()
-
-        input_dataset = self.dataset.from_dict(
-            {"type": "text_only", "instances": [{"text": context}]}
-        )
-
-        output_dataset = LLAMA_INFERENCER.inference(
-            model=LLAMA_MODEL,
-            dataset=input_dataset,
-            max_new_tokens=self.max_generation_token,
-            temperature=self.temperature,
-        )
-
-        response = output_dataset.to_dict()["instances"][0]["text"]
-        return response, len(response)
-
-    def get_answer_stream_iter(self):
-        context = self._get_llama_style_input()
-        partial_text = ""
-        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}]}
-            )
-            output_dataset = LLAMA_INFERENCER.inference(
-                model=LLAMA_MODEL,
-                dataset=input_dataset,
-                max_new_tokens=step,
-                temperature=self.temperature,
-            )
-            response = output_dataset.to_dict()["instances"][0]["text"]
-            if response == "" or response == self.end_string:
-                break
-            partial_text += response
-            yield partial_text
-
-
-class XMChat(BaseLLMModel):
-    def __init__(self, api_key):
-        super().__init__(model_name="xmchat")
-        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"
-        self.last_conv_id = None
-
-    def reset(self):
-        self.session_id = str(uuid.uuid4())
-        self.last_conv_id = None
-        return [], "已重置"
-
-    def image_to_base64(self, image_path):
-        # 打开并加载图片
-        img = Image.open(image_path)
-
-        # 获取图片的宽度和高度
-        width, height = img.size
-
-        # 计算压缩比例，以确保最长边小于4096像素
-        max_dimension = 2048
-        scale_ratio = min(max_dimension / width, max_dimension / height)
-
-        if scale_ratio < 1:
-            # 按压缩比例调整图片大小
-            new_width = int(width * scale_ratio)
-            new_height = int(height * scale_ratio)
-            img = img.resize((new_width, new_height), Image.ANTIALIAS)
-
-        # 将图片转换为jpg格式的二进制数据
-        buffer = BytesIO()
-        if img.mode == "RGBA":
-            img = img.convert("RGB")
-        img.save(buffer, format='JPEG')
-        binary_image = buffer.getvalue()
-
-        # 对二进制数据进行Base64编码
-        base64_image = base64.b64encode(binary_image).decode('utf-8')
-
-        return base64_image
-
-    def try_read_image(self, filepath):
-        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
-
-        if is_image_file(filepath):
-            logging.info(f"读取图片文件: {filepath}")
-            self.image_bytes = self.image_to_base64(filepath)
-            self.image_path = filepath
-        else:
-            self.image_bytes = None
-            self.image_path = None
-
-    def like(self):
-        if self.last_conv_id is None:
-            return "点赞失败，你还没发送过消息"
-        data = {
-            "uuid": self.last_conv_id,
-            "appraise": "good"
-        }
-        response = requests.post(self.url, json=data)
-        return "👍点赞成功，，感谢反馈～"
-
-    def dislike(self):
-        if self.last_conv_id is None:
-            return "点踩失败，你还没发送过消息"
-        data = {
-            "uuid": self.last_conv_id,
-            "appraise": "bad"
-        }
-        response = requests.post(self.url, json=data)
-        return "👎点踩成功，感谢反馈～"
-
-    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("使用图片作为输入")
-                # XMChat的一轮对话中实际上只能处理一张图片
-                self.reset()
-                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']}")
-        return None, chatbot, None
-
-    def get_answer_at_once(self):
-        question = self.history[-1]["content"]
-        conv_id = str(uuid.uuid4())
-        self.last_conv_id = conv_id
-        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)
-        try:
-            response = json.loads(response.text)
-            return response["data"], len(response["data"])
-        except Exception as e:
-            return response.text, len(response.text)
-
-
-
-
-def get_model(
-    model_name,
-    lora_model_path=None,
-    access_key=None,
-    temperature=None,
-    top_p=None,
-    system_prompt=None,
-) -> BaseLLMModel:
-    msg = i18n("模型设置为了：") + f" {model_name}"
-    model_type = ModelType.get_type(model_name)
-    lora_selector_visibility = False
-    lora_choices = []
-    dont_change_lora_selector = False
-    if model_type != ModelType.OpenAI:
-        config.local_embedding = True
-    # del current_model.model
-    model = None
-    try:
-        if model_type == ModelType.OpenAI:
-            logging.info(f"正在加载OpenAI模型: {model_name}")
-            model = OpenAIClient(
-                model_name=model_name,
-                api_key=access_key,
-                system_prompt=system_prompt,
-                temperature=temperature,
-                top_p=top_p,
-            )
-        elif model_type == ModelType.ChatGLM:
-            logging.info(f"正在加载ChatGLM模型: {model_name}")
-            model = ChatGLM_Client(model_name)
-        elif model_type == ModelType.LLaMA and lora_model_path == "":
-            msg = f"现在请为 {model_name} 选择LoRA模型"
-            logging.info(msg)
-            lora_selector_visibility = True
-            if os.path.isdir("lora"):
-                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}")
-            dont_change_lora_selector = True
-            if lora_model_path == "No LoRA":
-                lora_model_path = None
-                msg += " + No LoRA"
-            else:
-                msg += f" + {lora_model_path}"
-            model = LLaMA_Client(model_name, lora_model_path)
-        elif model_type == ModelType.XMChat:
-            if os.environ.get("XMCHAT_API_KEY") != "":
-                access_key = os.environ.get("XMCHAT_API_KEY")
-            model = XMChat(api_key=access_key)
-        elif model_type == ModelType.Unknown:
-            raise ValueError(f"未知模型: {model_name}")
-        logging.info(msg)
-    except Exception as e:
-        logging.error(e)
-        msg = f"{STANDARD_ERROR_MSG}: {e}"
-    if dont_change_lora_selector:
-        return model, msg
-    else:
-        return model, msg, gr.Dropdown.update(choices=lora_choices, visible=lora_selector_visibility)
-
-
-if __name__ == "__main__":
-    with open("config.json", "r") as f:
-        openai_api_key = cjson.load(f)["openai_api_key"]
-    # set logging level to debug
-    logging.basicConfig(level=logging.DEBUG)
-    # client = ModelManager(model_name="gpt-3.5-turbo", access_key=openai_api_key)
-    client = get_model(model_name="chatglm-6b-int4")
-    chatbot = []
-    stream = False
-    # 测试账单功能
-    logging.info(colorama.Back.GREEN + "测试账单功能" + colorama.Back.RESET)
-    logging.info(client.billing_info())
-    # 测试问答
-    logging.info(colorama.Back.GREEN + "测试问答" + colorama.Back.RESET)
-    question = "巴黎是中国的首都吗？"
-    for i in client.predict(inputs=question, chatbot=chatbot, stream=stream):
-        logging.info(i)
-    logging.info(f"测试问答后history : {client.history}")
-    # 测试记忆力
-    logging.info(colorama.Back.GREEN + "测试记忆力" + colorama.Back.RESET)
-    question = "我刚刚问了你什么问题？"
-    for i in client.predict(inputs=question, chatbot=chatbot, stream=stream):
-        logging.info(i)
-    logging.info(f"测试记忆力后history : {client.history}")
-    # 测试重试功能
-    logging.info(colorama.Back.GREEN + "测试重试功能" + colorama.Back.RESET)
-    for i in client.retry(chatbot=chatbot, stream=stream):
-        logging.info(i)
-    logging.info(f"重试后history : {client.history}")
-    # # 测试总结功能
-    # print(colorama.Back.GREEN + "测试总结功能" + colorama.Back.RESET)
-    # chatbot, msg = client.reduce_token_size(chatbot=chatbot)
-    # print(chatbot, msg)
-    # print(f"总结后history: {client.history}")