| import collections | |
| import os | |
| from typing import Dict, List | |
| import gradio as gr | |
| from index_list import read_index_list | |
| from protein_viz import get_pdb_title, render_html | |
| from search_engine import MilvusParams, ProteinSearchEngine | |
| model_repo = "ronig/protein_biencoder" | |
| available_indexes = read_index_list() | |
| engine = ProteinSearchEngine( | |
| milvus_params=MilvusParams( | |
| uri="https://in03-ddab8e9a5a09fcc.api.gcp-us-west1.zillizcloud.com", | |
| token=os.environ.get("MILVUS_TOKEN"), | |
| db_name="Protein", | |
| collection_name="Peptriever", | |
| ), | |
| model_repo=model_repo, | |
| ) | |
| max_results = 1000 | |
| choice_sep = " | " | |
| max_seq_length = 50 | |
| def search_and_display(seq, max_res, index_selection): | |
| n_search_res = 1024 | |
| _validate_sequence_length(seq) | |
| max_res = int(limit_n_results(max_res)) | |
| if index_selection == "All Species": | |
| index_selection = None | |
| search_res = engine.search_by_sequence( | |
| seq, n=n_search_res, organism=index_selection | |
| ) | |
| agg_search_results = aggregate_search_results(search_res, max_res) | |
| formatted_search_results = format_search_results(agg_search_results) | |
| results_options = update_dropdown_menu(agg_search_results) | |
| return formatted_search_results, results_options | |
| def _validate_sequence_length(seq): | |
| if len(seq) > max_seq_length: | |
| raise gr.Error("Only peptide input is currently supported") | |
| def limit_n_results(n): | |
| return max(min(n, max_results), 1) | |
| def aggregate_search_results(raw_results: List[dict], max_res: int) -> Dict[str, dict]: | |
| aggregated_by_uniprot = collections.defaultdict(list) | |
| for raw_result in raw_results: | |
| entry = select_keys( | |
| raw_result, | |
| keys=["pdb_name", "chain_id", "score", "organism", "uniprot_id", "genes"], | |
| ) | |
| uniprot_id = raw_result["uniprot_id"] | |
| if uniprot_id is not None: | |
| aggregated_by_uniprot[uniprot_id].append(entry) | |
| if len(aggregated_by_uniprot) >= max_res: | |
| return dict(aggregated_by_uniprot) | |
| return dict(aggregated_by_uniprot) | |
| def select_keys(d: dict, keys: List[str]): | |
| return {key: d[key] for key in keys} | |
| def format_search_results(agg_search_results): | |
| formatted_search_results = {} | |
| for uniprot_id, entries in agg_search_results.items(): | |
| entry = entries[0] | |
| organism = entry["organism"] | |
| score = entry["score"] | |
| genes = entry["genes"] | |
| key = f"Uniprot ID: {uniprot_id} | Organism: {organism} | Gene Names: {genes}" | |
| formatted_search_results[key] = score | |
| return formatted_search_results | |
| def update_dropdown_menu(agg_search_res): | |
| choices = [] | |
| for uniprot_id, entries in agg_search_res.items(): | |
| for entry in entries: | |
| choice = choice_sep.join( | |
| [ | |
| uniprot_id, | |
| entry["pdb_name"], | |
| entry["chain_id"], | |
| entry["genes"] or "", | |
| ] | |
| ) | |
| choices.append(choice) | |
| if choices: | |
| update = gr.update( | |
| gr.Dropdown.get_component_class_id(), | |
| choices=choices, | |
| interactive=True, | |
| value=choices[0], | |
| visible=True, | |
| ) | |
| else: | |
| update = gr.update( | |
| gr.Dropdown.get_component_class_id(), | |
| choices=choices, | |
| interactive=True, | |
| visible=False, | |
| value=None, | |
| ) | |
| return update | |
| def parse_pdb_search_result(raw_result): | |
| prot = raw_result["pdb_name"] | |
| chain = raw_result["chain_id"] | |
| value = raw_result["score"] | |
| gene_names = raw_result["genes"] | |
| species = raw_result["organism"] | |
| key = f"PDB: {prot}.{chain}" | |
| if gene_names is not None: | |
| key += f" | Genes: {gene_names} | Organism: {species}" | |
| return key, value | |
| def switch_viz(new_choice): | |
| if new_choice is None: | |
| html = "" | |
| title_update = gr.update(gr.Markdown.get_component_class_id(), visible=False) | |
| description_update = gr.update( | |
| gr.Markdown.get_component_class_id(), value=None, visible=False | |
| ) | |
| else: | |
| choice_parts = new_choice.split(choice_sep) | |
| pdb_id, chain = choice_parts[1:3] | |
| title_update = gr.update(gr.Markdown.get_component_class_id(), visible=True) | |
| pdb_title = get_pdb_title(pdb_id) | |
| new_value = f"""**PDB Title**: {pdb_title}""" | |
| description_update = gr.update( | |
| gr.Markdown.get_component_class_id(), 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: | |
| with gr.Column(): | |
| with gr.Column(): | |
| with gr.Row(): | |
| with gr.Column(): | |
| seq_input = gr.Textbox(value="APTMPPPLPP", label="Input Sequence") | |
| n_results = gr.Number(10, label="N Results") | |
| index_selector = gr.Dropdown( | |
| choices=available_indexes, | |
| value="All Species", | |
| multiselect=False, | |
| visible=True, | |
| label="Index", | |
| ) | |
| search_button = gr.Button("Search", variant="primary") | |
| search_results = gr.Label( | |
| num_top_classes=max_results, label="Search Results", scale=2 | |
| ) | |
| viz_header = gr.Markdown("## Visualization", visible=False) | |
| results_selector = gr.Dropdown( | |
| 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", | |
| ) | |
| gr.Examples( | |
| ["APTMPPPLPP", "KFLIYQMECSTMIFGL", "PHFAMPPIHEDHLE", "AEERIISLD"], | |
| inputs=[seq_input], | |
| ) | |
| search_button.click( | |
| search_and_display, | |
| inputs=[seq_input, n_results, index_selector], | |
| outputs=[search_results, results_selector], | |
| ) | |
| results_selector.change( | |
| switch_viz, inputs=results_selector, outputs=[protein_viz, viz_header, viz_body] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |