app.py

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


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

spam = pipeline("text-classification", model="valurank/distilroberta-spam-comments-detection")

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

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

nsfw = pipeline("text-classification", model = "michellejieli/NSFW_text_classifier")


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"]}

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

@app.get("/health")
def health_check():

    status = {
        "server": "running",
        "models": {}
    }

    models = {
        "spam": ("valurank/distilroberta-spam-comments-detection", spam),
        "toxic": ("s-nlp/roberta_toxicity_classifier", toxic),
        "sentiment": ("nlptown/bert-base-multilingual-uncased-sentiment", sentiment),
        "nsfw": ("michellejieli/NSFW_text_classifier", nsfw),
    }

    for key, (model_name, model_pipeline) in models.items():
        try:
            model_pipeline("test")
            status["models"][key] = {
                "model_name": model_name,
                "status": "running"
            }
        except Exception as e:
            status["models"][key] = {
                "model_name": model_name,
                "status": f"error: {str(e)}"
            }

    return status