demodeploy / app.py
Abineshkumar77
Add application file
d65dba3
from fastapi import FastAPI
<<<<<<< HEAD
import onnxruntime as ort
from transformers import AutoTokenizer
=======
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
>>>>>>> 86f452461e3f04611f1ee50ade207b1a64893e79
import time
import torch
<<<<<<< HEAD
# Load the tokenizer and ONNX model
tokenizer = AutoTokenizer.from_pretrained("cardiffnlp/twitter-roberta-base-sentiment")
onnx_model_path = "D:/demodeploy/sentiment_model.onnx"
onnx_session = ort.InferenceSession(onnx_model_path)
=======
# Load the tokenizer and model directly
tokenizer = AutoTokenizer.from_pretrained("cardiffnlp/twitter-roberta-base-sentiment")
model = AutoModelForSequenceClassification.from_pretrained("cardiffnlp/twitter-roberta-base-sentiment")
>>>>>>> 86f452461e3f04611f1ee50ade207b1a64893e79
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")
input_ids = inputs["input_ids"].numpy()
attention_mask = inputs["attention_mask"].numpy()
# Perform the inference using ONNX
onnx_inputs = {
'input_ids': input_ids,
'attention_mask': attention_mask
}
outputs = onnx_session.run(None, onnx_inputs)
# Calculate the inference time
inference_time = time.time() - start_time
# Get the probabilities from the logits
logits = outputs[0]
probabilities = torch.softmax(torch.tensor(logits), dim=1)
# Get the label with the highest probability
max_prob, max_index = torch.max(probabilities, dim=1)
# Map the labels to desired names
label_map = {
0: "Negative",
1: "Neutral",
2: "Positive"
}
# Get the highest label and its corresponding score
highest_label = label_map[max_index.item()]
highest_score = round(max_prob.item(), 4)
# 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
<<<<<<< HEAD
}
=======
}
>>>>>>> 86f452461e3f04611f1ee50ade207b1a64893e79