app.py

import os
import gradio as gr
from transformers import AutoTokenizer, TFAutoModelForSequenceClassification, TextClassificationPipeline

# -----------------------
# Config
# -----------------------
# Point this to your HF model repo in the Space settings (Variables → MODEL_ID).
MODEL_ID = os.getenv("MODEL_ID", "will-rads/distilbert-hatespeech-classifier")

# Expected label names (must match your model's id2label/label2id)
LABELS = ["hate speech", "offensive language", "neither"]  # fixed order for display
TOP_K = 3  # show all 3

# -----------------------
# Load model (TensorFlow)
# -----------------------
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = TFAutoModelForSequenceClassification.from_pretrained(MODEL_ID)
pipe = TextClassificationPipeline(
    model=model,
    tokenizer=tokenizer,
    framework="tf",
    return_all_scores=True
)

def format_scores(scores_list):
    # scores_list: [{'label': 'neither','score':0.12}, ...]
    score_map = {d["label"]: float(d["score"]) for d in scores_list}
    # Ensure every expected label is present (fallback to 0.0 if missing)
    ordered = [(lbl, score_map.get(lbl, 0.0)) for lbl in LABELS]
    # For Gradio Label top-k
    top_display = {lbl: s for lbl, s in ordered[:TOP_K]}
    # Detailed pretty JSON
    details = {lbl: round(s, 4) for lbl, s in ordered}
    return top_display, details

def moderate(text: str):
    text = (text or "").strip()
    if not text:
        return {"error": 1.0}, {"error": "Please enter some text."}
    try:
        scores = pipe(text)[0]  # list of dicts
        return format_scores(scores)
    except Exception as e:
        return {"error": 1.0}, {"exception": str(e)}

with gr.Blocks(title="Ethical Content Moderator") as demo:
    gr.Markdown(
        "# Ethical Content Moderator\n"
        "Paste text and get probabilities for **hate speech**, **offensive language**, or **neither**."
    )
    inp = gr.Textbox(label="Input Text", placeholder="Type or paste a sentence…", lines=5)
    btn = gr.Button("Analyze", variant="primary")
    topk = gr.Label(num_top_classes=TOP_K, label="Predictions (Top)")
    details = gr.JSON(label="Scores (All Labels)")
    btn.click(fn=moderate, inputs=inp, outputs=[topk, details])

    gr.Examples(
        examples=[
            ["I hate you and your stupid project."],
            ["Thanks for your help today, really appreciate it!"],
            ["Clouds exist. That's it."]
        ],
        inputs=inp,
        label="Examples"
    )

if __name__ == "__main__":
    demo.launch()