src/diffusion_lm/stories.py

"""TinyStories preparation: AR-pretraining stream and instruct data with thought slots.

Two products over one shared tokenizer:

- ``prepare-pretrain``: TinyStoriesV2 stories packed as a continuous EOS-separated
  uint16 stream for causal-LM pretraining.
- ``prepare-instruct``: TinyStories-Instruct records rendered with the reasoning
  slot geometry — prompt holds the writing instruction, thought slots hold the
  requirement restatements plus the story plan (the summary, which never appears
  in the prompt), and the answer region holds the story.
"""

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 (
    ExampleEncoder,
    LayoutSpec,
    ReasoningExample,
    REASONING_SPECIAL_TOKENS,
    _write_packed,
)
from diffusion_lm.tokenizer import (
    load_tokenizer,
    special_token_ids,
    train_tokenizer_from_iterator,
)

_FIELD_RE = re.compile(r'^(Features|Words|Summary|Random sentence):\s*(.+)$', re.MULTILINE)
_STORY_RE = re.compile(r'^Story:\s*$', re.MULTILINE)


def iter_stories(path: Path) -> Iterator[str]:
    """Yield individual stories from an ``<|endoftext|>``-separated text file."""

    buffer = ''
    with path.open('r', encoding='utf-8') as handle:
        while True:
            chunk = handle.read(1 << 24)
            if not chunk:
                break
            buffer += chunk
            *complete, buffer = buffer.split('<|endoftext|>')
            for piece in complete:
                story = piece.strip()
                if story:
                    yield story
    tail = buffer.strip()
    if tail:
        yield tail


def iter_instruct_records(path: Path) -> Iterator[tuple[dict[str, str], str]]:
    """Yield ``(fields, story)`` pairs from a TinyStories-Instruct dump."""

    for record in iter_stories(path):
        story_match = _STORY_RE.search(record)
        if not story_match:
            continue
        header = record[: story_match.start()]
        story = record[story_match.end():].strip()
        fields = dict(_FIELD_RE.findall(header))
        if story and fields:
            yield fields, story


def instruct_example(fields: dict[str, str], story: str) -> ReasoningExample | None:
    """Render one record as prompt, thought steps, and the story answer.

    The summary becomes plan thoughts and is deliberately excluded from the
    prompt, so planning is generative rather than copyable.
    """

    summary = ' '.join(fields.get('Summary', '').split())
    if not summary:
        return None
    prompt_parts = ['Write a short story.']
    thoughts: list[str] = []
    features = ' '.join(fields.get('Features', '').split())
    words = ' '.join(fields.get('Words', '').split())
    sentence = ' '.join(fields.get('Random sentence', '').split())
    if features:
        prompt_parts.append(f'It should feature: {features}.')
        thoughts.append(f'The story needs these elements: {features}.')
    if words:
        prompt_parts.append(f'Use the words: {words}.')
        thoughts.append(f'I have to work in the words {words}.')
    if sentence:
        prompt_parts.append(f'Include the sentence: {sentence}')
        thoughts.append(f'The sentence "{sentence}" must appear.')
    thoughts.append(f'Plan: {summary}')
    story = '\n'.join(line.strip() for line in story.splitlines() if line.strip())
    return ReasoningExample(
        problem=' '.join(prompt_parts),
        steps=tuple(thoughts),
        answer=story,
        expected_answer='',
    )


