28 lines
559 B
Python
28 lines
559 B
Python
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)
|