Add application file

app.py

@@ -1,17 +1,19 @@
 from fastapi import FastAPI
-import onnxruntime as ort
-import numpy as np
+from transformers import AutoTokenizer
+from optimum.onnxruntime import ORTModelForSequenceClassification
+from optimum.onnxruntime import ORTQuantizer
+from optimum.onnxruntime.configuration import AutoQuantizationConfig
 import time
-from transformers import RobertaTokenizer
 
 app = FastAPI()
 
-# Load the ONNX model
-onnx_model_path = "sentiment_model.onnx"
-session = ort.InferenceSession(onnx_model_path)
+# Load the ONNX model and tokenizer
+model_path = "./model_onnx/model_quantized.onnx"
+tokenizer_path = "./model_onnx"
 
-# Initialize the tokenizer
-tokenizer = RobertaTokenizer.from_pretrained("cardiffnlp/twitter-roberta-base-sentiment")
+model = ORTModelForSequenceClassification.from_pretrained(model_path)
+tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+pipe = pipeline("text-classification", model=model, tokenizer=tokenizer)
 
 def preprocess_tweet(tweet: str) -> str:
     tweet_words = []
@@ -29,36 +31,31 @@ def home():
 
 @app.get("/analyze")
 def analyze_sentiment(tweet: str):
+    # Preprocess the tweet
     tweet_proc = preprocess_tweet(tweet)
-    inputs = tokenizer(tweet_proc, return_tensors="np")
-    
+
     # Measure the time taken for the inference
     start_time = time.time()
     
-    # Perform inference with ONNX
-    ort_inputs = {k: v for k, v in inputs.items()}
-    ort_outputs = session.run(None, ort_inputs)
-    
-    # Extract the output
-    logits = ort_outputs[0]
+    # Use the pipeline to get the sentiment analysis result
+    results = pipe(tweet_proc, return_all_scores=True)
     
-    # Apply softmax to get probabilities
-    probs = np.exp(logits) / np.sum(np.exp(logits), axis=1, keepdims=True)
-    
-    # Get the predicted label
+    # Calculate the inference time
+    inference_time = time.time() - start_time
+
+    # Map the labels to desired names
     label_map = {
-        0: "Negative",
-        1: "Neutral",
-        2: "Positive"
+        "LABEL_0": "Negative",
+        "LABEL_1": "Neutral",
+        "LABEL_2": "Positive"
     }
-    
-    highest_score_index = np.argmax(probs[0])
-    highest_label = label_map[highest_score_index]
-    highest_score = round(probs[0][highest_score_index], 4)
 
-    # Calculate the inference time
-    inference_time = time.time() - start_time
+    # 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)
 
+    # Return the original tweet, the label with the highest score, and the inference time
     return {
         "text": tweet,
         "label": highest_label,

requirements.txt

@@ -3,10 +3,7 @@ uvicorn
 transformers
 torch
 scipy
-fastapi==0.99.0
-onnxruntime==1.14.0
-transformers==4.33.0
-numpy==1.25.2
-uvicorn==0.23.1
+optimum
+onnxruntime