| from fastapi import FastAPI, Request |
| from transformers import AutoTokenizer, AutoModel |
| import uvicorn, json, datetime |
| import torch |
|
|
| 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() |
|
|
|
|
| @app.post("/") |
| 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') |
| 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 |
|
|
|
|
| if __name__ == '__main__': |
| tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True) |
| model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True).half().cuda() |
| model.eval() |
| uvicorn.run(app, host='0.0.0.0', port=8000, workers=1) |
|
|
