Add application file

app.py

@@ -1,11 +1,18 @@
-from fastapi import FastAPI
-from transformers import pipeline
+from fastapi import FastAPI, Query
+from transformers import pipeline, AutoTokenizer
+import onnxruntime as ort
+import numpy as np
 import time
 
-# Initialize the sentiment analysis pipeline
+app = FastAPI()
+
+# Initialize the Hugging Face pipeline for sentiment analysis
 pipe = pipeline("text-classification", model="cardiffnlp/twitter-roberta-base-sentiment")
 
-app = FastAPI()
+# Load the ONNX model and tokenizer
+onnx_model_path = "sentiment_model.onnx"
+session = ort.InferenceSession(onnx_model_path)
+tokenizer = AutoTokenizer.from_pretrained("cardiffnlp/twitter-roberta-base-sentiment")
 
 def preprocess_tweet(tweet: str) -> str:
     tweet_words = []
@@ -17,35 +24,48 @@ def preprocess_tweet(tweet: str) -> str:
         tweet_words.append(word)
     return " ".join(tweet_words)
 
+def run_inference_onnx(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"}
 
 @app.get("/analyze")
-def analyze_sentiment(tweet: str):
+def analyze_sentiment(tweet: str, method: str = Query("pipeline", enum=["pipeline", "onnx"])):
     # Preprocess the tweet
     tweet_proc = preprocess_tweet(tweet)
 
     # 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)
-    
-    # Calculate the inference time
-    inference_time = time.time() - start_time
+    if method == "pipeline":
+        # Use the Hugging Face pipeline to get the sentiment analysis result
+        results = pipe(tweet_proc, return_all_scores=True)
 
-    # Map the labels to desired names
-    label_map = {
-        "LABEL_0": "Negative",
-        "LABEL_1": "Neutral",
-        "LABEL_2": "Positive"
-    }
+        # Find the label with the highest score
+        highest_score_result = max(results[0], key=lambda x: x['score'])
+        label_map = {
+            "LABEL_0": "Negative",
+            "LABEL_1": "Neutral",
+            "LABEL_2": "Positive"
+        }
+        highest_label = label_map[highest_score_result['label']]
+        highest_score = round(highest_score_result['score'], 4)
+
+    elif method == "onnx":
+        # Run inference using the ONNX model
+        logits = run_inference_onnx(tweet_proc)
+        label_map = ["Negative", "Neutral", "Positive"]
+        highest_label_idx = np.argmax(logits)
+        highest_label = label_map[highest_label_idx]
+        highest_score = round(float(np.max(logits)), 4)
 
-    # 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)
+    # Calculate the inference time
+    inference_time = time.time() - start_time
 
     # Return the original tweet, the label with the highest score, and the inference time
     return {
@@ -53,4 +73,4 @@ def analyze_sentiment(tweet: str):
         "label": highest_label,
         "score": highest_score,
         "inference_time": round(inference_time, 4)  # In seconds
-    }
+    }