Intoval/privateChatGPT
1
1from __future__ import annotations2from typing import TYPE_CHECKING, List3 4import logging5import json6import commentjson as cjson7import os8import sys9import requests10import urllib311import platform12 13from tqdm import tqdm14import colorama15from duckduckgo_search import ddg16import asyncio17import aiohttp18from enum import Enum19import uuid20 21from .presets import *22from .llama_func import *23from .utils import *24from . import shared25from .config import retrieve_proxy26from modules import config27from .base_model import BaseLLMModel, ModelType28 29 30class OpenAIClient(BaseLLMModel):31 def __init__(32 self,33 model_name,34 api_key,35 system_prompt=INITIAL_SYSTEM_PROMPT,36 temperature=1.0,37 top_p=1.0,38 ) -> None:39 super().__init__(40 model_name=model_name,41 temperature=temperature,42 top_p=top_p,43 system_prompt=system_prompt,44 )45 self.api_key = api_key46 self.need_api_key = True47 self._refresh_header()48 49 def get_answer_stream_iter(self):50 response = self._get_response(stream=True)51 if response is not None:52 iter = self._decode_chat_response(response)53 partial_text = ""54 for i in iter:55 partial_text += i56 yield partial_text57 else:58 yield STANDARD_ERROR_MSG + GENERAL_ERROR_MSG59 60 def get_answer_at_once(self):61 response = self._get_response()62 response = json.loads(response.text)63 content = response["choices"][0]["message"]["content"]64 total_token_count = response["usage"]["total_tokens"]65 return content, total_token_count66 67 def count_token(self, user_input):68 input_token_count = count_token(construct_user(user_input))69 if self.system_prompt is not None and len(self.all_token_counts) == 0:70 system_prompt_token_count = count_token(71 construct_system(self.system_prompt)72 )73 return input_token_count + system_prompt_token_count74 return input_token_count75 76 def billing_info(self):77 try:78 curr_time = datetime.datetime.now()79 last_day_of_month = get_last_day_of_month(80 curr_time).strftime("%Y-%m-%d")81 first_day_of_month = curr_time.replace(day=1).strftime("%Y-%m-%d")82 usage_url = f"{shared.state.usage_api_url}?start_date={first_day_of_month}&end_date={last_day_of_month}"83 try:84 usage_data = self._get_billing_data(usage_url)85 except Exception as e:86 logging.error(f"获取API使用情况失败:" + str(e))87 return i18n("**获取API使用情况失败**")88 rounded_usage = "{:.5f}".format(usage_data["total_usage"] / 100)89 return i18n("**本月使用金额** ") + f"\u3000 ${rounded_usage}"90 except requests.exceptions.ConnectTimeout:91 status_text = (92 STANDARD_ERROR_MSG + CONNECTION_TIMEOUT_MSG + ERROR_RETRIEVE_MSG93 )94 return status_text95 except requests.exceptions.ReadTimeout:96 status_text = STANDARD_ERROR_MSG + READ_TIMEOUT_MSG + ERROR_RETRIEVE_MSG97 return status_text98 except Exception as e:99 logging.error(i18n("获取API使用情况失败:") + str(e))100 return STANDARD_ERROR_MSG + ERROR_RETRIEVE_MSG101 102 def set_token_upper_limit(self, new_upper_limit):103 pass104 105 @shared.state.switching_api_key # 在不开启多账号模式的时候，这个装饰器不会起作用106 def _get_response(self, stream=False):107 openai_api_key = self.api_key108 system_prompt = self.system_prompt109 history = self.history110 logging.debug(colorama.Fore.YELLOW +111 f"{history}" + colorama.Fore.RESET)112 headers = {113 "Content-Type": "application/json",114 "Authorization": f"Bearer {openai_api_key}",115 }116 117 if system_prompt is not None:118 history = [construct_system(system_prompt), *history]119 120 payload = {121 "model": self.model_name,122 "messages": history,123 "temperature": self.temperature,124 "top_p": self.top_p,125 "n": self.n_choices,126 "stream": stream,127 "presence_penalty": self.presence_penalty,128 "frequency_penalty": self.frequency_penalty,129 }130 131 if self.max_generation_token is not None:132 payload["max_tokens"] = self.max_generation_token133 if self.stop_sequence is not None:134 payload["stop"] = self.stop_sequence135 if self.logit_bias is not None:136 payload["logit_bias"] = self.logit_bias137 if self.user_identifier is not None:138 payload["user"] = self.user_identifier139 140 if stream:141 timeout = TIMEOUT_STREAMING142 else:143 timeout = TIMEOUT_ALL144 145 # 如果有自定义的api-host，使用自定义host发送请求，否则使用默认设置发送请求146 if shared.state.completion_url != COMPLETION_URL:147 logging.info(f"使用自定义API URL: {shared.state.completion_url}")148 149 with retrieve_proxy():150 try:151 response = requests.post(152 shared.state.completion_url,153 headers=headers,154 json=payload,155 stream=stream,156 timeout=timeout,157 )158 except:159 return None160 return response161 162 def _refresh_header(self):163 self.headers = {164 "Content-Type": "application/json",165 "Authorization": f"Bearer {self.api_key}",166 }167 168 def _get_billing_data(self, billing_url):169 with retrieve_proxy():170 response = requests.get(171 billing_url,172 headers=self.headers,173 timeout=TIMEOUT_ALL,174 )175 176 if response.status_code == 200:177 data = response.json()178 return data179 else:180 raise Exception(181 f"API request failed with status code {response.status_code}: {response.text}"182 )183 184 def _decode_chat_response(self, response):185 error_msg = ""186 for chunk in response.iter_lines():187 if chunk:188 chunk = chunk.decode()189 chunk_length = len(chunk)190 try:191 chunk = json.loads(chunk[6:])192 except json.JSONDecodeError:193 print(i18n("JSON解析错误,收到的内容: ") + f"{chunk}")194 error_msg += chunk195 continue196 if chunk_length > 6 and "delta" in chunk["choices"][0]:197 if chunk["choices"][0]["finish_reason"] == "stop":198 break199 try:200 yield chunk["choices"][0]["delta"]["content"]201 except Exception as e:202 # logging.error(f"Error: {e}")203 continue204 if error_msg:205 raise Exception(error_msg)206 207 208class ChatGLM_Client(BaseLLMModel):209 def __init__(self, model_name) -> None:210 super().__init__(model_name=model_name)211 from transformers import AutoTokenizer, AutoModel212 import torch213 global CHATGLM_TOKENIZER, CHATGLM_MODEL214 if CHATGLM_TOKENIZER is None or CHATGLM_MODEL is None:215 system_name = platform.system()216 model_path = None217 if os.path.exists("models"):218 model_dirs = os.listdir("models")219 if model_name in model_dirs:220 model_path = f"models/{model_name}"221 if model_path is not None:222 model_source = model_path223 else:224 model_source = f"THUDM/{model_name}"225 CHATGLM_TOKENIZER = AutoTokenizer.from_pretrained(226 model_source, trust_remote_code=True227 )228 quantified = False229 if "int4" in model_name:230 quantified = True231 model = AutoModel.from_pretrained(232 model_source, trust_remote_code=True233 )234 if torch.cuda.is_available():235 # run on CUDA236 logging.info("CUDA is available, using CUDA")237 model = model.half().cuda()238 # mps加速还存在一些问题，暂时不使用239 elif system_name == "Darwin" and model_path is not None and not quantified:240 logging.info("Running on macOS, using MPS")241 # running on macOS and model already downloaded242 model = model.half().to("mps")243 else:244 logging.info("GPU is not available, using CPU")245 model = model.float()246 model = model.eval()247 CHATGLM_MODEL = model248 249 def _get_glm_style_input(self):250 history = [x["content"] for x in self.history]251 query = history.pop()252 logging.debug(colorama.Fore.YELLOW +253 f"{history}" + colorama.Fore.RESET)254 assert (255 len(history) % 2 == 0256 ), f"History should be even length. current history is: {history}"257 history = [[history[i], history[i + 1]]258 for i in range(0, len(history), 2)]259 return history, query260 261 def get_answer_at_once(self):262 history, query = self._get_glm_style_input()263 response, _ = CHATGLM_MODEL.chat(264 CHATGLM_TOKENIZER, query, history=history)265 return response, len(response)266 267 def get_answer_stream_iter(self):268 history, query = self._get_glm_style_input()269 for response, history in CHATGLM_MODEL.stream_chat(270 CHATGLM_TOKENIZER,271 query,272 history,273 max_length=self.token_upper_limit,274 top_p=self.top_p,275 temperature=self.temperature,276 ):277 yield response278 279 280class LLaMA_Client(BaseLLMModel):281 def __init__(282 self,283 model_name,284 lora_path=None,285 ) -> None:286 super().__init__(model_name=model_name)287 from lmflow.datasets.dataset import Dataset288 from lmflow.pipeline.auto_pipeline import AutoPipeline289 from lmflow.models.auto_model import AutoModel290 from lmflow.args import ModelArguments, DatasetArguments, InferencerArguments291 292 self.max_generation_token = 1000293 self.end_string = "\n\n"294 # We don't need input data295 data_args = DatasetArguments(dataset_path=None)296 self.dataset = Dataset(data_args)297 self.system_prompt = ""298 299 global LLAMA_MODEL, LLAMA_INFERENCER300 if LLAMA_MODEL is None or LLAMA_INFERENCER is None:301 model_path = None302 if os.path.exists("models"):303 model_dirs = os.listdir("models")304 if model_name in model_dirs:305 model_path = f"models/{model_name}"306 if model_path is not None:307 model_source = model_path308 else:309 model_source = f"decapoda-research/{model_name}"310 # raise Exception(f"models目录下没有这个模型: {model_name}")311 if lora_path is not None:312 lora_path = f"lora/{lora_path}"313 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,314 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)315 pipeline_args = InferencerArguments(316 local_rank=0, random_seed=1, deepspeed='configs/ds_config_chatbot.json', mixed_precision='bf16')317 318 with open(pipeline_args.deepspeed, "r") as f:319 ds_config = json.load(f)320 LLAMA_MODEL = AutoModel.get_model(321 model_args,322 tune_strategy="none",323 ds_config=ds_config,324 )325 LLAMA_INFERENCER = AutoPipeline.get_pipeline(326 pipeline_name="inferencer",327 model_args=model_args,328 data_args=data_args,329 pipeline_args=pipeline_args,330 )331 # Chats332 # model_name = model_args.model_name_or_path333 # if model_args.lora_model_path is not None:334 # model_name += f" + {model_args.lora_model_path}"335 336 # context = (337 # "You are a helpful assistant who follows the given instructions"338 # " unconditionally."339 # )340 341 def _get_llama_style_input(self):342 history = []343 instruction = ""344 if self.system_prompt:345 instruction = (f"Instruction: {self.system_prompt}\n")346 for x in self.history:347 if x["role"] == "user":348 history.append(f"{instruction}Input: {x['content']}")349 else:350 history.append(f"Output: {x['content']}")351 context = "\n\n".join(history)352 context += "\n\nOutput: "353 return context354 355 def get_answer_at_once(self):356 context = self._get_llama_style_input()357 358 input_dataset = self.dataset.from_dict(359 {"type": "text_only", "instances": [{"text": context}]}360 )361 362 output_dataset = LLAMA_INFERENCER.inference(363 model=LLAMA_MODEL,364 dataset=input_dataset,365 max_new_tokens=self.max_generation_token,366 temperature=self.temperature,367 )368 369 response = output_dataset.to_dict()["instances"][0]["text"]370 return response, len(response)371 372 def get_answer_stream_iter(self):373 context = self._get_llama_style_input()374 partial_text = ""375 step = 1376 for _ in range(0, self.max_generation_token, step):377 input_dataset = self.dataset.from_dict(378 {"type": "text_only", "instances": [379 {"text": context + partial_text}]}380 )381 output_dataset = LLAMA_INFERENCER.inference(382 model=LLAMA_MODEL,383 dataset=input_dataset,384 max_new_tokens=step,385 temperature=self.temperature,386 )387 response = output_dataset.to_dict()["instances"][0]["text"]388 if response == "" or response == self.end_string:389 break390 partial_text += response391 yield partial_text392 393 394class XMBot_Client(BaseLLMModel):395 def __init__(self, api_key):396 super().__init__(model_name="xmchat")397 self.api_key = api_key398 self.session_id = None399 self.reset()400 self.image_bytes = None401 self.image_path = None402 self.xm_history = []403 self.url = "https://xmbot.net/web"404 405 def reset(self):406 self.session_id = str(uuid.uuid4())407 return [], "已重置"408 409 def try_read_image(self, filepath):410 import base64411 412 def is_image_file(filepath):413 # 判断文件是否为图片414 valid_image_extensions = [".jpg", ".jpeg", ".png", ".bmp", ".gif", ".tiff"]415 file_extension = os.path.splitext(filepath)[1].lower()416 return file_extension in valid_image_extensions417 418 def read_image_as_bytes(filepath):419 # 读取图片文件并返回比特流420 with open(filepath, "rb") as f:421 image_bytes = f.read()422 return image_bytes423 424 if is_image_file(filepath):425 logging.info(f"读取图片文件: {filepath}")426 image_bytes = read_image_as_bytes(filepath)427 base64_encoded_image = base64.b64encode(image_bytes).decode()428 self.image_bytes = base64_encoded_image429 self.image_path = filepath430 else:431 self.image_bytes = None432 self.image_path = None433 434 def prepare_inputs(self, real_inputs, use_websearch, files, reply_language, chatbot):435 fake_inputs = real_inputs436 display_append = ""437 limited_context = False438 return limited_context, fake_inputs, display_append, real_inputs, chatbot439 440 def handle_file_upload(self, files, chatbot):441 """if the model accepts multi modal input, implement this function"""442 if files:443 for file in files:444 if file.name:445 logging.info(f"尝试读取图像: {file.name}")446 self.try_read_image(file.name)447 if self.image_path is not None:448 chatbot = chatbot + [((self.image_path,), None)]449 if self.image_bytes is not None:450 logging.info("使用图片作为输入")451 conv_id = str(uuid.uuid4())452 data = {453 "user_id": self.api_key,454 "session_id": self.session_id,455 "uuid": conv_id,456 "data_type": "imgbase64",457 "data": self.image_bytes458 }459 response = requests.post(self.url, json=data)460 response = json.loads(response.text)461 logging.info(f"图片回复: {response['data']}")462 return None, chatbot, None463 464 def get_answer_at_once(self):465 question = self.history[-1]["content"]466 conv_id = str(uuid.uuid4())467 data = {468 "user_id": self.api_key,469 "session_id": self.session_id,470 "uuid": conv_id,471 "data_type": "text",472 "data": question473 }474 response = requests.post(self.url, json=data)475 try:476 response = json.loads(response.text)477 return response["data"], len(response["data"])478 except Exception as e:479 return response.text, len(response.text)480 481 482 483 484def get_model(485 model_name,486 lora_model_path=None,487 access_key=None,488 temperature=None,489 top_p=None,490 system_prompt=None,491) -> BaseLLMModel:492 msg = i18n("模型设置为了：") + f" {model_name}"493 model_type = ModelType.get_type(model_name)494 lora_selector_visibility = False495 lora_choices = []496 dont_change_lora_selector = False497 if model_type != ModelType.OpenAI:498 config.local_embedding = True499 # del current_model.model500 model = None501 try:502 if model_type == ModelType.OpenAI:503 logging.info(f"正在加载OpenAI模型: {model_name}")504 model = OpenAIClient(505 model_name=model_name,506 api_key=access_key,507 system_prompt=system_prompt,508 temperature=temperature,509 top_p=top_p,510 )511 elif model_type == ModelType.ChatGLM:512 logging.info(f"正在加载ChatGLM模型: {model_name}")513 model = ChatGLM_Client(model_name)514 elif model_type == ModelType.LLaMA and lora_model_path == "":515 msg = f"现在请为 {model_name} 选择LoRA模型"516 logging.info(msg)517 lora_selector_visibility = True518 if os.path.isdir("lora"):519 lora_choices = get_file_names(520 "lora", plain=True, filetypes=[""])521 lora_choices = ["No LoRA"] + lora_choices522 elif model_type == ModelType.LLaMA and lora_model_path != "":523 logging.info(f"正在加载LLaMA模型: {model_name} + {lora_model_path}")524 dont_change_lora_selector = True525 if lora_model_path == "No LoRA":526 lora_model_path = None527 msg += " + No LoRA"528 else:529 msg += f" + {lora_model_path}"530 model = LLaMA_Client(model_name, lora_model_path)531 elif model_type == ModelType.XMBot:532 model = XMBot_Client(api_key=access_key)533 elif model_type == ModelType.Unknown:534 raise ValueError(f"未知模型: {model_name}")535 logging.info(msg)536 except Exception as e:537 logging.error(e)538 msg = f"{STANDARD_ERROR_MSG}: {e}"539 if dont_change_lora_selector:540 return model, msg541 else:542 return model, msg, gr.Dropdown.update(choices=lora_choices, visible=lora_selector_visibility)543 544 545if __name__ == "__main__":546 with open("config.json", "r") as f:547 openai_api_key = cjson.load(f)["openai_api_key"]548 # set logging level to debug549 logging.basicConfig(level=logging.DEBUG)550 # client = ModelManager(model_name="gpt-3.5-turbo", access_key=openai_api_key)551 client = get_model(model_name="chatglm-6b-int4")552 chatbot = []553 stream = False554 # 测试账单功能555 logging.info(colorama.Back.GREEN + "测试账单功能" + colorama.Back.RESET)556 logging.info(client.billing_info())557 # 测试问答558 logging.info(colorama.Back.GREEN + "测试问答" + colorama.Back.RESET)559 question = "巴黎是中国的首都吗？"560 for i in client.predict(inputs=question, chatbot=chatbot, stream=stream):561 logging.info(i)562 logging.info(f"测试问答后history : {client.history}")563 # 测试记忆力564 logging.info(colorama.Back.GREEN + "测试记忆力" + colorama.Back.RESET)565 question = "我刚刚问了你什么问题？"566 for i in client.predict(inputs=question, chatbot=chatbot, stream=stream):567 logging.info(i)568 logging.info(f"测试记忆力后history : {client.history}")569 # 测试重试功能570 logging.info(colorama.Back.GREEN + "测试重试功能" + colorama.Back.RESET)571 for i in client.retry(chatbot=chatbot, stream=stream):572 logging.info(i)573 logging.info(f"重试后history : {client.history}")574 # # 测试总结功能575 # print(colorama.Back.GREEN + "测试总结功能" + colorama.Back.RESET)576 # chatbot, msg = client.reduce_token_size(chatbot=chatbot)577 # print(chatbot, msg)578 # print(f"总结后history: {client.history}")579