BERT: Pre-training of Deep Bidirectional Transformers

def sentence_pair(
    document=0,
    sentence=0,
    is_random_next=False,
    random_document=1,
    random_sentence=0,
    max_length=10,
):
    """Pack an explicit corpus pair; label 0 means IsNext, 1 means NotNext."""
    if not isinstance(is_random_next, bool):
        raise ValueError("is_random_next must be boolean.")
    if not (0 <= document < len(CORPUS)):
        raise ValueError("Invalid document.")
    if not (0 <= sentence < len(CORPUS[document]) - 1):
        raise ValueError("Sentence A must have a successor.")
    a = CORPUS[document][sentence]
    if is_random_next:
        if (
            not (0 <= random_document < len(CORPUS))
            or random_document == document
        ):
            raise ValueError("Random B must come from another document.")
        if not (0 <= random_sentence < len(CORPUS[random_document])):
            raise ValueError("Invalid random sentence.")
        b = CORPUS[random_document][random_sentence]
    else:
        b = CORPUS[document][sentence + 1]
    tokens = [CLS, *a, SEP, *b, SEP]
    if max_length < len(tokens):
        raise ValueError("Pair does not fit; this fixture does not truncate.")
    segments = [0] * (len(a) + 2) + [1] * (len(b) + 1)
    padding = max_length - len(tokens)
    return {
        "ids": torch.tensor(tokens + [PAD] * padding),
        "segments": torch.tensor(segments + [0] * padding),
        "valid": torch.tensor([True] * len(tokens) + [False] * padding),
        "nsp_label": int(is_random_next),
    }

def select_and_corrupt(ids, valid, rng, probability=0.15, max_predictions=3):
    """Use original capped selection and two RNG draws for 80/10/10."""
    check_ids(ids, VOCAB_SIZE, "ids")
    if ids.ndim != 1 or valid.shape != ids.shape or valid.dtype != torch.bool:
        raise ValueError("Expected one sequence and a boolean validity mask.")
    length = int(valid.sum())
    if length == 0 or not valid[:length].all() or valid[length:].any():
        raise ValueError("Valid tokens must form a nonempty prefix.")
    if not 0 < probability <= 1 or max_predictions < 1:
        raise ValueError(
            "Probability and maximum predictions must be positive."
        )
    candidates = [i for i in range(length) if int(ids[i]) >= 4]
    if not candidates:
        raise ValueError("At least one eligible content token is required.")
    rng.shuffle(candidates)
    count = min(max_predictions, max(1, int(round(length * probability))))
    selected = candidates[:count]
    output = ids.clone()
    branches = {}
    for index in selected:
        if rng.random() < 0.8:
            output[index], branches[index] = MASK, "mask"
        elif rng.random() < 0.5:
            branches[index] = "unchanged"
        else:
            # Original replacement samples the ENTIRE vocabulary, specials too.
            output[index] = rng.randint(0, VOCAB_SIZE - 1)
            branches[index] = "random"
    selected.sort()
    slots = max_predictions - len(selected)
    positions = torch.tensor(selected + [0] * slots)
    targets = torch.tensor([int(ids[i]) for i in selected] + [PAD] * slots)
    weights = torch.tensor([1.0] * len(selected) + [0.0] * slots)
    return output, positions, targets, weights, [branches[i] for i in selected]

def input_embeddings(self, ids, segments):
    """Sum learned token, position and segment embeddings, then LN/dropout."""
    if (
        ids.ndim != 2
        or min(ids.shape) < 1
        or ids.shape[1] > self.max_length
    ):
        raise ValueError(
            "Expected nonempty [B, N] IDs within position capacity."
        )
    if segments.shape != ids.shape:
        raise ValueError("Segment shape must match tokens.")
    check_ids(ids, VOCAB_SIZE, "ids")
    check_ids(segments, 2, "segments")
    positions = torch.arange(ids.shape[1], device=ids.device)
    x = self.token(ids) + self.position(positions) + self.segment(segments)
    # CLS uses its learned token row, just like other vocabulary entries.
    return self.drop(self.embedding_norm(x))

