Add topN

main.py

@@ -144,7 +144,7 @@ def get_pdbid_from_filename(filename: str):
     # Assuming the filename would be of the standard form 11GS.pdb
     return filename.split(".")[0]
 
-def predict(pdb_code, pdb_file):
+def predict(pdb_code, pdb_file, topN):
     #path_to_pdb = get_pdb(pdb_code=pdb_code, filepath=pdb_file)
 
     #pdb = open(path_to_pdb, "r").read()
@@ -184,7 +184,6 @@ def predict(pdb_code, pdb_file):
     for i in range(adaptability.shape[0]):
         data.append([i, atom_mapping[atoms_protein.iloc[i, atoms_protein.columns.get_loc("element")] - 1], atoms_protein.iloc[i, atoms_protein.columns.get_loc("x")],atoms_protein.iloc[i, atoms_protein.columns.get_loc("y")],atoms_protein.iloc[i, atoms_protein.columns.get_loc("z")],adaptability[i]])
 
-    topN = 100
     topN_ind = np.argsort(adaptability)[::-1][:topN]    
 
     pdb = open(pdb_file.name, "r").read()
@@ -198,7 +197,7 @@ def predict(pdb_code, pdb_file):
     # view.setSpecular(0.5)  # Adjust the specular lighting effect
     # view.setAmbient(0.5)  # Adjust the ambient lighting effect
     
-    for i in range(10):
+    for i in range(topN):
         adaptability_value = adaptability[topN_ind[i]]
         color = 'orange'
         view.addSphere({
@@ -248,7 +247,7 @@ def run():
         with gr.Row():
             dataframe = gr.Dataframe()
                 
-        single_btn.click(fn=predict, inputs=[inp, pdb_file], outputs=[html, dataframe])
+        single_btn.click(fn=predict, inputs=[inp, pdb_file, topN], outputs=[html, dataframe])
 
 
     demo.launch(server_name="0.0.0.0", server_port=7860)