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

model = GLiNER.from_pretrained("knowledgator/gliner-multitask-large-v0.5").to('cpu')

text1 = """
"I recently purchased the Sony WH-1000XM4 Wireless Noise-Canceling Headphones from Amazon and I must say, I'm thoroughly impressed. The package arrived in New York within 2 days, thanks to Amazon Prime's expedited shipping.

The headphones themselves are remarkable. The noise-canceling feature works like a charm in the bustling city environment, and the 30-hour battery life means I don't have to charge them every day. Connecting them to my Samsung Galaxy S21 was a breeze, and the sound quality is second to none.

I also appreciated the customer service from Amazon when I had a question about the warranty. They responded within an hour and provided all the information I needed.

However, the headphones did not come with a hard case, which was listed in the product description. I contacted Amazon, and they offered a 10% discount on my next purchase as an apology.

Overall, I'd give these headphones a 4.5/5 rating and highly recommend them to anyone looking for top-notch quality in both product and service."""


text3 = """
Several studies have reported its pharmacological activities, including anti-inflammatory, antimicrobial, and antitumoral effects.
The effect of E-anethole was studied in the osteosarcoma MG-63 cell line, and the antiproliferative activity was evaluated by an MTT assay.
It showed a GI50 value of 60.25 μM with apoptosis induction through the mitochondrial-mediated pathway. Additionally, it induced cell cycle arrest at the G0/G1 phase, up-regulated the expression of p53, caspase-3, and caspase-9, and down-regulated Bcl-xL expression.
Moreover, the antitumoral activity of anethole was assessed against oral tumor Ca9-22 cells, and the cytotoxic effects were evaluated by MTT and LDH assays.
It demonstrated a LD50 value of 8 μM, and cellular proliferation was 42.7% and 5.2% at anethole concentrations of 3 μM and 30 μM, respectively.
It was reported that it could selectively and in a dose-dependent manner decrease cell proliferation and induce apoptosis, as well as induce autophagy, decrease ROS production, and increase glutathione activity. The cytotoxic effect was mediated through NF-kB, MAP kinases, Wnt, caspase-3 and -9, and PARP1 pathways. Additionally, treatment with anethole inhibited cyclin D1 oncogene expression, increased cyclin-dependent kinase inhibitor p21WAF1, up-regulated p53 expression, and inhibited the EMT markers.
"""

text5 = """
Dr. Paul Hammond, a renowned neurologist at Johns Hopkins University, has recently published a paper in the prestigious journal "Nature Neuroscience". His research focuses on a rare genetic mutation, found in less than 0.01% of the population, that appears to prevent the development of Alzheimer's disease. Collaborating with researchers at the University of California, San Francisco, the team is now working to understand the mechanism by which this mutation confers its protective effect. Funded by the National Institutes of Health, their research could potentially open new avenues for Alzheimer's treatment.
"""

ner_examples = [
    [
        text5,
        "neurologist, scientist, gene, disease, biological process, city, journal, university",
        0.5,
        False
    ],
    [
        text1,
        "product, brand, location, features, rating",
        0.5,
        False
    ],
    [
        text3,
        "cell line, protein, metric, substance",
        0.5,
        False
    ]]

def merge_entities(entities):
    if not entities:
        return []
    merged = []
    current = entities[0]
    for next_entity in entities[1:]:
        if next_entity['entity'] == current['entity'] and (next_entity['start'] == current['end'] + 1 or next_entity['start'] == current['end']):
            current['word'] += ' ' + next_entity['word']
            current['end'] = next_entity['end']
        else:
            merged.append(current)
            current = next_entity
    merged.append(current)
    return merged

def process(
    text, labels: str, threshold: float, nested_ner: bool
) -> Dict[str, Union[str, int, float]]:
    labels = [label.strip() for label in labels.split(",")]
    r = {
        "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
            )
        ],
    }
    r["entities"] =  merge_entities(r["entities"])
    return r

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

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


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