| """LightGBM (+ optional ExtraTrees and two-tower network blends) on enzyme + substrate features + substrate prior. Loaded by predict.py.""" |
| import json |
| import re |
| from pathlib import Path |
|
|
| import lightgbm as lgb |
| import numpy as np |
| import pandas as pd |
|
|
| from features import TARGETS, base_features, esm_embed, load_esm, prior_features |
|
|
| WEIGHTS = Path(__file__).parent / 'weights' |
|
|
|
|
| class Model: |
| def __init__(self, weights=WEIGHTS): |
| w = Path(weights) |
| self.meta = json.loads((w / 'meta.json').read_text()) |
| self.models = [lgb.Booster(model_file=str(w / f'lgb_{t}.txt')) for t in TARGETS] |
| self.prior = pd.read_parquet(w / 'prior.parquet').set_index('smiles') |
| self.global_mean = self.meta['global_mean'] |
| self.esm = None |
| if self.meta.get('esm'): |
| import torch |
| torch.set_num_threads(2) |
| e = self.meta['esm'] |
| self.esm = load_esm(str(w / e['dir']), quantize=e['quantize'], layer=e.get('layer')) |
| self.esm_cache = {} |
| self.et = None |
| if self.meta.get('et'): |
| import joblib |
| self.et = joblib.load(w / self.meta['et']['file']) |
| self.et.n_jobs = 2 |
| self.nets = [] |
| if self.meta.get('nn'): |
| import torch |
| from nn_model import TwoTower |
| n = self.meta['nn'] |
| self.q = dict(np.load(w / 'nn' / 'quantiles.npz')) |
| for info in n['seeds']: |
| net = TwoTower(n['de'], n['ds'], p=n.get('p', 0.3), act=n.get('act', 'gelu')) |
| net.load_state_dict(torch.load(w / 'nn' / f"net_{info['seed']}.pt", map_location='cpu')) |
| self.nets.append((net.eval(), np.array(info['mu']), np.array(info['sd']))) |
|
|
| def prepare(self, sequences): |
| """Embed every unique enzyme of the whole input once (better length batching than 128-row chunks).""" |
| if self.esm is not None: |
| new = [s for s in pd.unique(sequences) if s not in self.esm_cache] |
| if new: |
| self.esm_cache.update(zip(new, esm_embed(new, *self.esm, batch_tokens=1000))) |
|
|
| def predict(self, rows): |
| seqs = np.array([r['sequence'] for r in rows], dtype=object) |
| smis = np.array([r['smiles'] for r in rows], dtype=object) |
| base = base_features(seqs, smis) |
| blocks = [base] |
| if self.esm is not None: |
| self.prepare(seqs) |
| E = np.stack([self.esm_cache[s] for s in seqs]) |
| blocks.append(E) |
| P = prior_features(self.prior, smis, self.global_mean) |
| blocks.append(P) |
| X = np.hstack(blocks) |
| assert X.shape[1] == self.meta['n_features'], (X.shape, self.meta['n_features']) |
| p = np.column_stack([m.predict(X, num_threads=2) for m in self.models]) |
| if self.et is not None: |
| a = self.meta['et']['weight'] |
| p = (1 - a) * p + a * self.et.predict(np.nan_to_num(X, nan=-999).astype(np.float32)) |
| if self.nets: |
| p = self.meta['nn']['lgb_weight'] * p + (1 - self.meta['nn']['lgb_weight']) * self.predict_nn(base, E, P) |
| return np.nan_to_num(p, nan=0.0, posinf=0.0, neginf=0.0).tolist() |
|
|
| def predict_nn(self, base, E, P): |
| import torch |
| from nn_model import quantile_normal, split_blocks |
| enz, sub = split_blocks(base, E, P) |
| e = torch.tensor(quantile_normal(enz, self.q['enz_q'], self.q['enz_r'])) |
| s = torch.tensor(quantile_normal(sub, self.q['sub_q'], self.q['sub_r'])) |
| with torch.no_grad(): |
| return np.mean([net(e, s).numpy() * sd + mu for net, mu, sd in self.nets], axis=0) |
|
|
|
|
| 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 |
|
|
