Snippets

import torch, torch.nn.functional as F

# ─────────────────────────────────────────────────────────────────────────────
# FAMILY 1 — vocabulary readout.  (openjev, Bespoke Nimble, any LoRA-on-an-LLM)
#
# The options are TEXT IN THE PROMPT. The model only ever picks a letter.
# One forward pass, exactly like generation — you just never sample.
# ─────────────────────────────────────────────────────────────────────────────
def read_by_vocabulary(model, tok, evidence, question, options):
    letters = "ABCDEFGHIJKLMNOP"[: len(options)]          # hard cap: 16 options
    prompt = tok.apply_chat_template([
        {"role": "system", "content":
            "Apply the supplied criterion to the supplied evidence. Choose exactly "
            "one listed option. Respond with only its uppercase letter."},
        {"role": "user", "content": json.dumps({
            "evidence": evidence, "criterion": question,
            "options": [{"letter": L, "description": o}
                        for L, o in zip(letters, options)]})},
    ], tokenize=False, add_generation_prompt=True)

    # each letter must be ONE token that round-trips, or you cannot slice its row
    slots = [tok.encode(L, add_special_tokens=False)[0] for L in letters]

    ids = torch.tensor([tok.encode(prompt, add_special_tokens=False)])
    with torch.inference_mode():
        logits = model(ids).logits[:, -1, :]              # (1, 151936)  <- full vocab
    return F.softmax(logits[0, slots], dim=-1)            # (N,)  renormalised over N rows


# ─────────────────────────────────────────────────────────────────────────────
# FAMILY 2 — per-option scalar scorer.  (cua-s1-forms, open-jev-deberta, and —
# per the published disclosures — Jev itself)
#
# No letters. No vocabulary. The head's output dimension is ONE, and it runs
# once per option. N is a loop bound.
# ─────────────────────────────────────────────────────────────────────────────
class DecisionHead(torch.nn.Module):            # open-jev-deberta's real shapes:
    def __init__(self, d=1024):                 #   0.weight [1024, 3072]
        super().__init__()                      #   2.weight [1, 1024]
        self.net = torch.nn.Sequential(         #   = 3,147,777 params
            torch.nn.Linear(3 * d, d), torch.nn.GELU(), torch.nn.Linear(d, 1))

    def forward(self, q, o):                    # q, o: (B, 1024) pooled spans
        return self.net(torch.cat([q, o, q * o], dim=-1)).squeeze(-1)   # (B,)

def read_by_scoring(encoder, head, tok, evidence, question, options):
    ctx = pool(encoder, tok, f"{evidence}\n{question}")           # (1024,)
    scores = torch.stack([                                        # one scalar EACH
        head(ctx[None], pool(encoder, tok, o)[None])[0] for o in options
    ])                                                            # (N,)
    return F.softmax(scores, dim=-1)

model = model.to(device, memory_format=torch.channels_last)
scaler = torch.amp.GradScaler("cuda")

for x, y in loader:
    x = x.to(device, memory_format=torch.channels_last, non_blocking=True)
    y = y.to(device, non_blocking=True)

    optimizer.zero_grad(set_to_none=True)
    with torch.autocast("cuda", dtype=torch.float16):
        loss = criterion(model(x), y)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

__global__ void vadd(const float *a, const float *b, float *c, int n) {
  int stride = blockDim.x * gridDim.x;
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += stride) {
    c[i] = a[i] + b[i];
  }
}

// Launch: size the grid to the device, not to n.
int block = 256;
int grid = (n + block - 1) / block;        // cap this for very large n
vadd<<<grid, block>>>(d_a, d_b, d_c, n);