feat: changing model with a more finetuned respose

app.py

@@ -1,43 +1,85 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from typing import List
 from transformers import pipeline
+import torch
+from typing import List
 
+# Initialize FastAPI
 app = FastAPI()
 
-sentiment_analyzer = pipeline("sentiment-analysis", model="nlptown/bert-base-multilingual-uncased-sentiment")
+# Load RoBERTa sentiment analysis model
+sentiment_pipeline = pipeline(
+    "sentiment-analysis", 
+    model="cardiffnlp/twitter-roberta-base-sentiment"
+)
 
+# Request models
 class SentimentRequest(BaseModel):
+    content_id: str
     text: str
-    post_id: str
-
-class SentimentResponse(BaseModel):
-    post_id: str
-    label: str
-    score: float
 
-@app.post("/sentiment", response_model=SentimentResponse)
-async def analyze_sentiment(request: SentimentRequest):
-    result = sentiment_analyzer(request.text)
-    return {"post_id": request.post_id, "label": result[0]["label"], "score": result[0]["score"]}
+class BatchSentimentRequest(BaseModel):
+    posts: List[SentimentRequest]
 
-class SentimentPost(BaseModel):
-    post_id: str
-    text: str
+# Response model
+class SentimentResponse(BaseModel):
+    content_id: str
+    sentiment_score: float
+    sentiment_label: str
+    confidence: float
 
-class SentimentBatchRequest(BaseModel):
-    posts: List[SentimentPost]
+# Mapping RoBERTa labels to a floating-point scale
+LABEL_MAP = {
+    "LABEL_0": -1.0,  # Negative
+    "LABEL_1": 0.0,   # Neutral
+    "LABEL_2": 1.0    # Positive
+}
 
+@app.post("/analyze", response_model=SentimentResponse)
+def analyze_sentiment(request: SentimentRequest):
+    try:
+        # Get model prediction
+        result = sentiment_pipeline(request.text)[0]
+        label = result["label"]
+        score = result["score"]
+        
+        # Convert RoBERTa labels to floating-point scores
+        sentiment_score = LABEL_MAP[label]
+        confidence = round(score, 3)
+        
+        return SentimentResponse(
+            content_id=request.content_id,
+            sentiment_score=sentiment_score,
+            sentiment_label="positive" if sentiment_score > 0 else "neutral" if sentiment_score == 0 else "negative",
+            confidence=confidence
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
-@app.post("/sentiment/batch", response_model=List[SentimentResponse])
-async def analyze_sentiment_batch(request: SentimentBatchRequest):
-    texts = [post.text for post in request.posts]
-    results = sentiment_analyzer(texts)
-    return [
-        {"post_id": post.post_id, "label": res["label"], "score": res["score"]}
-        for post, res in zip(request.posts, results)
-    ]
+@app.post("/analyze_batch", response_model=List[SentimentResponse])
+def analyze_sentiment_batch(request: BatchSentimentRequest):
+    try:
+        responses = []
+        for post in request.posts:
+            result = sentiment_pipeline(post.text)[0]
+            label = result["label"]
+            score = result["score"]
+            
+            # Convert RoBERTa labels to floating-point scores
+            sentiment_score = LABEL_MAP[label]
+            confidence = round(score, 3)
+            
+            responses.append(SentimentResponse(
+                content_id=post.content_id,
+                sentiment_score=sentiment_score,
+                sentiment_label="positive" if sentiment_score > 0 else "neutral" if sentiment_score == 0 else "negative",
+                confidence=confidence
+            ))
+        return responses
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
+# Root endpoint
 @app.get("/")
-def greet_json():
-    return {"message": "BERT Sentiment Analysis API is running!"}
+def root():
+    return {"message": "RoBERTa Sentiment Analysis API is running!"}