修改模型位置

pull/547/head
Whitroom 2023-04-12 15:03:33 +08:00
parent be7e14ce45
commit 02d4d3dd2c
1 changed files with 3 additions and 3 deletions

6
api.py
View File

@ -16,11 +16,11 @@ models = {}
@asynccontextmanager @asynccontextmanager
async def lifespan(app: FastAPI): async def lifespan(app: FastAPI):
models['chat'] = AutoModel.from_pretrained( models['chat'] = AutoModel.from_pretrained(
"THUDM/models", "THUDM/chatglm-6b",
trust_remote_code=True).half().cuda() trust_remote_code=True).half().cuda()
models['chat'].eval() models['chat'].eval()
models['tokenizer'] = AutoTokenizer.from_pretrained( models['tokenizer'] = AutoTokenizer.from_pretrained(
"THUDM/models", "THUDM/chatglm-6b",
trust_remote_code=True) trust_remote_code=True)
yield yield
for model in models.values(): 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) return Answer(response=response, history=history, status=200, time=time)
if __name__ == '__main__': 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)