Update app.py

app.py

@@ -14,17 +14,29 @@ def get_model(model):
 @st.cache(allow_output_mutation=True)
 def loading_models(model='roberta-base'):
      return get_model(model), SentenceTransformer('all-MiniLM-L6-v2')
-     
-#@st.cache(allow_output_mutation=True)
-def main(nlp, semantic_model, data_load_state):
+
+def infer(text):
+    global nlp
+    return nlp(text+' '+nlp.tokenizer.mask_token)
+
+def sim(predicted_seq, sem_list):
+    return semantic_model.encode(predicted_seq, convert_to_tensor=True), \
+            semantic_model.encode(sem_list, convert_to_tensor=True)
+    
+def hash_func(inp):
+    #bypass hash function
+    return True
+    
+@st.cache(allow_output_mutation=True, hash_funcs={'tokenizers.Tokenizer': hash_func, 'tokenizers.AddedToken': hash_func})
+def main(text,semantic_text,history_keyword_text):
+    global semantic_model, data_load_state
     data_load_state.text('Inference from model...')
-    result = nlp(text+' '+nlp.tokenizer.mask_token)
+    result = infer(text)
     sem_list=[semantic_text.strip()]
     data_load_state.text('Checking similarity...')
     if len(semantic_text):
         predicted_seq=[rec['sequence'] for rec in result]
-        predicted_embeddings = semantic_model.encode(predicted_seq, convert_to_tensor=True)
-        semantic_history_embeddings = semantic_model.encode(sem_list, convert_to_tensor=True)
+        predicted_embeddings, semantic_history_embeddings = sim(predicted_seq, sem_list)
         cosine_scores = util.cos_sim(predicted_embeddings, semantic_history_embeddings)
     data_load_state.text('similarity check completed...')
     
@@ -39,11 +51,7 @@ def main(nlp, semantic_model, data_load_state):
             
     #sort the results        
     df=pd.DataFrame(result).sort_values(by='score', ascending=False)
-    
-#    show the results as a table
-    st.table(df)
-#    print(df)
-    data_load_state.text('')
+    return df
     
     
 if __name__ == '__main__':
@@ -54,21 +62,20 @@ This is an example of an auto-complete approach where the next token suggested b
 The 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 history
 """)
         history_keyword_text = st.text_input("Enter users's history <Keywords Match> (optional, i.e., 'Gates')", value="")
+        
         semantic_text = st.text_input("Enter users's history <Semantic> (optional, i.e., 'Microsoft' or 'President')", value="Microsoft")
         
         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"])
         
         data_load_state = st.text('1.Loading model ...')
 
-#        semantic_model = SentenceTransformer('all-MiniLM-L6-v2')
-#        nlp = get_model(model)
         nlp, semantic_model = loading_models(model)
         
-        main(nlp, semantic_model, data_load_state)
+        df=main(text,semantic_text,history_keyword_text)
+        #show the results as a table
+        st.table(df)
+        data_load_state.text('')
     else:
         sys.argv = ['streamlit', 'run', sys.argv[0]]
         sys.exit(stcli.main())