app.py

import os
os.environ["HF_HOME"] = "/tmp"  # ensure Hugging Face cache is writable

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

# Load the Hugging Face model
classifier = pipeline("text-classification", model="mrm8488/bert-tiny-finetuned-sms-spam-detection")

# Define FastAPI app
app = FastAPI()

# Health check route
@app.get("/")
def read_root():
    return {"status": "ok", "message": "API is running"}

# Define request schema
class Query(BaseModel):
    text: str

# Prediction route (POST /)
@app.post("/predict")
def predict(query: Query):
    result = classifier(query.text)[0]
    return {"label": result["label"], "score": result["score"]}