from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import re

app = FastAPI()

tokenizer = AutoTokenizer.from_pretrained("cardiffnlp/twitter-roberta-base-emotion-latest")
model = AutoModelForSequenceClassification.from_pretrained("cardiffnlp/twitter-roberta-base-emotion-latest")

class TextRequest(BaseModel):
    text: str

def clean_text(text: str) -> str:
    fillers = ["um", "uh", "like", "you know", "I mean", "sort of", "kind of", "hmm", "uhh"]
    text = re.sub(r'\b(?:' + '|'.join(fillers) + r')\b', '', text, flags=re.IGNORECASE)
    text = re.sub(r'\s+', ' ', text).strip()  
    return text

@app.get("/")
def home():
    return {"message": "Speak your mind emotion API is running"}

@app.post("/classify-emotion")
async def classify_emotion(request: TextRequest):
    try:
        text = request.text.strip()

        if not text:
            raise HTTPException(status_code=400, detail="Text cannot be empty")
        cleaned_text = clean_text(text)

        inputs = tokenizer(cleaned_text, return_tensors="pt", truncation=True, padding=True, max_length=512)

        with torch.no_grad():
            outputs = model(**inputs)

        logits = outputs.logits
        predicted_class_id = torch.argmax(logits, dim=-1).item()
        predicted_emotion = model.config.id2label[predicted_class_id]

        return {
            "original_text": text,
            "cleaned_text": cleaned_text,
            "predicted_emotion": predicted_emotion
        }

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error processing text: {str(e)}")