26 lines
770 B
Python
26 lines
770 B
Python
|
import argparse
|
|||
|
|
import uvicorn
|
|||
|
|
from fastapi_qa import app
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
parser = argparse.ArgumentParser()
|
|||
|
|
parser.add_argument("--model_dir", type=str, default="/model",
|
|||
|
|
help="模型目录（挂载到容器内的路径）")
|
|||
|
|
parser.add_argument("--use_gpu", action="store_true", default=True,
|
|||
|
|
help="是否使用 GPU（CUDA）")
|
|||
|
|
parser.add_argument("--port", type=int, default=8000,
|
|||
|
|
help="FastAPI 服务端口，默认 8000")
|
|||
|
|
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
|
|||
|
|
app.state.config = {
|
|||
|
|
"model_dir": args.model_dir,
|
|||
|
|
"use_gpu": args.use_gpu,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
uvicorn.run("fastapi_qa:app",
|
|||
|
|
host="0.0.0.0",
|
|||
|
|
port=args.port,
|
|||
|
|
workers=1,
|
|||
|
|
)
|