# Import gradio and transformers libraries
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# Load the small deberta models for hate and offensive speech detection
hate_model = AutoModelForSequenceClassification.from_pretrained("KoalaAI/HateSpeechDetector")
hate_tokenizer = AutoTokenizer.from_pretrained("KoalaAI/HateSpeechDetector")

# Define a function that takes an input text and returns the scores from the models
def get_scores(text):
  # Tokenize and encode the input text
  hate_input = hate_tokenizer(text, return_tensors="pt")

  # Get the logits from the models
  hate_logits = hate_model(**hate_input).logits

  # Apply softmax to get probabilities
  hate_probs = hate_logits.softmax(dim=1)

  # Get the labels from the models
  hate_labels = hate_model.config.id2label

  # Format the output as a dictionary of scores
  output = {}
  output["Hate speech"] = {hate_labels[i]: round(p.item(), 4) for i, p in enumerate(hate_probs[0])}

  return output

# Create a gradio interface with a text input and a json output
iface = gr.Interface(fn=get_scores, inputs="text", outputs="json")

# Launch the interface
iface.launch()