import re
import gradio as gr

# A simulated GLiNER output mapping for demonstration:
SIMULATED_ENTITIES = {
    "Household Survey Data": "named dataset",
    "Ethiopia": "data geography",
    "2020": "reference year",
    "World Bank": "publisher"
}

def highlight_text(text):
    """
    Simulate GLiNER entity extraction by looking for a few known phrases
    and returning start/end spans + labels.
    """
    entities = []
    for phrase, label in SIMULATED_ENTITIES.items():
        for match in re.finditer(re.escape(phrase), text):
            entities.append({
                "start": match.start(),
                "end": match.end(),
                "label": label
            })
    # Return in the format gr.HighlightedText expects:
    return {"text": text, "entities": entities}

# Build the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## Simulated GLiNER Entity Highlighter\n"
                "Enter any sentence (or use the default) and see simulated entity highlights.")
    text_input = gr.Textbox(
        label="Input Text",
        lines=3,
        value="The Household Survey Data for Ethiopia in 2020 was published by the World Bank."
    )
    highlighted = gr.HighlightedText(label="Highlighted Entities")
    # Trigger highlighting whenever the text changes
    text_input.change(fn=highlight_text, inputs=text_input, outputs=highlighted)
    
    gr.Markdown("""
**Entity legend**  
- **named dataset**: Household Survey Data  
- **data geography**: Ethiopia  
- **reference year**: 2020  
- **publisher**: World Bank
""")

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