supporting multiple indexes

app.py

@@ -1,16 +1,18 @@
 import gradio as gr
 
 from concurrency import execute_multithread
-from get_index import get_engine
+from get_index import get_engines
 from protein_viz import get_gene_name, get_protein_name, render_html
 
 index_repo = "ronig/siamese_protein_index"
 model_repo = "ronig/protein_search_engine"
-engine = get_engine(index_repo, model_repo)
+engines = get_engines(index_repo, model_repo)
+available_indexes = list(engines.keys())
 
 
-def search_and_display(seq, n_res):
+def search_and_display(seq, n_res, index_selection):
     n_res = int(limit_n_results(n_res))
+    engine = engines[index_selection]
     search_res = engine.search_by_sequence(seq, n=n_res)
     results_options = update_dropdown_menu(search_res)
     formatted_search_results = format_search_results(search_res)
@@ -24,12 +26,18 @@ def limit_n_results(n):
 def update_dropdown_menu(search_res):
     choices = []
     for row in search_res:
-        choice = ".".join([row["pdb_name"], row["chain_id"]])
-        choices.append(choice)
-
-    return gr.Dropdown.update(
-        choices=choices, interactive=True, value=choices[0], visible=True
-    )
+        if "pdb_name" in row and "chain_id" in row:
+            choice = ".".join([row["pdb_name"], row["chain_id"]])
+            choices.append(choice)
+    if choices:
+        update = gr.Dropdown.update(
+            choices=choices, interactive=True, value=choices[0], visible=True
+        )
+    else:
+        update = gr.Dropdown.update(
+            choices=choices, interactive=True, visible=False, value=None
+        )
+    return update
 
 
 def format_search_results(raw_search_results):
@@ -44,6 +52,22 @@ def format_search_results(raw_search_results):
 
 
 def format_search_result(raw_result):
+    is_pdb = "pdb_name" in raw_result
+    if is_pdb:
+        key, value = parse_pdb_search_result(raw_result)
+    else:
+        key, value = parse_fasta_search_result(raw_result)
+    return key, value
+
+
+def parse_fasta_search_result(raw_result):
+    gene = parse_gene_from_fasta_entry(raw_result["description"])
+    key = f"Gene: {gene}"
+    value = raw_result["score"]
+    return key, value
+
+
+def parse_pdb_search_result(raw_result):
     prot = raw_result["pdb_name"]
     chain = raw_result["chain_id"]
     value = raw_result["score"]
@@ -54,16 +78,28 @@ def format_search_result(raw_result):
     return key, value
 
 
+def parse_gene_from_fasta_entry(description):
+    after = description.split("GN=")[1]
+    gene = after.split(" ")[0]
+    return gene
+
+
 def switch_viz(new_choice):
-    choice_parts = new_choice.split(".")
-    pdb_id, chain = choice_parts[0], choice_parts[1]
-    title_update = gr.Markdown.update(visible=True)
-    protein_name = get_protein_name(pdb_id)
+    if new_choice is None:
+        html = ""
+        title_update = gr.Markdown.update(visible=False)
+        description_update = gr.Markdown.update(value=None, visible=False)
+    else:
+        choice_parts = new_choice.split(".")
+        pdb_id, chain = choice_parts[0], choice_parts[1]
+        title_update = gr.Markdown.update(visible=True)
+        protein_name = get_protein_name(pdb_id)
 
-    new_value = f"""**PDB Title**: {protein_name}"""
+        new_value = f"""**PDB Title**: {protein_name}"""
 
-    description_update = gr.Markdown.update(value=new_value, visible=True)
-    return render_html(pdb_id=pdb_id, chain=chain), title_update, description_update
+        description_update = gr.Markdown.update(value=new_value, visible=True)
+        html = render_html(pdb_id=pdb_id, chain=chain)
+    return html, title_update, description_update
 
 
 with gr.Blocks() as demo:
@@ -79,10 +115,15 @@ with gr.Blocks() as demo:
         with gr.Column():
             with gr.Row():
                 with gr.Column():
-                    seq_input = gr.Textbox(
-                        value="APTMPPPLPP", label="Input Sequence"
-                    )
+                    seq_input = gr.Textbox(value="APTMPPPLPP", label="Input Sequence")
                     n_results = gr.Number(5, label="N Results")
+                    index_selector = gr.Dropdown(
+                        choices=available_indexes,
+                        value=available_indexes[0],
+                        multiselect=False,
+                        visible=True,
+                        label="Index",
+                    )
                     search_button = gr.Button("Search", variant="primary")
                 search_results = gr.Label(num_top_classes=20, label="Search Results")
             viz_header = gr.Markdown("## Visualization", visible=False)
@@ -103,7 +144,7 @@ with gr.Blocks() as demo:
             )
     search_button.click(
         search_and_display,
-        inputs=[seq_input, n_results],
+        inputs=[seq_input, n_results, index_selector],
         outputs=[search_results, results_selector],
     )
     results_selector.change(

get_index.py

@@ -1,4 +1,6 @@
+import os.path
 import sys
+from glob import glob
 from pathlib import Path
 
 from huggingface_hub import snapshot_download
@@ -6,10 +8,11 @@ from huggingface_hub import snapshot_download
 from credentials import get_token
 
 
-def get_engine(index_repo: str, model_repo: str):
+def get_engines(index_repo: str, model_repo: str):
     index_path = Path(
         snapshot_download(index_repo, use_auth_token=get_token(), repo_type="dataset")
     )
+
     local_arch_path = Path(
         snapshot_download(model_repo, use_auth_token=get_token(), repo_type="model")
     )
@@ -18,4 +21,9 @@ def get_engine(index_repo: str, model_repo: str):
         ProteinSearchEngine,
     )
 
-    return ProteinSearchEngine(data_path=index_path)
+    subindex_paths = glob(str(index_path / "*/"))
+    engines = {}
+    for subindex_path in subindex_paths:
+        subindex_name = os.path.basename(subindex_path).title()
+        engines[subindex_name] = ProteinSearchEngine(data_path=Path(subindex_path))
+    return engines