alberti / app.py
alvp's picture
Initial commit
92627cb
raw
history blame
909 Bytes
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)