import gradio as gr
from transformers import pipeline
import os

# 1. Load the classifier from your Hugging Face Repo
# This replaces the /content/drive path
model_repo = "MakD1227/afriberta-hsd-model"
classifier = pipeline("text-classification", model=model_repo)

# 2. Prediction function
def predict_speech(text): 
    results = classifier(text)
    # Mapping: LABEL_0 -> Free, LABEL_1 -> Offensive, LABEL_2 -> Hate
    label_map = {"LABEL_0": "Free (Neutral)", "LABEL_1": "Offensive", "LABEL_2": "Hate"}

    label = results[0]['label']
    score = results[0]['score']

    return label_map.get(label, label), f"{round(score * 100, 2)}%"

# 3. Gradio Interface
interface = gr.Interface(
    fn=predict_speech,
    inputs=gr.Textbox(
        lines=2,
        label="Input Text",
        placeholder="Enter Amharic or Afan Oromo text..."
    ),
    outputs=[
        gr.Label(label="Classification"),
        gr.Text(label="Confidence")
    ],
    title="Amharic & Afan Oromo Hate Speech Detector",
    description="Classify text into Free, Offensive, or Hate Speech",
    article="<p style='text-align: center;'>@2025 Mequanent Degu Belete </p><p style='text-align: center;'>mekuanentde@gmail.com</p><p style='text-align: center;'>SNHCC, Academia Sinica, Taiwan</p>",
    examples=[
        ["ኢትዮጵያ ለዘላለም ትኑር"],
        ["haatee sali shamtuu situ nuu beekaa waa ee baalee"]
    ]
)

# Launch (No 'share=True' needed on Hugging Face Spaces)
if __name__ == "__main__":
    interface.launch()