Add application file

app.py

@@ -1,15 +1,9 @@
-import onnxruntime as ort
-from transformers import AutoTokenizer
-import numpy as np
-import time
 from fastapi import FastAPI
+from transformers import pipeline
+import time
 
-# Load the ONNX model
-onnx_model_path = "model_quantized.onnx"
-session = ort.InferenceSession(onnx_model_path)
-
-# Load the tokenizer
-tokenizer = AutoTokenizer.from_pretrained("cardiffnlp/twitter-roberta-base-sentiment")
+# Initialize the sentiment analysis pipeline
+pipe = pipeline("text-classification", model="cardiffnlp/twitter-roberta-base-sentiment")
 
 app = FastAPI()
 
@@ -23,12 +17,6 @@ def preprocess_tweet(tweet: str) -> str:
         tweet_words.append(word)
     return " ".join(tweet_words)
 
-def run_inference(tweet: str):
-    inputs = tokenizer(tweet, return_tensors="np", padding=True, truncation=True)
-    ort_inputs = {k: v for k, v in inputs.items()}
-    ort_outs = session.run(None, ort_inputs)
-    return ort_outs[0]
-
 @app.get("/")
 def home():
     return {"message": "Welcome to the sentiment analysis API"}
@@ -41,21 +29,28 @@ def analyze_sentiment(tweet: str):
     # Measure the time taken for the inference
     start_time = time.time()
     
-    # Run inference using ONNX model
-    logits = run_inference(tweet_proc)
-
+    # Use the pipeline to get the sentiment analysis result
+    results = pipe(tweet_proc, return_all_scores=True)
+    
     # Calculate the inference time
     inference_time = time.time() - start_time
 
+    # Map the labels to desired names
+    label_map = {
+        "LABEL_0": "Negative",
+        "LABEL_1": "Neutral",
+        "LABEL_2": "Positive"
+    }
+
     # Find the label with the highest score
-    label_map = ["Negative", "Neutral", "Positive"]
-    highest_label_idx = np.argmax(logits)
-    highest_score = np.max(logits)
+    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)
 
     # Return the original tweet, the label with the highest score, and the inference time
     return {
         "text": tweet,
-        "label": label_map[highest_label_idx],
-        "score": round(float(highest_score), 4),
+        "label": highest_label,
+        "score": highest_score,
         "inference_time": round(inference_time, 4)  # In seconds
-    }
+    }

requirements.txt

@@ -5,6 +5,8 @@ torch
 scipy
 onnx
 onnxruntime
-
+onnxruntime-gpu
+numpy
+time