masked-probs / app.py
Rik Bose
oops
a5ff044
raw
history blame contribute delete
626 Bytes
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)