修复gitignore导致的文件缺失

.gitignore

@@ -140,7 +140,7 @@ dmypy.json
 api_key.txt
 config.json
 auth.json
-models/
+.models/
 lora/
 .idea
 templates/*

modules/models/StableLM.py

@@ -0,0 +1,95 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, StoppingCriteria, StoppingCriteriaList
+import time
+import numpy as np
+from torch.nn import functional as F
+import os
+from base_model import BaseLLMModel
+
+class StopOnTokens(StoppingCriteria):
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        stop_ids = [50278, 50279, 50277, 1, 0]
+        for stop_id in stop_ids:
+            if input_ids[0][-1] == stop_id:
+                return True
+        return False
+
+class StableLM_Client(BaseLLMModel):
+    def __init__(self, model_name) -> None:
+        super().__init__(model_name=model_name)
+        print(f"Starting to load StableLM to memory")
+        self.model = AutoModelForCausalLM.from_pretrained(
+            "stabilityai/stablelm-tuned-alpha-7b", torch_dtype=torch.float16).cuda()
+        self.tokenizer = AutoTokenizer.from_pretrained("stabilityai/stablelm-tuned-alpha-7b")
+        self.generator = pipeline('text-generation', model=self.model, tokenizer=self.tokenizer, device=0)
+        print(f"Sucessfully loaded StableLM to the memory")
+        self.system_prompt = """StableAssistant
+- StableAssistant is A helpful and harmless Open Source AI Language Model developed by Stability and CarperAI.
+- StableAssistant is excited to be able to help the user, but will refuse to do anything that could be considered harmful to the user.
+- StableAssistant is more than just an information source, StableAssistant is also able to write poetry, short stories, and make jokes.
+- StableAssistant will refuse to participate in anything that could harm a human."""
+
+    def user(self, user_message, history):
+        history = history + [[user_message, ""]]
+        return "", history, history
+
+
+    def bot(self, history, curr_system_message):
+        messages = f"<|SYSTEM|># {self.system_prompt}" + \
+            "".join(["".join(["<|USER|>"+item[0], "<|ASSISTANT|>"+item[1]])
+                    for item in history])
+        output = self.generate(messages)
+        history[-1][1] = output
+        time.sleep(1)
+        return history, history
+
+    def _get_stablelm_style_input(self):
+        messages = self.system_prompt + \
+            "".join(["".join(["<|USER|>"+self.history[i]["content"], "<|ASSISTANT|>"+self.history[i + 1]["content"]])
+                    for i in range(0, len(self.history), 2)])
+        return messages
+
+    def generate(self, text, bad_text=None):
+        stop = StopOnTokens()
+        result = self.generator(text, max_new_tokens=1024, num_return_sequences=1, num_beams=1, do_sample=True,
+                        temperature=1.0, top_p=0.95, top_k=1000, stopping_criteria=StoppingCriteriaList([stop]))
+        return result[0]["generated_text"].replace(text, "")
+
+    def contrastive_generate(self, text, bad_text):
+        with torch.no_grad():
+            tokens = self.tokenizer(text, return_tensors="pt")[
+                'input_ids'].cuda()[:, :4096-1024]
+            bad_tokens = self.tokenizer(bad_text, return_tensors="pt")[
+                'input_ids'].cuda()[:, :4096-1024]
+            history = None
+            bad_history = None
+            curr_output = list()
+            for i in range(1024):
+                out = self.model(tokens, past_key_values=history, use_cache=True)
+                logits = out.logits
+                history = out.past_key_values
+                bad_out = self.model(bad_tokens, past_key_values=bad_history,
+                            use_cache=True)
+                bad_logits = bad_out.logits
+                bad_history = bad_out.past_key_values
+                probs = F.softmax(logits.float(), dim=-1)[0][-1].cpu()
+                bad_probs = F.softmax(bad_logits.float(), dim=-1)[0][-1].cpu()
+                logits = torch.log(probs)
+                bad_logits = torch.log(bad_probs)
+                logits[probs > 0.1] = logits[probs > 0.1] - bad_logits[probs > 0.1]
+                probs = F.softmax(logits)
+                out = int(torch.multinomial(probs, 1))
+                if out in [50278, 50279, 50277, 1, 0]:
+                    break
+                else:
+                    curr_output.append(out)
+                out = np.array([out])
+                tokens = torch.from_numpy(np.array([out])).to(
+                    tokens.device)
+                bad_tokens = torch.from_numpy(np.array([out])).to(
+                    tokens.device)
+            return self.tokenizer.decode(curr_output)
+
+    def get_answer_at_once(self):
+        messages = self._get_stablelm_style_input()
+        return self.generate(messages)

modules/models/base_model.py

@@ -0,0 +1,561 @@
+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/models.py

@@ -0,0 +1,625 @@
+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}")