zero-shot-classification-zh/main.py

import uvicorn
from fastapi import FastAPI
from pydantic import BaseModel
from classification import ClassifyResult, classify as classifier
from config import HTTP_ARGS

app = FastAPI()


class TextClassificationRequest(BaseModel):
    text: str
    labels: list[str]


# class TextClassificationResponse(BaseModel):
#     prediction: str
#     ranks: list[str]


@app.post("/classify")
def classify(req: TextClassificationRequest) -> ClassifyResult:

    return classifier(req.text, req.labels)


if __name__ == "__main__":
    uvicorn.run(app, **HTTP_ARGS)