Update app.py

app.py

@@ -327,7 +327,7 @@ if 'xq' not in st.session_state:
                 expander.text("""Contact prediction is based on a logistic regression over the model's attention maps. \
                 This methodology is based on ICLR 2021 paper, Transformer protein language models are unsupervised structure learners. 
                 (Rao et al. 2020) The MSA Transformer (ESM-MSA-1) takes a multiple sequence alignment (MSA) as input, and uses the tied row self-attention maps in the same way.""")
-            st.session_state['xq'] = model
+            st.session_state['xq'] = st.session_state.model
         elif option == function_list[1]:
             sequence = st.text_input('protein sequence', '')
             st.write('Try an example:')
@@ -353,7 +353,7 @@ if 'xq' not in st.session_state:
                     print(result_temp_seq[4])
 
                 start[2] = st.pyplot(visualize_3D_Coordinates(result_temp_coords).figure)
-            st.session_state['xq'] = model
+            st.session_state['xq'] = st.session_state.model
         elif option == function_list[2]:
             st.text('we predict the biological activity of mutations of a protein, using fixed embeddings from ESM.')
             sequence = st.text_input('protein sequence', '')
@@ -370,7 +370,7 @@ if 'xq' not in st.session_state:
                 start[3] = showmol(render_pdb_resn(viewer = render_pdb(id = id_PDB),resn_lst = [residues_marker]))
             else:
                 start[3] = showmol(render_pdb(id = id_PDB))
-            st.session_state['xq'] = model
+            st.session_state['xq'] = st.session_state.model
             
 else:
     if st.session_state.query_num < len(messages):