import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
from typing import Dict, List, Any
import numpy as np

# Load the model and tokenizer
model = AutoModelForSequenceClassification.from_pretrained('kietnt0603/bertweet-base-hate-speech-offensive')
tokenizer = AutoTokenizer.from_pretrained('kietnt0603/bertweet-base-hate-speech-offensive')

# Define the labels
labels = ["Hate", "Offensive", "Neither"]

# Function for prediction
def predict(inputs: str) -> Dict[str, Any]:
    # Tokenize input text
    inputs_dict = tokenizer(inputs, return_tensors="pt")

    # Forward pass
    with torch.no_grad():
        outputs = model(**inputs_dict)

    # Softmax to get probabilities
    probabilities = torch.nn.functional.softmax(outputs.logits, dim=1)

    # Get probabilities for each label
    label_probabilities = {label: round(prob, 4) for label, prob in zip(labels, probabilities[0].tolist())}

    # Return the result
    return label_probabilities

# Create title and description for the task
title = "Text Classification Demo"
description = "Classify text into categories: Hate, Offensive, Neither"
article = "Model loaded from https://huggingface.co/kietnt0603/bertweet-base-hate-speech-offensive"

# Create the Gradio interface
iface = gr.Interface(fn=predict,
                     inputs="textbox",
                     outputs=gr.Label(num_top_classes=3, label='Predictions'),
                     title=title,
                     description=description,
                     article=article)

# Launch the interface
iface.launch()