Merge pull request #37 from Vedasheersh/codex/catpred-webkinpred-adapter

Vedasheersh · web-flow · commit b4f5bc31bacc · 2026-03-28T22:23:35.000-04:00
Add webKinPred adapter for local CatPred inference
diff --git a/catpred/data/utils.py b/catpred/data/utils.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from collections import OrderedDict, defaultdict
+from pathlib import Path
 import sys
 import csv
 import ctypes
@@ -22,7 +23,7 @@
 from .scaffold import log_scaffold_stats, scaffold_split
 from catpred.features import load_features, load_valid_atom_or_bond_features, is_mol
 from catpred.rdkit import make_mol
-from catpred.security import load_index_artifact, load_pickle_artifact
+from catpred.security import load_index_artifact, load_pickle_artifact, load_torch_artifact
 
 if TYPE_CHECKING:
     from catpred.args import PredictArgs, TrainArgs
@@ -556,13 +557,26 @@ def get_data(path: str,
                         'This usually means multiple sequences share the same pdbpath identifier.'
                     )
 
+                if 'esm2_feats' in protein_record and not isinstance(protein_record['esm2_feats'], torch.Tensor):
+                    protein_record['esm2_feats'] = torch.as_tensor(protein_record['esm2_feats'])
+
                 if 'esm2_feats' not in protein_record:
-                    sequence_features, _ = sequence_feat_getter(
-                        protein_record['seq'],
-                        name=pdbname,
-                        device='cpu'
-                    )
-                    protein_record['esm2_feats'] = sequence_features[0]  # batch dim
+                    esm2_feats_path = protein_record.get('esm2_feats_path')
+                    if esm2_feats_path:
+                        esm2_feats_path = Path(esm2_feats_path).resolve()
+                        protein_record['esm2_feats'] = load_torch_artifact(
+                            esm2_feats_path,
+                            purpose="protein ESM2 features",
+                            map_location='cpu',
+                            roots=[esm2_feats_path.parent],
+                        )
+                    else:
+                        sequence_features, _ = sequence_feat_getter(
+                            protein_record['seq'],
+                            name=pdbname,
+                            device='cpu'
+                        )
+                        protein_record['esm2_feats'] = sequence_features[0]  # batch dim
                 
             targets, atom_targets, bond_targets = [], [], []
             for column in target_columns:
diff --git a/catpred/inference/service.py b/catpred/inference/service.py
@@ -44,6 +44,15 @@ def _resolve_input_path(input_file: str, repo_root: Path) -> Path:
     return path
 
 
+def _resolve_existing_path(path_str: str, repo_root: Path, purpose: str) -> Path:
+    path = Path(path_str)
+    if not path.is_absolute():
+        path = (repo_root / path).resolve()
+    if not path.exists():
+        raise FileNotFoundError(f'{purpose} not found: "{path}"')
+    return path
+
+
 def _validate_and_prepare_dataframe(parameter: str, df: pd.DataFrame, input_csv: Path) -> pd.DataFrame:
     required_columns = {"SMILES", "sequence", "pdbpath"}
     missing = required_columns.difference(df.columns)
@@ -152,7 +161,19 @@ def run_raw_prediction(request: PredictionRequest, paths: PreparedInputPaths) ->
     env = os.environ.copy()
     env["PROTEIN_EMBED_USE_CPU"] = "0" if request.use_gpu else "1"
 
-    subprocess.run(create_records_cmd, cwd=str(root), env=env, check=True)
+    protein_records_path = paths.records_file
+    if request.protein_records_file:
+        protein_records_path = str(
+            _resolve_existing_path(
+                request.protein_records_file,
+                repo_root=root,
+                purpose="Protein records file",
+            )
+        )
+    else:
+        subprocess.run(create_records_cmd, cwd=str(root), env=env, check=True)
+
+    predict_cmd[-1] = protein_records_path
     subprocess.run(predict_cmd, cwd=str(root), env=env, check=True)
 
     if not os.path.exists(paths.output_csv):
diff --git a/catpred/inference/types.py b/catpred/inference/types.py
@@ -11,6 +11,7 @@ class PredictionRequest:
     use_gpu: bool = False
     repo_root: str | None = None
     python_executable: str = "python"
+    protein_records_file: str | None = None
 
 
 @dataclass(frozen=True)
