Add application file

app.py

@@ -1,9 +1,11 @@
 from fastapi import FastAPI
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
 import time
 
-# Initialize the sentiment analysis pipeline
-pipe = pipeline("text-classification", model="cardiffnlp/twitter-roberta-base-sentiment")
+# Load the tokenizer and model directly
+tokenizer = AutoTokenizer.from_pretrained("cardiffnlp/twitter-roberta-base-sentiment")
+model = AutoModelForSequenceClassification.from_pretrained("cardiffnlp/twitter-roberta-base-sentiment")
 
 app = FastAPI()
 
@@ -28,24 +30,33 @@ def analyze_sentiment(tweet: str):
 
     # Measure the time taken for the inference
     start_time = time.time()
-    
-    # Use the pipeline to get the sentiment analysis result
-    results = pipe(tweet_proc, return_all_scores=True)
+
+    # Tokenize the input tweet
+    inputs = tokenizer(tweet_proc, return_tensors="pt")
+
+    # Perform the inference
+    with torch.no_grad():
+        outputs = model(**inputs)
     
     # Calculate the inference time
     inference_time = time.time() - start_time
 
+    # Get the probabilities from the logits
+    probabilities = torch.softmax(outputs.logits, dim=1)
+
+    # Get the label with the highest probability
+    max_prob, max_index = torch.max(probabilities, dim=1)
+
     # Map the labels to desired names
     label_map = {
-        "LABEL_0": "Negative",
-        "LABEL_1": "Neutral",
-        "LABEL_2": "Positive"
+        0: "Negative",
+        1: "Neutral",
+        2: "Positive"
     }
 
-    # Find the label with the highest score
-    highest_score_result = max(results[0], key=lambda x: x['score'])
-    highest_label = label_map[highest_score_result['label']]
-    highest_score = round(highest_score_result['score'], 4)
+    # Get the highest label and its corresponding score
+    highest_label = label_map[max_index.item()]
+    highest_score = round(max_prob.item(), 4)
 
     # Return the original tweet, the label with the highest score, and the inference time
     return {
@@ -53,4 +64,4 @@ def analyze_sentiment(tweet: str):
         "label": highest_label,
         "score": highest_score,
         "inference_time": round(inference_time, 4)  # In seconds
-    }
+    }