mbahrami's picture
Update app.py
40c9d2b
raw
history blame
1.12 kB
import streamlit as st
import pandas as pd
from transformers import pipeline
@st.cache(allow_output_mutation=True)
def get_model(model):
return pipeline("fill-mask", model=model, top_k=100)
HISTORY_WEIGHT = 100 # set history weight (if found any keyword from history, it will priorities based on its weight)
history_keyword_text = st.text_input("Enter users's history keywords (optional, i.e., 'Gates')", value="Gates")
text = st.text_input("Enter a text for auto completion...", value='Where is Bill')
model = st.selectbox("choose a model", ["roberta-base", "bert-base-uncased", "gpt2", "t5"])
data_load_state = st.text('Loading model...')
nlp = get_model(model)
if text:
data_load_state = st.text('Inference to model...')
result = nlp(text+' '+nlp.tokenizer.mask_token)
data_load_state.text('')
for index, r in enumerate(result):
if r['token_str'].lower().strip() in history_keyword_text.lower().strip():
result[index]['score']*=HISTORY_WEIGHT
df=pd.DataFrame(result)
#result={k: v for k, v in sorted(result.items(), key=lambda item: item[0])}
st.table(df)