def encode(self, ids, segments, valid):
    """Attend to every valid key in both directions, then post-norm FFN."""
    x = self.input_embeddings(ids, segments)
    if valid.shape != ids.shape or valid.dtype != torch.bool:
        raise ValueError("Expected boolean validity [B, N].")
    if not valid.any(dim=1).all():
        raise ValueError("Every sequence needs at least one valid key.")
    batch, length, width = x.shape
    head_width = width // self.heads
    q, k, v = [
        projection(x)
        .reshape(batch, length, self.heads, head_width)
        .transpose(1, 2)
        for projection in (self.q, self.k, self.v)
    ]
    scores = q @ k.transpose(-1, -2) / math.sqrt(head_width)
    scores = scores + (~valid[:, None, None, :]).to(x.dtype) * -10000.0
    attention = torch.softmax(scores, dim=-1)  # No causal triangle.
    context = (self.drop(attention) @ v).transpose(1, 2).reshape(x.shape)
    z = self.attention_norm(x + self.drop(self.attention_output(context)))
    # Original modeling.gelu uses this tanh approximation, not exact erf.
    ffn = self.output(F.gelu(self.intermediate(z), approximate="tanh"))
    return self.output_norm(z + self.drop(ffn)), attention

def gather_selected_states(encoded, positions):
    """Gather [B, M, D]; padded slots remain present but have zero loss weight."""
    if encoded.ndim != 3 or positions.ndim != 2 or min(positions.shape) < 1:
        raise ValueError(
            "Expected encoded [B, N, D] and nonempty positions [B, M]."
        )
    if positions.shape[0] != encoded.shape[0]:
        raise ValueError("Position batch must match encoded batch.")
    check_ids(positions, encoded.shape[1], "positions")
    return encoded.gather(
        1, positions[..., None].expand(-1, -1, encoded.shape[-1])
    )

def mlm_logits(self, states):
    """Dense/GELU/LN, then tied token matrix and output-only vocab bias."""
    h = self.mlm_norm(F.gelu(self.mlm_dense(states), approximate="tanh"))
    return F.linear(h, self.token.weight, self.vocab_bias)

def nsp_logits(self, encoded):
    """Pool final CLS through dense+tanh, then predict IsNext/NotNext."""
    pooled_cls = torch.tanh(self.pooler(encoded[:, 0]))
    return self.nsp(pooled_cls)

def joint_loss(logits, targets, weights, nsp_logits, labels):
    """Normalize selected MLM slots and batch NSP separately, then add."""
    validate_objectives(logits, targets, weights, nsp_logits, labels)
    per_slot = (
        -F.log_softmax(logits, dim=-1)
        .gather(-1, targets[..., None])
        .squeeze(-1)
    )
    mlm_loss = (per_slot * weights).sum() / (weights.sum() + 1e-5)
    nsp_loss = F.cross_entropy(nsp_logits, labels, reduction="mean")
    return mlm_loss + nsp_loss, mlm_loss, nsp_loss


def joint_update(model, batch, learning_rate=0.05):
    """Take one illustrative SGD step; both objectives reach shared parameters."""
    if not math.isfinite(learning_rate) or learning_rate <= 0:
        raise ValueError("Learning rate must be finite and positive.")
    optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
    optimizer.zero_grad(set_to_none=True)
    encoded, logits, nsp_logits, attention = model(
        batch["ids"], batch["segments"], batch["valid"], batch["positions"]
    )
    losses = joint_loss(
        logits, batch["targets"], batch["weights"], nsp_logits, batch["labels"]
    )
    losses[0].backward()
    norms = {name: float(p.grad.norm()) for name, p in model.named_parameters()}
    if not all(math.isfinite(value) for value in norms.values()):
        raise ValueError("Nonfinite gradients.")
    optimizer.step()
    return encoded, logits, nsp_logits, attention, losses, norms