app.py

from fastapi import FastAPI, Request
from pydantic import BaseModel
from setfit import SetFitModel

app = FastAPI()

class TextInput(BaseModel):
    text: str

class LabelScore(BaseModel):
    label: str
    score: float

@app.post("/predict", response_model=LabelScore)
async def predict(input: TextInput):
    model = SetFitModel.from_pretrained("assets")
    preds = model.predict(input)
    return preds

@app.get("/")
def read_root():
    return {"message": "Welcome to the FastAPI prediction service"}