src/diffusion_lm/reasoning.py

"""OpenMathInstruct-2 preparation and packed datasets for reasoning experiments.

Each example is rendered into two layouts over one shared tokenizer:

- ``flat``:    problem <think> steps... </think> answer <eos>
- ``slotted``: problem <think> [fixed-size thought slots, one step each, <tpad>-filled]
               [terminal slot: </think> <tpad>...] answer <eos>

The slotted layout gives every thought a fixed-geometry block so a block-diffusion
model can denoise one thought at a time; the flat layout serves the autoregressive
and pure-diffusion baselines on identical example sets.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterator

import numpy as np
import torch
from torch import Tensor
from torch.utils.data import Dataset

from diffusion_lm.tokenizer import (
    RoleAwareTokenizer,
    load_tokenizer,
    special_token_ids,
    train_tokenizer_from_iterator,
)

REASONING_PACKED_FORMAT = 'mini-diffusion-lm-reasoning-packed-v1'

_REASONING_NAMESPACE = 'mdlm-r-8f4e2a6c'
THINK_TOKEN = f'<|{_REASONING_NAMESPACE}:think|>'
END_THINK_TOKEN = f'<|{_REASONING_NAMESPACE}:endthink|>'
THOUGHT_PAD_TOKEN = f'<|{_REASONING_NAMESPACE}:tpad|>'
REASONING_SPECIAL_TOKENS = (THINK_TOKEN, END_THINK_TOKEN, THOUGHT_PAD_TOKEN)

# Block sizes a thought may occupy in the adaptive layout. The model emits one
# ``<szN>`` control token per thought to choose the block length; the full menu is
# baked into the tokenizer so the vocabulary stays stable across experiments even
# when a run activates only a subset of the sizes.
BLOCK_SIZES = (32, 64, 128, 256, 512)
SIZE_TOKENS = tuple(f'<|{_REASONING_NAMESPACE}:sz{size}|>' for size in BLOCK_SIZES)
ADAPTIVE_SPECIAL_TOKENS = REASONING_SPECIAL_TOKENS + SIZE_TOKENS

# Native spellings honored when the tokenizer comes from a pretrained backbone whose vocabulary
# already carries trained think delimiters (e.g. Qwen3). The namespaced sentinel wins when both
# spellings exist so project tokenizers keep their exact historical ids.
_NATIVE_REASONING_TOKENS = {'think': '<think>', 'end_think': '</think>'}


def size_token_ids(
    tokenizer: RoleAwareTokenizer, sizes: tuple[int, ...] | None = None
) -> dict[int, int]:
    """Map block sizes to ``<szN>`` control-token ids.

    Without ``sizes`` the full menu is scanned and absent tokens are skipped, so tokenizers
    trained before a menu extension keep resolving. Explicitly requested sizes must exist.
    """

    ids: dict[int, int] = {}
    for size, token in zip(BLOCK_SIZES, SIZE_TOKENS):
        if sizes is not None and size not in sizes:
            continue
        token_id = tokenizer.token_to_id(token)
        if token_id is None:
            if sizes is None:
                continue
            raise ValueError(f'tokenizer is missing the size token {token!r}')
        ids[size] = token_id
    if sizes is not None and (missing := set(sizes) - set(ids)):
        raise ValueError(f'sizes {sorted(missing)} are not in the block-size menu {BLOCK_SIZES}')
    if not ids:
        raise ValueError('tokenizer carries no <szN> size tokens')
    return ids

_BOILERPLATE_STEP = re.compile(
    r"^let'?s\s+(solve|answer|tackle)\s+(the\s+)?(new\s+)?(question|problem)\b.*$",
    re.IGNORECASE,
)
_MATH_SPAN = re.compile(r'\$[^$]*\$')
_SENTENCE_BREAK = re.compile(r'(?<=[.!?])\s+(?=[A-Z\d(])')
_BOXED = re.compile(r'\\boxed\{([^{}]*)\}')


def reasoning_token_ids(tokenizer: RoleAwareTokenizer) -> dict[str, int]:
    """Resolve the reasoning control-token ids, failing on non-reasoning tokenizers.

    Pretrained-backbone tokenizers resolve ``think``/``end_think`` through their native
    trained spellings when the namespaced sentinels are absent.
    """

    ids: dict[str, int] = {}
    for name, token in (
        ('think', THINK_TOKEN),
        ('end_think', END_THINK_TOKEN),
        ('thought_pad', THOUGHT_PAD_TOKEN),
    ):
        token_id = tokenizer.token_to_id(token)
        if token_id is None and name in _NATIVE_REASONING_TOKENS:
            token_id = tokenizer.token_to_id(_NATIVE_REASONING_TOKENS[name])
        if token_id is None:
            raise ValueError(f'tokenizer is missing the reasoning token {token!r}')
        ids[name] = token_id
    return ids


def extract_boxed_answer(text: str) -> str | None:
    matches = _BOXED.findall(text)
    if not matches:
        return None
    return matches[-1].strip()


def _normalize_answer(answer: str) -> str:
    cleaned = answer.strip().replace(',', '').replace('$', '').rstrip('.')
    try:
        value = float(cleaned)
    except ValueError:
        return cleaned
    return str(int(value)) if value == int(value) else str(value)


def split_solution_steps(solution: str) -> list[str]:
    """Split a natural-language CoT solution into sentence-level steps."""

    protected: list[str] = []

    def _protect(match: re.Match[str]) -> str:
        protected.append(match.group(0))
        return f'\x00{len(protected) - 1}\x00'

    masked = _MATH_SPAN.sub(_protect, solution)
    parts: list[str] = []
    for line in masked.splitlines():
        line = line.strip()
        if not line:
            continue
        parts.extend(piece.strip() for piece in _SENTENCE_BREAK.split(line) if piece.strip())

    def _restore(text: str) -> str:
        return re.sub(r'\x00(\d+)\x00', lambda m: protected[int(m.group(1))], text)

    steps = [_restore(part) for part in parts]
    return [step for step in steps if not _BOILERPLATE_STEP.match(step)]


@dataclass(frozen=True)
class ReasoningExample:
    problem: str
    steps: tuple[str, ...]
    answer: str
    expected_answer: str


def parse_example(row: dict[str, Any]) -> ReasoningExample | None:
    """Clean one dataset row into problem, thought steps, and an answer sentence."""

    problem = ' '.join(str(row['problem']).split())
    expected = _normalize_answer(str(row['expected_answer']))
    steps = split_solution_steps(str(row['generated_solution']))
    answer_index = None
    for index in range(len(steps) - 1, -1, -1):
        if _BOXED.search(steps[index]):
            answer_index = index
            break
    if answer_index is None or not problem or not expected:
        return None
    answer = steps[answer_index]
    thoughts = tuple(steps[:answer_index])
    if not thoughts:
        return None
    boxed = extract_boxed_answer(answer)
    if boxed is None or _normalize_answer(boxed) != expected:
        return None
    return ReasoningExample(problem, thoughts, answer, expected)


@dataclass(frozen=True)
class LayoutSpec:
    seq_len: int
    block: int
    max_slots: int
    sizes: tuple[int, ...] = ()

    def __post_init__(self) -> None:
        if self.seq_len <= 0 or self.block <= 1 or self.max_slots <= 1:
            raise ValueError('seq_len, block, and max_slots must be meaningfully positive')
        sizes = tuple(self.sizes)
        if sizes:
            if any(size not in BLOCK_SIZES for size in sizes):
                raise ValueError(f'adaptive sizes must be drawn from {BLOCK_SIZES}')
            if list(sizes) != sorted(set(sizes)):
                raise ValueError('adaptive sizes must be unique and ascending')
        object.__setattr__(self, 'sizes', sizes)


@dataclass(frozen=True)
class EncodedExample:
    flat: np.ndarray
    flat_regions: np.ndarray
    slotted: np.ndarray
    slotted_regions: np.ndarray


@dataclass(frozen=True)
class EncodedAdaptive:
    tokens: np.ndarray
    regions: np.ndarray
    block_sizes: tuple[int, ...]


class ExampleEncoder:
    """Render cleaned examples into the flat and slotted token layouts."""

    def __init__(self, tokenizer: RoleAwareTokenizer, spec: LayoutSpec) -> None:
        self.tokenizer = tokenizer
        self.spec = spec
        roles = special_token_ids(tokenizer)
        reasoning = reasoning_token_ids(tokenizer)
        self.pad_id = roles['pad']
        self.eos_id = roles['eos']
        self.think_id = reasoning['think']
        self.end_think_id = reasoning['end_think']
        self.tpad_id = reasoning['thought_pad']
        self.size_ids = size_token_ids(tokenizer, spec.sizes) if spec.sizes else {}
        vocab_size = tokenizer.get_vocab_size(with_added_tokens=True)
        self.token_dtype = np.dtype(
            'uint16' if vocab_size <= np.iinfo(np.uint16).max else 'uint32'
        )

    def _encode(self, text: str) -> list[int]:
        return self.tokenizer.encode(text, add_special_tokens=False).ids

    def encode_adaptive(self, example: ReasoningExample) -> EncodedAdaptive | None:
        """Render an example into the adaptive layout with per-thought size tokens.

        Each thought is placed in the smallest active block that fits it; thoughts
        longer than the largest block continue across consecutive maximum blocks.
        The ``<szN>`` control token precedes every block, giving the model an
        autoregressive target for the block-length decision, and the remainder of
        each block is filled with ``<tpad>``.
        """

        spec = self.spec
        sizes = spec.sizes
        if not sizes:
            raise ValueError('encode_adaptive requires a LayoutSpec with sizes set')
        max_size = sizes[-1]
        problem_ids = self._encode(example.problem)
        answer_ids = self._encode(example.answer)
        if not answer_ids or not problem_ids:
            return None

        blocks: list[tuple[int, list[int]]] = []
        for step in example.steps:
            ids = self._encode(step)
            for start in range(0, len(ids), max_size):
                chunk = ids[start : start + max_size]
                size = next(candidate for candidate in sizes if candidate >= len(chunk))
                blocks.append((size, chunk))
        # An empty chain is a legitimate example: it teaches the controller to answer without
        # thinking, which a conversation needs on every trivial turn.
        if len(blocks) > spec.max_slots:
            return None

        tokens = [*problem_ids, self.think_id]
        for size, chunk in blocks:
            tokens.append(self.size_ids[size])
            tokens.extend(chunk)
            tokens.extend([self.tpad_id] * (size - len(chunk)))
        tokens.append(self.end_think_id)
        answer_start = len(tokens)
        tokens.extend(answer_ids)
        tokens.append(self.eos_id)
        answer_end = len(tokens)
        if answer_end > spec.seq_len:
            return None
        tokens.extend([self.pad_id] * (spec.seq_len - len(tokens)))

        regions = np.asarray(
            [len(problem_ids), len(blocks), answer_start, answer_end], dtype=np.int32
        )
        return EncodedAdaptive(
            tokens=np.asarray(tokens, dtype=self.token_dtype),
            regions=regions,
            block_sizes=tuple(size for size, _ in blocks),
        )

    def encode_example(self, example: ReasoningExample) -> EncodedExample | None:
        spec = self.spec
        problem_ids = self._encode(example.problem)
        step_ids = [self._encode(step) for step in example.steps]
        answer_ids = self._encode(example.answer)

        if not answer_ids or not problem_ids:
            return None

        # A thought that exceeds one slot continues into the next, so long steps
        # cost extra slots instead of dropping the example.
        chunked: list[list[int]] = []
        for ids in step_ids:
            for start in range(0, len(ids), spec.block):
                chunked.append(ids[start : start + spec.block])
        step_ids = chunked
        if len(step_ids) > spec.max_slots - 1:
            return None

        slots = len(step_ids) + 1
        slotted_len = len(problem_ids) + 1 + slots * spec.block + len(answer_ids) + 1
        if slotted_len > spec.seq_len:
            return None

        slotted = [*problem_ids, self.think_id]
        for ids in step_ids:
            slotted.extend(ids)
            slotted.extend([self.tpad_id] * (spec.block - len(ids)))
        slotted.append(self.end_think_id)
        slotted.extend([self.tpad_id] * (spec.block - 1))
        answer_start_slotted = len(slotted)
        slotted.extend(answer_ids)
        slotted.append(self.eos_id)
        answer_end_slotted = len(slotted)
        slotted.extend([self.pad_id] * (spec.seq_len - len(slotted)))

        flat_think_ids = self._encode(' '.join(example.steps))
        flat = [*problem_ids, self.think_id, *flat_think_ids, self.end_think_id]
        answer_start_flat = len(flat)
        flat.extend(answer_ids)
        flat.append(self.eos_id)
        answer_end_flat = len(flat)
        if answer_end_flat > spec.seq_len:
            return None
        flat.extend([self.pad_id] * (spec.seq_len - len(flat)))

        def _regions(problem_len: int, n_slots: int, start: int, end: int) -> np.ndarray:
            return np.asarray([problem_len, n_slots, start, end], dtype=np.int32)

        return EncodedExample(
            flat=np.asarray(flat, dtype=self.token_dtype),
            flat_regions=_regions(len(problem_ids), 0, answer_start_flat, answer_end_flat),
            slotted=np.asarray(slotted, dtype=self.token_dtype),
            slotted_regions=_regions(
                len(problem_ids), slots, answer_start_slotted, answer_end_slotted
            ),
        )


def regions_path(token_path: str | Path) -> Path:
    path = Path(token_path)
    return path.with_suffix('.regions.npy')


class ReasoningTokenDataset(Dataset[tuple[Tensor, Tensor]]):
    """Fixed-length reasoning sequences with per-example region annotations.

    Regions hold ``[problem_len, n_slots, answer_start, answer_end]``. ``n_slots``
    is zero for the flat layout and counts thought slots including the terminal
    ``</think>`` slot for the slotted layout.
    """

    def __init__(self, path: str | Path) -> None:
        self.path = Path(path)
        metadata_path = self.path.with_suffix(self.path.suffix + '.json')
        if not metadata_path.is_file():
            raise FileNotFoundError(f'reasoning metadata not found: {metadata_path}')
        with metadata_path.open('r', encoding='utf-8') as handle:
            self.metadata = json.load(handle)
        if self.metadata.get('format') != REASONING_PACKED_FORMAT:
            raise ValueError(f'unsupported reasoning data format in {metadata_path}')
        self.seq_len = int(self.metadata['seq_len'])
        count = int(self.metadata['example_count'])
        # Datasets packed before the dtype key existed are always uint16.
        self.token_dtype = np.dtype(self.metadata.get('dtype', 'uint16'))
        self._tokens = np.memmap(
            self.path, mode='r', dtype=self.token_dtype, shape=(count, self.seq_len)
        )
        self._regions = np.load(regions_path(self.path))
        if self._regions.shape != (count, 4):
            raise ValueError(f'regions shape {self._regions.shape} does not match example count')

    def __len__(self) -> int:
        return int(self._tokens.shape[0])

    def __getitem__(self, index: int) -> tuple[Tensor, Tensor]:
        tokens = torch.from_numpy(np.asarray(self._tokens[index], dtype=np.int64))
        regions = torch.from_numpy(np.asarray(self._regions[index], dtype=np.int64))
        return tokens, regions

    def __getstate__(self) -> dict[str, Any]:
        state = self.__dict__.copy()
        state['_tokens'] = None
        return state

    def __setstate__(self, state: dict[str, Any]) -> None:
        self.__dict__.update(state)
        count = int(self.metadata['example_count'])
        self._tokens = np.memmap(
            self.path, mode='r', dtype=self.token_dtype, shape=(count, self.seq_len)
        )


def _write_packed(
    path: Path,
    tokens: np.ndarray,
    regions: np.ndarray,
    *,
    layout: str,
    spec: LayoutSpec,
    tokenizer_path: Path,
    tokenizer: RoleAwareTokenizer,
    extra_metadata: dict[str, Any] | None = None,
) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tokens.tofile(path)
    np.save(regions_path(path), regions)
    if tokens.dtype not in (np.dtype('uint16'), np.dtype('uint32')):
        raise ValueError(f'packed reasoning tokens must be uint16 or uint32, got {tokens.dtype}')
    metadata = {
        'format': REASONING_PACKED_FORMAT,
        'layout': layout,
        'seq_len': spec.seq_len,
        'block': spec.block,
        'max_slots': spec.max_slots,
        'dtype': tokens.dtype.name,
        'example_count': int(tokens.shape[0]),
        'vocab_size': tokenizer.get_vocab_size(with_added_tokens=True),
        'special_token_ids': special_token_ids(tokenizer),
        'reasoning_token_ids': reasoning_token_ids(tokenizer),
        'tokenizer_sha256': hashlib.sha256(tokenizer_path.read_bytes()).hexdigest(),
    }
    if extra_metadata:
        metadata.update(extra_metadata)
    metadata_path = path.with_suffix(path.suffix + '.json')
    with metadata_path.open('w', encoding='utf-8') as handle:
        json.dump(metadata, handle, indent=2)
        handle.write('\n')


def _iter_rows(parquet_paths: list[Path], sources: set[str]) -> Iterator[dict[str, Any]]:
    import pyarrow.parquet as pq

    for path in parquet_paths:
        table = pq.read_table(
            path, columns=['problem', 'generated_solution', 'expected_answer', 'problem_source']
        )
        for row in table.to_pylist():
            if row['problem_source'] in sources:
                yield row


def _is_validation(problem: str, val_fraction: float) -> bool:
    digest = hashlib.sha256(problem.encode('utf-8')).digest()
    return int.from_bytes(digest[:4], 'big') / 2**32 < val_fraction


def prepare(args: argparse.Namespace) -> None:
    sources = set(args.sources.split(','))
    parquet_paths = [Path(path) for path in args.parquet]
    spec = LayoutSpec(seq_len=args.seq_len, block=args.block, max_slots=args.max_slots)
    output_dir = Path(args.output_dir)
    tokenizer_path = Path(args.tokenizer)

    examples: list[ReasoningExample] = []
    dropped_parse = 0
    seen: set[str] = set()
    for row in _iter_rows(parquet_paths, sources):
        example = parse_example(row)
        if example is None:
            dropped_parse += 1
            continue
        key = hashlib.sha256(
            (example.problem + '\x1f' + ' '.join(example.steps)).encode('utf-8')
        ).hexdigest()
        if key in seen:
            continue
        seen.add(key)
        examples.append(example)
    print(f'parsed {len(examples):,} unique examples ({dropped_parse:,} dropped at parse)')

    if tokenizer_path.is_file():
        tokenizer = load_tokenizer(tokenizer_path)
        reasoning_token_ids(tokenizer)
        print(f'reusing tokenizer {tokenizer_path}')
    else:
        def _texts() -> Iterator[str]:
            for example in examples:
                yield f'{example.problem}\n{" ".join(example.steps)}\n{example.answer}'

        tokenizer = train_tokenizer_from_iterator(
            _texts(),
            tokenizer_path,
            vocab_size=args.vocab_size,
            min_frequency=4,
            length=len(examples),
            extra_special_tokens=REASONING_SPECIAL_TOKENS,
        )
        print(f'trained tokenizer {tokenizer_path}')

    encoder = ExampleEncoder(tokenizer, spec)
    split: dict[str, dict[str, list[np.ndarray]]] = {
        'train': {'flat': [], 'flat_regions': [], 'slotted': [], 'slotted_regions': []},
        'validation': {'flat': [], 'flat_regions': [], 'slotted': [], 'slotted_regions': []},
    }
    val_problems: list[dict[str, str]] = []
    dropped_encode = 0
    for example in examples:
        encoded = encoder.encode_example(example)
        if encoded is None:
            dropped_encode += 1
            continue
        is_val = _is_validation(example.problem, args.val_fraction)
        bucket = split['validation' if is_val else 'train']
        bucket['flat'].append(encoded.flat)
        bucket['flat_regions'].append(encoded.flat_regions)
        bucket['slotted'].append(encoded.slotted)
        bucket['slotted_regions'].append(encoded.slotted_regions)
        if is_val:
            val_problems.append(
                {'problem': example.problem, 'expected_answer': example.expected_answer}
            )

    total = sum(len(bucket['flat']) for bucket in split.values())
    print(f'encoded {total:,} examples ({dropped_encode:,} dropped at encode)')
    for split_name, bucket in split.items():
        if not bucket['flat']:
            raise ValueError(f'no examples in the {split_name} split; adjust filters')
        for layout in ('flat', 'slotted'):
            _write_packed(
                output_dir / f'{split_name}-{layout}.bin',
                np.stack(bucket[layout]),
                np.stack(bucket[f'{layout}_regions']),
                layout=layout,
                spec=spec,
                tokenizer_path=tokenizer_path,
                tokenizer=tokenizer,
            )
        print(f'{split_name}: {len(bucket["flat"]):,} examples -> {output_dir}')

    problems_path = output_dir / 'validation-problems.jsonl'
    with problems_path.open('w', encoding='utf-8') as handle:
        for record in val_problems:
            handle.write(json.dumps(record, ensure_ascii=False) + '\n')
    print(f'wrote {len(val_problems):,} validation problems to {problems_path}')


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    subparsers = parser.add_subparsers(dest='command', required=True)
    prepare_parser = subparsers.add_parser(
        'prepare', help='filter, tokenize, and pack OpenMathInstruct-2 parquet shards'
    )
    prepare_parser.add_argument('--parquet', type=Path, nargs='+', required=True)
    prepare_parser.add_argument('--output-dir', type=Path, required=True)
    prepare_parser.add_argument('--tokenizer', type=Path, required=True)
    prepare_parser.add_argument('--vocab-size', type=int, default=8192)
    prepare_parser.add_argument('--seq-len', type=int, default=512)
    prepare_parser.add_argument('--block', type=int, default=32)
    prepare_parser.add_argument('--max-slots', type=int, default=11)
    prepare_parser.add_argument('--val-fraction', type=float, default=0.02)
    prepare_parser.add_argument('--sources', default='gsm8k,augmented_gsm8k')
    return parser


def main() -> None:
    args = _build_parser().parse_args()
    if args.command == 'prepare':
        prepare(args)


if __name__ == '__main__':
    main()