| """ESM-2 and hierarchical-head scoring, run as a separate process. |
| |
| Isolation is deliberate. PyTorch and LightGBM each bring their own OpenMP runtime, and loading both |
| into one interpreter aborts the process on some platforms. A child process that never imports |
| LightGBM removes the conflict, and frees the model's memory when it exits. |
| |
| Reads a JSON job from argv[1] and writes {query: {ec: [cls, l3, cen, knn]}} to argv[2]. Exits 0 with |
| whatever it finished when the deadline passes, so a slow machine degrades rather than fails. |
| """ |
| import json |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
|
|
| sys.path.insert(0, str(Path(__file__).parent)) |
| from embmodel import evidence, load_esm, load_head |
|
|
|
|
| def main(): |
| job = json.loads(Path(sys.argv[1]).read_text()) |
| out_path = Path(sys.argv[2]) |
| cfg, weights = job["cfg"], Path(job["weights"]) |
| seqs, order, cands = job["seqs"], job["order"], job["cands"] |
| budget = job["budget"] |
| start = time.time() |
|
|
| torch.set_num_threads(job.get("threads", 2)) |
| model, alphabet = load_esm(weights / cfg["esm_file"]) |
| net, ecs, l3 = load_head(weights / cfg["head_file"]) |
| C = torch.tensor(np.load(weights / cfg["centroids_file"])) |
| Z = F.normalize(torch.tensor(np.load(weights / cfg["z_file"]).astype(np.float32)), dim=-1) |
| ref_ecs = json.loads((weights / cfg["z_ecs_file"]).read_text()) |
|
|
| bc = alphabet.get_batch_converter() |
| layer = model.num_layers |
| max_len, max_tok = cfg["max_len"], cfg["max_tokens"] |
|
|
| out = {} |
| p = 0 |
| with torch.inference_mode(): |
| while p < len(order) and time.time() - start < budget: |
| q = p + 1 |
| top = min(len(seqs[str(order[p])]), max_len) |
| while q < len(order): |
| top2 = max(top, min(len(seqs[str(order[q])]), max_len)) |
| if (q - p + 1) * top2 > max_tok: |
| break |
| top = top2 |
| q += 1 |
| batch = order[p:q] |
| _, _, toks = bc([(str(i), seqs[str(i)][:max_len]) for i in batch]) |
| rep = model(toks, repr_layers=[layer])["representations"][layer] |
| mask = (toks != alphabet.padding_idx) & (toks != alphabet.cls_idx) & (toks != alphabet.eos_idx) |
| x = (rep * mask.unsqueeze(-1)).sum(1) / mask.sum(1, keepdim=True) |
| ev = evidence(net, x, C, ecs, l3, None, Z, ref_ecs, [cands[str(i)] for i in batch]) |
| for i, d in zip(batch, ev): |
| out[str(i)] = d |
| p = q |
| out_path.write_text(json.dumps(out)) |
| print(f"embedded {len(out)}/{len(order)} queries in {time.time() - start:.1f}s", file=sys.stderr) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
