import streamlit as st

from transformers import pipeline

@st.cache(allow_output_mutation=True)
def get_model(model):
	return pipeline("fill-mask", model=model)


text = st.text_input("Enter a sentence. Use a * for a mask.")

model = st.selectbox("choose a model", ["roberta-base", "bert-base-uncased", "gpt2", "t5"])

# Create a text element and let the reader know the data is loading.

if text:
	data_load_state = st.text('Loading data...')
	nlp = get_model(model)
	result = nlp(text.replace("*", nlp.tokenizer.mask_token))
	data_load_state.text('Loading data...done!')
	for c in result:
		del c["token"]
	st.table(result)