Spaces:
Sleeping
Sleeping
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 | |
def home(): | |
return {"message": "Speak your mind emotion API is running"} | |
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)}") | |