File size: 10,108 Bytes
685e018 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 | """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()
|