diff --git a/api.py b/api.py index 3e0e533..8997712 100644 --- a/api.py +++ b/api.py @@ -16,11 +16,11 @@ models = {} @asynccontextmanager async def lifespan(app: FastAPI): models['chat'] = AutoModel.from_pretrained( - "THUDM/models", + "THUDM/chatglm-6b", trust_remote_code=True).half().cuda() models['chat'].eval() models['tokenizer'] = AutoTokenizer.from_pretrained( - "THUDM/models", + "THUDM/chatglm-6b", trust_remote_code=True) yield for model in models.values(): @@ -63,4 +63,4 @@ async def create_item(item: Item): return Answer(response=response, history=history, status=200, time=time) if __name__ == '__main__': - uvicorn.run(app, host='0.0.0.0', port=8000, workers=1) + uvicorn.run(app, host='0.0.0.0', port=8010, workers=1)