Spaces:
Sleeping
Sleeping
| # importing libraries | |
| import streamlit as st | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| import tensorflow as tf | |
| from helper import processed_text | |
| # Load your pre-trained model | |
| def load_pretrained_model(): | |
| model = tf.keras.models.load_model('toxic.keras') | |
| return model | |
| # Load the pre-trained model | |
| model = load_pretrained_model() | |
| # Define the classes | |
| classes = ['toxic', 'severe_toxic', 'obscene', 'threat', 'insult', 'identity_hate'] | |
| # Function to predict classes and probabilities | |
| def predict_toxicity(text): | |
| vectorized_text = processed_text(text) # Modify this as per your preprocessing | |
| # Predict probabilities | |
| probabilities = model.predict(np.expand_dims(vectorized_text, 0))[0] | |
| # Determine predicted classes (those with probability > 0.5) | |
| predicted_classes = [classes[i] for i, prob in enumerate(probabilities) if prob >= 0.5] | |
| return probabilities, predicted_classes | |
| # Streamlit Front End | |
| st.title('Comment Toxicity Detection') | |
| # Input field for the comment | |
| user_input = st.text_area("Enter a comment:", "") | |
| if st.button('Predict'): | |
| if user_input: | |
| # Make predictions | |
| probabilities, predicted_classes = predict_toxicity(user_input) | |
| # Create two columns | |
| col1, col2 = st.columns([2, 3]) # Adjust column width as needed | |
| with col2: # Right Column | |
| # Display the original text | |
| st.subheader("Original Comment:") | |
| st.write(user_input) | |
| # Display the predicted classes | |
| st.subheader("Predicted Toxicity Classes:") | |
| if predicted_classes: | |
| # Capitalize, colorize, and bold each class | |
| for i, cls in enumerate(predicted_classes): | |
| color = f'<span style="color:{plt.cm.Set1(i % 9)}; font-weight:bold;">{cls.capitalize()}</span>' | |
| st.markdown(color, unsafe_allow_html=True) | |
| else: | |
| st.write("No toxic content detected.") | |
| with col1: # Left Column | |
| # Plot the vertical bar graph of probabilities with different colors | |
| st.subheader("Class Probabilities:") | |
| # Define colors for each bar (one for each class) | |
| colors = ['red', 'green', 'blue', 'orange', 'purple', 'cyan'] | |
| fig, ax = plt.subplots() | |
| ax.bar(classes, probabilities, color=colors) | |
| ax.set_ylim(0, 1) | |
| ax.set_ylabel('Probability') | |
| ax.set_title('Toxicity Class Probabilities') | |
| # Display the bar graph | |
| st.pyplot(fig) | |
| else: | |
| st.write("Please enter a comment to predict.") | |
| # print('Text vectorization for input is', processed_text('Yor are a bad person')) |