import re
import gradio as gr

# Simulated GLiNER‐style output
SIM_ENTITIES = {
    "Household Survey Data": "named dataset",
    "Ethiopia":             "data geography",
    "2020":                 "reference year",
    "World Bank":           "publisher"
}

def annotate_text(text):
    entities = []
    for phrase, label in SIM_ENTITIES.items():
        for m in re.finditer(re.escape(phrase), text):
            entities.append({
                "entity": label,    # note: 'entity' key
                "start":  m.start(),
                "end":    m.end(),
            })
    return {"text": text, "entities": entities}

with gr.Blocks() as demo:
    gr.Markdown("## Simulated GLiNER Highlighter\n"
                "Type or paste a sentence and click **Highlight**.")

    txt = gr.Textbox(
        lines=3,
        value="The Household Survey Data for Ethiopia in 2020 was published by the World Bank.",
        label="Input Text",
        placeholder="Enter any sentence here…"
    )
    btn = gr.Button("Highlight")
    out = gr.HighlightedText(label="Annotated Entities")

    # fire on button click or Enter
    btn.click(fn=annotate_text, inputs=txt, outputs=out)
    txt.submit(fn=annotate_text, inputs=txt, outputs=out)
    # initial render
    demo.load(fn=annotate_text, inputs=txt, outputs=out)

    gr.Markdown("""
**Legend**  
- **named dataset** → Household Survey Data  
- **data geography** → Ethiopia  
- **reference year** → 2020  
- **publisher** → World Bank
""")

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