from typing import Dict, Union
from gliner import GLiNER
import gradio as gr

model = GLiNER.from_pretrained("BSC-NLP4BIA/chagas-ner",
                               revision="NuNER_Zero_2025-02-16_22-10-21",
                               load_tokenizer=True
                              )

examples = [
    [
        """Sexo: Femenino
País de Nacimiento: Paraguay
Fecha de Nacimiento: XXXX
Resultado Chagas: Positivo
Informe: Seguimiento de consultas externas
Visita: Primera visita
Fecha: 20 de julio XXXX
Paciente paraguaya de 32 años diagnosticada con enfermedad de Chagas tras cribado comunitario. Sin antecedentes transfusionales ni de familiares afectados. Refiere buen estado general.
EF:
Buen estado general.
AC: rítmico, sin soplos.
AP: MVC sin hallazgos.
Abdomen blando, sin masas.
PLAN:
Solicitar estudio inicial de Chagas con ECG, ETT y serología.
Informe: Seguimiento de consultas externas
Visita: Seguimiento
Fecha: 10 de diciembre XXXX
Asintomática. PCR para T.cruzi negativa. ECG muestra ritmo sinusal con bloqueo de rama derecha. ETT sin alteraciones funcionales.
EF:
Sin cambios en exploración física.
PLAN:
Seguimiento anual con ECG y serología.""",
        "ENFERMEDAD, SINTOMA",
        0.5,
        False,
    ],
]


def ner(
        text, labels: str, threshold: float, nested_ner: bool
) -> Dict[str, Union[str, int, float]]:
    labels = labels.split(",")
    return {
        "text": text,
        "entities": [
            {
                "entity": entity["label"],
                "word": entity["text"],
                "start": entity["start"],
                "end": entity["end"],
                "score": 0,
            }
            for entity in model.predict_entities(
                text, labels, flat_ner=not nested_ner, threshold=threshold
            )
        ],
    }


with gr.Blocks(title="GLiNER for Chagas detection - NER ") as demo:

    input_text = gr.Textbox(
        value=examples[0][0], label="Text input", placeholder="Enter your text here"
    )
    with gr.Row() as row:
        labels = gr.Textbox(
            value=examples[0][1],
            label="Labels",
            placeholder="Enter your labels here (comma separated)",
            scale=2,
        )
        threshold = gr.Slider(
            0,
            1,
            value=0.5,
            step=0.01,
            label="Threshold",
            info="Lower the threshold to increase how many entities get predicted.",
            scale=1,
        )
        nested_ner = gr.Checkbox(
            value=examples[0][2],
            label="Nested NER",
            info="Allow for nested NER?",
            scale=0,
        )
    output = gr.HighlightedText(label="Predicted Entities")
    submit_btn = gr.Button("Submit")
    examples = gr.Examples(
        examples,
        fn=ner,
        inputs=[input_text, labels, threshold, nested_ner],
        outputs=output,
        cache_examples=True,
    )

    # Submitting
    input_text.submit(
        fn=ner, inputs=[input_text, labels, threshold, nested_ner], outputs=output
    )
    labels.submit(
        fn=ner, inputs=[input_text, labels, threshold, nested_ner], outputs=output
    )
    threshold.release(
        fn=ner, inputs=[input_text, labels, threshold, nested_ner], outputs=output
    )
    submit_btn.click(
        fn=ner, inputs=[input_text, labels, threshold, nested_ner], outputs=output
    )
    nested_ner.change(
        fn=ner, inputs=[input_text, labels, threshold, nested_ner], outputs=output
    )

demo.queue()
demo.launch(debug=True)