Update app.py

app.py

@@ -288,7 +288,7 @@ with st.spinner("Connecting DB..."):
 with st.spinner("Loading Models..."):
     # Initialize SAGE model
     if 'xq' not in st.session_state:
-        model, alphabet = init_esm()
+        st.session_state.model, st.session_state.alphabet = init_esm()
         batch_converter = alphabet.get_batch_converter()
         st.session_state['batch'] = batch_converter
         st.session_state.query_num = 0
@@ -383,7 +383,10 @@ else:
         st.title("Evolutionary Scale Modeling")
         start = [st.empty(), st.empty(), st.empty(), st.empty(), st.empty(), st.empty(), st.empty(), st.empty(), st.empty()]      
         start[0].info(msg)
-        option = st.selectbox('Application options', ('self-contact prediction', 'search the database', 'activity prediction','PDB viewer'))
+        option = st.selectbox('Application options', ('self-contact prediction',
+                         'search the database for similar proteins',
+                         'activity prediction with similar proteins',
+                         'PDB viewer'))
 
         st.session_state.db_name_ref = 'default.esm_protein'
         if option == 'self-contact prediction':
@@ -411,9 +414,7 @@ else:
             
             if sequence:
                 st.write('you have entered: ', sequence)
-                st.session_state['xq'] = model
-                batch_converter = alphabet.get_batch_converter()
-                result_temp_seq = esm_search(st.session_state['xq'], sequence, batch_converter ,top_k=10)
+                result_temp_seq = esm_search(st.session_state.model, sequence, st.session_state.batch_converter ,top_k=10)
                 st.text('search result (top 5): ')
                 # tab1, tab2, tab3, tab4, = st.tabs(["Cat", "Dog", "Owl"])
                 tab1, tab2, tab3 , tab4, tab5 = st.tabs(['1','2','3','4','5'])