import gradio as gr
from transformers import pipeline


ner = pipeline("token-classification", model="chuuhtetnaing/myanmar-ner-model", grouped_entities=True)


def tag_text(text):
    """Run NER tagging on input text."""
    ner_results = ner(text)
    highlighted_text = []
    last_idx = 0
    for entity in ner_results:
        # Skip empty words or overlapping entities (model artifacts)
        if not entity["word"].strip():
            continue

        start = entity["start"]
        end = entity["end"]
        label = entity["entity_group"]
        if start > last_idx:
            highlighted_text.append((text[last_idx:start], None))
        highlighted_text.append((text[start:end], label))
        last_idx = end
    if last_idx < len(text):
        highlighted_text.append((text[last_idx:], None))
    return highlighted_text


css = """
#col-container {
    margin: 0 auto;
    max-width: 1200px;
    padding: 1rem;
    gap: 2.5rem;
}
#col-container h1 {
    text-align: center;
    margin-bottom: 1rem;
}
#input-output-row {
    flex-direction: row;
    gap: 1rem;
}
#input-output-row label {
    text-align: center;
}
@media (max-width: 768px) {
    #input-output-row {
        flex-direction: column;
    }
}
"""

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# Myanmar Named Entity Recognition")

        with gr.Row(elem_id="input-output-row", equal_height=True):
            text_input = gr.Textbox(label="Input Text", placeholder="Enter Myanmar text here...", lines=8)
            output = gr.HighlightedText(label="NER Result", combine_adjacent=True)

        run_button = gr.Button("Recognize Entities", variant="primary")

        gr.Examples(
            examples=[
                ["ကုလသမဂ္ဂသည်နယူးယောက်မြို့တွင်အခြေစိုက်သည်။"],
                ["ကိုမောင်သည်ရန်ကုန်မြို့သို့သွားသည်။"],
                ["ဘတ်စ်ကားပေါ်မှာခရီးသည်၄၈ယောက်လောက်ထိုင်နေကြတယ်။"]
            ],
            inputs=text_input,
        )

        with gr.Accordion("NER Tag Reference", open=False):
            gr.Markdown("""
| Tag | Description |
|-----|-------------|
| DATE | Date |
| LOC | Location |
| NUM | Number |
| ORG | Organization |
| PER | Person |
| TIME | Time |
""")

    run_button.click(fn=tag_text, inputs=text_input, outputs=output)
    text_input.submit(fn=tag_text, inputs=text_input, outputs=output)

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