import gradio as gr
import tensorflow as tf
from tensorflow.keras.layers import TextVectorization


# Load the saved model
model = tf.keras.models.load_model('toxicity.h5')

# Columns names
columns = ['toxic', 'severe_toxic', 'obscene', 'threat','insult', 'identity_hate']

# Specifying the maximum number of words in the vocabulary
MAX_FEATURES=200000

# Creating a TextVectorization layer with the specified parameters
vectorizer = TextVectorization(max_tokens=MAX_FEATURES,output_sequence_length=1800,output_mode="int")

# Define a function to score a comment
def score_comment(comment):
  # Vectorize the comment using the vectorizer
  vectorized_comment = vectorizer([comment])
    
  # Get the prediction results from the model
  results = model.predict(vectorized_comment)
    
  # Create a string to return the prediction results for each class
  text = ''
  for idx,col in enumerate(columns):
    text += '{}: {}\n'.format(col,results[0][idx]>0.5)
  
  return text


# Create a Gradio interface for the score_comment function
interface = gr.Interface(fn=score_comment,inputs=gr.inputs.Textbox(lines=2,placeholder="Comment to score"),outputs="text",title='Comment Toxicity Classifier')

# Launch the Gradio interface
interface.launch()