Spaces:
Sleeping
Sleeping
File size: 1,697 Bytes
5c50bbf f54c8a8 5c50bbf f54c8a8 5c50bbf f54c8a8 5c50bbf f54c8a8 5c50bbf f54c8a8 5c50bbf f54c8a8 5c50bbf f54c8a8 5c50bbf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
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)}")
|