Merge pull request #42 from Vedasheersh/codex/inference-speedup

Vedasheersh · web-flow · commit 2e7013532f75 · 2026-06-14T10:04:22.000-04:00
Batch uncached ESM embeddings
diff --git a/catpred/data/cache_utils.py b/catpred/data/cache_utils.py
@@ -28,6 +28,37 @@ def md5_hash_fn(s):
     encoded = s.encode('utf-8')
     return hashlib.md5(encoded).hexdigest()
 
+
+def cache_entry_path(path='', cache_key=None, hash_fn=md5_hash_fn, name=None):
+    (CACHE_PATH / path).mkdir(parents=True, exist_ok=True)
+    if name is None:
+        key = hash_fn(cache_key)
+    else:
+        key = name
+    return CACHE_PATH / path / f'{key}.pt'
+
+
+def load_cache_value(path='', cache_key=None, *, purpose="cache entry", hash_fn=md5_hash_fn, name=None, map_location=None):
+    entry_path = cache_entry_path(path=path, cache_key=cache_key, hash_fn=hash_fn, name=name)
+    if not entry_path.exists():
+        return None
+
+    log(f'cache hit: fetching {cache_key} from {str(entry_path)}')
+    return load_torch_artifact(
+        str(entry_path),
+        purpose=purpose,
+        map_location=map_location,
+        roots=[CACHE_PATH],
+    )
+
+
+def save_cache_value(value, path='', cache_key=None, *, hash_fn=md5_hash_fn, name=None):
+    entry_path = cache_entry_path(path=path, cache_key=cache_key, hash_fn=hash_fn, name=name)
+    log(f'saving: {cache_key} to {str(entry_path)}')
+    cache_value = value.detach().cpu() if isinstance(value, torch.Tensor) else value
+    torch.save(cache_value, str(entry_path))
+    return entry_path
+
 # run once function
 
 GLOBAL_RUN_RECORDS = dict()
@@ -83,27 +114,20 @@ def inner(t, *args, __cache_key = None, **kwargs):
         if clear:
             clear_cache_folder_()
 
-        if name is None: 
-            cache_str = __cache_key if exists(__cache_key) else t
-            key = hash_fn(cache_str)
-        else:
-            key = name
-
-        entry_path = CACHE_PATH / path / f'{key}.pt'
-
-        if entry_path.exists():
-            log(f'cache hit: fetching {t} from {str(entry_path)}')
-            return load_torch_artifact(
-                str(entry_path),
-                purpose="esm cache entry",
-                roots=[CACHE_PATH],
-            )
+        cache_str = __cache_key if exists(__cache_key) else t
+        cached = load_cache_value(
+            path=path,
+            cache_key=cache_str,
+            purpose="esm cache entry",
+            hash_fn=hash_fn,
+            name=name,
+        )
+        if cached is not None:
+            return cached
 
         out = fn(t, *args, **kwargs)
 
-        log(f'saving: {t} to {str(entry_path)}')
-        cache_out = out.detach().cpu() if isinstance(out, torch.Tensor) else out
-        torch.save(cache_out, str(entry_path))
+        save_cache_value(out, path=path, cache_key=cache_str, hash_fn=hash_fn, name=name)
         return out
         
     return inner
diff --git a/catpred/data/esm_utils.py b/catpred/data/esm_utils.py
@@ -3,7 +3,7 @@
 from functools import partial
 import esm
 from torch.nn.utils.rnn import pad_sequence
-from .cache_utils import cache_fn, run_once
+from .cache_utils import cache_fn, load_cache_value, run_once, save_cache_value
 
 def exists(val):
     return val is not None
@@ -63,6 +63,11 @@ def calc_protein_representations_with_subunits(proteins, get_repr_fn, *, device)
 
 ESM_MAX_LENGTH = 2048
 ESM_EMBED_DIM = 1280
