src/diffusion_lm/distill.py

"""Distilled-trace preparation: reasoning traces packed with chained thought slots.

Streams a distilled-reasoning dataset (traces carrying a ``<think>...</think>`` region
followed by a final answer), segments each think region into a chain of thought units, and
renders every example through the shared slot geometry so a block-diffusion model denoises
one thought at a time. The number of thoughts varies per example, so slot count stands in
for how much the model chose to think.

Segmentation splits on paragraph breaks first, then on reasoning-marker sentence boundaries
inside long paragraphs, so a thought is a coherent reasoning move rather than a fixed-size
cut.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import re
from pathlib import Path
from typing import Iterator

import numpy as np

from diffusion_lm.reasoning import (
    ADAPTIVE_SPECIAL_TOKENS,
    ExampleEncoder,
    LayoutSpec,
    REASONING_SPECIAL_TOKENS,
    ReasoningExample,
    size_token_ids,
    _write_packed,
)
from diffusion_lm.tokenizer import load_tokenizer, train_tokenizer_from_iterator

_THINK_RE = re.compile(r'<think>(.*?)</think>\s*', re.DOTALL)
_MARKER_BREAK = re.compile(
    r'(?<=[.!?])\s+(?=(?:Wait|Alternatively|Hmm|Okay|Now|But wait|Let me|First|Second|'
    r'Next|Then|Finally|So,|Also|Another|Actually|Therefore|Thus)\b)'
)
_LONG_PARAGRAPH_CHARS = 900


def segment_thoughts(think: str) -> list[str]:
    """Split a think region into a chain of coherent thought units."""

    thoughts: list[str] = []
    for paragraph in re.split(r'\n\s*\n', think):
        paragraph = ' '.join(paragraph.split())
        if not paragraph:
            continue
        if len(paragraph) > _LONG_PARAGRAPH_CHARS:
            thoughts.extend(part.strip() for part in _MARKER_BREAK.split(paragraph) if part.strip())
        else:
            thoughts.append(paragraph)
    return thoughts


def parse_glaive(row: dict[str, object]) -> ReasoningExample | None:
    """Render one glaive ``reasoning-v1`` row into a chained-thought example."""

    prompt = ' '.join(str(row.get('prompt') or '').split())
    response = str(row.get('response') or '')
    match = _THINK_RE.search(response)
    if not prompt or not match:
        return None
    answer = ' '.join(response[match.end():].split())
    thoughts = segment_thoughts(match.group(1))
    if not answer or not thoughts:
        return None
    return ReasoningExample(prompt, tuple(thoughts), answer, expected_answer='')


PARSERS = {'glaive': parse_glaive}


def _stream_examples(
    dataset: str, split: str, parser_name: str, limit: int
) -> Iterator[ReasoningExample]:
    from datasets import load_dataset

    parser = PARSERS[parser_name]
    kept = 0
    for row in load_dataset(dataset, split=split, streaming=True):
        example = parser(row)
        if example is None:
            continue
        yield example
        kept += 1
        if limit and kept >= limit:
            return


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 _resolve_tokenizer(
    examples: list[ReasoningExample],
    args: argparse.Namespace,
    special_tokens: tuple[str, ...] = REASONING_SPECIAL_TOKENS,
):
    from pathlib import Path

    tokenizer_path = Path(args.tokenizer)
    if tokenizer_path.is_file():
        print(f'reusing tokenizer {tokenizer_path}')
        return load_tokenizer(tokenizer_path), tokenizer_path

    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=special_tokens,
    )
    print(f'trained tokenizer {tokenizer_path} (vocab {args.vocab_size}) on {len(examples):,} traces')
    return tokenizer, tokenizer_path


def prepare(args: argparse.Namespace) -> None:
    adaptive = bool(getattr(args, 'sizes', None))
    spec = LayoutSpec(
        seq_len=args.seq_len,
        block=args.block,
        max_slots=args.max_slots,
        sizes=tuple(args.sizes) if adaptive else (),
    )
    output_dir = Path(args.output_dir)

    # Buffer unique parsed examples so a native tokenizer can be trained in the same
    # streaming pass the corpus is packed from, avoiding a second dataset download.
    examples: list[ReasoningExample] = []
    seen: set[str] = set()
    scanned = 0
    for example in _stream_examples(args.dataset, args.split, args.parser, args.scan_limit):
        scanned += 1
        key = hashlib.sha256(example.problem.encode('utf-8')).hexdigest()
        if key in seen:
            continue
        seen.add(key)
        examples.append(example)
        if args.limit and len(examples) >= args.limit:
            break
    if not examples:
        raise ValueError('no usable examples; check dataset, parser, and seq_len')
    print(f'scanned {scanned:,}, buffered {len(examples):,} unique examples')

    special_tokens = ADAPTIVE_SPECIAL_TOKENS if adaptive else REASONING_SPECIAL_TOKENS
    tokenizer, tokenizer_path = _resolve_tokenizer(examples, args, special_tokens)
    encoder = ExampleEncoder(tokenizer, spec)

    if adaptive:
        _pack_adaptive(examples, encoder, spec, tokenizer, tokenizer_path, output_dir, args)
        return

    split: dict[str, dict[str, list[np.ndarray]]] = {
        'train': {'flat': [], 'flat_regions': [], 'slotted': [], 'slotted_regions': []},
        'validation': {'flat': [], 'flat_regions': [], 'slotted': [], 'slotted_regions': []},
    }
    val_prompts: list[dict[str, str]] = []
    dropped = 0
    for example in examples:
        encoded = encoder.encode_example(example)
        if encoded is None:
            dropped += 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 and len(val_prompts) < 500:
            val_prompts.append({'problem': example.problem, 'expected_answer': ''})

    total = sum(len(bucket['flat']) for bucket in split.values())
    if not total:
        raise ValueError('no examples survived encoding; raise seq_len or max_slots')
    print(f'encoded {total:,} examples ({dropped:,} dropped at encode)')
    for split_name, bucket in split.items():
        if not bucket['flat']:
            raise ValueError(f'no examples in the {split_name} split; lower val_fraction or scan more')
        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_prompts:
            handle.write(json.dumps(record, ensure_ascii=False) + '\n')
    print(f'wrote {len(val_prompts):,} validation prompts to {problems_path}')


def _pack_adaptive(
    examples: list[ReasoningExample],
    encoder: ExampleEncoder,
    spec: LayoutSpec,
    tokenizer,
    tokenizer_path: Path,
    output_dir: Path,
    args: argparse.Namespace,
) -> None:
    """Encode the adaptive layout and report the block-size distribution it produces."""

    split: dict[str, dict[str, list[np.ndarray]]] = {
        'train': {'tokens': [], 'regions': []},
        'validation': {'tokens': [], 'regions': []},
    }
    val_prompts: list[dict[str, str]] = []
    size_histogram: dict[int, int] = {size: 0 for size in spec.sizes}
    blocks_per_example: list[int] = []
    dropped = 0
    for example in examples:
        encoded = encoder.encode_adaptive(example)
        if encoded is None:
            dropped += 1
            continue
        is_val = _is_validation(example.problem, args.val_fraction)
        bucket = split['validation' if is_val else 'train']
        bucket['tokens'].append(encoded.tokens)
        bucket['regions'].append(encoded.regions)
        blocks_per_example.append(len(encoded.block_sizes))
        for size in encoded.block_sizes:
            size_histogram[size] += 1
        if is_val and len(val_prompts) < 500:
            val_prompts.append({'problem': example.problem, 'expected_answer': ''})

    total = sum(len(bucket['tokens']) for bucket in split.values())
    if not total:
        raise ValueError('no examples survived encoding; raise seq_len, max_slots, or sizes')

    total_blocks = sum(size_histogram.values())
    fractions = {
        size: round(count / max(1, total_blocks), 4) for size, count in size_histogram.items()
    }
    sorted_blocks = sorted(blocks_per_example)

    def percentile(fraction: float) -> int:
        return sorted_blocks[min(len(sorted_blocks) - 1, int(len(sorted_blocks) * fraction))]

    print(f'encoded {total:,} examples ({dropped:,} dropped at encode)')
    print(f'block-size counts {size_histogram} fractions {fractions}')
    print(
        f'blocks/example p50 {percentile(0.5)} p90 {percentile(0.9)} '
        f'max {sorted_blocks[-1]}'
    )

    extra_metadata = {
        'sizes': list(spec.sizes),
        'size_token_ids': size_token_ids(tokenizer),
    }
    for split_name, bucket in split.items():
        if not bucket['tokens']:
            raise ValueError(f'no examples in the {split_name} split; lower val_fraction')
        _write_packed(
            output_dir / f'{split_name}-adaptive.bin',
            np.stack(bucket['tokens']),
            np.stack(bucket['regions']),
            layout='adaptive',
            spec=spec,
            tokenizer_path=tokenizer_path,
            tokenizer=tokenizer,
            extra_metadata=extra_metadata,
        )
        print(f'{split_name}: {len(bucket["tokens"]):,} examples -> {output_dir}')

    problems_path = output_dir / 'validation-problems.jsonl'
    with problems_path.open('w', encoding='utf-8') as handle:
        for record in val_prompts:
            handle.write(json.dumps(record, ensure_ascii=False) + '\n')
    print(f'wrote {len(val_prompts):,} validation prompts 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='stream a distilled-reasoning dataset into chained-thought slots'
    )
    prepare_parser.add_argument('--dataset', default='glaiveai/reasoning-v1-20m')
    prepare_parser.add_argument('--split', default='train')
    prepare_parser.add_argument('--parser', choices=sorted(PARSERS), default='glaive')
    prepare_parser.add_argument(
        '--tokenizer', required=True,
        help='tokenizer path; trained from the buffered traces when the file is absent'
    )
    prepare_parser.add_argument('--vocab-size', type=int, default=16384)
    prepare_parser.add_argument('--output-dir', required=True)
    prepare_parser.add_argument('--seq-len', type=int, default=2048)
    prepare_parser.add_argument('--block', type=int, default=32)
    prepare_parser.add_argument('--max-slots', type=int, default=64)
    prepare_parser.add_argument(
        '--sizes', type=int, nargs='+',
        help='activate the adaptive layout with these ascending block sizes (e.g. 64 256)'
    )
    prepare_parser.add_argument('--val-fraction', type=float, default=0.02)
    prepare_parser.add_argument(
        '--limit', type=int, default=0, help='stop after this many unique kept examples'
    )
    prepare_parser.add_argument(
        '--scan-limit', type=int, default=0, help='stop streaming after this many parsed rows'
    )
    return parser


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


if __name__ == '__main__':
    main()