| """Hierarchical CLEAN-style head (shared by training and inference) and offline ESM-2 loader. |
| |
| v1 used a flat softmax. HIT-EC's central idea is that an EC number is a path, not a label, so this |
| head predicts all four levels from one trunk. That keeps level-3 evidence usable when the exact |
| fourth digit is rare, which is precisely the remote-homology case the benchmark tests. |
| """ |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| def prefix(ec, level): |
| return ".".join(ec.split(".")[:level]) |
|
|
|
|
| class HierNet(nn.Module): |
| def __init__(self, d_in, sizes, d_hid=1024, d_proj=256, drop=0.3): |
| """sizes = [number of L1, L2, L3, L4 classes].""" |
| super().__init__() |
| self.body = nn.Sequential(nn.LayerNorm(d_in), nn.Linear(d_in, d_hid), nn.GELU(), nn.Dropout(drop), |
| nn.Linear(d_hid, d_hid), nn.GELU(), nn.Dropout(drop)) |
| self.proj = nn.Linear(d_hid, d_proj) |
| self.cls = nn.ModuleList(nn.Linear(d_hid, n) for n in sizes) |
|
|
| def forward(self, x): |
| h = self.body(x) |
| return F.normalize(self.proj(h), dim=-1), [c(h) for c in self.cls] |
|
|
|
|
| def load_head(path): |
| ck = torch.load(path, map_location="cpu", weights_only=False) |
| net = HierNet(ck["d_in"], ck["sizes"]) |
| net.load_state_dict(ck["state"]) |
| return net.eval(), ck["ecs"], ck["l3"] |
|
|
|
|
| @torch.no_grad() |
| def evidence(net, x, C, ecs, l3, l3_of, Zdb, db_ecs, cands, topk=8, knn=20, cen_temp=0.05, batch=2048): |
| """Per-query {ec: (cls, l3, cen, knn)} for the head's top ECs plus the given candidates. |
| |
| `cands[i]` are the ECs homology already proposed: they are always scored exactly, where v1 |
| could only score an EC that also happened to reach the head's own top 15. |
| """ |
| index = {e: j for j, e in enumerate(ecs)} |
| out = [] |
| for s in range(0, len(x), batch): |
| z, logits = net(x[s:s + batch]) |
| p4 = torch.softmax(logits[3], -1) |
| p3 = torch.softmax(logits[2], -1) |
| pc = torch.softmax(z @ C.T / cen_temp, -1) |
| sim = z @ Zdb.T |
| near = torch.topk(sim, knn, -1) |
| top = torch.topk(p4 + pc, topk, -1).indices.tolist() |
| for r in range(len(z)): |
| nb = {} |
| for v, j in zip(near.values[r].tolist(), near.indices[r].tolist()): |
| for e in db_ecs[j]: |
| nb[e] = nb.get(e, 0.0) + max(v, 0.0) ** 8 / len(db_ecs[j]) |
| tot = sum(nb.values()) or 1.0 |
| d = {} |
| for e in set(cands[s + r]) | {ecs[j] for j in top[r]} | set(sorted(nb, key=nb.get)[-3:]): |
| j = index.get(e) |
| k = l3.get(prefix(e, 3)) |
| d[e] = (float(p4[r, j]) if j is not None else 0.0, |
| float(p3[r, k]) if k is not None else 0.0, |
| float(pc[r, j]) if j is not None else 0.0, |
| nb.get(e, 0.0) / tot) |
| out.append(d) |
| return out |
|
|
|
|
| def load_esm(path): |
| """Load an ESM-2 checkpoint from a local file (no network, contact head not needed).""" |
| import warnings |
|
|
| import esm |
| data = torch.load(str(path), map_location="cpu", weights_only=False) |
| with warnings.catch_warnings(): |
| warnings.simplefilter("ignore") |
| model, alphabet = esm.pretrained.load_model_and_alphabet_core(path.stem, data, None) |
| return model.eval(), alphabet |
|
|
