initial protein search engine implementation

.gitignore

@@ -0,0 +1,4 @@
+.idea
+venv
+__pycache__
+scratch

Makefile

@@ -0,0 +1,15 @@
+setup:
+	python -m venv venv
+	venv/bin/pip install -r requirements.txt
+	venv/bin/pip install -r requirements-dev.txt
+
+format:
+	venv/bin/black .
+
+check-code: lint-python check-formatting
+
+check-formatting:
+	venv/bin/black --check .
+
+lint-python:
+	venv/bin/pylint --rcfile=pylintrc .

app.py

@@ -0,0 +1,36 @@
+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()

credentials.py

@@ -0,0 +1,5 @@
+import os
+
+
+def get_token():
+    return os.environ.get("TOKEN", True)

get_index.py

@@ -0,0 +1,23 @@
+import sys
+from pathlib import Path
+
+from huggingface_hub import snapshot_download
+
+from credentials import get_token
+
+
+def get_engine(index_repo: str, model_repo: str):
+    index_path = Path(
+        snapshot_download(
+            index_repo, use_auth_token=get_token(), repo_type="dataset"
+        )
+    )
+    local_arch_path = Path(
+        snapshot_download(
+            model_repo, use_auth_token=get_token(), repo_type="model"
+        )
+    )
+    sys.path.append(str(local_arch_path))
+    from protein_index import ProteinSearchEngine  # pylint: disable=import-error,import-outside-toplevel
+
+    return ProteinSearchEngine(data_path=index_path)

pylintrc

@@ -0,0 +1,20 @@
+[MESSAGES CONTROL]
+disable=missing-docstring,invalid-name,logging-fstring-interpolation
+
+[DESIGN]
+min-public-methods=1
+
+[FORMAT]
+max-line-length=88
+
+[SIMILARITIES]
+min-similarity-lines=10
+
+[TYPECHECK]
+
+[MASTER]
+init-hook=import sys; sys.path.append(".")
+extension-pkg-whitelist=pydantic,cassandra
+generated-members=torch.*,cv2.*,np.random.*
+ignore-patterns=setup,py,tasks.py
+max-args=6

requirements-dev.txt

@@ -0,0 +1,5 @@
+pytest
+pylint
+black
+mypy
+huggingface_hub

requirements.txt

@@ -0,0 +1,3 @@
+torch
+transformers
+annoy