Spaces:
Running
Running
| import gradio as gr | |
| from get_index import get_engine | |
| from protein_viz import render_html | |
| index_repo = "ronig/siamese_protein_index" | |
| model_repo = "ronig/protein_search_engine" | |
| engine = get_engine(index_repo, model_repo) | |
| def search_and_display(seq, n_res): | |
| n_res = int(limit_n_results(n_res)) | |
| 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) | |
| return formatted_search_results, results_options | |
| def limit_n_results(n): | |
| return max(min(n, 20), 1) | |
| def update_dropdown_menu(search_res): | |
| choices = [ | |
| ','.join([row['pdb_name'], row['chain_id']]) | |
| for row in search_res | |
| ] | |
| return gr.Dropdown.update( | |
| choices=choices, interactive=True, value=choices[0], visible=True | |
| ) | |
| def format_search_results(raw_search_results): | |
| formatted_search_results = {} | |
| for res in raw_search_results: | |
| prot = res["pdb_name"] | |
| chain = res["chain_id"] | |
| value = res["score"] | |
| genes = ",".join(res["gene_ids"]) | |
| key = f"PDB: {prot} | Chain: {chain}" | |
| if genes != "Unknown": | |
| key += f" | Gene: {genes}" | |
| formatted_search_results[key] = value | |
| return formatted_search_results | |
| def switch_viz(new_choice): | |
| pdb_id, chain = new_choice.split(',') | |
| title_update = gr.Markdown.update(visible=True) | |
| return render_html(pdb_id=pdb_id, chain=chain), title_update | |
| with gr.Blocks() as demo: | |
| with gr.Column(): | |
| 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. | |
| 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" | |
| ) | |
| n_results = gr.Number(5, label="N Results") | |
| 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" | |
| ) | |
| protein_viz = gr.HTML( | |
| value=render_html(pdb_id=None, chain=None), | |
| label="Protein Visualization" | |
| ) | |
| gr.Examples([ | |
| "KFLIYQMECSTMIFGL", | |
| "PHFAMPPIHEDHLE", | |
| "AEERIISLD" | |
| ], inputs=[seq_input]) | |
| search_button.click( | |
| search_and_display, | |
| inputs=[seq_input, n_results], | |
| outputs=[search_results, results_selector] | |
| ) | |
| results_selector.change( | |
| switch_viz, inputs=results_selector, outputs=[protein_viz, viz_header] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |