Create app.py

app.py

@@ -0,0 +1,67 @@
+import streamlit as st
+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])
+
+# Define the classes and their corresponding labels
+classes = {
+    0: 'Non-Toxic',
+    1: 'Toxic',
+    2: 'Severely Toxic',
+    3: 'Obscene',
+    4: 'Threat',
+    5: 'Insult',
+    6: 'Identity Hate'
+}
+
+# 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):
+    data = {'Tweet': [], 'Highest Toxicity Class': [], 'Probability': []}
+    for tweet, prediction in predictions.items():
+        data['Tweet'].append(tweet)
+        data['Highest Toxicity Class'].append(prediction[0])
+        data['Probability'].append(prediction[1])
+    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)
+
+    # Display the toxicity predictions in a table
+    predictions = {tweet_input: (predicted_class_label, predicted_prob)}
+    table = create_table(predictions)
+    st.table(table)