from transformers import RobertaTokenizer, RobertaForSequenceClassification, Trainer, TrainingArguments
from transformers import DataCollatorWithPadding
import torch
import numpy as np
import gradio as gr


load_tokenizer = RobertaTokenizer.from_pretrained("./saved_model")
load_model = RobertaForSequenceClassification.from_pretrained("./saved_model")

# Define the prediction function
def predict(text):
    inputs = load_tokenizer(text, return_tensors='pt', truncation=True, padding=True)
    with torch.no_grad():
        outputs = load_model(**inputs)
    logits = outputs.logits
    prediction = torch.argmax(logits, dim=1).item()
    return "Hate speech detected" if prediction == 1 else "Hate speech not detected"

# Create Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs="text",
    outputs="text",
    title="Hate Speech Detection System",
    description="Enter a Pidgin or English text to check if it contains hate speech.",
        examples = [
        ["Yoruba men dey craze"],
        ["Yoruba men are crazy"],
        ["How una dey"],
        ["How are you"],
        ["All these Christians dey mad"],
        ["All Christians are mad"]
    ]
)

# Launch the Gradio interface
iface.launch()