import gradio as gr
from gliner import GLiNER

MODEL_NAME = "cutaa/gliner-au-pii-v1"
LABELS = ["AU_ORGANISATION", "AU_GOV_AGENCY", "AU_LOCATION"]

model = GLiNER.from_pretrained(MODEL_NAME)


def detect_entities(text: str, threshold: float):
    if not text or not text.strip():
        return "Please enter some text to analyze."

    entities = model.predict_entities(text, LABELS, threshold=threshold)

    if not entities:
        return "No entities found."

    lines = []
    for entity in entities:
        lines.append(
            f"[{entity['label']}] '{entity['text']}' ({entity['score']:.2f})"
        )

    return "\n".join(lines)


with gr.Blocks(title="GLiNER AU PII Detector") as demo:
    gr.Markdown("# GLiNER AU Entity Detector")
    gr.Markdown(
        f"Model: `{MODEL_NAME}`\n\n"
        f"Labels: `{', '.join(LABELS)}`"
    )

    text_input = gr.Textbox(
        label="Text",
        placeholder="Enter text to detect AU entities...",
        lines=6,
    )
    threshold_input = gr.Slider(
        minimum=0.0,
        maximum=1.0,
        value=0.5,
        step=0.01,
        label="Threshold",
    )
    output = gr.Textbox(label="Detected Entities", lines=10)

    run_button = gr.Button("Detect")
    run_button.click(
        fn=detect_entities,
        inputs=[text_input, threshold_input],
        outputs=output,
    )


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