diff --git a/api.py b/api.py index 4ad1db6..693c70a 100644 --- a/api.py +++ b/api.py @@ -50,8 +50,7 @@ async def create_item(request: Request): 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) - -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()