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["protein_name"][0]
        chain = res["chain_id"][0]
        value = res["score"]
        key = f"Protein: {prot} | Chain: {chain}"
        outputs[key] = value
    return outputs


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


with gr.Blocks() as demo:
    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)

demo.launch()