ChadWong/CharttGLM2-6B

1from fastapi import FastAPI, Request2from transformers import AutoTokenizer, AutoModel3import uvicorn, json, datetime4import torch5 6DEVICE = "cuda"7DEVICE_ID = "0"8CUDA_DEVICE = f"{DEVICE}:{DEVICE_ID}" if DEVICE_ID else DEVICE9 10 11def torch_gc():12    if torch.cuda.is_available():13        with torch.cuda.device(CUDA_DEVICE):14            torch.cuda.empty_cache()15            torch.cuda.ipc_collect()16 17 18app = FastAPI()19 20 21@app.post("/")22async def create_item(request: Request):23    global model, tokenizer24    json_post_raw = await request.json()25    json_post = json.dumps(json_post_raw)26    json_post_list = json.loads(json_post)27    prompt = json_post_list.get('prompt')28    history = json_post_list.get('history')29    max_length = json_post_list.get('max_length')30    top_p = json_post_list.get('top_p')31    temperature = json_post_list.get('temperature')32    response, history = model.chat(tokenizer,33                                   prompt,34                                   history=history,35                                   max_length=max_length if max_length else 2048,36                                   top_p=top_p if top_p else 0.7,37                                   temperature=temperature if temperature else 0.95)38    now = datetime.datetime.now()39    time = now.strftime("%Y-%m-%d %H:%M:%S")40    answer = {41        "response": response,42        "history": history,43        "status": 200,44        "time": time45    }46    log = "[" + time + "] " + '", prompt:"' + prompt + '", response:"' + repr(response) + '"'47    print(log)48    torch_gc()49    return answer50 51 52if __name__ == '__main__':53    tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)54    model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda()55    # 多显卡支持，使用下面三行代替上面两行，将num_gpus改为你实际的显卡数量56    # model_path = "THUDM/chatglm2-6b"57    # tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)58    # model = load_model_on_gpus(model_path, num_gpus=2)59    model.eval()60    uvicorn.run(app, host='0.0.0.0', port=8000, workers=1)61