Update demo/modules/search.py

demo/modules/search.py

@@ -33,13 +33,6 @@ samples = {
     ],
 }
 
-# Databases for different modalities
-now_db = {
-    "sequence": list(all_index["sequence"].keys())[0],
-    "structure": list(all_index["structure"].keys())[0],
-    "text": list(all_index["text"].keys())[0]
-}
-
 
 def clear_results():
     return "", gr.update(visible=False), gr.update(visible=False)
@@ -75,6 +68,8 @@ def plot(scores) -> None:
 
 # Search from database
 def search(input: str, nprobe: int, topk: int, input_type: str, query_type: str, subsection_type: str, db: str):
+    print(f"Input type: {input_type}\n Output type: {query_type}\nDatabase: {db}\nSubsection: {subsection_type}")
+    
     input_modality = input_type.replace("sequence", "protein")
     with torch.no_grad():
         input_embedding = getattr(model, f"get_{input_modality}_repr")([input]).cpu().numpy()
@@ -169,17 +164,18 @@ def load_example(example_id):
  
 # Change the visibility of subsection type
 def change_output_type(query_type: str, subsection_type: str):
-    nprobe_visible = check_index_ivf(query_type, subsection_type)
+    db_type = list(all_index[query_type].keys())[0]
+    nprobe_visible = check_index_ivf(query_type, db_type, subsection_type)
     subsection_visible = True if query_type == "text" else False
-    
+
     return (
-        gr.update(visible=subsection_visible), 
+        gr.update(visible=subsection_visible),
         gr.update(visible=nprobe_visible),
-        gr.update(choices=list(all_index[query_type].keys()), value=now_db[query_type])
+        gr.update(choices=list(all_index[query_type].keys()), value=db_type)
     )
 
 
-def check_index_ivf(index_type: str, subsection_type: str = None) -> bool:
+def check_index_ivf(index_type: str, db: str, subsection_type: str = None) -> bool:
     """
     Check if the index is of IVF type.
     Args:
@@ -189,7 +185,6 @@ def check_index_ivf(index_type: str, subsection_type: str = None) -> bool:
     Returns:
         Whether the index is of IVF type or not.
     """
-    db = now_db[index_type]
     if index_type == "sequence":
         index = all_index["sequence"][db]["index"]
     
@@ -211,14 +206,12 @@ def change_db_type(query_type: str, subsection_type: str, db_type: str):
         query_type: The output type.
         db_type: The database to search.
     """
-    now_db[query_type] = db_type
-    
     if query_type == "text":
-        subsection_update = gr.update(choices=list(valid_subsections[now_db["text"]]), value="Function")
+        subsection_update = gr.update(choices=list(valid_subsections[db_type]), value="Function")
     else:
         subsection_update = gr.update(visible=False)
     
-    nprobe_visible = check_index_ivf(query_type, subsection_type)
+    nprobe_visible = check_index_ivf(query_type, db_type, subsection_type)
     return subsection_update, gr.update(visible=nprobe_visible)
 
 
@@ -240,10 +233,12 @@ def build_search_module():
                 )
             
                 # If the output type is "text", provide an option to choose the subsection of text
-                subsection_type = gr.Dropdown(valid_subsections[now_db["text"]], label="Subsection of text", value="Function",
+                text_db = list(all_index["text"].keys())[0]
+                sequence_db = list(all_index["sequence"].keys())[0]
+                subsection_type = gr.Dropdown(valid_subsections[text_db], label="Subsection of text", value="Function",
                                               interactive=True, visible=False, scale=0)
                 
-                db_type = gr.Dropdown(all_index["sequence"].keys(), label="Database", value=now_db["sequence"],
+                db_type = gr.Dropdown(all_index["sequence"].keys(), label="Database", value=sequence_db,
                                               interactive=True, visible=True, scale=0)
 
             with gr.Row():
@@ -256,7 +251,7 @@ def build_search_module():
             
             
             # If the index is of IVF type, provide an option to choose the number of clusters.
-            nprobe_visible = check_index_ivf(query_type.value)
+            nprobe_visible = check_index_ivf(query_type.value, db_type.value)
             nprobe = gr.Slider(1, 1000000, 1000,  step=1, visible=nprobe_visible,
                                label="Number of clusters to search (lower value for faster search and higher value for more accurate search)")