Update app.py

app.py

@@ -3,19 +3,9 @@ import pandas as pd
 import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
-# Define the available models to choose from
-models = {
-    'BERT': 'bert-base-uncased',
-    'RoBERTa': 'roberta-base',
-    'DistilBERT': 'distilbert-base-uncased'
-}
-
-# Create a drop-down menu to select the model
-model_name = st.sidebar.selectbox('Select Model', list(models.keys()))
 
-# Load the tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained(models[model_name])
-model = AutoModelForSequenceClassification.from_pretrained(models[model_name])
+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 = {
@@ -30,17 +20,6 @@ classes = {
 
 # Create a function to generate the toxicity predictions
 @st.cache(allow_output_mutation=True)
-def predict_toxicity(tweet, model, tokenizer):
-    # Preprocess the text
-    inputs = tokenizer(tweet, padding=True, truncation=True, return_tensors='pt')
-    # Get the predictions from the model
-    outputs = model(**inputs)
-    predictions = torch.nn.functional.softmax(outputs.logits, dim=1).detach().numpy()
-    # Get the class with the highest probability
-    predicted_class = int(predictions.argmax())
-    predicted_class_label = classes[predicted_class]
-    predicted_prob = predictions[0][predicted_class]
-    return predicted_class_label, predicted_prob
 
 # Create a table to display the toxicity predictions
 def create_table(predictions):
@@ -52,11 +31,9 @@ def create_table(predictions):
     df = pd.DataFrame(data)
     return df
 
-# Create the user interface
 st.title('Toxicity Prediction App')
 tweet_input = st.text_input('Enter a tweet:')
 if st.button('Predict'):
-    # Generate the toxicity prediction for the tweet using the selected model
     predicted_class_label, predicted_prob = predict_toxicity(tweet_input, model, tokenizer)
     prediction_text = f'Prediction: {predicted_class_label} ({predicted_prob:.2f})'
     st.write(prediction_text)