mbahrami/Auto-Complete_Semantic

1import streamlit as st2import pandas as pd3from streamlit import cli as stcli4from transformers import pipeline5from sentence_transformers import SentenceTransformer, util6import sys 7 8HISTORY_WEIGHT = 100 # set history weight (if found any keyword from history, it will priorities based on its weight)9 10@st.cache(allow_output_mutation=True, suppress_st_warning=True)11def get_model(model):12	return pipeline("fill-mask", model=model, top_k=10)#set the maximum of tokens to be retrieved after each inference to model13 14def hash_func(inp):15    return True16 17@st.cache(allow_output_mutation=True, suppress_st_warning=True)18def loading_models(model='roberta-base'):19     return get_model(model), SentenceTransformer('all-MiniLM-L6-v2')20 21@st.cache(allow_output_mutation=True, 22          suppress_st_warning=True,23          hash_funcs={'tokenizers.Tokenizer': hash_func, 'tokenizers.AddedToken': hash_func})24def infer(text):25#    global nlp 26    return nlp(text+' '+nlp.tokenizer.mask_token)27 28 29@st.cache(allow_output_mutation=True, 30          suppress_st_warning=True,31          hash_funcs={'tokenizers.Tokenizer': hash_func, 'tokenizers.AddedToken': hash_func})32def sim(predicted_seq, sem_list):33    return semantic_model.encode(predicted_seq, convert_to_tensor=True), \34            semantic_model.encode(sem_list, convert_to_tensor=True)35    36@st.cache(allow_output_mutation=True, 37          suppress_st_warning=True,38          hash_funcs={'tokenizers.Tokenizer': hash_func, 'tokenizers.AddedToken': hash_func})39def main(text,semantic_text,history_keyword_text):40    global semantic_model, data_load_state41    data_load_state.text('Inference from model...')42    result = infer(text)43    sem_list=[semantic_text.strip()]44    data_load_state.text('Checking similarity...')45    if len(semantic_text):46        predicted_seq=[rec['sequence'] for rec in result]47        predicted_embeddings, semantic_history_embeddings = sim(predicted_seq, sem_list)48        cosine_scores = util.cos_sim(predicted_embeddings, semantic_history_embeddings)49    data_load_state.text('similarity check completed...')50    51    for index, r in enumerate(result):52        if len(semantic_text):53                if len(r['token_str'])>2: #skip spcial chars such as "?"54                    result[index]['score']+=float(sum(cosine_scores[index]))*HISTORY_WEIGHT55        if r['token_str'].lower().strip() in history_keyword_text.lower().strip() and len(r['token_str'].lower().strip())>1:56            #found from history, then increase the score of tokens57            result[index]['score']*=HISTORY_WEIGHT58    data_load_state.text('Score updated...')59            60    #sort the results        61    df=pd.DataFrame(result).sort_values(by='score', ascending=False)62    return df63    64    65if __name__ == '__main__':66    if st._is_running_with_streamlit:67        st.markdown("""68# Auto-Complete69This is an example of an auto-complete approach where the next token suggested based on users's history Keyword match & Semantic similarity of users's history (log).70The next token is predicted per probability and a weight if it is appeared in keyword user's history or there is a similarity to semantic user's history71""")72        history_keyword_text = st.text_input("Enter users's history <Keywords Match> (optional, i.e., 'Gates')", value="")73        74        semantic_text = st.text_input("Enter users's history <Semantic> (optional, i.e., 'Microsoft' or 'President')", value="Microsoft")75        76        text = st.text_input("Enter a text for auto completion...", value='Where is Bill')77        model = st.selectbox("Choose a model", ["roberta-base", "bert-base-uncased"])78        79        data_load_state = st.text('1.Loading model ...')80 81        nlp, semantic_model = loading_models(model)82        83        df=main(text,semantic_text,history_keyword_text)84        #show the results as a table85        st.table(df)86        data_load_state.text('')87    else:88        sys.argv = ['streamlit', 'run', sys.argv[0]]89        sys.exit(stcli.main())