File size: 675 Bytes
da676c8
 
 
 
 
 
 
 
 
 
 
d37b190
da676c8
 
 
 
d37b190
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
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 text for auto completion...")

history_keyword_text = st.text_input("Enter users's history keywords")

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

if text:
    data_load_state = st.text('Loading...')
    nlp = get_model(model)
    result = nlp(text+nlp.tokenizer.mask_token)
    data_load_state.text('Loading data...done!')
    for c in result:
        if c in history_keyword_text:
            c['score']*=0.10
    st.table(result)