showing gene info

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 
 from get_index import get_engine
-from protein_viz import get_protein_name, render_html
+from protein_viz import get_gene_names, get_protein_name, render_html
 
 index_repo = "ronig/siamese_protein_index"
 model_repo = "ronig/protein_search_engine"
@@ -21,10 +21,16 @@ def limit_n_results(n):
 
 
 def update_dropdown_menu(search_res):
-    choices = [
-        ','.join([row['pdb_name'], row['chain_id']])
-        for row in search_res
-    ]
+    choices = []
+    for row in search_res:
+        gene = row["gene_ids"][0]
+        if gene != "Unknown":
+            choice_parts = [row["pdb_name"], row["chain_id"], gene]
+        else:
+            choice_parts = [row["pdb_name"], row["chain_id"]]
+        choice = ",".join(choice_parts)
+        choices.append(choice)
+
     return gr.Dropdown.update(
         choices=choices, interactive=True, value=choices[0], visible=True
     )
@@ -45,57 +51,63 @@ def format_search_results(raw_search_results):
 
 
 def switch_viz(new_choice):
-    pdb_id, chain = new_choice.split(',')
+    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)
-    description_update = gr.Markdown.update(
-        value=f"""**PDB Title**: {protein_name}""",
-        visible=True
-    )
+
+    new_value = f"""**PDB Title**: {protein_name}"""
+    if len(choice_parts) > 2:
+        gene = choice_parts[2]
+        gene_name = get_gene_names([gene])[0]
+        new_value += f"""\n\n**Gene Name**: {gene_name.title()}"""
+
+    description_update = gr.Markdown.update(value=new_value, visible=True)
     return render_html(pdb_id=pdb_id, chain=chain), title_update, description_update
 
 
 with gr.Blocks() as demo:
     with gr.Column():
-        gr.Markdown("""
+        gr.Markdown(
+            """
         # Protein Binding Search Engine
-        This application examines all files uploaded to [PDB](https://www.rcsb.org/) to find the chains with which a given protein sequence is most likely to bind.
+        This application examines all files uploaded to [PDB](https://www.rcsb.org/) 
+        to find the chains with which a given protein sequence is most likely to bind.
         If the results are linked to a specific gene, their IDs will also be displayed.
-        """)
+        """
+        )
         with gr.Column():
             with gr.Row():
                 with gr.Column():
                     seq_input = gr.Textbox(
-                        value="KFLIYQMECSTMIFGL",
-                        label="Input Sequence"
+                        value="KFLIYQMECSTMIFGL", label="Input Sequence"
                     )
                     n_results = gr.Number(5, label="N Results")
-                    search_button = gr.Button("Search", variant='primary')
+                    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)
             results_selector = gr.Dropdown(
-                choices=[], multiselect=False, visible=False,
-                label="Visualized Search Result"
+                choices=[],
+                multiselect=False,
+                visible=False,
+                label="Visualized Search Result",
             )
             viz_body = gr.Markdown("", visible=False)
             protein_viz = gr.HTML(
                 value=render_html(pdb_id=None, chain=None),
-                label="Protein Visualization"
+                label="Protein Visualization",
+            )
+            gr.Examples(
+                ["KFLIYQMECSTMIFGL", "PHFAMPPIHEDHLE", "AEERIISLD"], inputs=[seq_input]
             )
-            gr.Examples([
-                "KFLIYQMECSTMIFGL",
-                "PHFAMPPIHEDHLE",
-                "AEERIISLD"
-            ], inputs=[seq_input])
     search_button.click(
         search_and_display,
         inputs=[seq_input, n_results],
-        outputs=[search_results, results_selector]
+        outputs=[search_results, results_selector],
     )
     results_selector.change(
-        switch_viz,
-        inputs=results_selector,
-        outputs=[protein_viz, viz_header, viz_body]
+        switch_viz, inputs=results_selector, outputs=[protein_viz, viz_header, viz_body]
     )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()

get_index.py

@@ -8,16 +8,14 @@ from credentials import get_token
 
 def get_engine(index_repo: str, model_repo: str):
     index_path = Path(
-        snapshot_download(
-            index_repo, use_auth_token=get_token(), repo_type="dataset"
-        )
+        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"
-        )
+        snapshot_download(model_repo, use_auth_token=get_token(), repo_type="model")
     )
     sys.path.append(str(local_arch_path))
-    from protein_index import ProteinSearchEngine  # pylint: disable=import-error,import-outside-toplevel
+    from protein_index import (  # pylint: disable=import-error,import-outside-toplevel
+        ProteinSearchEngine,
+    )
 
     return ProteinSearchEngine(data_path=index_path)

protein_viz.py

@@ -35,8 +35,8 @@ def render_html(pdb_id, chain):
 
 def get_gene_names(genes: List[str]):
     mg = mygene.MyGeneInfo()
-    ginfo = mg.querymany(genes, scopes='ensembl.gene')
-    gene_names = [gene['name'] for gene in ginfo]
+    ginfo = mg.querymany(genes, scopes="ensembl.gene", verbose=False, fields="name")
+    gene_names = [gene["name"] for gene in ginfo]
     return gene_names
 
 
@@ -45,7 +45,7 @@ def get_protein_name(pdb_id: str):
     response = requests.get(url, timeout=1)
     if response.ok:
         data = response.json()
-        protein_name = data['struct']['title']
+        protein_name = data["struct"]["title"]
     else:
-        protein_name = 'Unknown'
+        protein_name = "Unknown"
     return protein_name