File size: 6,149 Bytes
307cc2f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | """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)
|