423A35C7/LLaMA-Factory
1
0
Fork 0
mirror of https://github.com/hiyouga/LLaMA-Factory.git synced 2025-08-23 22:32:54 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix api

Browse Source 
Former-commit-id: 8333795fcb25728dc7146596869f06e5e9533e58
This commit is contained in:
hiyouga 2023-07-20 22:14:54 +08:00
parent 38109c3da3
commit f1010ff65f
2 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
src/api_demo.py
View File

@ -5,11 +5,14 @@
import uvicorn
from llmtuner import ChatModel
from llmtuner.api.app import create_app
from llmtuner.tuner import get_infer_args
def main():
app = create_app()
chat_model = ChatModel(*get_infer_args())
app = create_app(chat_model)
uvicorn.run(app, host="0.0.0.0", port=8000, workers=1)

2
src/llmtuner/api/app.py
View File

@ -30,7 +30,7 @@ async def lifespan(app: FastAPI): # collects GPU memory
torch_gc()
def create_app(chat_model: ChatModel):
def create_app(chat_model: ChatModel) -> FastAPI:
app = FastAPI(lifespan=lifespan)
app.add_middleware(