Intoval/privateChatGPT
1
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 XMBot = 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.XMBot46 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())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 )250 reference_results = add_source_numbers(reference_results)251 display_append = "\n\n" + "".join(display_append)252 real_inputs = (253 replace_today(WEBSEARCH_PTOMPT_TEMPLATE)254 .replace("{query}", real_inputs)255 .replace("{web_results}", "\n\n".join(reference_results))256 .replace("{reply_language}", reply_language)257 )258 else:259 display_append = ""260 return limited_context, fake_inputs, display_append, real_inputs, chatbot261 262 def predict(263 self,264 inputs,265 chatbot,266 stream=False,267 use_websearch=False,268 files=None,269 reply_language="中文",270 should_check_token_count=True,271 ): # repetition_penalty, top_k272 273 status_text = "开始生成回答……"274 logging.info(275 "输入为：" + colorama.Fore.BLUE + f"{inputs}" + colorama.Style.RESET_ALL276 )277 if should_check_token_count:278 yield chatbot + [(inputs, "")], status_text279 if reply_language == "跟随问题语言（不稳定）":280 reply_language = "the same language as the question, such as English, 中文, 日本語, Español, Français, or Deutsch."281 282 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)283 yield chatbot + [(fake_inputs, "")], status_text284 285 if (286 self.need_api_key and287 self.api_key is None288 and not shared.state.multi_api_key289 ):290 status_text = STANDARD_ERROR_MSG + NO_APIKEY_MSG291 logging.info(status_text)292 chatbot.append((inputs, ""))293 if len(self.history) == 0:294 self.history.append(construct_user(inputs))295 self.history.append("")296 self.all_token_counts.append(0)297 else:298 self.history[-2] = construct_user(inputs)299 yield chatbot + [(inputs, "")], status_text300 return301 elif len(inputs.strip()) == 0:302 status_text = STANDARD_ERROR_MSG + NO_INPUT_MSG303 logging.info(status_text)304 yield chatbot + [(inputs, "")], status_text305 return306 307 if self.single_turn:308 self.history = []309 self.all_token_counts = []310 self.history.append(construct_user(inputs))311 312 try:313 if stream:314 logging.debug("使用流式传输")315 iter = self.stream_next_chatbot(316 inputs,317 chatbot,318 fake_input=fake_inputs,319 display_append=display_append,320 )321 for chatbot, status_text in iter:322 yield chatbot, status_text323 else:324 logging.debug("不使用流式传输")325 chatbot, status_text = self.next_chatbot_at_once(326 inputs,327 chatbot,328 fake_input=fake_inputs,329 display_append=display_append,330 )331 yield chatbot, status_text332 except Exception as e:333 traceback.print_exc()334 status_text = STANDARD_ERROR_MSG + str(e)335 yield chatbot, status_text336 337 if len(self.history) > 1 and self.history[-1]["content"] != inputs:338 logging.info(339 "回答为："340 + colorama.Fore.BLUE341 + f"{self.history[-1]['content']}"342 + colorama.Style.RESET_ALL343 )344 345 if limited_context:346 # self.history = self.history[-4:]347 # self.all_token_counts = self.all_token_counts[-2:]348 self.history = []349 self.all_token_counts = []350 351 max_token = self.token_upper_limit - TOKEN_OFFSET352 353 if sum(self.all_token_counts) > max_token and should_check_token_count:354 count = 0355 while (356 sum(self.all_token_counts)357 > self.token_upper_limit * REDUCE_TOKEN_FACTOR358 and sum(self.all_token_counts) > 0359 ):360 count += 1361 del self.all_token_counts[0]362 del self.history[:2]363 logging.info(status_text)364 status_text = f"为了防止token超限，模型忘记了早期的 {count} 轮对话"365 yield chatbot, status_text366 367 def retry(368 self,369 chatbot,370 stream=False,371 use_websearch=False,372 files=None,373 reply_language="中文",374 ):375 logging.debug("重试中……")376 if len(self.history) > 0:377 inputs = self.history[-2]["content"]378 del self.history[-2:]379 self.all_token_counts.pop()380 elif len(chatbot) > 0:381 inputs = chatbot[-1][0]382 else:383 yield chatbot, f"{STANDARD_ERROR_MSG}上下文是空的"384 return385 386 iter = self.predict(387 inputs,388 chatbot,389 stream=stream,390 use_websearch=use_websearch,391 files=files,392 reply_language=reply_language,393 )394 for x in iter:395 yield x396 logging.debug("重试完毕")397 398 # def reduce_token_size(self, chatbot):399 # logging.info("开始减少token数量……")400 # chatbot, status_text = self.next_chatbot_at_once(401 # summarize_prompt,402 # chatbot403 # )404 # max_token_count = self.token_upper_limit * REDUCE_TOKEN_FACTOR405 # num_chat = find_n(self.all_token_counts, max_token_count)406 # logging.info(f"previous_token_count: {self.all_token_counts}, keeping {num_chat} chats")407 # chatbot = chatbot[:-1]408 # self.history = self.history[-2*num_chat:] if num_chat > 0 else []409 # self.all_token_counts = self.all_token_counts[-num_chat:] if num_chat > 0 else []410 # msg = f"保留了最近{num_chat}轮对话"411 # logging.info(msg)412 # logging.info("减少token数量完毕")413 # return chatbot, msg + "，" + self.token_message(self.all_token_counts if len(self.all_token_counts) > 0 else [0])414 415 def interrupt(self):416 self.interrupted = True417 418 def recover(self):419 self.interrupted = False420 421 def set_token_upper_limit(self, new_upper_limit):422 self.token_upper_limit = new_upper_limit423 print(f"token上限设置为{new_upper_limit}")424 425 def set_temperature(self, new_temperature):426 self.temperature = new_temperature427 428 def set_top_p(self, new_top_p):429 self.top_p = new_top_p430 431 def set_n_choices(self, new_n_choices):432 self.n_choices = new_n_choices433 434 def set_stop_sequence(self, new_stop_sequence: str):435 new_stop_sequence = new_stop_sequence.split(",")436 self.stop_sequence = new_stop_sequence437 438 def set_max_tokens(self, new_max_tokens):439 self.max_generation_token = new_max_tokens440 441 def set_presence_penalty(self, new_presence_penalty):442 self.presence_penalty = new_presence_penalty443 444 def set_frequency_penalty(self, new_frequency_penalty):445 self.frequency_penalty = new_frequency_penalty446 447 def set_logit_bias(self, logit_bias):448 logit_bias = logit_bias.split()449 bias_map = {}450 encoding = tiktoken.get_encoding("cl100k_base")451 for line in logit_bias:452 word, bias_amount = line.split(":")453 if word:454 for token in encoding.encode(word):455 bias_map[token] = float(bias_amount)456 self.logit_bias = bias_map457 458 def set_user_identifier(self, new_user_identifier):459 self.user_identifier = new_user_identifier460 461 def set_system_prompt(self, new_system_prompt):462 self.system_prompt = new_system_prompt463 464 def set_key(self, new_access_key):465 self.api_key = new_access_key.strip()466 msg = f"API密钥更改为了{hide_middle_chars(self.api_key)}"467 logging.info(msg)468 return new_access_key, msg469 470 def set_single_turn(self, new_single_turn):471 self.single_turn = new_single_turn472 473 def reset(self):474 self.history = []475 self.all_token_counts = []476 self.interrupted = False477 return [], self.token_message([0])478 479 def delete_first_conversation(self):480 if self.history:481 del self.history[:2]482 del self.all_token_counts[0]483 return self.token_message()484 485 def delete_last_conversation(self, chatbot):486 if len(chatbot) > 0 and STANDARD_ERROR_MSG in chatbot[-1][1]:487 msg = "由于包含报错信息，只删除chatbot记录"488 chatbot.pop()489 return chatbot, self.history490 if len(self.history) > 0:491 self.history.pop()492 self.history.pop()493 if len(chatbot) > 0:494 msg = "删除了一组chatbot对话"495 chatbot.pop()496 if len(self.all_token_counts) > 0:497 msg = "删除了一组对话的token计数记录"498 self.all_token_counts.pop()499 msg = "删除了一组对话"500 return chatbot, msg501 502 def token_message(self, token_lst=None):503 if token_lst is None:504 token_lst = self.all_token_counts505 token_sum = 0506 for i in range(len(token_lst)):507 token_sum += sum(token_lst[: i + 1])508 return i18n("Token 计数: ") + f"{sum(token_lst)}" + i18n("，本次对话累计消耗了 ") + f"{token_sum} tokens"509 510 def save_chat_history(self, filename, chatbot, user_name):511 if filename == "":512 return513 if not filename.endswith(".json"):514 filename += ".json"515 return save_file(filename, self.system_prompt, self.history, chatbot, user_name)516 517 def export_markdown(self, filename, chatbot, user_name):518 if filename == "":519 return520 if not filename.endswith(".md"):521 filename += ".md"522 return save_file(filename, self.system_prompt, self.history, chatbot, user_name)523 524 def load_chat_history(self, filename, chatbot, user_name):525 logging.debug(f"{user_name} 加载对话历史中……")526 if type(filename) != str:527 filename = filename.name528 try:529 with open(os.path.join(HISTORY_DIR, user_name, filename), "r") as f:530 json_s = json.load(f)531 try:532 if type(json_s["history"][0]) == str:533 logging.info("历史记录格式为旧版，正在转换……")534 new_history = []535 for index, item in enumerate(json_s["history"]):536 if index % 2 == 0:537 new_history.append(construct_user(item))538 else:539 new_history.append(construct_assistant(item))540 json_s["history"] = new_history541 logging.info(new_history)542 except:543 # 没有对话历史544 pass545 logging.debug(f"{user_name} 加载对话历史完毕")546 self.history = json_s["history"]547 return filename, json_s["system"], json_s["chatbot"]548 except FileNotFoundError:549 logging.warning(f"{user_name} 没有找到对话历史文件，不执行任何操作")550 return filename, self.system_prompt, chatbot551