composer_replication/batch.py

"""real_batch.py — build a real, tokenized 3-channel batch from a HF tokenizer.

Used by Spike 006's smoke to generate inputs for `compose_loss` from a real
chat-template-formatted conversation, NOT random ints.
"""
from __future__ import annotations

from typing import Any

import torch


def build_batch(
    tokenizer: Any,
    *,
    device: torch.device | str = "cpu",
    seed: int = 42,
    variant: str = "factorial",
    align_sdpo_shapes: bool = False,
) -> dict[str, torch.Tensor]:
    """Construct a full 3-channel input batch from a real tokenizer.

    Returns a dict with all keys `compose_loss` may consume:
        input_ids, response_mask
        ctx_teacher_input_ids, sdpo_loss_mask
        dpo_chosen_input_ids, dpo_chosen_response_mask
        dpo_rejected_input_ids, dpo_rejected_response_mask
        dpo_chosen_ref_logprobs, dpo_rejected_ref_logprobs

    The DPO ref logprobs are dummy tensors (not from a real reference policy
    forward); the smoke is verifying the loss composition wires together,
    not the reference-policy precompute pipeline.

    Args:
        tokenizer: real HF tokenizer
        device: torch device for the returned tensors
        seed: reproducibility — fixes torch.manual_seed before any random
            tensor (only the dummy logprobs use random; the chat-template
            text is deterministic)
        variant: "factorial" or "binary_search" — pick which canned
            conversation. Used by Spike 006-strict to alternate batches
            so the loss-decrease isn't memorization of a single sample.
        align_sdpo_shapes: if True, truncate ctx_teacher_input_ids to
            match input_ids length so the SDPO channel actually fires
            (no shape-mismatch fallback). Used by Spike 006-strict to
            exercise the SDPO loss on a real model.
    """
    torch.manual_seed(seed)

    # ------------------------------------------------------------------
    # Conversation 1: student rollout (variants for non-tautological tests)
    # ------------------------------------------------------------------
    if variant == "factorial":
        student_msgs = [
            {"role": "system", "content": "You are a careful coding assistant."},
            {"role": "user", "content": "Write a Python function to compute the factorial of n."},
            {"role": "assistant", "content": "def factorial(n):\n    if n <= 1: return 1\n    return n * factorial(n - 1)"},
        ]
        teacher_msgs = [
            {"role": "system", "content": "You are a careful coding assistant."},
            {"role": "user", "content": "Write a Python function to compute the factorial of n."},
            {"role": "user", "content": "[HINT] Recursion overflows for n>1000. Use an iterative loop."},
            {"role": "assistant", "content": "def factorial(n):\n    result = 1\n    for i in range(2, n + 1):\n        result *= i\n    return result"},
        ]
    elif variant == "binary_search":
        student_msgs = [
            {"role": "system", "content": "You are a careful coding assistant."},
            {"role": "user", "content": "Implement binary search in Python."},
            {"role": "assistant", "content": "def bsearch(a, t):\n    l, r = 0, len(a)\n    while l < r:\n        m = (l + r) // 2\n        if a[m] < t: l = m + 1\n        else: r = m\n    return l"},
        ]
        teacher_msgs = [
            {"role": "system", "content": "You are a careful coding assistant."},
            {"role": "user", "content": "Implement binary search in Python."},
            {"role": "user", "content": "[HINT] Use right = len(a) - 1 with inclusive upper bound is more standard."},
            {"role": "assistant", "content": "def bsearch(a, t):\n    l, r = 0, len(a) - 1\n    while l <= r:\n        m = (l + r) // 2\n        if a[m] == t: return m\n        if a[m] < t: l = m + 1\n        else: r = m - 1\n    return -1"},
        ]
    else:
        raise ValueError(f"unknown variant: {variant!r}")

    student_text = tokenizer.apply_chat_template(student_msgs, tokenize=False, add_generation_prompt=False)
    student_enc = tokenizer(student_text, return_tensors="pt", add_special_tokens=False)
    input_ids = student_enc["input_ids"].to(device)

    T = input_ids.shape[1]
    response_mask = torch.zeros_like(input_ids)
    response_mask[:, int(T * 0.7):] = 1

    # ------------------------------------------------------------------
    # Conversation 2: hint-conditioned teacher context (SDPO)
    # ------------------------------------------------------------------
    teacher_text = tokenizer.apply_chat_template(teacher_msgs, tokenize=False, add_generation_prompt=False)
    teacher_enc = tokenizer(teacher_text, return_tensors="pt", add_special_tokens=False)
    ctx_teacher_input_ids = teacher_enc["input_ids"].to(device)

    if align_sdpo_shapes:
        # Truncate the teacher context to the student length so SDPO actually fires
        # (compose_loss falls back to zero when shapes mismatch). This is a
        # correctness-relaxing test mode — production will pad/align via the
        # real data collator, but for the smoke we just need the SDPO loss
        # to exercise the generalized_jsd_loss code path on a real HF model.
        T_t = ctx_teacher_input_ids.shape[1]
        if T_t > T:
            ctx_teacher_input_ids = ctx_teacher_input_ids[:, :T]
        elif T_t < T:
            pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id
            pad = torch.full((1, T - T_t), pad_id, dtype=ctx_teacher_input_ids.dtype, device=device)
            ctx_teacher_input_ids = torch.cat([ctx_teacher_input_ids, pad], dim=1)

    T_t = ctx_teacher_input_ids.shape[1]
    sdpo_loss_mask = torch.zeros_like(ctx_teacher_input_ids)
    sdpo_loss_mask[:, int(T_t * 0.7):] = 1

    # ------------------------------------------------------------------
    # Conversation 3 + 4: DPO chosen / rejected pairs
    # ------------------------------------------------------------------
    dpo_chosen_msgs = [
        {"role": "system", "content": "You are a careful coding assistant."},
        {"role": "user", "content": "What's the time complexity of binary search?"},
        {"role": "assistant", "content": "Binary search is O(log n) because each comparison halves the search space."},
    ]
    dpo_rejected_msgs = [
        {"role": "system", "content": "You are a careful coding assistant."},
        {"role": "user", "content": "What's the time complexity of binary search?"},
        {"role": "assistant", "content": "It's O(n) I think, you have to look at every element."},
    ]
    chosen_text = tokenizer.apply_chat_template(dpo_chosen_msgs, tokenize=False, add_generation_prompt=False)
    rejected_text = tokenizer.apply_chat_template(dpo_rejected_msgs, tokenize=False, add_generation_prompt=False)

    # Pad both sequences to the same length so we can stack them
    chosen_enc = tokenizer(chosen_text, return_tensors="pt", add_special_tokens=False, padding=False)
    rejected_enc = tokenizer(rejected_text, return_tensors="pt", add_special_tokens=False, padding=False)

    pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id

    chosen_ids = chosen_enc["input_ids"]
    rejected_ids = rejected_enc["input_ids"]
    L = max(chosen_ids.shape[1], rejected_ids.shape[1])

    def _pad(ids: torch.Tensor, length: int) -> torch.Tensor:
        cur = ids.shape[1]
        if cur >= length:
            return ids[:, :length]
        return torch.cat([ids, torch.full((1, length - cur), pad_id, dtype=ids.dtype)], dim=1)

    dpo_chosen_input_ids = _pad(chosen_ids, L).to(device)
    dpo_rejected_input_ids = _pad(rejected_ids, L).to(device)

    chosen_resp_mask = torch.zeros_like(dpo_chosen_input_ids)
    chosen_resp_mask[:, int(L * 0.6):chosen_ids.shape[1]] = 1
    rejected_resp_mask = torch.zeros_like(dpo_rejected_input_ids)
    rejected_resp_mask[:, int(L * 0.6):rejected_ids.shape[1]] = 1

    # Dummy reference-policy logprobs (in production: precomputed by data collator)
    dpo_chosen_ref_logprobs = torch.tensor([-30.0], device=device)
    dpo_rejected_ref_logprobs = torch.tensor([-35.0], device=device)

    return {
        "input_ids": input_ids,
        "response_mask": response_mask,
        "ctx_teacher_input_ids": ctx_teacher_input_ids,
        "sdpo_loss_mask": sdpo_loss_mask,
        "dpo_chosen_input_ids": dpo_chosen_input_ids,
        "dpo_chosen_response_mask": chosen_resp_mask,
        "dpo_rejected_input_ids": dpo_rejected_input_ids,
        "dpo_rejected_response_mask": rejected_resp_mask,
        "dpo_chosen_ref_logprobs": dpo_chosen_ref_logprobs,
        "dpo_rejected_ref_logprobs": dpo_rejected_ref_logprobs,
    }


__all__ = ["build_batch"]