| """Offline CPU inference: MMseqs2 search -> (time-budgeted) ESM-2 head on least-confident queries -> LightGBM-style reranker (numpy).""" |
| import argparse, gzip, json, os, pickle, shutil, subprocess, sys, tempfile, time |
| from pathlib import Path |
| import numpy as np |
| import pandas as pd |
|
|
| T0 = time.time() |
| ROOT = Path(__file__).resolve().parent |
| WDIR = Path(os.environ.get('EC_WEIGHTS', str(ROOT / 'weights'))) |
| sys.path.insert(0, str(ROOT)) |
| import ecpipe |
|
|
|
|
| def log(*a): |
| print(f'[{time.time() - T0:6.1f}s]', *a, file=sys.stderr, flush=True) |
|
|
|
|
| |
| def diagnostics(): |
| try: |
| import platform |
| cpu = open('/proc/cpuinfo').read() if os.path.exists('/proc/cpuinfo') else '' |
| flags = [f for f in ('avx2', 'sse4_1', 'avx512f') if (' ' + f + ' ') in cpu or (' ' + f + '\n') in cpu] |
| log('python', sys.version.split()[0], '|', platform.platform(), '| cpus', os.cpu_count(), '| uid', os.getuid(), '| cpu flags', flags) |
| for line in open('/proc/mounts').read().splitlines(): |
| p = line.split() |
| if len(p) > 3 and p[1] in ('/', '/tmp', '/dev/shm', '/app', '/var/tmp', '/home'): log('mount', p[1], p[2], p[3][:80]) |
| for d in ('/tmp', '/dev/shm', '/var/tmp', str(Path.home()), str(ROOT), os.getcwd()): |
| try: log('dir', d, 'writable' if os.access(d, os.W_OK) else 'NOT writable', 'free GB', round(shutil.disk_usage(d).free / 1e9, 1)) |
| except Exception as e: log('dir', d, 'err', repr(e)[:80]) |
| for f in ('/sys/fs/cgroup/memory.max', '/sys/fs/cgroup/memory/memory.limit_in_bytes'): |
| if os.path.exists(f): log('cgroup mem', open(f).read().strip()); break |
| mi = open('/proc/meminfo').read().splitlines(); log('meminfo', [x for x in mi if x.startswith(('MemTotal', 'MemAvailable'))]) |
| for p in sorted(WDIR.rglob('*')): |
| if p.is_file() and p.stat().st_size > 1e6: log('weights', str(p.relative_to(WDIR)), round(p.stat().st_size / 1e6), 'MB', oct(p.stat().st_mode & 0o777)) |
| except Exception as e: |
| log('diagnostics failed', repr(e)[:200]) |
|
|
|
|
| |
| |
| |
| MM_CHUNK = int(os.environ.get('EC_MM_CHUNK', 400)) |
| MM_DEADLINE_S = 260 |
| MM2_DEADLINE_S = 330 |
| EMBED_DEADLINE_S = 460 |
| KNN_DEADLINE_S = 530 |
| MOTIF_DEADLINE_S = 520 |
|
|
|
|
| class Runner: |
| """Finds a way to run the bundled static mmseqs binary: in place (chmod if needed) first, copies only as a last resort.""" |
| def __init__(self): |
| self.cmd0, self.fds, self.how = None, (), None |
| try: cpu = open('/proc/cpuinfo').read() |
| except Exception: cpu = '' |
| names = ['mmseqs-avx2', 'mmseqs-sse41'] if ' avx2' in cpu else ['mmseqs-sse41', 'mmseqs-avx2'] |
| for name in names: |
| src = WDIR / 'bin' / name |
| if not src.exists(): continue |
| if self._try(lambda: self._inplace(src), f'{name} in place'): return |
| for d in ('/tmp', '/dev/shm', '/var/tmp', str(Path.home()), str(ROOT), os.getcwd()): |
| if self._try(lambda d=d: self._copy(src, d), f'{name} copied to {d}'): return |
| log('NO RUNNABLE mmseqs binary found') |
|
|
| def _inplace(self, src): |
| if not os.access(src, os.X_OK): os.chmod(src, 0o755) |
| return [str(src)], () |
| def _copy(self, src, d): |
| dst = Path(tempfile.mkdtemp(dir=d)) / 'mmseqs'; shutil.copy(src, dst); os.chmod(dst, 0o755); return [str(dst)], () |
|
|
| def _try(self, make, label): |
| try: |
| cmd0, fds = make() |
| r = subprocess.run(cmd0 + ['version'], capture_output=True, text=True, timeout=60, pass_fds=fds) |
| if r.returncode == 0: |
| self.cmd0, self.fds, self.how = cmd0, fds, label; log('mmseqs runnable:', label, r.stdout.strip()[:30]); return True |
| log('mmseqs not runnable:', label, 'rc', r.returncode, (r.stdout + r.stderr)[-200:].strip()) |
| except Exception as e: |
| log('mmseqs not runnable:', label, repr(e)[:200]) |
| return False |
|
|
| def run(self, args, timeout): |
| r = subprocess.run(self.cmd0 + args, capture_output=True, text=True, timeout=max(20, timeout), pass_fds=self.fds) |
| if r.returncode: raise RuntimeError(f'rc {r.returncode}: ' + (r.stdout + r.stderr)[-700:]) |
| return r |
|
|
|
|
| def pick_workdir(need_gb=0.3): |
| for d in ('/tmp', '/dev/shm', '/var/tmp', str(Path.home()), str(ROOT), os.getcwd()): |
| try: |
| if os.access(d, os.W_OK) and shutil.disk_usage(d).free > need_gb * 1e9: return tempfile.mkdtemp(dir=d) |
| except Exception: pass |
| return tempfile.mkdtemp() |
|
|
|
|
| RUNNER = None |
|
|
|
|
| def parse_m8(path): |
| if not os.path.exists(path) or os.path.getsize(path) == 0: return {} |
| m8 = pd.read_csv(path, sep='\t', header=None, names=['q', 't', 'bits', 'pid', 'qc', 'tc']) |
| if os.environ.get('EC_HIT_CAP'): m8 = m8[m8.pid < float(os.environ['EC_HIT_CAP'])] |
| m8 = m8.sort_values(['q', 'bits'], ascending=[True, False], kind='stable') |
| q = m8.q.values; cut = np.flatnonzero(np.diff(q)) + 1; starts = np.r_[0, cut]; ends = np.r_[cut, len(q)] |
| cols = [m8[c].values for c in ('t', 'bits', 'pid', 'qc', 'tc')] |
| return {int(q[s]): tuple(c[s:e] for c in cols) for s, e in zip(starts, ends)} |
|
|
|
|
| def mm_hits(seqs, meta, deadline, idxs=None, s=None, max_seqs=None, chunk=None, tag='mmseqs', budget=None): |
| """-> {query index: (target idx, bits, pident, qcov, tcov) sorted by bits desc}. Partial results are kept if later chunks fail (reasons are logged).""" |
| hits = {} |
| idxs = list(range(len(seqs))) if idxs is None else list(idxs) |
| s = meta['mm_s'] if s is None else s; max_seqs = meta['mm_max_seqs'] if max_seqs is None else max_seqs; chunk = chunk or MM_CHUNK |
| try: |
| rn = RUNNER if RUNNER is not None else Runner() |
| globals()['RUNNER'] = rn |
| if rn.cmd0 is None: return {} |
| db = str(WDIR / 'mm' / 'targetDB'); fails = 0; t_all = time.time(); degraded = False |
| for st in range(0, len(idxs), chunk): |
| if time.time() > deadline: log(f'{tag} deadline: stopping at {st}/{len(idxs)}'); break |
| if budget and st >= 0.25 * len(idxs) and not degraded and (time.time() - t_all) / (st / len(idxs)) > budget: |
| degraded = True; s = min(float(s), 2.5); max_seqs = min(int(max_seqs), 50) |
| log(f'{tag}: projected {(time.time() - t_all) / (st / len(idxs)):.0f}s > budget {budget}s, switching to s={s} max-seqs {max_seqs} for the rest') |
| work = pick_workdir(); t = time.time(); rem = lambda: deadline - time.time() + 60 |
| try: |
| with open(work + '/q.fa', 'w') as f: |
| for i in idxs[st:st + chunk]: f.write('>%d\n%s\n' % (i, seqs[i])) |
| rn.run(['createdb', work + '/q.fa', work + '/qdb', '-v', '1'], rem()) |
| rn.run(['search', work + '/qdb', db, work + '/res', work + '/tmp', '-s', str(s), '--max-seqs', str(max_seqs), |
| '--threads', '2', '--split-memory-limit', '10G', '-v', '1'], rem()) |
| rn.run(['convertalis', work + '/qdb', db, work + '/res', work + '/res.m8', '--format-output', 'query,target,bits,pident,qcov,tcov', '-v', '1'], rem()) |
| hits.update(parse_m8(work + '/res.m8')); fails = 0 |
| except Exception as e: |
| fails += 1; log(f'MMSEQS chunk at query {st} FAILED ({fails}):', repr(e)[:800]) |
| if fails >= 3: log('giving up on MMseqs2'); break |
| finally: |
| shutil.rmtree(work, ignore_errors=True) |
| if (st // chunk) % 4 == 0: log(f'{tag} chunk {st // chunk} took {time.time() - t:.1f}s; hits so far {len(hits)}') |
| log(f'{tag} total {time.time() - t_all:.0f}s') |
| except Exception as e: |
| log('MMSEQS FAILED, continuing with what we have:', repr(e)[:900]) |
| return hits |
|
|
|
|
| |
| def embed_queries(seqs, order, meta, deadline, sess=None): |
| """Mean-pooled ESM embeddings (raw, not normalised) for queries in `order` until `deadline`. -> {query index: vector}""" |
| if sess is None: |
| import onnxruntime as ort |
| so = ort.SessionOptions(); so.intra_op_num_threads = 2; so.inter_op_num_threads = 1 |
| sess = ort.InferenceSession(str(WDIR / 'esm' / 'model.onnx'), so, providers=['CPUExecutionProvider']) |
| v = meta['vocab']; unk, cls, eos, pad = v['<unk>'], v['<cls>'], v['<eos>'], v['<pad>'] |
| maxlen, budget = meta['maxlen'], meta['tok_budget'] |
| out = {} |
| for st in range(0, len(order), 96): |
| if time.time() > deadline: break |
| blk = sorted(order[st:st + 96], key=lambda i: len(seqs[i])) |
| i = 0 |
| while i < len(blk): |
| j = i + 1 |
| while j < len(blk) and (j - i + 1) * (min(len(seqs[blk[j]]), maxlen) + 2) <= budget: j += 1 |
| lst = [[cls] + [v.get(c, unk) for c in seqs[k][:maxlen]] + [eos] for k in blk[i:j]] |
| L = max(len(x) for x in lst); ids = np.full((len(lst), L), pad, np.int64); am = np.zeros((len(lst), L), np.int64) |
| for r, x in enumerate(lst): ids[r, :len(x)] = x; am[r, :len(x)] = 1 |
| hid = sess.run(None, {'input_ids': ids, 'attention_mask': am})[0] |
| for r, x in enumerate(lst): |
| if len(x) > 2: out[blk[i + r]] = hid[r, 1:len(x) - 1].mean(0) |
| i = j |
| return out |
|
|
|
|
| def norm_emb(E, hw): |
| X = E.astype(np.float32) - hw['MU']; return X / np.maximum(np.linalg.norm(X, axis=1, keepdims=True), 1e-8) |
|
|
|
|
| def knn_search(E, keys, hw, path, K=20): |
| """Cosine nearest neighbours of embedded queries among all training embeddings. -> {query index: (positions, sims)}""" |
| out = {} |
| try: |
| T = np.load(path).astype(np.float32) |
| for st in range(0, len(keys), 256): |
| if time.time() > T0 + KNN_DEADLINE_S: log(f'kNN deadline: {st}/{len(keys)} done'); break |
| ks = keys[st:st + 256]; S = norm_emb(np.stack([E[k] for k in ks]), hw) @ T.T |
| ix = np.argpartition(-S, K, axis=1)[:, :K]; sv = np.take_along_axis(S, ix, 1); o = np.argsort(-sv, axis=1) |
| for r, k in enumerate(ks): out[k] = (ix[r][o[r]], sv[r][o[r]]) |
| except Exception as e: |
| log('kNN failed, continuing without it:', repr(e)[:400]); out = {} |
| return out |
|
|
|
|
| def head2_and_protos(E, hw2, pw, ctx): |
| """-> (within-family probabilities, prototype cosine similarities), each (n, n_ECs).""" |
| X = norm_emb(E, hw2) |
| Z = ecpipe.gelu(X @ hw2['W1'].T + hw2['b1']) @ hw2['W2'].T + hw2['b2'] |
| H2 = ecpipe.family_softmax(Z, ctx.l3_id, ctx.nl3).astype(np.float32) |
| Y = ecpipe.gelu(X @ pw['W1'].T + pw['b1']) @ pw['W2'].T + pw['b2'] |
| Y = Y / np.maximum(np.linalg.norm(Y, axis=1, keepdims=True), 1e-8) |
| PS = Y @ pw['P'].astype(np.float32).T; PS[:, ~pw['has']] = np.nan |
| return H2, PS |
|
|
|
|
| def head_probs(E, hw): |
| X = norm_emb(E, hw) |
| return ecpipe.softmax(ecpipe.gelu(X @ hw['W1'].T + hw['b1']) @ hw['W2'].T + hw['b2']) |
|
|
|
|
| |
| def main(input_path, output_path): |
| from inputs import prepare_inputs |
| diagnostics() |
| meta = json.load(open(WDIR / 'meta.json')) |
| raw = pd.read_parquet(input_path) |
| frame = prepare_inputs(raw) |
| if 'sequence' not in frame.columns: |
| frame = frame.merge(raw[['id', 'sequence']].drop_duplicates('id'), on='id') |
| ids, seqs = frame.id.tolist(), frame.sequence.tolist(); n = len(seqs) |
| log(f'{n} unique queries') |
| with gzip.open(WDIR / 'labels.pkl.gz', 'rb') as f: d = pickle.load(f) |
| ctx = ecpipe.Ctx(d['labels'], d['ecs'], d['prior']) |
| ctx.use_unseen = bool(meta.get('use_unseen')); log(f'unseen-EC candidates: {ctx.use_unseen} ({sum(len(v) for v in ctx.unseen_by_l3.values())} vocabulary ECs without training sequences)') |
| if meta.get('use_nb'): |
| z = np.load(WDIR / 'nb.npz'); ctx.set_nb(z['E'], z['S']); log(f'neighbourhood table {z["E"].shape}') |
| sj = json.load(open(WDIR / 'stacker.json')) |
| forests = [ecpipe.Forest(t) for t in sj['forests']] if isinstance(sj, dict) else [ecpipe.Forest(sj)] |
|
|
| hits = mm_hits(seqs, meta, deadline=T0 + meta.get('mm_deadline_s', MM_DEADLINE_S), chunk=(MM_CHUNK if meta['mm_max_seqs'] <= 100 else 200), budget=meta.get('mm_budget_s', 230)) |
| log(f'mmseqs hits for {len(hits)}/{n} queries') |
| m2 = meta.get('mm2') |
| if m2: |
| weak = [i for i in range(n) if (hits[i][1][0] if i in hits else 0.0) < m2['tb']] |
| log(f'second pass for {len(weak)} weak queries (s={m2["s"]}, max-seqs {m2["max_seqs"]})') |
| h2 = mm_hits(seqs, meta, T0 + meta.get('mm2_deadline_s', MM2_DEADLINE_S), idxs=weak, s=m2['s'], max_seqs=m2['max_seqs'], chunk=m2.get('chunk', 200), tag='mmseqs2') |
| hits.update(h2); log(f'second pass gave hits for {len(h2)}/{len(weak)}; total hits {len(hits)}/{n}') |
|
|
| H, H2, PS = {}, {}, {} |
| if meta.get('use_plm') and (WDIR / 'esm' / 'model.onnx').exists(): |
| try: |
| topbits = np.array([hits[i][1][0] if i in hits else 0.0 for i in range(n)]) |
| order = [int(i) for i in np.argsort(topbits, kind='stable') if topbits[i] < meta['embed_max_topbits']] |
| E = embed_queries(seqs, order, meta, deadline=T0 + EMBED_DEADLINE_S) |
| log(f'embedded {len(E)}/{len(order)} candidate queries') |
| hw = dict(np.load(WDIR / 'head.npz')) |
| keys = list(E) |
| for st in range(0, len(keys), 500): |
| ks = keys[st:st + 500]; P = head_probs(np.stack([E[k] for k in ks]), hw) |
| for k, p in zip(ks, P): H[k] = p |
| if meta.get('use_hier'): |
| hw2 = dict(np.load(WDIR / 'head2.npz')); pw = dict(np.load(WDIR / 'proto.npz')); pw['has'] = pw['has'].astype(bool) |
| for st in range(0, len(keys), 500): |
| ks = keys[st:st + 500]; A2, A3 = head2_and_protos(np.stack([E[k] for k in ks]), hw2, pw, ctx) |
| for k, a2, a3 in zip(ks, A2, A3): H2[k] = a2; PS[k] = a3 |
| ctx.use_hier_cands = bool(meta.get('use_hier_cands')) |
| except Exception as e: |
| log('PLM path failed, continuing without it:', repr(e)[:500]); H = {}; H2 = {}; PS = {}; E = {}; hw = None |
| log(f'head probabilities for {len(H)} queries | within-family head and prototypes for {len(H2)}') |
| KNN = {} |
| if H and (WDIR / 'knn_emb.npy').exists(): |
| KNN = knn_search(E, list(H), hw, WDIR / 'knn_emb.npy'); log(f'kNN for {len(KNN)} queries') |
|
|
| MOT = {} |
| if meta.get('use_motif') and (WDIR / 'motif.npz').exists(): |
| try: |
| mfs = meta.get('motif_files') or [dict(file='motif.npz', ks=[3, 4, 5, 6])] |
| mws = [(dict(np.load(WDIR / m['file'])), tuple(m['ks'])) for m in mfs]; t_m = time.time() |
| for st in range(0, n, 1000): |
| if time.time() - T0 > MOTIF_DEADLINE_S: log(f'motif deadline: stopping at {st}/{n}'); break |
| P = np.mean([ecpipe.motif_probs(seqs[st:st + 1000], mw, ks) for mw, ks in mws], axis=0) |
| for k, p in enumerate(P): MOT[st + k] = p |
| ctx.use_motif_cands = bool(meta.get('use_motif_cands')) |
| log(f'motif scores for {len(MOT)}/{n} queries in {time.time() - t_m:.0f}s') |
| except Exception as e: |
| log('motif path failed, continuing without it:', repr(e)[:500]); MOT = {} |
| rows = [] |
| n2c = 3 if meta.get('nb_cands', True) else 0 |
| for i in range(n): rows += ecpipe.query_rows(i, hits.get(i), len(seqs[i]), H.get(i), ctx, knn=KNN.get(i), h2=H2.get(i), ps=PS.get(i), mot=MOT.get(i), TOPN2C=n2c) |
| log(f'{len(rows)} candidate rows') |
| best = {} |
| if rows: |
| A = np.array(rows, dtype=np.float64) |
| cols = [ecpipe.COLS.index(f) for f in meta['feats']]; pred = np.mean([f.predict(A[:, cols]) for f in forests], axis=0) |
| df = pd.DataFrame(dict(q=A[:, 0].astype(int), e=A[:, 1].astype(int), p=pred)) |
| best = df.loc[df.groupby('q').p.idxmax()].set_index('q').e.to_dict() |
| vp = ROOT / 'ec_vocabulary.json' |
| vocab = set(json.loads(vp.read_text())) if vp.exists() else None |
| with Path(output_path).open('w') as f: |
| for i, qid in enumerate(ids): |
| ec = ctx.ecs[int(best[i])] if i in best else ctx.top_ec |
| if vocab is not None and ec not in vocab: ec = ctx.top_ec if ctx.top_ec in vocab else sorted(vocab)[0] |
| f.write(json.dumps({'id': qid, 'prediction': [ec]}) + '\n') |
| log('fallback (no candidates) for', sum(1 for i in range(n) if i not in best), 'of', n, 'queries; done') |
|
|
|
|
| if __name__ == '__main__': |
| ap = argparse.ArgumentParser(description=__doc__) |
| ap.add_argument('--input', required=True); ap.add_argument('--output', required=True) |
| a = ap.parse_args(); main(a.input, a.output) |
|
|
