import gradio as gr
from transformers import pipeline
import os
import nltk

# Check if the resource is available; if not, download it
nltk.download("averaged_perceptron_tagger")

# Your existing app code follows

token = os.getenv("HF_TOKEN")
get_completion = pipeline(
    "newsagency-ner",
    model="impresso-project/bert-newsagency-ner-fr",
    trust_remote_code=True,
    revision="main",
    token=token,
)

#
# PIPELINE_REGISTRY.register_pipeline(
#     "newsagency-ner",
#     pipeline_class=NewsAgencyModelPipeline,
#     pt_model=AutoModelForTokenClassification,
# )
# model.config.custom_pipelines = {
#     "newsagency-ner": {
#         "impl": "newsagency_ner.NewsAgencyModelPipeline",
#         "pt": ["BertForTokenClassification"],
#         "tf": [],
#     }
# }
#
# classifier = pipeline("newsagency-ner", model=model, tokenizer=tokenizer)


def ner(input):
    output = get_completion(input)
    return {"text": input, "entities": output}


demo = gr.Interface(
    fn=ner,
    inputs=[
        gr.Textbox(
            label="Find the news agency mentions in the following text in French:",
            lines=2,
        )
    ],
    outputs=[gr.HighlightedText(label="Text with news agency mentions")],
    title="News Agency Recognition with impresso-project/bert-newsagency-ner-fr",
    description="Find entities using the `impresso-project/bert-newsagency-ner-fr` model under the hood!",
    allow_flagging="never",
    # Here we introduce a new tag, examples, easy to use examples for your application
    examples=[
        "Des chercheurs de l'Université de Cambridge ont développé une nouvelle technique de calcul quantique qui "
        "promet d'augmenter exponentiellement les vitesses de calcul. Cette percée, décrite comme un 'bond quantique' "
        "dans la technologie informatique, pourrait ouvrir la voie à des capacités de traitement de données "
        "ultra-rapides et sécurisées. Le rapport complet sur ces découvertes a été publié dans la "
        "prestigieuse revue 'Nature Physics'. (Reuters)",
        "Les tensions continuent de monter alors que les récentes élections parlementaires en Suède ont conduit à un "
        "changement inattendu dans le paysage politique. Le parti d'extrême droite, connu pour ses politiques "
        "strictes en matière d'immigration et ses vues eurosceptiques, a gagné un nombre substantiel de sièges, "
        "remettant en question l'équilibre traditionnel des pouvoirs. Alors que les partis négocient pour former un nouveau gouvernement, "
        "la communauté internationale observe attentivement pour voir comment ces développements influenceront les "
        "politiques étrangères et domestiques de la Suède. (AFP)",
        "United Press - On the home front, the British populace remains steadfast in the face of ongoing air "
        "raids. In London, despite the destruction, the spirit of the people is unbroken, with volunteers and civil "
        "defense units working tirelessly to support the war effort. Reports from BUP correspondents highlight the "
        "nationwide push for increased production in factories, essential for supplying the front lines with the "
        "materials needed for victory.",
    ],
)
demo.launch()