File size: 2,731 Bytes
7ca06ca | 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 | """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 # noqa: E402
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()
|