+ESM_CACHE_PATH = 'esm/proteins'
+DEFAULT_ESM_BATCH_SIZE = max(
+    int(os.getenv("CATPRED_ESM_BATCH_SIZE", "1" if PROTEIN_EMBED_USE_CPU else "4")),
+    1,
+)
 
 INT_TO_AA_STR_MAP = {
     0: '<cls>',
@@ -154,14 +159,94 @@ def get_single_esm_repr(protein_str):
     representation = token_representations[0][1 : len(protein_str) + 1]
     return representation
 
+
+def _run_esm_batch(protein_strs):
+    init_esm()
+    model, batch_converter = GLOBAL_VARIABLES['model']
+
+    data = [(f'protein_{index}', protein_str) for index, protein_str in enumerate(protein_strs)]
+    batch_labels, batch_strs, batch_tokens = batch_converter(data)
+
+    if batch_tokens.shape[1] > ESM_MAX_LENGTH:
+        print(f'warning max length protein esm')
+
+    batch_tokens = batch_tokens[:, :ESM_MAX_LENGTH]
+
+    if not PROTEIN_EMBED_USE_CPU:
+        batch_tokens = batch_tokens.to(next(model.parameters()).device)
+
+    with torch.no_grad():
+        results = model(batch_tokens, repr_layers=[33])
+
+    token_representations = results['representations'][33]
+    representations = []
+    max_residue_tokens = ESM_MAX_LENGTH - 1
+    for index, protein_str in enumerate(protein_strs):
+        representation_length = min(len(protein_str), max_residue_tokens)
+        representations.append(
+            token_representations[index][1 : representation_length + 1].detach()
+        )
+    return representations
+
+
+def _run_esm_batch_with_fallback(protein_strs):
+    try:
+        return _run_esm_batch(protein_strs)
+    except RuntimeError as e:
+        if 'out of memory' not in str(e) or len(protein_strs) == 1:
+            raise e
+        print('| WARNING: ran out of memory, retrying smaller ESM batches')
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        midpoint = len(protein_strs) // 2
+        return (
+            _run_esm_batch_with_fallback(protein_strs[:midpoint])
+            + _run_esm_batch_with_fallback(protein_strs[midpoint:])
+        )
+
+
+def get_many_esm_reprs(proteins, device='cpu', batch_size=None):
+    if isinstance(proteins, torch.Tensor):
+        proteins = tensor_to_aa_str(proteins)
+
+    batch_size = max(int(batch_size or DEFAULT_ESM_BATCH_SIZE), 1)
+    ordered_unique_proteins = list(dict.fromkeys(proteins))
+    representations_by_sequence = {}
+    uncached_proteins = []
+
+    for protein_str in ordered_unique_proteins:
+        cached = load_cache_value(
+            path=ESM_CACHE_PATH,
+            cache_key=protein_str,
+            purpose="esm cache entry",
+            map_location='cpu',
+        )
+        if cached is None:
+            uncached_proteins.append(protein_str)
+        else:
+            representations_by_sequence[protein_str] = cached
+
+    for start in range(0, len(uncached_proteins), batch_size):
+        batch = uncached_proteins[start : start + batch_size]
+        batch_representations = _run_esm_batch_with_fallback(batch)
+        for protein_str, representation in zip(batch, batch_representations):
+            save_cache_value(representation, path=ESM_CACHE_PATH, cache_key=protein_str)
+            representations_by_sequence[protein_str] = representation
+
+    return {
+        protein_str: representations_by_sequence[protein_str].to(device)
+        for protein_str in ordered_unique_proteins
+    }
+
+
 def get_esm_repr(proteins, name, device):
     if isinstance(proteins, torch.Tensor):
         proteins = tensor_to_aa_str(proteins)
 
     # Cache by sequence content to avoid collisions when different proteins
     # are accidentally given the same pdb/name identifier.
     _ = name
-    get_protein_repr_fn = cache_fn(get_single_esm_repr, path='esm/proteins')
+    get_protein_repr_fn = cache_fn(get_single_esm_repr, path=ESM_CACHE_PATH)
 
     return calc_protein_representations_with_subunits([proteins], get_protein_repr_fn, device=device)
 
@@ -208,6 +293,7 @@ def get_esm_tokens(protein_str, device):
     'esm': {
         'dim': ESM_EMBED_DIM,
         'fn': get_esm_repr,
+        'batch_fn': get_many_esm_reprs,
         'tokenizer': get_esm_tokens,
     }
 }
