import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

tokenizer = AutoTokenizer.from_pretrained("impresso-project/nel-hipe-multilingual")
model = AutoModelForSeq2SeqLM.from_pretrained(
    "impresso-project/nel-hipe-multilingual"
).eval()


def disambiguate_sentences(sentences):
    results = []
    for sentence in sentences:
        outputs = model.generate(
            **tokenizer([sentence], return_tensors="pt"),
            num_beams=5,
            num_return_sequences=5
        )
        decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True)
        results.append(decoded)
    return results


demo = gr.Interface(
    fn=disambiguate_sentences,
    inputs=[
        gr.Textbox(
            label="Input Sentences:",
            lines=5,
            placeholder="Enter your sentence here in the following format: \\  `It is reported in [START] Paris [END], "
            "that the opening of the chambers will take place on the 27th January.' \\ "
            "This format ensures that the model knows which entities to disambiguate, more exactly the entity should "
            "be surrounded by `[START]` and `[END]`.",
        )
    ],
    outputs=[gr.outputs.Textbox(label="Predictions")],
    title="Entity Linking with impresso-project/nel-hipe-multilingual",
    description="Link entities using the `impresso-project/nel-hipe-multilingual` 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' [START] Université de Cambridge [END] 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'. ([START] Reuters [END])",
        "In the [START] year 1789 [END], [START] King Louis XVI, ruler of France [END], convened the Estates-General at the [START] Palace of Versailles [END], "
        "where Marie Antoinette, the Queen of France, alongside Maximilien Robespierre, a leading member of the National Assembly, "
        "debated with [START] Jean-Jacques Rousseau, the famous philosopher [END], and [START] Charles de Talleyrand, the Bishop of Autun [END], "
        "regarding the future of the French monarchy. At the same time, across the Atlantic in Philadelphia, "
        "[START] George Washington, the first President of the United States [END], and [START] Thomas Jefferson, the nation's Secretary of State [END], "
        "were drafting policies for the newly established American government following the signing of the Constitution.",
    ],
)
demo.launch()