| """LightGBM on RDKit substrate features, sequence composition, and ESM-2 (8M) embeddings.""" |
| import json |
| from pathlib import Path |
| import re |
|
|
| import numpy as np |
|
|
| from featurize import composition_features, esm_embeddings, smiles_features |
|
|
| WEIGHTS = Path(__file__).parent / 'weights' |
|
|
|
|
| class Model: |
| def __init__(self, weights=WEIGHTS, threads=2): |
| import lightgbm as lgb |
|
|
| self.meta = json.loads((Path(weights) / 'meta.json').read_text()) |
| self.threads = threads |
| self.boosters = [lgb.Booster(model_file=str(Path(weights) / f'lgbm_{t}.txt')) for t in self.meta['targets']] |
|
|
| def predict(self, frame): |
| """Return an (n, 2) array of log10_kcat, log10_km for a frame with smiles and sequence.""" |
| blocks = [] |
| groups = self.meta['groups'] |
| if 'smiles' in groups: |
| blocks.append(self._per_unique(frame.smiles, smiles_features)) |
| if 'comp' in groups: |
| blocks.append(self._per_unique(frame.sequence, composition_features)) |
| if 'esm' in groups: |
| blocks.append(self._per_unique(frame.sequence, lambda s: esm_embeddings(s, threads=self.threads))) |
| x = np.hstack(blocks) |
| return np.column_stack([b.predict(x, num_threads=self.threads) for b in self.boosters]) |
|
|
| @staticmethod |
| def _per_unique(values, fn): |
| keys, inverse = np.unique(values.to_numpy(), return_inverse=True) |
| return fn(list(keys))[inverse] |
|
|
|
|
| def validate_rows(rows, max_rows=256): |
| if not isinstance(rows, list) or not 1 <= len(rows) <= max_rows: |
| raise ValueError(f'Provide between 1 and {max_rows} rows.') |
| ids = set() |
| for row in rows: |
| if not isinstance(row, dict) or set(row) != {'id', 'smiles', 'sequence'}: |
| raise ValueError('Each row must contain exactly id, smiles, sequence.') |
| if not isinstance(row['id'], str) or not re.fullmatch(r'[A-Za-z0-9_-]{1,80}', row['id']): |
| raise ValueError('Invalid row ID.') |
| if row['id'] in ids: |
| raise ValueError('Duplicate row ID.') |
| ids.add(row['id']) |
| for key, limit in [('sequence', 50000), ('smiles', 20000)]: |
| if not isinstance(row[key], str) or not 1 <= len(row[key]) <= limit: |
| raise ValueError(f'Invalid {key}.') |
| return rows |
|
|
