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)}") | |