from typing import Union from fastapi import FastAPI, Header, BackgroundTasks,Request,Body from urllib.parse import unquote import re import pickle import os from fastapi.responses import JSONResponse import asyncio from pydantic import BaseModel import time from transformers import AutoTokenizer, AutoModel import uvicorn, json, datetime import torch import json import base64 import requests from tentcentSMS import sendSms #使用int8的量化模型 tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b-int4", trust_remote_code=True) #model = AutoModel.from_pretrained("THUDM/chatglm2-6b-int4", trust_remote_code=True).half().quantize(8).cuda() model = AutoModel.from_pretrained("THUDM/chatglm2-6b-int4",trust_remote_code=True).cuda() model.eval() #显存满配全开 #tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True) #model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda() #model.eval() def submit_post(url: str, data: dict): """ Submit a POST request to the given URL with the given data. """ return requests.post(url, data=json.dumps(data)) def save_encoded_image(b64_image: str, output_path: str): """ Save the given image to the given output path. """ with open(output_path, "wb") as image_file: image_file.write(base64.b64decode(b64_image)) # 正则查找11位手机号 def findPhoneNum(text): pattern = re.compile(r"1[356789]\d{9}") # strs = '小明的手机号是13987692110，你明天打给他' result = pattern.findall(text) print(result) # ['13987692110'] if len(result) > 0: # 通过腾讯云-短信发送api sendSms(result[0]) return result[0] else: return None # bot发消息到tg群组 def sendTg(message): baseUrl="https://api.telegram.org/bot" toekn="558659722:AAENA3v8gKq7R7_nsr8V58Iu-_Z6BWx2SCw" chat_id="-861577609" text=message url=f'{baseUrl}{toekn}/sendMessage?chat_id={chat_id}&text={text}' try: r = requests.get(url) except: print('请检查网络代理') # https://api.telegram.org/bot558659722:AAENA3v8gKq7R7_nsr8V58Iu-_Z6BWx2SCw/sendMessage?chat_id=-861577609&text=1234 DEVICE = "cuda" DEVICE_ID = "0" CUDA_DEVICE = f"{DEVICE}:{DEVICE_ID}" if DEVICE_ID else DEVICE def torch_gc(): if torch.cuda.is_available(): with torch.cuda.device(CUDA_DEVICE): torch.cuda.empty_cache() torch.cuda.ipc_collect() app = FastAPI() class Item(BaseModel): query: Union[str, None] = None @app.get("/") def index(): return "hello" @app.get("/home") # def read_root(): # return {"Hello": "World"} def read_items( appid: Union[str, None] = Header(default=None), bid: Union[str, None] = Header(default=None), requestid: Union[str, None] = Header(default=None), uid: Union[str, None] = Header(default=None), ): print("console :,", "appid", appid, "bid", bid, "requestid", requestid, "uid", uid) return {"appid": appid, "bid": bid, "requestid": requestid, "uid": uid} @app.post("/glm") async def create_item(request: Request): # 获取请求参数 global model, tokenizer json_post_raw = await request.json() json_post = json.dumps(json_post_raw) json_post_list = json.loads(json_post) prompt = json_post_list.get('prompt') #问题 history = json_post_list.get('history') #默认空数组 max_length = json_post_list.get('max_length') top_p = json_post_list.get('top_p') temperature = json_post_list.get('temperature') # 请求ai模型 response, history = model.chat(tokenizer, prompt, history=history, max_length=max_length if max_length else 2048, top_p=top_p if top_p else 0.7, temperature=temperature if temperature else 0.95) now = datetime.datetime.now() time = now.strftime("%Y-%m-%d %H:%M:%S") answer = { "response": response, "history": history, "status": 200, "time": time } log = "[" + time + "] " + '", prompt:"' + prompt + '", response:"' + repr(response) + '"' print(log) torch_gc() return answer # AI智能对话 @app.post("/callback") async def callback( # item: Item, # inputArguments: Union[str, None] = None, query=Body(None), appid: Union[str, None] = Header(default=None), bid: Union[str, None] = Header(default=None), requestid: Union[str, None] = Header(default=None), uid: Union[str, None] = Header(default=None), ): print('query',query) print('query[query]',query["query"]) # 获取请求参数 global model, tokenizer # time.sleep(10) # print("item", item) history = [] #默认空数组 prompt='' try: query = unquote(query["query"], "utf-8") prompt = query #问题 except: prompt = "" max_length=None top_p=None temperature=None # 请求ai模型 response, history = model.chat(tokenizer, prompt, history=history, max_length=max_length if max_length else 2048, top_p=top_p if top_p else 0.7, temperature=temperature if temperature else 0.95) # log日志 now = datetime.datetime.now() time = now.strftime("%Y-%m-%d %H:%M:%S") log = "[" + time + "] " + '", prompt:"' + prompt + '", response:"' + repr(response) + '"' print(log) torch_gc() # return answer content = { "answer": response, "answer_type": "text", } headers = {"Content-Type": "application/json"} return JSONResponse(content=content, headers=headers) # AI生成图片 @app.post("/img") async def img( # item: Item, # inputArguments: Union[str, None] = None, query=Body(None), appid: Union[str, None] = Header(default=None), bid: Union[str, None] = Header(default=None), requestid: Union[str, None] = Header(default=None), uid: Union[str, None] = Header(default=None), ): # print("console :,", "appid", appid, "bid", bid, "requestid", requestid, "uid", uid) print('query',query) print('query[query]',query["query"]) # sendTg(query["query"]) # 获取请求参数 global model, tokenizer # time.sleep(10) # print("item", item) history = [] #默认空数组 prompt='' max_length=None top_p=None temperature=None response=None try: query = unquote(query["query"], "utf-8") # 判断有没有'/img'标识符 if ('/img' in query) is True: query=query.replace('/img', '') prompt=query txt2img_url = 'http://127.0.0.1:7860/sdapi/v1/txt2img' data = {'prompt': prompt} res = submit_post(txt2img_url, data) # 以时间戳命名 now =str(time.time()) save_encoded_image(res.json()['images'][0], f'img/{now}.png') # 网络图片 response=f'#web_img http://aiimg.hackrobot.cn/{now}.png' print('response',response) # return responseimg elif ('手机号' in query) is True: phonenumber=findPhoneNum(query) if phonenumber is not None: response="请查收短信,并在微信回复短信中的数字,自动邀请加群" else: response="手机号格式不正确,或其他错误" else: prompt = query #问题 # 请求ai模型 response, history = model.chat(tokenizer, prompt, history=history, max_length=max_length if max_length else 2048, top_p=top_p if top_p else 0.7, temperature=temperature if temperature else 0.95) except(ValueError, ArithmeticError): print('ValueError',ValueError) print('ArithmeticError',ArithmeticError) prompt = "" # log日志 # now = datetime.datetime.now() # time = now.strftime("%Y-%m-%d %H:%M:%S") # log = "[" + time + "] " + '", prompt:"' + prompt + '", response:"' + repr(response) + '"' # print(log) content = { "answer": response, "answer_type": "text", } headers = {"Content-Type": "application/json"} return JSONResponse(content=content, headers=headers) #if __name__ == '__main__': #tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True) #model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda() # 多显卡支持，使用下面三行代替上面两行，将num_gpus改为你实际的显卡数量 # model_path = "THUDM/chatglm2-6b" # tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) # model = load_model_on_gpus(model_path, num_gpus=2) # model.eval() # uvicorn.run(app, host='0.0.0.0', port=8000, workers=1)