"""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()