File size: 1,118 Bytes
da676c8
40c9d2b
da676c8
 
 
 
 
70c19cc
da676c8
f089045
da676c8
 
6dd0ae0
005c6a4
6dd0ae0
005c6a4
da676c8
 
 
1ef9e65
 
 
da676c8
1ef9e65
6dd0ae0
 
 
ac5b8a7
fa97b6f
f089045
40c9d2b
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
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)