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)