Add application file

app.py

@@ -1,22 +1,68 @@
 from fastapi import FastAPI
-from pydantic import BaseModel
-from transformers import pipeline
+from optimum.onnxruntime import ORTModelForSequenceClassification
+from transformers import AutoTokenizer
+import time
+
+# Load the tokenizer
+tokenizer = AutoTokenizer.from_pretrained("cardiffnlp/twitter-roberta-base-sentiment")
+
+# Load the quantized ONNX model from Hugging Face
+model = ORTModelForSequenceClassification.from_pretrained("minhdang/model_onnx", file_name="quantized_model.onnx")
 
 app = FastAPI()
 
-# Create a pipeline for text classification using the ONNX model
-pipe = pipeline("text-classification", model="minhdang/model_onnx")
+def preprocess_tweet(tweet: str) -> str:
+    tweet_words = []
+    for word in tweet.split(' '):
+        if word.startswith('@') and len(word) > 1:
+            word = '@user'
+        elif word.startswith('http'):
+            word = "http"
+        tweet_words.append(word)
+    return " ".join(tweet_words)
+
+@app.get("/")
+def home():
+    return {"message": "Welcome to the sentiment analysis API"}
+
+@app.get("/analyze")
+def analyze_sentiment(tweet: str):
+    # Preprocess the tweet
+    tweet_proc = preprocess_tweet(tweet)
+
+    # Measure the time taken for the inference
+    start_time = time.time()
+
+    # Tokenize the input tweet
+    inputs = tokenizer(tweet_proc, return_tensors="pt")
+
+    # Perform the inference with the ONNX model
+    outputs = model(**inputs)
+    
+    # Calculate the inference time
+    inference_time = time.time() - start_time
+
+    # Get the probabilities from the logits
+    probabilities = outputs.logits.softmax(dim=1)
+
+    # Get the label with the highest probability
+    max_prob, max_index = probabilities.max(dim=1)
 
-# Define a Pydantic model for input data
-class TextRequest(BaseModel):
-    text: str
+    # Map the labels to desired names
+    label_map = {
+        0: "Negative",
+        1: "Neutral",
+        2: "Positive"
+    }
 
-@app.post("/classify")
-async def classify(request: TextRequest):
-    text = request.text
-    # Use the pipeline to classify the text
-    result = pipe(text)
-    # Return the result as a JSON response
-    return {"result": result}
+    # Get the highest label and its corresponding score
+    highest_label = label_map[max_index.item()]
+    highest_score = round(max_prob.item(), 4)
 
-# Run the app with `uvicorn main:app --reload` in your terminal
+    # Return the original tweet, the label with the highest score, and the inference time
+    return {
+        "text": tweet,
+        "label": highest_label,
+        "score": highest_score,
+        "inference_time": round(inference_time, 4)  # In seconds
+    }