mirror of https://github.com/THUDM/ChatGLM-6B
修改模型位置
parent
be7e14ce45
commit
02d4d3dd2c
6
api.py
6
api.py
|
@ -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)
|
||||||
|
|
Loading…
Reference in New Issue