from transformers import pipeline
import gradio as gr
import base64

model_checkpoint2 = "dccuchile/bert-base-spanish-wwm-cased"

mask_filler = pipeline(
    "fill-mask", model=model_checkpoint2,tokenizer=model_checkpoint2
)

def fill_mask_interface(sentence):
    results = mask_filler(sentence)
    #suggestions = [f"{result['token_str']} (confidence: {result['score']:.4f})" for result in results]
    dictt ={}
    for text,score in zip([d['token_str'] for d in results],[d['score'] for d in results]):
      dictt[text] = score
    return dictt


with open("Iso_Logotipo_Ceibal.png", "rb") as image_file:
    encoded_image = base64.b64encode(image_file.read()).decode()

description = """
<p>
<center>
<img src='data:image/jpg;base64,{}' width=200px>
Interactúa con este espacio para probar la predicción de palabras en una frase con una palabra enmascarada. La palabra enmascarada debe ser [MASK] y sólo  una por frase.
</center>
</p>
""".format(encoded_image)


textbox = gr.Textbox(label="Agrega tu frase con una palabra enmascarada aquí!", placeholder="Hola, [MASK] estás?", lines=2)

gr.Interface(fn=fill_mask_interface, 
             inputs=textbox, 
             outputs="label",
             title = "Uso de AI para la predicción de palabras enmascaradas.",
             description = description,
             examples=[["Hola, cómo te lleva el [MASK]?"], ["Dónde deberíamos [MASK]?"]]
             ).launch()