import gradio as gr
from transformers import pipeline

pipe = pipeline("token-classification", model="cogniveon/nlpcw_bert-base-uncased-abbr", grouped_entities=True)

def predict(input) -> list[tuple[str, str | float | None]] | dict | None:
    output = pipe(input)
    return {"text": input, "entities": output}  


demo = gr.Interface(
    predict,
    gr.Textbox(
        label="Input",
        lines=3,
    ),
    gr.HighlightedText(
        label="Output",
        combine_adjacent=True,
        show_legend=True
    ),
    examples=[
        ["We developed a variant of gene set enrichment analysis (GSEA) to determine whether a genetic pathway shows evidence for age regulation [23]."],
    ],
    allow_flagging="manual",
    flagging_options=["Correct", "Incorrect", "Ambiguous"],
    flagging_callback=gr.CSVLogger(),
).launch()