| """Standalone loader for the UCDP mini event encoder. Requires torch>=2.0.""" |
| import json, torch, torch.nn as nn |
|
|
| class UcdpEncoder(nn.Module): |
| def __init__(self, vsm, fields, d, n_heads, n_layers, dropout=0.1): |
| super().__init__() |
| self.fields = list(fields); self.embed_dim = d |
| self.emb = nn.ModuleDict({f: nn.Embedding(vsm[f], d, padding_idx=0) for f in fields}) |
| self.field_proj = nn.Linear(len(fields) * d, d) |
| self.proj_dropout = nn.Dropout(dropout) |
| self.cls_token = nn.Parameter(torch.zeros(1, 1, d)) |
| layer = nn.TransformerEncoderLayer(d_model=d, nhead=n_heads, dim_feedforward=d * 4, |
| dropout=dropout, batch_first=True) |
| self.transformer = nn.TransformerEncoder(layer, num_layers=n_layers, |
| enable_nested_tensor=False) |
| self.ln = nn.LayerNorm(d) |
| self.heads = nn.ModuleDict({f: nn.Linear(d, vsm[f] - 1) for f in fields}) |
|
|
| def encode_sequence(self, seq, mask): |
| B = seq.size(0) |
| e = torch.cat([self.emb[f](seq[:, :, i]) for i, f in enumerate(self.fields)], dim=-1) |
| e = self.proj_dropout(self.field_proj(e)) |
| e = torch.cat([self.cls_token.expand(B, -1, -1), e], dim=1) |
| cm = torch.ones(B, 1, dtype=torch.bool, device=mask.device) |
| out = self.transformer(e, src_key_padding_mask=~torch.cat([cm, mask], dim=1)) |
| z = self.ln(out[:, 0, :]) |
| pad = (mask.sum(dim=1) == 0).unsqueeze(1) |
| return torch.where(pad, torch.zeros_like(z), z) |
|
|
| def load_encoder(release_dir="."): |
| cfg = json.load(open(f"{release_dir}/config.json")) |
| vsm = {f: s + 1 for f, s in cfg["vocab_sizes_no_mask"].items()} |
| m = UcdpEncoder(vsm, cfg["fields"], cfg["embed_dim"], cfg["n_heads"], cfg["n_layers"]) |
| m.load_state_dict(torch.load(f"{release_dir}/pytorch_model.pt", map_location="cpu", |
| weights_only=True)) |
| m.eval() |
| return m, cfg |
|
|
