model.py

"""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 = {}   # predict.py calls us in batches; the same enzyme can recur across batches
        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