| """Offline CPU inference: ensemble of fine-tuned reaction BERTs (int8) blended with a TF-IDF reaction kNN. | |
| Output: top-5 distinct ECs per unique input, best first (only the first is scored).""" | |
| import argparse | |
| import json | |
| import re | |
| from pathlib import Path | |
| import numpy as np | |
| import pandas as pd | |
| import scipy.sparse as sp | |
| import torch | |
| import torch.nn as nn | |
| from transformers import BertConfig, BertModel | |
| import featurize as F | |
| from inputs import prepare_inputs | |
| SMI_RE = re.compile(r"(\[[^\]]+]|Br?|Cl?|N|O|S|P|F|I|b|c|n|o|s|p|\(|\)|\.|=|#|-|\+|\\|/|:|~|@|\?|>>?|\*|\$|%[0-9]{2}|[0-9])") | |
| MAX_LEN = 256 | |
| TOP = 5 | |
| class Net(nn.Module): | |
| def __init__(self, config, n_cls, n_aux): | |
| super().__init__() | |
| self.bert = BertModel(config, add_pooling_layer=False) | |
| self.head = nn.Linear(config.hidden_size, n_cls) | |
| self.aux = nn.ModuleList([nn.Linear(config.hidden_size, n) for n in n_aux]) | |
| def forward(self, ids, mask): | |
| x = self.bert(input_ids=ids, attention_mask=mask).last_hidden_state | |
| m = mask.unsqueeze(-1).to(x.dtype) | |
| return self.head((x * m).sum(1) / m.sum(1)) | |
| def bert_probs(wdir, rxns, all_idx, quantize): | |
| """Softmax probabilities over the shared class list (all_idx maps this model's classes into it).""" | |
| classes = json.loads((wdir / 'classes.json').read_text()) | |
| tokens = (wdir / 'vocab.txt').read_text().splitlines() | |
| idx = {t: i for i, t in enumerate(tokens)} | |
| state = torch.load(wdir / 'best.pt', map_location='cpu', weights_only=True) | |
| n_aux = [state[f'aux.{i}.weight'].shape[0] for i in range(3)] if 'aux.0.weight' in state else [] | |
| net = Net(BertConfig.from_pretrained(wdir), len(classes), n_aux) | |
| net.load_state_dict(state, strict=False) | |
| net.eval() | |
| if quantize: | |
| net = torch.quantization.quantize_dynamic(net, {nn.Linear}, dtype=torch.qint8) | |
| enc = [[idx['[CLS]']] + [idx.get(t, 1) for t in SMI_RE.findall(r)][:MAX_LEN - 2] + [idx['[SEP]']] for r in rxns] | |
| order = sorted(range(len(enc)), key=lambda i: len(enc[i])) | |
| out = np.zeros((len(rxns), len(all_idx[1])), np.float32) | |
| with torch.inference_mode(): | |
| for a in range(0, len(order), 64): | |
| b = order[a:a + 64] | |
| n = max(len(enc[i]) for i in b) | |
| ids = torch.zeros(len(b), n, dtype=torch.long) | |
| mask = torch.zeros(len(b), n, dtype=torch.long) | |
| for j, i in enumerate(b): | |
| ids[j, :len(enc[i])] = torch.tensor(enc[i]); mask[j, :len(enc[i])] = 1 | |
| p = torch.softmax(net(ids, mask).float(), -1).numpy() | |
| for j, i in enumerate(b): | |
| out[i, all_idx[0]] = p[j] | |
| return out | |
| def knn_probs(kdir, Q, n, all_cls, k=10, power=8): | |
| X = sp.load_npz(kdir / 'X.npz').tocsr() | |
| Y = sp.load_npz(kdir / 'Y.npz').tocsr() | |
| kc = json.loads((kdir / 'classes.json').read_text()) | |
| XT = X.T.tocsr() | |
| out = np.zeros((n, len(kc)), np.float32) | |
| for a in range(0, n, 250): | |
| S = (Q[a:a + 250] @ XT).toarray() | |
| ik = np.argpartition(-S, k, axis=1)[:, :k] | |
| sk = np.maximum(np.take_along_axis(S, ik, 1), 0) ** power | |
| for j in range(k): | |
| out[a:a + 250] += sk[:, j:j + 1] * Y[ik[:, j]].toarray() | |
| out /= np.maximum(out.sum(1, keepdims=True), 1e-9) | |
| pos = {e: i for i, e in enumerate(all_cls)} | |
| full = np.zeros((n, len(all_cls)), np.float32) | |
| for i, e in enumerate(kc): | |
| if e in pos: | |
| full[:, pos[e]] = out[:, i] | |
| return full | |
| def mlp_probs(mdir, Q, all_cls): | |
| """Fingerprint MLP (Linear-ReLU-Linear) on the same TF-IDF features as the kNN.""" | |
| st = torch.load(mdir / 'mlp.pt', map_location='cpu', weights_only=True) | |
| cl = json.loads((mdir / 'classes.json').read_text()) | |
| net = nn.Sequential(nn.Linear(st['0.weight'].shape[1], st['0.weight'].shape[0]), nn.ReLU(), nn.Dropout(0.0), | |
| nn.Linear(st['3.weight'].shape[1], st['3.weight'].shape[0])) | |
| net.load_state_dict({k: v.float() for k, v in st.items()}) | |
| net.eval() | |
| pos = {e: i for i, e in enumerate(all_cls)} | |
| idx = [pos[e] for e in cl] | |
| full = np.zeros((Q.shape[0], len(all_cls)), np.float32) | |
| with torch.inference_mode(): | |
| for a in range(0, Q.shape[0], 512): | |
| full[a:a + 512][:, idx] = torch.softmax(net(torch.from_numpy(Q[a:a + 512].toarray())), -1).numpy() | |
| return full | |
| def predict_file(input_path, output_path): | |
| torch.set_num_threads(2) | |
| root = Path(__file__).parent | |
| w = root / 'weights' | |
| cfg = json.loads((w / 'blend.json').read_text()) | |
| frame = prepare_inputs(pd.read_parquet(input_path)) | |
| rxns = frame.rxn_smiles.tolist() | |
| valid = set(json.loads((root / 'ec_vocabulary.json').read_text())) | |
| kdir = w / 'knn' | |
| all_cls = json.loads((kdir / 'classes.json').read_text()) | |
| pos = {e: i for i, e in enumerate(all_cls)} | |
| probs = 0 | |
| models = sorted(d for d in w.iterdir() if d.is_dir() and (d / 'best.pt').exists()) | |
| for d in models: | |
| cl = json.loads((d / 'classes.json').read_text()) | |
| probs = probs + bert_probs(d, rxns, ([pos[e] for e in cl], all_cls), cfg['quantize']) | |
| probs = probs / len(models) | |
| kw, mw = cfg['knn_weight'], cfg.get('mlp_weight', 0.0) | |
| idf = np.load(kdir / 'idf.npz') | |
| whole, diff = F.raw_features(rxns) | |
| Q = F.combine(whole, diff, idf['whole'], idf['diff']) | |
| blend = (1 - kw - mw) * probs + kw * knn_probs(kdir, Q, len(rxns), all_cls) | |
| if mw: | |
| blend = blend + mw * mlp_probs(w / 'mlp', Q, all_cls) | |
| probs = blend | |
| top = np.argsort(-probs, axis=1)[:, :TOP] | |
| with Path(output_path).open('w') as out: | |
| for key, row in zip(frame.id, top): | |
| p = [all_cls[t] for t in row if all_cls[t] in valid] | |
| out.write(json.dumps({'id': key, 'prediction': p or [all_cls[0]]}) + '\n') | |
| if __name__ == '__main__': | |
| ap = argparse.ArgumentParser(description=__doc__) | |
| ap.add_argument('--input', required=True) | |
| ap.add_argument('--output', required=True) | |
| a = ap.parse_args() | |
| predict_file(a.input, a.output) | |