from transformers import pipeline
import streamlit as st
import pandas as pd



def filter_candidates(candidates):
    df = pd.DataFrame(columns=["Candidates", "Probability"])
    cand_list = []
    score_list = []
    for candidate in candidates:
        if candidate["token_str"][:2] != "##":
            cand = candidate["sequence"]
            score = candidate["score"]
            cand_list.append(cand)
            score_list.append('{0:.5f}'.format(score))
        if len(score_list) == 5:
            break
    df["Candidates"] = cand_list
    df["Probability"] = score_list
    
    df.index = [1,2,3,4,5]

    return df

nlp = pipeline("fill-mask", model="flax-community/alberti-bert-base-multilingual-cased")

user_input = st.text_input("Mask token: [MASK]", "Me encanta escribir [MASK].")

if st.button("Guess!"):
    results = filter_candidates(nlp(user_input, top_k=20))
    st.table(results)