import gradio as gr
import os
import sys
import subprocess

# ---- SAFETY: ensure transformers is installed ----
try:
    from transformers import pipeline
except ModuleNotFoundError:
    subprocess.check_call(
        [sys.executable, "-m", "pip", "install", "transformers", "torch", "sentencepiece"]
    )
    from transformers import pipeline

# 1. Load the classifier from Hugging Face Hub
MODEL_REPO = "MakD1227/afriberta-hsd-model"

classifier = pipeline(
    task="text-classification",
    model=MODEL_REPO,
    tokenizer=MODEL_REPO
)

# 2. Prediction function
def predict_speech(text):
    if not text or not text.strip():
        return "Free (Neutral)", "0.00%"

    results = classifier(text)

    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"{score * 100:.2f}%"

# 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"]
    ]
)

# 4. Launch
if __name__ == "__main__":
    interface.launch()