| from contextlib import asynccontextmanager |
|
|
| from fastapi import FastAPI, HTTPException |
| from pydantic import BaseModel |
|
|
| from predictor import load_benchmark, load_teacher, load_student |
|
|
| _models: dict = {} |
|
|
|
|
| @asynccontextmanager |
| async def lifespan(app: FastAPI): |
| _models["benchmark"] = load_benchmark() |
| _models["teacher"] = load_teacher() |
| _models["student"] = load_student() |
| yield |
| _models.clear() |
|
|
|
|
| app = FastAPI(title="AMANE Toxicity API", version="1.0.0", lifespan=lifespan) |
|
|
|
|
| class PredictRequest(BaseModel): |
| text: str |
| model: str = "student" |
|
|
|
|
| class BatchRequest(BaseModel): |
| texts: list[str] |
| model: str = "student" |
|
|
|
|
| @app.get("/health") |
| def health(): |
| return { |
| name: {"loaded": p.loaded, "error": p.error} |
| for name, p in _models.items() |
| } |
|
|
|
|
| @app.post("/predict") |
| def predict(req: PredictRequest): |
| p = _models.get(req.model) |
| if p is None: |
| raise HTTPException(status_code=404, detail=f"Unknown model: {req.model}") |
| if not p.loaded: |
| raise HTTPException(status_code=503, detail=f"Model not available: {p.error}") |
| return p.predict(req.text) |
|
|
|
|
| @app.post("/predict_batch") |
| def predict_batch(req: BatchRequest): |
| p = _models.get(req.model) |
| if p is None: |
| raise HTTPException(status_code=404, detail=f"Unknown model: {req.model}") |
| if not p.loaded: |
| raise HTTPException(status_code=503, detail=f"Model not available: {p.error}") |
| return p.predict_batch(req.texts) |
|
|