diff --git a/catpred/integration/__init__.py b/catpred/integration/__init__.py
@@ -0,0 +1 @@
+"""Integration helpers for external runtimes."""
diff --git a/catpred/integration/webkinpred_adapter.py b/catpred/integration/webkinpred_adapter.py
@@ -0,0 +1,285 @@
+from __future__ import annotations
+
+import argparse
+import gzip
+import hashlib
+import json
+import os
+from pathlib import Path
+import subprocess
+import sys
+import tempfile
+from typing import Any
+
+import pandas as pd
+import torch
+
+_THIS_FILE = Path(__file__).resolve()
+_REPO_ROOT = _THIS_FILE.parents[2]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+from catpred.inference import PredictionRequest, run_prediction_pipeline
+from catpred.security import load_torch_artifact
+
+_VALID_AAS = set("ACDEFGHIKLMNPQRSTVWY")
+_TARGET_TO_PARAMETER = {
+    "kcat": "kcat",
+    "Km": "km",
+    "km": "km",
+    "ki": "ki",
+    "Ki": "ki",
+}
+_PREDICTION_COLUMN = {
+    "kcat": "Prediction_(s^(-1))",
+    "km": "Prediction_(mM)",
+    "ki": "Prediction_(mM)",
+}
+
+
+def _repo_root() -> Path:
+    return _REPO_ROOT
+
+
+def _contains_model_checkpoints(path: Path) -> bool:
+    return path.exists() and path.is_dir() and any(path.rglob("model.pt"))
+
+
+def _discover_checkpoint_root(repo_root: Path) -> Path:
+    production_root = (repo_root / ".e2e-assets" / "pretrained" / "production").resolve()
+    checkpoints_root = (repo_root / "checkpoints").resolve()
+    if _contains_model_checkpoints(production_root):
+        return production_root
+    return checkpoints_root
+
+
+def _env_path(name: str, default: Path) -> Path:
+    value = os.environ.get(name)
+    return Path(value).resolve() if value else default.resolve()
+
+
+def _resolve_parameter(payload: dict[str, Any]) -> str:
+    target = payload.get("target")
+    if isinstance(target, str) and target in _TARGET_TO_PARAMETER:
+        return _TARGET_TO_PARAMETER[target]
+
+    params = payload.get("params") or {}
+    kinetics_type = str(params.get("kinetics_type", "")).strip().upper()
+    if kinetics_type == "KCAT":
+        return "kcat"
+    if kinetics_type == "KM":
+        return "km"
+    if kinetics_type == "KI":
+        return "ki"
+    raise RuntimeError(f"Unsupported target in payload: {target!r}")
+
+
+def _stable_seq_id(sequence: str) -> str:
+    digest = hashlib.sha1(sequence.encode("utf-8")).hexdigest()[:16]
+    return f"seq_{digest}"
+
+
+def _resolve_seq_ids(sequences: list[str], tools_path: Path, media_path: Path) -> list[str]:
+    seqmap_cli = tools_path / "seqmap" / "main.py"
+    seqmap_db = media_path / "sequence_info" / "seqmap.sqlite3"
+    if not seqmap_cli.exists() or not seqmap_db.exists():
+        return [_stable_seq_id(sequence) for sequence in sequences]
+
+    seqmap_python = os.environ.get("CATPRED_SEQMAP_PYTHON", sys.executable)
+    payload = "\n".join(sequences) + "\n"
+    cmd = [
+        seqmap_python,
+        str(seqmap_cli),
+        "--db",
+        str(seqmap_db),
+        "batch-get-or-create",
+        "--stdin",
+    ]
+    proc = subprocess.run(
+        cmd,
+        input=payload,
+        text=True,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+    )
+    if proc.returncode != 0:
+        raise RuntimeError(
+            f"seqmap failed (rc={proc.returncode})\nSTDOUT:\n{proc.stdout}\nSTDERR:\n{proc.stderr}"
+        )
+    seq_ids = [line.strip() for line in proc.stdout.splitlines() if line.strip()]
+    if len(seq_ids) != len(sequences):
+        raise RuntimeError(f"seqmap returned {len(seq_ids)} ids for {len(sequences)} sequences")
+    return seq_ids
+
+
+def _load_or_compute_esm2_feature(sequence: str, cache_file: Path) -> Path:
+    cache_file.parent.mkdir(parents=True, exist_ok=True)
+    if cache_file.exists():
+        load_torch_artifact(
+            cache_file,
+            purpose="CatPred ESM2 cache entry",
+            map_location="cpu",
+            roots=[cache_file.parent],
+        )
+        return cache_file
+
+    os.environ.setdefault("PROTEIN_EMBED_USE_CPU", "1")
+    from catpred.data.esm_utils import get_single_esm_repr
+
+    embedding = get_single_esm_repr(sequence).cpu()
+    torch.save(embedding, cache_file)
+    return cache_file
+
+
+def _build_input_dataframe(rows: list[dict[str, Any]], seq_ids: list[str]) -> pd.DataFrame:
+    formatted_rows = []
+    for row, seq_id in zip(rows, seq_ids):
+        substrate = row.get("substrates", row.get("substrate", row.get("Substrate", "")))
+        if isinstance(substrate, list):
+            if len(substrate) != 1:
+                raise RuntimeError("CatPred expects exactly one substrate per row.")
+            substrate = substrate[0]
+        substrate = str(substrate).strip()
+        sequence = str(row.get("sequence", "")).strip()
+        formatted_rows.append(
+            {
+                "SMILES": substrate,
+                "sequence": sequence,
+                "pdbpath": seq_id,
+            }
+        )
+    return pd.DataFrame(formatted_rows)
+
+
+def _write_protein_records(
+    rows: list[dict[str, Any]],
+    seq_ids: list[str],
+    parameter: str,
+    media_path: Path,
+    out_path: Path,
+) -> None:
+    records: dict[str, dict[str, Any]] = {}
+    needs_esm = parameter in {"kcat", "km"}
+    esm_cache_dir = media_path / "sequence_info" / "esm2_last" / "per_residue"
+
+    for row, seq_id in zip(rows, seq_ids):
+        sequence = str(row.get("sequence", "")).strip()
+        record: dict[str, Any] = {"name": seq_id, "seq": sequence}
+        if needs_esm:
+            cache_file = _load_or_compute_esm2_feature(sequence, esm_cache_dir / f"{seq_id}.pt")
+            record["esm2_feats_path"] = str(cache_file.resolve())
+        records[seq_id] = record
+
+    with gzip.open(out_path, "wt", encoding="utf-8") as handle:
+        json.dump(records, handle)
+
+
+def run_from_payload(payload: dict[str, Any]) -> dict[str, Any]:
+    rows = payload.get("rows") or []
+    if not isinstance(rows, list):
+        raise RuntimeError("'rows' must be a list in input payload.")
+
+    valid_rows: list[dict[str, Any]] = []
+    invalid_indices: list[int] = []
+    for idx, row in enumerate(rows):
+        sequence = str(row.get("sequence", "")).strip()
+        substrate = row.get("substrates", row.get("substrate", row.get("Substrate", "")))
+        if isinstance(substrate, list):
+            substrate = substrate[0] if len(substrate) == 1 else ""
+        substrate = str(substrate).strip()
+        if not sequence or not substrate or not set(sequence).issubset(_VALID_AAS):
+            invalid_indices.append(idx)
+            continue
+        valid_rows.append(row)
+
+    predictions: list[float | None] = [None] * len(rows)
+    if not valid_rows:
+        for idx in range(len(rows)):
+            print(f"Progress: {idx + 1}/{len(rows)}", flush=True)
+        return {"predictions": predictions, "invalid_indices": invalid_indices}
+
+    repo_root = _env_path("CATPRED_REPO_ROOT", _repo_root())
+    media_path = _env_path("CATPRED_MEDIA_PATH", repo_root / "media")
+    tools_path = _env_path("CATPRED_TOOLS_PATH", repo_root / "tools")
+    checkpoint_root = _env_path("CATPRED_CHECKPOINT_ROOT", _discover_checkpoint_root(repo_root))
+    parameter = _resolve_parameter(payload)
+
+    seq_ids = _resolve_seq_ids(
+        [str(row.get("sequence", "")).strip() for row in valid_rows],
+        tools_path=tools_path,
+        media_path=media_path,
+    )
+
+    with tempfile.TemporaryDirectory(prefix="catpred_webkinpred_") as tmp_dir_str:
+        tmp_dir = Path(tmp_dir_str).resolve()
+        input_csv = tmp_dir / "input.csv"
+        protein_records = tmp_dir / "protein_records.json.gz"
+        results_dir = tmp_dir / "results"
+
+        _build_input_dataframe(valid_rows, seq_ids).to_csv(input_csv, index=False)
+        _write_protein_records(
+            rows=valid_rows,
+            seq_ids=seq_ids,
+            parameter=parameter,
+            media_path=media_path,
+            out_path=protein_records,
+        )
+
+        request = PredictionRequest(
+            parameter=parameter,
+            input_file=str(input_csv),
+            checkpoint_dir=str((checkpoint_root / parameter).resolve()),
+            use_gpu=False,
+            repo_root=str(repo_root),
+            python_executable=sys.executable,
+            protein_records_file=str(protein_records),
+        )
+        output_file = run_prediction_pipeline(request=request, results_dir=str(results_dir))
+        output_df = pd.read_csv(output_file)
+        value_col = _PREDICTION_COLUMN[parameter]
+        if value_col not in output_df.columns:
+            raise RuntimeError(f"CatPred output is missing expected column: {value_col}")
+
+        valid_predictions = output_df[value_col].tolist()
+        if len(valid_predictions) != len(valid_rows):
+            raise RuntimeError(
+                f"CatPred produced {len(valid_predictions)} predictions for {len(valid_rows)} rows."
+            )
+
+    valid_iter = iter(valid_predictions)
+    for idx in range(len(rows)):
+        if idx in invalid_indices:
+            continue
+        predictions[idx] = float(next(valid_iter))
+
+    total = len(rows)
+    for idx in range(total):
+        print(f"Progress: {idx + 1}/{total}", flush=True)
+
+    return {
+        "predictions": predictions,
+        "invalid_indices": sorted(set(invalid_indices)),
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Run CatPred via the webKinPred subprocess contract.")
+    parser.add_argument("--input", required=True, help="Input JSON path.")
+    parser.add_argument("--output", required=True, help="Output JSON path.")
+    args = parser.parse_args()
+
+    with open(args.input, "r", encoding="utf-8") as handle:
+        payload = json.load(handle)
+
+    result = run_from_payload(payload)
+
+    with open(args.output, "w", encoding="utf-8") as handle:
+        json.dump(result, handle)
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except Exception as exc:
+        print(f"[CatPred] ERROR: {exc}", file=sys.stderr, flush=True)
+        raise

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+"""Integration helpers for external runtimes."""`