def prepare_pretrain(args: argparse.Namespace) -> None:
    tokenizer_path = Path(args.tokenizer)
    if tokenizer_path.is_file():
        tokenizer = load_tokenizer(tokenizer_path)
        print(f'reusing tokenizer {tokenizer_path}')
    else:
        def _texts() -> Iterator[str]:
            for index, story in enumerate(iter_stories(Path(args.train))):
                if index >= args.tokenizer_sample:
                    break
                yield story

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

    eos_id = special_token_ids(tokenizer)['eos']
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    for split, source in (('train', args.train), ('validation', args.validation)):
        source_path = Path(source)
        out_path = output_dir / f'{split}.bin'
        token_count = 0
        document_count = 0
        batch: list[str] = []
        with out_path.open('wb') as destination:
            def _flush(batch: list[str]) -> tuple[int, int]:
                encodings = tokenizer.encode_batch(batch, add_special_tokens=False)
                tokens = 0
                for encoding in encodings:
                    ids = encoding.ids + [eos_id]
                    np.asarray(ids, dtype=np.uint16).tofile(destination)
                    tokens += len(ids)
                return tokens, len(encodings)

            for story in iter_stories(source_path):
                batch.append(story)
                if len(batch) >= 2048:
                    tokens, docs = _flush(batch)
                    token_count += tokens
                    document_count += docs
                    batch = []
            if batch:
                tokens, docs = _flush(batch)
                token_count += tokens
                document_count += docs

        metadata = {
            'format': 'mini-diffusion-lm-packed-tokens-v1',
            'dtype': 'uint16',
            'token_count': token_count,
            'document_count': document_count,
            'vocab_size': tokenizer.get_vocab_size(with_added_tokens=True),
            'mask_token_id': special_token_ids(tokenizer)['mask'],
            'eos_token_id': eos_id,
            'special_token_ids': special_token_ids(tokenizer),
            'tokenizer_sha256': hashlib.sha256(tokenizer_path.read_bytes()).hexdigest(),
            'source_files': [str(source_path)],
        }
        with (out_path.parent / f'{out_path.name}.json').open('w', encoding='utf-8') as handle:
            json.dump(metadata, handle, indent=2)
            handle.write('\n')
        print(f'{split}: {document_count:,} stories, {token_count:,} tokens -> {out_path}')


def prepare_instruct(args: argparse.Namespace) -> None:
    tokenizer = load_tokenizer(Path(args.tokenizer))
    spec = LayoutSpec(seq_len=args.seq_len, block=args.block, max_slots=args.max_slots)
    encoder = ExampleEncoder(tokenizer, spec)
    output_dir = Path(args.output_dir)

    for split, source in (('train', args.train), ('validation', args.validation)):
        flat, flat_regions, slotted, slotted_regions = [], [], [], []
        prompts: list[dict[str, str]] = []
        dropped = 0
        count = 0
        for fields, story in iter_instruct_records(Path(source)):
            if args.limit and count >= args.limit:
                break
            example = instruct_example(fields, story)
            if example is None:
                dropped += 1
                continue
            encoded = encoder.encode_example(example)
            if encoded is None:
                dropped += 1
                continue
            count += 1
            flat.append(encoded.flat)
            flat_regions.append(encoded.flat_regions)
            slotted.append(encoded.slotted)
            slotted_regions.append(encoded.slotted_regions)
            if split == 'validation' and len(prompts) < 500:
                prompts.append({'problem': example.problem, 'expected_answer': ''})
        if not flat:
            raise ValueError(f'no usable records in {source}')
        for layout, tokens, regions in (
            ('flat', flat, flat_regions),
            ('slotted', slotted, slotted_regions),
        ):
            _write_packed(
                output_dir / f'{split}-{layout}.bin',
                np.stack(tokens),
                np.stack(regions),
                layout=layout,
                spec=spec,
                tokenizer_path=Path(args.tokenizer),
                tokenizer=tokenizer,
            )
        print(f'{split}: {count:,} examples ({dropped:,} dropped) -> {output_dir}')
        if split == 'validation':
            with (output_dir / 'validation-problems.jsonl').open('w', encoding='utf-8') as fh:
                for record in prompts:
                    fh.write(json.dumps(record, ensure_ascii=False) + '\n')


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    subparsers = parser.add_subparsers(dest='command', required=True)

    pretrain = subparsers.add_parser('prepare-pretrain')
    pretrain.add_argument('--train', required=True)
    pretrain.add_argument('--validation', required=True)
    pretrain.add_argument('--output-dir', required=True)
    pretrain.add_argument('--tokenizer', required=True)
    pretrain.add_argument('--vocab-size', type=int, default=8192)
    pretrain.add_argument('--tokenizer-sample', type=int, default=400_000)

    instruct = subparsers.add_parser('prepare-instruct')
    instruct.add_argument('--train', required=True)
    instruct.add_argument('--validation', required=True)
    instruct.add_argument('--output-dir', required=True)
    instruct.add_argument('--tokenizer', required=True)
    instruct.add_argument('--seq-len', type=int, default=768)
    instruct.add_argument('--block', type=int, default=32)
    instruct.add_argument('--max-slots', type=int, default=9)
    instruct.add_argument('--limit', type=int, default=0)
    return parser


def main() -> None:
    args = _build_parser().parse_args()
    if args.command == 'prepare-pretrain':
        prepare_pretrain(args)
    elif args.command == 'prepare-instruct':
        prepare_instruct(args)


if __name__ == '__main__':
    main()