main.py

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

# Set custom cache directory to avoid permission issues
os.environ["TRANSFORMERS_CACHE"] = "/app/cache"

app = FastAPI()

# Explicitly specify a model (avoid default selection)
MODEL_NAME = "distilbert-base-uncased-finetuned-sst-2-english"
sentiment_pipeline = pipeline("sentiment-analysis", model=MODEL_NAME)

class SentimentRequest(BaseModel):
    text: str

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

@app.get("/")
def home():
    return {"message": "Sentiment Analysis API is running!"}

@app.post("/predict/", response_model=SentimentResponse)
def predict(request: SentimentRequest):
    result = sentiment_pipeline(request.text)
    return SentimentResponse(label=result[0]['label'], score=result[0]['score'])