app.py

import gradio as gr

from get_index import get_engine

index_repo = "ronig/siamese_protein_index"
model_repo = "ronig/protein_search_engine"
engine = get_engine(index_repo, model_repo)


def search(seq, n_res):
    n_res = int(limit_n_results(n_res))
    search_results = engine.search_by_sequence(seq, n=n_res)
    outputs = {}
    for res in 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" | Genes: {genes}"
        outputs[key] = value
    return outputs


def limit_n_results(n):
    return max(min(n, 20), 1)


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 specific genes, their IDs will also be displayed.
        """)
        with gr.Row():
            with gr.Column():
                seq_input = gr.Textbox("KFLIYQMECSTMIFGL", label="Input Sequence")
                n_results = gr.Number(5, label="N Results")
                search_button = gr.Button("Search")
            output = gr.Label(num_top_classes=20, label="Search Results")
    search_button.click(search, inputs=[seq_input, n_results], outputs=output)

if __name__ == '__main__':
    demo.launch()