Danielzero/GPT3.5

1from __future__ import annotations2from typing import TYPE_CHECKING, List3 4import logging5import json6import commentjson as cjson7import os8import sys9import requests10import urllib311import traceback12 13from tqdm import tqdm14import colorama15from duckduckgo_search import ddg16import asyncio17import aiohttp18from enum import Enum19 20from .presets import *21from .llama_func import *22from .utils import *23from . import shared24from .config import retrieve_proxy25 26 27class ModelType(Enum):28    Unknown = -129    OpenAI = 030    ChatGLM = 131    LLaMA = 232    XMChat = 333 34    @classmethod35    def get_type(cls, model_name: str):36        model_type = None37        model_name_lower = model_name.lower()38        if "gpt" in model_name_lower:39            model_type = ModelType.OpenAI40        elif "chatglm" in model_name_lower:41            model_type = ModelType.ChatGLM42        elif "llama" in model_name_lower or "alpaca" in model_name_lower:43            model_type = ModelType.LLaMA44        elif "xmchat" in model_name_lower:45            model_type = ModelType.XMChat46        else:47            model_type = ModelType.Unknown48        return model_type49 50 51class BaseLLMModel:52    def __init__(53        self,54        model_name,55        system_prompt="",56        temperature=1.0,57        top_p=1.0,58        n_choices=1,59        stop=None,60        max_generation_token=None,61        presence_penalty=0,62        frequency_penalty=0,63        logit_bias=None,64        user="",65    ) -> None:66        self.history = []67        self.all_token_counts = []68        self.model_name = model_name69        self.model_type = ModelType.get_type(model_name)70        try:71            self.token_upper_limit = MODEL_TOKEN_LIMIT[model_name]72        except KeyError:73            self.token_upper_limit = DEFAULT_TOKEN_LIMIT74        self.interrupted = False75        self.system_prompt = system_prompt76        self.api_key = None77        self.need_api_key = False78        self.single_turn = False79 80        self.temperature = temperature81        self.top_p = top_p82        self.n_choices = n_choices83        self.stop_sequence = stop84        self.max_generation_token = None85        self.presence_penalty = presence_penalty86        self.frequency_penalty = frequency_penalty87        self.logit_bias = logit_bias88        self.user_identifier = user89 90    def get_answer_stream_iter(self):91        """stream predict, need to be implemented92        conversations are stored in self.history, with the most recent question, in OpenAI format93        should return a generator, each time give the next word (str) in the answer94        """95        logging.warning("stream predict not implemented, using at once predict instead")96        response, _ = self.get_answer_at_once()97        yield response98 99    def get_answer_at_once(self):100        """predict at once, need to be implemented101        conversations are stored in self.history, with the most recent question, in OpenAI format102        Should return:103        the answer (str)104        total token count (int)105        """106        logging.warning("at once predict not implemented, using stream predict instead")107        response_iter = self.get_answer_stream_iter()108        count = 0109        for response in response_iter:110            count += 1111        return response, sum(self.all_token_counts) + count112 113    def billing_info(self):114        """get billing infomation, inplement if needed"""115        logging.warning("billing info not implemented, using default")116        return BILLING_NOT_APPLICABLE_MSG117 118    def count_token(self, user_input):119        """get token count from input, implement if needed"""120        logging.warning("token count not implemented, using default")121        return len(user_input)122 123    def stream_next_chatbot(self, inputs, chatbot, fake_input=None, display_append=""):124        def get_return_value():125            return chatbot, status_text126 127        status_text = i18n("开始实时传输回答……")128        if fake_input:129            chatbot.append((fake_input, ""))130        else:131            chatbot.append((inputs, ""))132 133        user_token_count = self.count_token(inputs)134        self.all_token_counts.append(user_token_count)135        logging.debug(f"输入token计数: {user_token_count}")136 137        stream_iter = self.get_answer_stream_iter()138 139        for partial_text in stream_iter:140            chatbot[-1] = (chatbot[-1][0], partial_text + display_append)141            self.all_token_counts[-1] += 1142            status_text = self.token_message()143            yield get_return_value()144            if self.interrupted:145                self.recover()146                break147        self.history.append(construct_assistant(partial_text))148 149    def next_chatbot_at_once(self, inputs, chatbot, fake_input=None, display_append=""):150        if fake_input:151            chatbot.append((fake_input, ""))152        else:153            chatbot.append((inputs, ""))154        if fake_input is not None:155            user_token_count = self.count_token(fake_input)156        else:157            user_token_count = self.count_token(inputs)158        self.all_token_counts.append(user_token_count)159        ai_reply, total_token_count = self.get_answer_at_once()160        self.history.append(construct_assistant(ai_reply))161        if fake_input is not None:162            self.history[-2] = construct_user(fake_input)163        chatbot[-1] = (chatbot[-1][0], ai_reply + display_append)164        if fake_input is not None:165            self.all_token_counts[-1] += count_token(construct_assistant(ai_reply))166        else:167            self.all_token_counts[-1] = total_token_count - sum(self.all_token_counts)168        status_text = self.token_message()169        return chatbot, status_text170 171    def handle_file_upload(self, files, chatbot):172        """if the model accepts multi modal input, implement this function"""173        status = gr.Markdown.update()174        if files:175            construct_index(self.api_key, file_src=files)176            status = "索引构建完成"177        return gr.Files.update(), chatbot, status178 179    def prepare_inputs(self, real_inputs, use_websearch, files, reply_language, chatbot):180        fake_inputs = None181        display_append = []182        limited_context = False183        fake_inputs = real_inputs184        if files:185            from llama_index.indices.vector_store.base_query import GPTVectorStoreIndexQuery186            from llama_index.indices.query.schema import QueryBundle187            from langchain.embeddings.huggingface import HuggingFaceEmbeddings188            from langchain.chat_models import ChatOpenAI189            from llama_index import (190                GPTSimpleVectorIndex,191                ServiceContext,192                LangchainEmbedding,193                OpenAIEmbedding,194            )195            limited_context = True196            msg = "加载索引中……"197            logging.info(msg)198            # yield chatbot + [(inputs, "")], msg199            index = construct_index(self.api_key, file_src=files)200            assert index is not None, "获取索引失败"201            msg = "索引获取成功，生成回答中……"202            logging.info(msg)203            if local_embedding or self.model_type != ModelType.OpenAI:204                embed_model = LangchainEmbedding(HuggingFaceEmbeddings(model_name = "sentence-transformers/distiluse-base-multilingual-cased-v2"))205            else:206                embed_model = OpenAIEmbedding()207            # yield chatbot + [(inputs, "")], msg208            with retrieve_proxy():209                prompt_helper = PromptHelper(210                    max_input_size=4096,211                    num_output=5,212                    max_chunk_overlap=20,213                    chunk_size_limit=600,214                )215                from llama_index import ServiceContext216 217                service_context = ServiceContext.from_defaults(218                    prompt_helper=prompt_helper, embed_model=embed_model219                )220                query_object = GPTVectorStoreIndexQuery(221                    index.index_struct,222                    service_context=service_context,223                    similarity_top_k=5,224                    vector_store=index._vector_store,225                    docstore=index._docstore,226                )227                query_bundle = QueryBundle(real_inputs)228                nodes = query_object.retrieve(query_bundle)229            reference_results = [n.node.text for n in nodes]230            reference_results = add_source_numbers(reference_results, use_source=False)231            display_append = add_details(reference_results)232            display_append = "\n\n" + "".join(display_append)233            real_inputs = (234                replace_today(PROMPT_TEMPLATE)235                .replace("{query_str}", real_inputs)236                .replace("{context_str}", "\n\n".join(reference_results))237                .replace("{reply_language}", reply_language)238            )239        elif use_websearch:240            limited_context = True241            search_results = ddg(real_inputs, max_results=5)242            reference_results = []243            for idx, result in enumerate(search_results):244                logging.debug(f"搜索结果{idx + 1}：{result}")245                domain_name = urllib3.util.parse_url(result["href"]).host246                reference_results.append([result["body"], result["href"]])247                display_append.append(248                    # f"{idx+1}. [{domain_name}]({result['href']})\n"249                    f"<li><a href=\"{result['href']}\" target=\"_blank\">{domain_name}</a></li>\n"250                )251            reference_results = add_source_numbers(reference_results)252            display_append = "<ol>\n\n" + "".join(display_append) + "</ol>"253            real_inputs = (254                replace_today(WEBSEARCH_PTOMPT_TEMPLATE)255                .replace("{query}", real_inputs)256                .replace("{web_results}", "\n\n".join(reference_results))257                .replace("{reply_language}", reply_language)258            )259        else:260            display_append = ""261        return limited_context, fake_inputs, display_append, real_inputs, chatbot262 263    def predict(264        self,265        inputs,266        chatbot,267        stream=False,268        use_websearch=False,269        files=None,270        reply_language="中文",271        should_check_token_count=True,272    ):  # repetition_penalty, top_k273 274        status_text = "开始生成回答……"275        logging.info(276            "输入为：" + colorama.Fore.BLUE + f"{inputs}" + colorama.Style.RESET_ALL277        )278        if should_check_token_count:279            yield chatbot + [(inputs, "")], status_text280        if reply_language == "跟随问题语言（不稳定）":281            reply_language = "the same language as the question, such as English, 中文, 日本語, Español, Français, or Deutsch."282 283        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)284        yield chatbot + [(fake_inputs, "")], status_text285 286        if (287            self.need_api_key and288            self.api_key is None289            and not shared.state.multi_api_key290        ):291            status_text = STANDARD_ERROR_MSG + NO_APIKEY_MSG292            logging.info(status_text)293            chatbot.append((inputs, ""))294            if len(self.history) == 0:295                self.history.append(construct_user(inputs))296                self.history.append("")297                self.all_token_counts.append(0)298            else:299                self.history[-2] = construct_user(inputs)300            yield chatbot + [(inputs, "")], status_text301            return302        elif len(inputs.strip()) == 0:303            status_text = STANDARD_ERROR_MSG + NO_INPUT_MSG304            logging.info(status_text)305            yield chatbot + [(inputs, "")], status_text306            return307 308        if self.single_turn:309            self.history = []310            self.all_token_counts = []311        self.history.append(construct_user(inputs))312 313        try:314            if stream:315                logging.debug("使用流式传输")316                iter = self.stream_next_chatbot(317                    inputs,318                    chatbot,319                    fake_input=fake_inputs,320                    display_append=display_append,321                )322                for chatbot, status_text in iter:323                    yield chatbot, status_text324            else:325                logging.debug("不使用流式传输")326                chatbot, status_text = self.next_chatbot_at_once(327                    inputs,328                    chatbot,329                    fake_input=fake_inputs,330                    display_append=display_append,331                )332                yield chatbot, status_text333        except Exception as e:334            traceback.print_exc()335            status_text = STANDARD_ERROR_MSG + str(e)336            yield chatbot, status_text337 338        if len(self.history) > 1 and self.history[-1]["content"] != inputs:339            logging.info(340                "回答为："341                + colorama.Fore.BLUE342                + f"{self.history[-1]['content']}"343                + colorama.Style.RESET_ALL344            )345 346        if limited_context:347            # self.history = self.history[-4:]348            # self.all_token_counts = self.all_token_counts[-2:]349            self.history = []350            self.all_token_counts = []351 352        max_token = self.token_upper_limit - TOKEN_OFFSET353 354        if sum(self.all_token_counts) > max_token and should_check_token_count:355            count = 0356            while (357                sum(self.all_token_counts)358                > self.token_upper_limit * REDUCE_TOKEN_FACTOR359                and sum(self.all_token_counts) > 0360            ):361                count += 1362                del self.all_token_counts[0]363                del self.history[:2]364            logging.info(status_text)365            status_text = f"为了防止token超限，模型忘记了早期的 {count} 轮对话"366            yield chatbot, status_text367 368    def retry(369        self,370        chatbot,371        stream=False,372        use_websearch=False,373        files=None,374        reply_language="中文",375    ):376        logging.debug("重试中……")377        if len(self.history) > 0:378            inputs = self.history[-2]["content"]379            del self.history[-2:]380            self.all_token_counts.pop()381        elif len(chatbot) > 0:382            inputs = chatbot[-1][0]383        else:384            yield chatbot, f"{STANDARD_ERROR_MSG}上下文是空的"385            return386 387        iter = self.predict(388            inputs,389            chatbot,390            stream=stream,391            use_websearch=use_websearch,392            files=files,393            reply_language=reply_language,394        )395        for x in iter:396            yield x397        logging.debug("重试完毕")398 399    # def reduce_token_size(self, chatbot):400    #     logging.info("开始减少token数量……")401    #     chatbot, status_text = self.next_chatbot_at_once(402    #         summarize_prompt,403    #         chatbot404    #     )405    #     max_token_count = self.token_upper_limit * REDUCE_TOKEN_FACTOR406    #     num_chat = find_n(self.all_token_counts, max_token_count)407    #     logging.info(f"previous_token_count: {self.all_token_counts}, keeping {num_chat} chats")408    #     chatbot = chatbot[:-1]409    #     self.history = self.history[-2*num_chat:] if num_chat > 0 else []410    #     self.all_token_counts = self.all_token_counts[-num_chat:] if num_chat > 0 else []411    #     msg = f"保留了最近{num_chat}轮对话"412    #     logging.info(msg)413    #     logging.info("减少token数量完毕")414    #     return chatbot, msg + "，" + self.token_message(self.all_token_counts if len(self.all_token_counts) > 0 else [0])415 416    def interrupt(self):417        self.interrupted = True418 419    def recover(self):420        self.interrupted = False421 422    def set_token_upper_limit(self, new_upper_limit):423        self.token_upper_limit = new_upper_limit424        print(f"token上限设置为{new_upper_limit}")425 426    def set_temperature(self, new_temperature):427        self.temperature = new_temperature428 429    def set_top_p(self, new_top_p):430        self.top_p = new_top_p431 432    def set_n_choices(self, new_n_choices):433        self.n_choices = new_n_choices434 435    def set_stop_sequence(self, new_stop_sequence: str):436        new_stop_sequence = new_stop_sequence.split(",")437        self.stop_sequence = new_stop_sequence438 439    def set_max_tokens(self, new_max_tokens):440        self.max_generation_token = new_max_tokens441 442    def set_presence_penalty(self, new_presence_penalty):443        self.presence_penalty = new_presence_penalty444 445    def set_frequency_penalty(self, new_frequency_penalty):446        self.frequency_penalty = new_frequency_penalty447 448    def set_logit_bias(self, logit_bias):449        logit_bias = logit_bias.split()450        bias_map = {}451        encoding = tiktoken.get_encoding("cl100k_base")452        for line in logit_bias:453            word, bias_amount = line.split(":")454            if word:455                for token in encoding.encode(word):456                    bias_map[token] = float(bias_amount)457        self.logit_bias = bias_map458 459    def set_user_identifier(self, new_user_identifier):460        self.user_identifier = new_user_identifier461 462    def set_system_prompt(self, new_system_prompt):463        self.system_prompt = new_system_prompt464 465    def set_key(self, new_access_key):466        self.api_key = new_access_key.strip()467        msg = i18n("API密钥更改为了") + hide_middle_chars(self.api_key)468        logging.info(msg)469        return self.api_key, msg470 471    def set_single_turn(self, new_single_turn):472        self.single_turn = new_single_turn473 474    def reset(self):475        self.history = []476        self.all_token_counts = []477        self.interrupted = False478        return [], self.token_message([0])479 480    def delete_first_conversation(self):481        if self.history:482            del self.history[:2]483            del self.all_token_counts[0]484        return self.token_message()485 486    def delete_last_conversation(self, chatbot):487        if len(chatbot) > 0 and STANDARD_ERROR_MSG in chatbot[-1][1]:488            msg = "由于包含报错信息，只删除chatbot记录"489            chatbot.pop()490            return chatbot, self.history491        if len(self.history) > 0:492            self.history.pop()493            self.history.pop()494        if len(chatbot) > 0:495            msg = "删除了一组chatbot对话"496            chatbot.pop()497        if len(self.all_token_counts) > 0:498            msg = "删除了一组对话的token计数记录"499            self.all_token_counts.pop()500        msg = "删除了一组对话"501        return chatbot, msg502 503    def token_message(self, token_lst=None):504        if token_lst is None:505            token_lst = self.all_token_counts506        token_sum = 0507        for i in range(len(token_lst)):508            token_sum += sum(token_lst[: i + 1])509        return i18n("Token 计数: ") + f"{sum(token_lst)}" + i18n("，本次对话累计消耗了 ") + f"{token_sum} tokens"510 511    def save_chat_history(self, filename, chatbot, user_name):512        if filename == "":513            return514        if not filename.endswith(".json"):515            filename += ".json"516        return save_file(filename, self.system_prompt, self.history, chatbot, user_name)517 518    def export_markdown(self, filename, chatbot, user_name):519        if filename == "":520            return521        if not filename.endswith(".md"):522            filename += ".md"523        return save_file(filename, self.system_prompt, self.history, chatbot, user_name)524 525    def load_chat_history(self, filename, chatbot, user_name):526        logging.debug(f"{user_name} 加载对话历史中……")527        if type(filename) != str:528            filename = filename.name529        try:530            with open(os.path.join(HISTORY_DIR, user_name, filename), "r") as f:531                json_s = json.load(f)532            try:533                if type(json_s["history"][0]) == str:534                    logging.info("历史记录格式为旧版，正在转换……")535                    new_history = []536                    for index, item in enumerate(json_s["history"]):537                        if index % 2 == 0:538                            new_history.append(construct_user(item))539                        else:540                            new_history.append(construct_assistant(item))541                    json_s["history"] = new_history542                    logging.info(new_history)543            except:544                # 没有对话历史545                pass546            logging.debug(f"{user_name} 加载对话历史完毕")547            self.history = json_s["history"]548            return filename, json_s["system"], json_s["chatbot"]549        except FileNotFoundError:550            logging.warning(f"{user_name} 没有找到对话历史文件，不执行任何操作")551            return filename, self.system_prompt, chatbot552 553    def like(self):554        """like the last response, implement if needed555        """556        return gr.update()557 558    def dislike(self):559        """dislike the last response, implement if needed560        """561        return gr.update()562