import gradio as gr
import numpy as np

from transformers import AutoTokenizer, TFAutoModelForSequenceClassification
  
tokenizer = AutoTokenizer.from_pretrained("risingodegua/hate-speech-detector")

model = TFAutoModelForSequenceClassification.from_pretrained("risingodegua/hate-speech-detector")

def make_prediction(text):
    '''
    This function takes a string as input and returns a prediction for the hate speech class.
    Hate speech class labels are: Normal(0), Offensive(1), and Hate speech(2).

    Parameters:
        text (str): The text to be classified.

    Returns:
        str: The predicted class label.   
    '''
    input_ids = tokenizer.encode(text)
    input_ids = np.array(input_ids)
    input_ids = np.expand_dims(input_ids, axis=0)
    prediction_arr = model.predict(input_ids)[0][0]

    labels = ["Normal", "Offensive", "Hate Speech"]
    prediction = labels[np.argmax(prediction_arr)]
    return prediction



iface = gr.Interface(
  fn=make_prediction, 
  inputs=gr.inputs.Textbox(lines=3, placeholder="Enter your text here..."), 
  outputs="text",
  title="Hate Speech Detector",
  description="A model for detecting if a given text is an hate speech.",
  )
iface.launch()