zero-shot-classification-zh/main.py

import uvicorn
from fastapi import FastAPI
from pydantic import BaseModel
import classification

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) -> TextClassificationResponse:
    result = classification.classify(req.text, req.labels)

    return TextClassificationResponse(prediction=result.prediction, ranks=result.rank)


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)