File size: 1,346 Bytes
a559d4f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | import time
from fastapi import APIRouter, HTTPException, Query, Request
from src.schemas.requests import SampleInput
from src.schemas.responses import SampleOutput
from src.logger import log_to_betterstack
router = APIRouter()
MAX_LEN = 256
BATCH_SIZE = 32
@router.post("/predict-all-models", response_model=dict[str, list[SampleOutput]])
def predict_all_models_endpoint(
request: Request,
samples: list[SampleInput],
deduplicate: bool = Query(False),
) -> dict[str, list[SampleOutput]]:
pipeline = getattr(request.app.state, "all_models_pipeline", None)
if pipeline is None or not pipeline.models:
raise HTTPException(
status_code=503,
detail="No models are available.",
)
samples_raw = [s.model_dump() for s in samples]
start = time.perf_counter()
results = pipeline.run(samples_raw, MAX_LEN, BATCH_SIZE, deduplicate=deduplicate)
elapsed = time.perf_counter() - start
output = {
mode: [SampleOutput(**r) for r in preds]
for mode, preds in results.items()
}
log_to_betterstack(
endpoint_name="/predict-all-models",
original_text=samples_raw,
formatted_text={m: [o.model_dump() for o in preds] for m, preds in output.items()},
model_name="all",
time_elapsed=elapsed,
)
return output
|