Update app.py

app.py

@@ -13,7 +13,6 @@ tweet = df['comment_text'][r.randint(0,1000)]
 tokenizer = AutoTokenizer.from_pretrained("APJ23/MultiHeaded_Sentiment_Analysis_Model")
 model = AutoModelForSequenceClassification.from_pretrained("APJ23/MultiHeaded_Sentiment_Analysis_Model")
 
-# Define the classes and their corresponding labels
 classes = {
     0: 'Non-Toxic',
     1: 'Toxic',
@@ -23,11 +22,14 @@ classes = {
     5: 'Insult',
     6: 'Identity Hate'
 }
-
-# Create a function to generate the toxicity predictions
 @st.cache(allow_output_mutation=True)
-
-# Create a table to display the toxicity predictions
+def prediction(tweet,model,tokenizer):
+    inputs = tokenizer(tweet, return_tensors="pt", padding=True, truncation=True)
+    outputs = model(**inputs)
+    predicted_class = torch.argmax(outputs.logits, dim=1)
+    predicted_prob = torch.softmax(outputs.logits, dim=1)[0][predicted_class].item()
+    return classes[predicted_class], predicted_prob
+    
 def create_table(predictions):
     data = {'Tweet': [], 'Highest Toxicity Class': [], 'Probability': []}
     for tweet, prediction in predictions.items():
@@ -38,13 +40,11 @@ def create_table(predictions):
     return df
 
 st.title('Toxicity Prediction App')
-tweet_input = st.text_input('Enter a tweet:')
+st.write(f'The random tweet select is {tweet}',tweet)
 if st.button('Predict'):
-    predicted_class_label, predicted_prob = predict_toxicity(tweet_input, model, tokenizer)
+    predicted_class_label, predicted_prob = predict_toxicity(tweet, model, tokenizer)
     prediction_text = f'Prediction: {predicted_class_label} ({predicted_prob:.2f})'
     st.write(prediction_text)
-
-    # Display the toxicity predictions in a table
     predictions = {tweet_input: (predicted_class_label, predicted_prob)}
     table = create_table(predictions)
     st.table(table)