from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline
import os


os.environ["HF_HOME"] = "/tmp" 

spam = pipeline("text-classification", model="mrm8488/bert-tiny-finetuned-sms-spam-detection")

toxic = pipeline("text-classification", model="s-nlp/roberta_toxicity_classifier")

sentiment = pipeline("text-classification", model = "nlptown/bert-base-multilingual-uncased-sentiment")



app = FastAPI()

@app.get("/")
def root():
    return {"status": "ok"}

class Query(BaseModel):
    text: str

@app.post("/spam")
def predict_spam(query: Query):
    result = spam(query.text)[0]
    return {"label": result["label"], "score": result["score"]}

@app.post("/toxic")
def predict_toxic(query: Query):
    result = toxic(query.text)[0]
    return {"label": result["label"], "score": result["score"]}

@app.post("/sentiment")
def predict_sentiment(query: Query):
    result = sentiment(query.text)[0]
    return {"label": result["label"], "score": result["score"]}