bug in retrieval

app.py

@@ -70,7 +70,6 @@ def predict_dti():
             selected_encoder = st.selectbox(
                 'Select encoder for drug compound',('None', 'CDDD', 'MolBERT')
             )
-            st.image('molecule_encoder.png')
             if smiles: 
                 if selected_encoder == 'CDDD':
                     from cddd.inference import InferenceModel
@@ -92,11 +91,13 @@ def predict_dti():
                     molbert_model = MolBertFeaturizer(checkpoint_path, max_seq_len=500, embedding_type='average-1-cat-pooled')
                     embedding = molbert_model.transform([smiles])
                 else: 
-                    st.write('No pre-trained version of HyperPCM is available for the chosen encoder.')
+                    #st.write('No pre-trained version of HyperPCM is available for the chosen encoder.')
                     embedding = None
+                    st.image('molecule_encoder.png')
                 if embedding is not None:
-                    st.write(f'{selected_encoder} embedding')
-                    st.write(embedding)
+                    #st.write(f'{selected_encoder} embedding')
+                    #st.write(embedding)
+                    st.image('molecule_encoder_done.png')
 
     with col2:
         st.markdown('### Target')
@@ -113,7 +114,6 @@ def predict_dti():
             selected_encoder = st.selectbox(
                 'Select encoder for protein target',('None', 'SeqVec', 'UniRep', 'ESM-1b', 'ProtT5')
             )
-            st.image('protein_encoder.png')
             if sequence: 
                 if selected_encoder == 'SeqVec':
                     from bio_embeddings.embed import SeqVecEmbedder
@@ -143,11 +143,13 @@ def predict_dti():
                         embedding = encoder.reduce_per_protein(emb)
                         break
                 else: 
-                    st.write('No pre-trained version of HyperPCM is available for the chosen encoder.')
+                    #st.write('No pre-trained version of HyperPCM is available for the chosen encoder.')
                     embedding = None
+                    st.image('protein_encoder.png')
                 if embedding is not None:
-                    st.write(f'{selected_encoder} embedding')
-                    st.write(embedding)
+                    #st.write(f'{selected_encoder} embedding')
+                    #st.write(embedding)
+                    st.image('protein_encoder_done.png')
 
     st.write('TODO run inference with HyperPCM on the given drug compound and protein target.')
     
@@ -167,7 +169,7 @@ def retrieval():
         
         with col2:
             st.write('Encoding with SecVec')
-            st.image('protein_encoder.png')
+            st.image('protein_encoder_done.png')
             
             from bio_embeddings.embed import SeqVecEmbedder
             encoder = SeqVecEmbedder()
@@ -175,9 +177,6 @@ def retrieval():
             for emb in embeddings:
                 embedding = encoder.reduce_per_protein(emb)
                 break
-            st.write('SeqVec embedding')
-            st.write(embedding)
-            st.write(np.transpose(embedding))
         
     st.markdown('### Retrieval')
     st.write('TODO HyperPCM predicts the QSAR model for the given protein target.')
@@ -201,7 +200,7 @@ def retrieval():
                     'CC(=O)Nc1ccc(O[C@@H]2O[C@H](C(=O)O)[C@@H](O)[C@H](O)[C@H]2O)cc1']]
     cols = st.columns(5)
     for j, col in enumerate(cols):
-        with cols:
+        with col:
             for i in range(int(selected_k/5)):
                 mol = Chem.MolFromSmiles(dummy_smiles[i,j])
                 mol_img = Chem.Draw.MolToImage(mol)