Update app.py

app.py

@@ -2,12 +2,15 @@ import streamlit as st
 import pandas as pd
 import transformers
 import torch
+import seaborn as sns
+import matplotlib.pyplot as plt
 
 # Load the pre-trained BERT model and tokenizer
 tokenizer = transformers.BertTokenizer.from_pretrained('bert-base-uncased')
 model = transformers.BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=6)
 
 # Set up the Streamlit app
+st.set_page_config(layout="wide")
 st.title('Toxicity Classification App')
 
 # Create a text input for the user to enter their text
@@ -57,3 +60,13 @@ if st.button('Classify'):
 
 # Display the persistent DataFrame
 st.write('Classification Results:', st.session_state.get('results', pd.DataFrame()))
+
+# Plot the distribution of probabilities for each category
+if len(st.session_state.get('results', pd.DataFrame())) > 0:
+    df = st.session_state['results']
+    fig, axes = plt.subplots(ncols=2, figsize=(12, 6))
+    sns.histplot(data=df, x='Toxic', kde=True, ax=axes[0])
+    axes[0].set_title('Toxic Probability Distribution')
+    sns.histplot(data=df, x='Severe Toxic', kde=True, ax=axes[1])
+    axes[1].set_title('Severe Toxic Probability Distribution')
+    st.pyplot(fig)