from typing import Dict, Union
import sys
sys.path.extend(["./GLiNER"])
from GLiNER.model import GLiNER
import gradio as gr
import os
from transformers import AutoModel, AutoTokenizer, logging
from huggingface_hub import HfFolder

# Set the Hugging Face logging to error only to avoid sensitive info in logs
logging.set_verbosity_error()

# Retrieve the HF API key from the environment
hf_api_key = os.getenv('HF_KEY')

# If the key is found, use it to authenticate
if hf_api_key:
    HfFolder.save_token(hf_api_key)  # This authenticates you for this session
else:
    print("No HF_KEY found. Please make sure you've set up your Hugging Face API key as an environment variable.")


model = GLiNER.from_pretrained("DeepMount00/universal_ner_ita").eval()


examples = [
    [
        "La ditta Giardini Belli S.p.A., con partita IVA 01234567890, ha emesso una fattura di 500€ per la consulenza giuridica fornita dall'avvocato Giulia Bianchi, il cui studio si trova in Piazza del Duomo, Milano, dal giorno 01/04/2024.",
        "ragione sociale, partita iva, importo, avvocato, indirizzo, data",
        True,
    ],
    [
        "Il ristorante 'Sapori d'Arte', identificato dalla partita IVA 09876543210, ha commissionato la realizzazione di un nuovo sito web alla ditta Web Innovations S.r.l., per un valore di 2.500€, con conclusione dei lavori prevista per il 15/06/2024, presso la loro sede in Via Roma 123, Torino.",
        "ristorante, partita iva, importo, ditta, indirizzo, data di conclusione",
        True,
    ],
    [
        "Il comune di Castelrosso, con codice fiscale 80012345678, ha approvato il finanziamento di 15.000€ destinati alla ristrutturazione del parco giochi cittadino, affidando l'incarico alla società 'Verde Vivo Società Cooperativa', con sede legale in Corso della Libertà 45, Verona, da completarsi entro il 30/09/2024.",
        "comune, codice fiscale, importo, società, indirizzo, data di completamento",
        True,
    ],
    [
        "Durante la visita del 12/03/2024, il dottor Marco Verdi ha diagnosticato al paziente Luca Rossi, nato il 05/05/1985 a Roma e residente in Via della Salute 10, Milano, una bronchite asmatica, prescrivendo un trattamento di 10 giorni con Ventolin 2 mg.",
        "data visita, medico, diagnosi, paziente, data di nascita, luogo di nascita, indirizzo, prescrizione, durata trattamento, farmaco",
        True,
    ],
    [
        "Il notaio Anna Bianchi, con studio in Corso Impero 45, Torino, ha autenticato il testamento di Giovanni Neri, residente in Via Roma 123, Firenze, il 15/07/2024, assegnando come eredi i suoi due figli, Maria e Luca Neri, e stabilendo una donazione di 50.000€ all'associazione 'Aiuti Umanitari'.",
        "notaio, indirizzo studio, atto, data, testatore, indirizzo testatore, eredi, donazione, beneficiario",
        True,
    ],
    [
        "La Banca del Futuro S.p.A. ha approvato il 20/04/2024 un prestito ipotecario di 250.000€ a favore di Chiara Verdi, impiegata presso la società 'Tecnologie Innovative S.r.l.', con rimborso in 20 anni, garantito sull'immobile sito in Via Lago di Garda 123, Verona.",
        "banca, data approvazione, tipo di prestito, importo, beneficiario, datore di lavoro, durata rimborso, garanzia, indirizzo immobile",
        True,
    ],
    [
        "Il 10/06/2024, il giornalista Marco Bianchi, collaboratore del quotidiano 'La Voce del Popolo', ha pubblicato un articolo sulla crescita economica della regione Lombardia, evidenziando l'aumento del PIL del 3,5% nell'ultimo anno e intervistando l'assessore all'economia, Laura Rossi.",
        "data pubblicazione, giornalista, giornale, tema articolo, dato economico, tasso di crescita, intervistato, carica",
        True,
    ],
    [
        "Il 25/08/2024, l'Università degli Studi di Arcadia ha inaugurato il nuovo dipartimento di Robotica Applicata, diretto dal Prof. Andrea Verdi, con un investimento iniziale di 5 milioni di euro, raccolti grazie a contributi statali e donazioni private, tra cui spicca quella dell'imprenditore tech Gianluca Bianchi.",
        "data inaugurazione, istituzione, nuovo dipartimento, direttore, investimento iniziale, fonti finanziamento, donatore principale",
        True,
    ],
    [
        "L'avvocato difensore Maria Rossi, del Foro di Napoli, ha ottenuto il 30/09/2024 un'importante vittoria legale per il suo assistito, Franco Verdi, accusato ingiustamente di frode fiscale, dimostrando la sua totale estraneità ai fatti grazie all'analisi delle transazioni bancarie effettuate nel periodo contestato.",
        "data, avvocato, foro di appartenenza, esito causa, assistito, accusa, difesa, periodo contestato",
        True,
    ],
    [
        "Il 15/07/2024, l'Ospedale San Luigi di Torino ha introdotto un innovativo trattamento per la cura della sclerosi multipla, sviluppato dal team di ricerca guidato dal Dr. Elena Bianchi. Il trattamento, che prevede l'uso di una nuova molecola denominata SM-204, ha mostrato una riduzione dei sintomi nel 70% dei pazienti partecipanti alla sperimentazione clinica, finanziata con fondi dell'Unione Europea e donazioni della Fondazione Verdi per la Ricerca Scientifica.",
        "data introduzione, struttura ospedaliera, trattamento, malattia, responsabile ricerca, nuova molecola, efficacia, finanziamento, percentale efficacia",
        True,
    ],
]

def ner(text, labels: str, 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)
        ],
    }

with gr.Blocks(title="GLiNER-base") as demo:
    gr.Markdown(
        """
        # Universal Ner: A GliNer-Based Named Entity Recognition Model
        
        **Universal Ner** is a state-of-the-art Named Entity Recognition (NER) model that leverages the GliNer architecture. This model stands out by its capability to identify any entity type utilizing a bidirectional transformer encoder. Unlike traditional NER models constrained to predefined entity types, **Universal Ner** offers a versatile and practical solution. Moreover, it presents an efficient alternative to Large Language Models (LLMs) which, despite their adaptability, may not be suitable for resource-constrained environments due to their substantial size and associated costs. Currently, the model runs on CPU environments.
        
        ### Performance
        
        - **Inference Time on GPU**: Approximately 0.01 seconds. This showcases the model's efficiency and suitability for real-time applications, even when deployed on GPU resources.
        
        ### Useful Links
        - **Research Paper**: [Exploring Universal Ner](https://arxiv.org/abs/2311.08526)
        - **Model Page**: [universal_ner_ita](https://huggingface.co/DeepMount00/universal_ner_ita) on Hugging Face. This specific implementation is designed for Italian language named entity recognition tasks, making use of the advanced capabilities of the Universal Ner architecture.
        """
    )

    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=1,
        )
    output = gr.HighlightedText(label="Predicted Entities")
    submit_btn = gr.Button("Submit")
    examples = gr.Examples(
        examples,
        fn=ner,
        inputs=[input_text, labels],
        outputs=output,
        cache_examples=True,
    )

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


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