ADK09's picture
Added clean_text function to remove filler words and extra spaces
f54c8a8
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)}")