Merge upstream CatPred updates

Author: theproteinbot

catpred/args.py

@@ -943,6 +943,8 @@ class PredictArgs(CommonArgs):
     """Deprecated. Whether to calculate the variance of ensembles as a measure of epistemic uncertainty. If True, the variance is saved as an additional column for each target in the preds_path."""
     individual_ensemble_predictions: bool = False
     """Whether to return the predictions made by each of the individual models rather than the average of the ensemble"""
+    save_uncertainty_components: bool = False
+    """Whether to save aleatoric and epistemic uncertainty variance components when available."""
     # Uncertainty arguments
     uncertainty_method: Literal[
         'mve',

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,26 +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)}')
-        torch.save(out, str(entry_path))
+        save_cache_value(out, path=path, cache_key=cache_str, hash_fn=hash_fn, name=name)
         return out
 
     return inner

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,
     }
 }

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:

catpred/inference/__init__.py

@@ -23,17 +23,21 @@
     "PreparedInputPaths",
     "prepare_prediction_inputs",
     "run_raw_prediction",
+    "run_inprocess_prediction",
     "postprocess_predictions",
     "run_prediction_pipeline",
+    "run_inprocess_prediction_pipeline",
 ]
 
 
 def __getattr__(name: str):
     if name in {
         "prepare_prediction_inputs",
         "run_raw_prediction",
+        "run_inprocess_prediction",
         "postprocess_predictions",
         "run_prediction_pipeline",
+        "run_inprocess_prediction_pipeline",
     }:
         from . import service
 

catpred/inference/backends.py

@@ -37,8 +37,17 @@ def predict(self, request_obj: PredictionRequest, results_dir: str) -> BackendPr
 class LocalInferenceBackend(InferenceBackend):
     name = "local"
 
-    def __init__(self, repo_root: str | None = None) -> None:
+    def __init__(
+        self,
+        repo_root: str | None = None,
+        use_subprocess: bool | None = None,
+    ) -> None:
         self._repo_root = repo_root
+        self._use_subprocess = (
+            _env_flag("CATPRED_LOCAL_SUBPROCESS", default=False)
+            if use_subprocess is None
+            else use_subprocess
+        )
 
     def readiness(self) -> dict[str, Any]:
         root = Path(self._repo_root) if self._repo_root else Path.cwd()
@@ -53,17 +62,27 @@ def readiness(self) -> dict[str, Any]:
             "ready": len(missing) == 0,
             "missing_files": missing,
             "repo_root": str(root),
+            "mode": "subprocess" if self._use_subprocess else "in_process",
         }
 
     def predict(self, request_obj: PredictionRequest, results_dir: str) -> BackendPredictionResult:
-        from .service import run_prediction_pipeline
+        from .service import run_inprocess_prediction_pipeline, run_prediction_pipeline
 
         effective_request = request_obj
         if not request_obj.repo_root and self._repo_root:
             effective_request = replace(request_obj, repo_root=self._repo_root)
 
-        output_file = run_prediction_pipeline(effective_request, results_dir=results_dir)
-        return BackendPredictionResult(backend_name=self.name, output_file=output_file)
+        if self._use_subprocess:
+            output_file = run_prediction_pipeline(effective_request, results_dir=results_dir)
+            mode = "subprocess"
+        else:
+            output_file = run_inprocess_prediction_pipeline(effective_request, results_dir=results_dir)
+            mode = "in_process"
+        return BackendPredictionResult(
+            backend_name=self.name,
+            output_file=output_file,
+            metadata={"mode": mode},
+        )
 
 
 class ModalHTTPInferenceBackend(InferenceBackend):