JulianHame's picture
removed comments
8bc0e72
raw
history blame
No virus
1.75 kB
import streamlit as st
from transformers import pipeline
import tensorflow as tf
import numpy as np
import pandas as pd
from tensorflow.keras.layers import TextVectorization
from tensorflow import keras
model = tf.keras.models.load_model('toxicity_model.h5')
df = pd.read_csv('train.csv')
X = df['comment_text']
y = df[df.columns[2:]].values
MAX_FEATURES = 200000
vectorizer = TextVectorization(max_tokens=MAX_FEATURES,
output_sequence_length=1800,
output_mode='int')
vectorizer.adapt(X.values)
st.title('Toxicity Classifier')
st.header('Write a message here:')
text = st.text_area('The toxicity of the message will be evaluated.',
value = "I hate your guts and you should die alone.")
input_str = vectorizer(text)
res = model.predict(np.expand_dims(input_str,0))
classification = res[0].tolist()
toxicity = classification[0]
toxicity_severe = classification[1]
obscene = classification[2]
threat = classification[3]
insult = classification[4]
identity_hate = classification[5]
highest_class = "Severe toxicity"
highest_class_rating = toxicity_severe
if(obscene > highest_class_rating):
highest_class = "Obscenity"
highest_class_rating = obscene
if(threat > highest_class_rating):
highest_class = "Threat"
highest_class_rating = threat
if(insult > highest_class_rating):
highest_class = "Insult"
highest_class_rating = insult
if(identity_hate > highest_class_rating):
highest_class = "Identity hate"
highest_class_rating = identity_hate
st.write("---")
st.write("Toxicity rating:")
st.write(toxicity)
st.write("---")
st.write("Highest Toxicity Class:")
st.write(highest_class)
st.write("Rating:")
st.write(highest_class_rating)