| """Offline CPU inference for Leak CURBER enzyme classification, version 2. |
| |
| Pipeline |
| -------- |
| 1. Safety net: a quick DIAMOND search, so every input has homology evidence early on. |
| 2. Main search: DIAMOND `--very-sensitive` with full alignments, read through a pipe query by query |
| and turned into features on arrival, including EC-specific residue profiles. |
| 3. ESM-2 and the hierarchical head, in a child process, for queries with weak homology. |
| 4. A bag of LightGBM LambdaRank models ranks each (query, candidate EC) pair. |
| 5. Anything with no evidence falls back to the most frequent training EC. |
| |
| Running inside the rules |
| ------------------------ |
| The task says no file written by predict.py may exceed 16 MiB, and a limit like that may be enforced |
| by the OS, killing the process at the first large write. So the DIAMOND index is built when the |
| image is built (see Dockerfile), `--hit-membuf` keeps DIAMOND's intermediate hits in memory, and |
| results arrive through stdout. Nothing large touches the disk. |
| |
| Every stage has a deadline and a failure path. A slow or failing stage costs accuracy on the queries |
| it did not reach; it never costs the run, because a valid output file is always written. |
| """ |
| import argparse |
| import gzip |
| import json |
| import os |
| import shutil |
| import subprocess |
| import sys |
| import tempfile |
| import threading |
| import time |
| from pathlib import Path |
|
|
| import lightgbm as lgb |
| import numpy as np |
| import pandas as pd |
|
|
| sys.path.insert(0, str(Path(__file__).parent)) |
| from features import OUTFMT, assemble, query_features, read_hits |
| from inputs import prepare_inputs |
|
|
| ROOT = Path(__file__).parent |
| WEIGHTS = ROOT / "weights" |
| T0 = time.time() |
| |
| BUDGET = float(os.environ.get("EC_TIME_BUDGET", 540)) |
| THREADS = int(os.environ.get("EC_THREADS", 2)) |
|
|
|
|
| def elapsed(): |
| return time.time() - T0 |
|
|
|
|
| def log(msg): |
| print(f"[{elapsed():6.1f}s] {msg}", flush=True) |
|
|
|
|
| def diamond_bin(): |
| return shutil.which("diamond") or str(ROOT / "bin" / "diamond") |
|
|
|
|
| def database(tmp): |
| """The image carries a prebuilt index. Building here is only for runs outside the image.""" |
| built = Path(os.environ.get("EC_DMND", WEIGHTS / "train.dmnd")) |
| if built.exists(): |
| return built.with_suffix("") |
| fa = Path(tmp) / "db.fa" |
| with gzip.open(WEIGHTS / "train_db.fa.gz", "rb") as src, open(fa, "wb") as dst: |
| shutil.copyfileobj(src, dst) |
| db = Path(tmp) / "db" |
| subprocess.run([diamond_bin(), "makedb", "--in", str(fa), "-d", str(db), "--threads", str(THREADS), |
| "--quiet"], check=True) |
| fa.unlink() |
| return db |
|
|
|
|
| def blast(query, db, tmp, columns, k, mode, evalue): |
| cmd = [diamond_bin(), "blastp", "-q", str(query), "-d", str(db), "--outfmt", "6", *columns, |
| "-k", str(k), "--max-hsps", "1", "--evalue", str(evalue), "--threads", str(THREADS), |
| "--quiet", "--tmpdir", str(tmp), "--hit-membuf"] |
| if mode: |
| cmd.append(f"--{mode}") |
| return subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, |
| bufsize=1 << 20) |
|
|
|
|
| def stream(proc, deadline, k, keep_aln, on_query): |
| """Feed hits to `on_query` until DIAMOND ends or the deadline kills it. Partial results stay.""" |
| timer = threading.Timer(max(deadline - elapsed(), 0.1), proc.kill) |
| timer.start() |
| done = 0 |
| try: |
| for q, rows in read_hits(proc.stdout, k, keep_aln): |
| on_query(q, rows) |
| done += 1 |
| finally: |
| timer.cancel() |
| proc.stdout.close() |
| proc.wait() |
| return done, proc.returncode |
|
|
|
|
| def embed(seqs, order, cands, cfg, tmp, budget): |
| """Head evidence from an isolated child process; see embed_worker.py for why.""" |
| job = Path(tmp) / "emb_job.json" |
| res = Path(tmp) / "emb_out.json" |
| job.write_text(json.dumps({"seqs": {str(i): seqs[i] for i in order}, |
| "order": [int(i) for i in order], |
| "cands": {str(i): cands[i] for i in order}, "cfg": cfg, |
| "weights": str(WEIGHTS), "budget": budget, "threads": THREADS})) |
| try: |
| proc = subprocess.run([sys.executable, str(ROOT / "embed_worker.py"), str(job), str(res)], |
| capture_output=True, text=True, timeout=budget + 15) |
| except subprocess.TimeoutExpired: |
| log("embedding worker exceeded its time limit") |
| return {} |
| if proc.returncode != 0 or not res.exists(): |
| log(f"embedding worker failed (rc={proc.returncode}): {proc.stderr.strip()[-300:]}") |
| return {} |
| log(proc.stderr.strip() or "embedding worker finished") |
| raw = json.loads(res.read_text()) |
| return {int(k): {ec: tuple(v) for ec, v in d.items()} for k, d in raw.items()} |
|
|
|
|
| def rank(F, cfg): |
| scores = np.zeros(len(F)) |
| for name in cfg["rerankers"]: |
| booster = lgb.Booster(model_file=str(WEIGHTS / name)) |
| p = booster.predict(F, num_threads=THREADS) |
| scores += (p - p.mean()) / (p.std() + 1e-9) |
| return scores |
|
|
|
|
| def evidence_and_rank(seqs, cfg, db_ecs, ec_count, ec_prior, vocab, tmp): |
| n = len(seqs) |
| scanned = [None] * n |
| fa = Path(tmp) / "q.fa" |
| with open(fa, "w") as f: |
| for i, s in enumerate(seqs): |
| f.write(f">{i}\n{s}\n") |
| db = database(tmp) |
| log("DIAMOND database ready") |
|
|
| def on_query(q, rows): |
| scanned[q] = query_features(rows, seqs[q], db_ecs, ec_count, ec_prior, cfg) |
|
|
| plain = OUTFMT[:9] |
| try: |
| proc = blast(fa, db, tmp, plain, cfg["k_quick"], "fast", cfg["evalue"]) |
| got, rc = stream(proc, BUDGET * 0.35, cfg["k_quick"], 0, on_query) |
| log(f"safety-net search covered {got}/{n} queries (rc={rc})") |
| except Exception as exc: |
| log(f"safety-net search failed: {exc}") |
|
|
| try: |
| proc = blast(fa, db, tmp, OUTFMT, cfg["k"], cfg["sensitivity"], cfg["evalue"]) |
| got, rc = stream(proc, BUDGET - cfg["reserve"], cfg["k"], cfg["prof_hits"], on_query) |
| log(f"main search covered {got}/{n} queries (rc={rc})") |
| except Exception as exc: |
| log(f"main search failed: {exc}") |
|
|
| |
| best = np.array([r[2]["q_level"][0] if r else 0.0 for r in scanned]) |
| weak = [i for i in np.where(best < cfg["emb_bits_gate"])[0].tolist()] |
| weak.sort(key=lambda i: (best[i] >= 80, len(seqs[i]))) |
| log(f"{len(weak)} queries below the embedding gate") |
| emb = {} |
| room = BUDGET - 20 - elapsed() |
| if weak and room > 30: |
| try: |
| emb = embed(seqs, weak, {i: (scanned[i][0] if scanned[i] else []) for i in weak}, |
| cfg, tmp, room) |
| except Exception as exc: |
| log(f"embedding stage skipped: {exc}") |
|
|
| q, ec, F = assemble(scanned, n, emb, ec_prior, ec_count, seqs, cfg["extra_from_emb"]) |
| log(f"{len(q)} candidate rows") |
| scores = rank(F, cfg) if len(F) else np.zeros(0) |
| best_ec = {} |
| for qi, e, sc in zip(q.tolist(), ec.tolist(), scores.tolist()): |
| if e in vocab and (qi not in best_ec or sc > best_ec[qi][1]): |
| best_ec[qi] = (e, sc) |
| log(f"ranked predictions for {len(best_ec)} queries") |
| return best_ec |
|
|
|
|
| def predict_file(input_path, output_path): |
| cfg = json.loads((WEIGHTS / "config.json").read_text()) |
| vocab = set(json.loads((ROOT / "ec_vocabulary.json").read_text())) |
| frame = prepare_inputs(pd.read_parquet(input_path)) |
| seqs = frame.sequence.tolist() |
| log(f"{len(seqs)} unique inputs") |
|
|
| db_ecs = json.loads((WEIGHTS / "db_ecs.json").read_text()) |
| ec_prior = json.loads((WEIGHTS / "ec_prior.json").read_text()) |
| ec_count = json.loads((WEIGHTS / "ec_count.json").read_text()) |
| fallback = [e for e in json.loads((WEIGHTS / "prior.json").read_text()) if e in vocab][:1] |
|
|
| tmp = tempfile.mkdtemp(dir="/tmp") |
| best_ec = {} |
| try: |
| best_ec = evidence_and_rank(seqs, cfg, db_ecs, ec_count, ec_prior, vocab, tmp) |
| except Exception as exc: |
| log(f"pipeline failed, writing fallback predictions: {exc!r}") |
| finally: |
| shutil.rmtree(tmp, ignore_errors=True) |
| with open(output_path, "w") as f: |
| for i, key in enumerate(frame.id): |
| pred = [best_ec[i][0]] if i in best_ec else fallback |
| f.write(json.dumps({"id": key, "prediction": pred}) + "\n") |
| log(f"wrote {len(seqs)} predictions to {output_path}") |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--input", required=True) |
| parser.add_argument("--output", required=True) |
| args = parser.parse_args() |
| predict_file(args.input, args.output) |
|
|
