FastAPIExample / main.py
Akbartus's picture
Update main.py
9f2c4c2 verified
raw
history blame
541 Bytes
from fastapi import FastAPI
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from transformers import pipeline
app = FastAPI(debug=True, title="Text classifier")
class Payload(BaseModel):
text: str
async def classify_text(text):
pipe = pipeline("text-classification",
model="SamLowe/roberta-base-go_emotions")
return pipe(text)
@app.post("/classify")
async def classify(payload: Payload):
result = await classify_text(payload.text)
return JSONResponse({"data": result})