diff --git a/catpred/data/utils.py b/catpred/data/utils.py
@@ -58,6 +58,30 @@ def _load_protein_records(protein_records_path: str):
                 "Expected a JSON mapping in .json or .json.gz format."
             ) from plain_err
 
+
+def _populate_missing_esm2_features(
+        protein_records,
+        sequence_feat_getter,
+        sequence_batch_getter=None):
+    if not protein_records:
+        return
+
+    unique_sequences = list(dict.fromkeys(record['seq'] for record in protein_records))
+    if sequence_batch_getter is not None:
+        features_by_sequence = sequence_batch_getter(unique_sequences, device='cpu')
+        for record in protein_records:
+            record['esm2_feats'] = features_by_sequence[record['seq']]
+        return
+
+    features_by_sequence = {}
+    for record in protein_records:
+        sequence = record['seq']
+        if sequence not in features_by_sequence:
+            sequence_features, _ = sequence_feat_getter(sequence, name=record['name'], device='cpu')
+            features_by_sequence[sequence] = sequence_features[0]  # batch dim
+        record['esm2_feats'] = features_by_sequence[sequence]
+
+
 def get_header(path: str) -> List[str]:
     """
     Returns the header of a data CSV file.
@@ -509,7 +533,9 @@ def get_data(path: str,
 
     # Load sequence features
     if not protein_records is None:
-        sequence_feat_getter = get_protein_embedder('esm')['fn']
+        protein_embedder = get_protein_embedder('esm')
+        sequence_feat_getter = protein_embedder['fn']
+        sequence_batch_getter = protein_embedder.get('batch_fn')
                  
     # Load data
     smoke_test_counter = 0
@@ -523,6 +549,7 @@ def get_data(path: str,
 
         all_smiles, all_sequences, all_targets, all_atom_targets, all_bond_targets, all_rows, all_features, all_phase_features, all_constraints_data, all_raw_constraints_data, all_weights, all_gt, all_lt = [], [], [], [], [], [], [], [], [], [], [], [], []
         all_protein_records = []
+        missing_esm2_records = []
         for i, row in enumerate(tqdm(reader)):
             smoke_test_counter+=1
             if args.smoke_test: 
@@ -571,12 +598,7 @@ def get_data(path: str,
                             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
+                        missing_esm2_records.append(protein_record)
                 
             targets, atom_targets, bond_targets = [], [], []
             for column in target_columns:
@@ -650,6 +672,13 @@ def get_data(path: str,
             if len(all_smiles) >= max_data_size:
                 break
 
+        if protein_records is not None:
+            _populate_missing_esm2_features(
+                missing_esm2_records,
+                sequence_feat_getter=sequence_feat_getter,
+                sequence_batch_getter=sequence_batch_getter,
+            )
+
         atom_features = None
         atom_descriptors = None
         if args is not None and args.atom_descriptors is not None:
diff --git a/tests/test_esm_batching.py b/tests/test_esm_batching.py
@@ -0,0 +1,118 @@
+import importlib.util
+import unittest
+from unittest.mock import patch
+
+
+def _has_module(name: str) -> bool:
+    try:
+        return importlib.util.find_spec(name) is not None
+    except ValueError:
+        return False
+
+
+HAS_DATA_DEPS = all(_has_module(name) for name in ("numpy", "pandas", "rdkit", "torch"))
+HAS_ESM_DEPS = HAS_DATA_DEPS and _has_module("esm")
+
+
+@unittest.skipUnless(HAS_DATA_DEPS, "CatPred data dependencies are required")
+class PopulateMissingEsmFeaturesTests(unittest.TestCase):
+    def test_uses_batch_getter_once_for_unique_missing_sequences(self) -> None:
+        from catpred.data import utils
+
+        records = [
+            {"name": "protein_a", "seq": "AAA"},
+            {"name": "protein_b", "seq": "BBB"},
+            {"name": "protein_c", "seq": "AAA"},
+        ]
+        calls = []
+
+        def batch_getter(sequences, device):
+            calls.append((list(sequences), device))
+            return {sequence: f"features-{sequence}" for sequence in sequences}
+
+        utils._populate_missing_esm2_features(
+            records,
+            sequence_feat_getter=None,
+            sequence_batch_getter=batch_getter,
+        )
+
+        self.assertEqual(calls, [(["AAA", "BBB"], "cpu")])
+        self.assertEqual(records[0]["esm2_feats"], "features-AAA")
+        self.assertEqual(records[1]["esm2_feats"], "features-BBB")
+        self.assertEqual(records[2]["esm2_feats"], "features-AAA")
+
+    def test_fallback_getter_deduplicates_sequences(self) -> None:
+        from catpred.data import utils
+
+        records = [
+            {"name": "protein_a", "seq": "AAA"},
+            {"name": "protein_b", "seq": "BBB"},
+            {"name": "protein_c", "seq": "AAA"},
+        ]
+        calls = []
+
+        def single_getter(sequence, name, device):
+            calls.append((sequence, name, device))
+            return ([f"features-{sequence}"], None)
+
+        utils._populate_missing_esm2_features(
+            records,
+            sequence_feat_getter=single_getter,
+            sequence_batch_getter=None,
+        )
+
+        self.assertEqual(
+            calls,
+            [
+                ("AAA", "protein_a", "cpu"),
+                ("BBB", "protein_b", "cpu"),
+            ],
+        )
+        self.assertEqual(records[0]["esm2_feats"], "features-AAA")
+        self.assertEqual(records[1]["esm2_feats"], "features-BBB")
+        self.assertEqual(records[2]["esm2_feats"], "features-AAA")
+
+
+@unittest.skipUnless(HAS_ESM_DEPS, "ESM and torch dependencies are required")
+class BatchedEsmCacheTests(unittest.TestCase):
+    def test_get_many_esm_reprs_skips_cached_and_batches_unique_sequences(self) -> None:
+        import torch
+
+        from catpred.data import esm_utils
+
+        cached = {"AAA": torch.tensor([[1.0]])}
+        saved = []
+        batch_calls = []
+
+        def fake_load(path, cache_key, purpose, map_location):
+            return cached.get(cache_key)
+
+        def fake_save(value, path, cache_key):
+            saved.append((cache_key, value.detach().cpu().clone()))
+
+        def fake_batch(sequences):
+            batch_calls.append(list(sequences))
+            return [torch.tensor([[float(index + 2)]]) for index, _ in enumerate(sequences)]
+
+        with patch("catpred.data.esm_utils.load_cache_value", side_effect=fake_load), patch(
+            "catpred.data.esm_utils.save_cache_value", side_effect=fake_save
+        ), patch(
+            "catpred.data.esm_utils._run_esm_batch_with_fallback",
+            side_effect=fake_batch,
+        ):
+            result = esm_utils.get_many_esm_reprs(
+                ["AAA", "BBB", "CCC", "BBB"],
+                device="cpu",
+                batch_size=2,
+            )
+
+        self.assertEqual(batch_calls, [["BBB", "CCC"]])
+        self.assertEqual([key for key, _ in saved], ["BBB", "CCC"])
+        self.assertEqual(set(result), {"AAA", "BBB", "CCC"})
+        self.assertTrue(torch.equal(result["AAA"], torch.tensor([[1.0]])))
+        self.assertTrue(torch.equal(result["BBB"], torch.tensor([[2.0]])))
+        self.assertTrue(torch.equal(result["CCC"], torch.tensor([[3.0]])))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_postprocess_predictions.py b/tests/test_postprocess_predictions.py
