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()