Add application file

app.py

@@ -1,68 +1,22 @@
 from fastapi import FastAPI
-from transformers import AutoTokenizer
-from optimum.onnxruntime import ORTModelForSequenceClassification
-import time
-import torch
-
-# Load the tokenizer and ONNX model directly
-tokenizer = AutoTokenizer.from_pretrained("minhdang/model_onnx")
-model = ORTModelForSequenceClassification.from_pretrained("minhdang/model_onnx", file_name="model_quantized.onnx")
+from pydantic import BaseModel
+from transformers import pipeline
 
 app = FastAPI()
 
-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 ONNX 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)
+# Create a pipeline for text classification using the ONNX model
+pipe = pipeline("text-classification", model="minhdang/model_onnx")
 
-    # Map the labels to desired names
-    label_map = {
-        0: "Negative",
-        1: "Neutral",
-        2: "Positive"
-    }
+# Define a Pydantic model for input data
+class TextRequest(BaseModel):
+    text: str
 
-    # Get the highest label and its corresponding score
-    highest_label = label_map[max_index.item()]
-    highest_score = round(max_prob.item(), 4)
+@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}
 
-    # 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
-    }
+# Run the app with `uvicorn main:app --reload` in your terminal