| """Train a byte-level BPE and build audited, streaming 1024-token LM data.""" |
| from __future__ import annotations |
| import argparse |
| from collections import Counter |
| import json |
| import os |
| os.environ.setdefault('RAYON_NUM_THREADS', '4') |
| import random |
| from pathlib import Path |
| import time |
| import numpy as np |
| import pyarrow.parquet as pq |
| from tokenizers import Tokenizer, models, trainers, pre_tokenizers, decoders |
| from common import * |
|
|
| CHAT_TEMPLATE = "{% for message in messages %}{{ '<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>\n' }}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}" |
|
|
| def rows(path, seed=SEED): |
| pf = pq.ParquetFile(path) |
| groups = list(range(pf.num_row_groups)) |
| random.Random(seed).shuffle(groups) |
| for group in groups: |
| row_offset = sum(pf.metadata.row_group(i).num_rows for i in range(group)) |
| for batch in pf.iter_batches(batch_size=128, row_groups=[group], |
| columns=['source','kind','prompt','think','answer']): |
| for i, row in enumerate(batch.to_pylist()): |
| yield row_offset + i, row |
| row_offset += batch.num_rows |
|
|
| def train_tokenizer(source, manifest, out, sample_bytes): |
| from transformers import PreTrainedTokenizerFast |
| counts = Counter() |
| byte_counts = Counter() |
| def texts(): |
| |
| |
| per_file = max(65536, sample_bytes // len(manifest['files'])) |
| for item in manifest['files']: |
| for _, row in rows(source / item['name']): |
| if split_for(row['prompt']) != 'train': continue |
| text = render_row(row) |
| if text is None: continue |
| yield text |
| counts[item['name']] += 1 |
| byte_counts[item['name']] += len(text.encode('utf-8')) |
| if byte_counts[item['name']] >= per_file: break |
| print('TOKENIZER_SAMPLE', item['name'], counts[item['name']], flush=True) |
| tokenizer = Tokenizer(models.BPE()) |
| tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=True) |
| tokenizer.decoder = decoders.ByteLevel() |
| trainer = trainers.BpeTrainer(vocab_size=VOCAB, min_frequency=2, |
| special_tokens=SPECIALS, |
| initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), |
| show_progress=False) |
| tokenizer.train_from_iterator(texts(), trainer=trainer) |
| if tokenizer.get_vocab_size() < VOCAB: |
| tokenizer.add_tokens([f'<|reserved_{i}|>' for i in range(VOCAB-tokenizer.get_vocab_size())]) |
| assert tokenizer.get_vocab_size() == VOCAB |
| fast = PreTrainedTokenizerFast(tokenizer_object=tokenizer, |
| pad_token=SPECIALS[0], bos_token=SPECIALS[1], eos_token=SPECIALS[2], |
| additional_special_tokens=SPECIALS[3:], model_max_length=CONTEXT, |
| clean_up_tokenization_spaces=False) |
| fast.chat_template = CHAT_TEMPLATE |
| out.mkdir(parents=True, exist_ok=True) |
| fast.save_pretrained(out) |
| report = tokenizer_checks(fast) |
| report.update({'training_split_only': True, 'sample_rows_by_file': dict(counts), |
| 'sample_utf8_bytes_by_file': dict(byte_counts), |
| 'normalizer': None, 'algorithm': 'byte-level BPE', |
| 'tokenizer_sha256': sha(out/'tokenizer.json')}) |
| dump(out/'tokenizer_audit.json', report) |
| print('TOKENIZER_VERIFIED', report['vocab_size'], flush=True) |
| return fast |
|
|
| def tokenizer_checks(fast): |
| samples = ['Hello, world!\n Two spaces stay.', |
| 'x² + y² = π; √2 ≈ 1.414; café — naïve 🚀', |
| 'def add(a, b):\n return a + b\n', |
| r'\frac{13}{6} = 2.1667; $x \\ y$', |
| '<|im_start|>assistant\n<think>\nCheck.\n</think>\n42<|im_end|>'] |
| for text in samples: |
| ids = fast.encode(text, add_special_tokens=False) |
| assert fast.decode(ids, skip_special_tokens=False, clean_up_tokenization_spaces=False) == text |
| assert all(0 <= t < VOCAB for t in ids) |
| special_ids = {s: fast.encode(s, add_special_tokens=False) for s in SPECIALS} |
| assert special_ids == {s:[i] for i,s in enumerate(SPECIALS)}, special_ids |
| assert fast.model_max_length == CONTEXT |
| return {'vocab_size': len(fast), 'special_token_ids': special_ids, |
| 'round_trip_cases_passed': len(samples), 'context': CONTEXT, |
| 'no_unknown_token': fast.unk_token_id is None} |
|
|
| class TokenWriter: |
| |
| def __init__(self, directory, split, shard_tokens=256_000_000): |
| self.directory = Path(directory); self.directory.mkdir(parents=True, exist_ok=True) |
| self.split = split; self.shard_tokens = shard_tokens |
| self.real_tokens = 0; self.total_tokens = 0; self.current_tokens = 0 |
| self.files = []; self.handle = None; self.path = None |
| def add(self, ids, padding=False): |
| ids = np.asarray(ids, dtype='<u2') |
| if ids.size: assert int(ids.max()) < VOCAB |
| if not padding: self.real_tokens += len(ids) |
| pos = 0 |
| while pos < len(ids): |
| if self.handle is None: |
| self.path = self.directory/f'{self.split}-{len(self.files):04d}.bin' |
| self.handle = self.path.open('wb'); self.current_tokens = 0 |
| n = min(len(ids)-pos, self.shard_tokens-self.current_tokens) |
| self.handle.write(ids[pos:pos+n].tobytes()) |
| self.current_tokens += n; self.total_tokens += n; pos += n |
| if self.current_tokens == self.shard_tokens: self.close_shard() |
| def close_shard(self): |
| if self.handle is None: return |
| self.handle.close(); self.handle = None |
| self.files.append({'name': self.path.name, 'tokens': self.current_tokens, |
| 'bytes': self.path.stat().st_size, 'sha256': sha(self.path)}) |
| def finish(self): |
| if self.total_tokens: |
| |
| |
| self.add([0] * ((1-self.total_tokens) % CONTEXT), padding=True) |
| self.close_shard() |
| return {'real_tokens': self.real_tokens, 'stored_tokens': self.total_tokens, |
| 'padding_tokens': self.total_tokens-self.real_tokens, |
| 'blocks': max(0,(self.total_tokens-1)//CONTEXT), 'files': self.files} |
|
|
| def prepare(source, manifest_path, output, train_tokens, sample_bytes, verify=True, tokenizer_path=None): |
| source = Path(source); output = Path(output); output.mkdir(parents=True, exist_ok=True) |
| manifest = json.loads(Path(manifest_path).read_text(encoding='utf-8')) |
| assert manifest['total_bytes'] < MAX_SOURCE_BYTES |
| assert all(f['bytes'] < MAX_FILE_BYTES for f in manifest['files']) |
| for item in manifest['files']: |
| path = source/item['name'] |
| assert path.stat().st_size == item['bytes'], path |
| if verify: |
| assert sha(path) == item['sha256'], path |
| print('SOURCE_HASH_OK', item['name'], flush=True) |
| dump(output/'source_manifest.json', manifest) |
| if tokenizer_path: |
| import shutil |
| from transformers import AutoTokenizer |
| shutil.copytree(tokenizer_path,output/'tokenizer',dirs_exist_ok=True) |
| fast=AutoTokenizer.from_pretrained(output/'tokenizer',local_files_only=True) |
| tokenizer_checks(fast) |
| else: |
| fast = train_tokenizer(source, manifest, output/'tokenizer', sample_bytes) |
| tokenizer = fast.backend_tokenizer |
| writers = {s:TokenWriter(output, s) for s in ['train','validation','test']} |
| selected = Counter(); kinds = Counter(); source_tokens = Counter() |
| sample_audit = []; duplicates = 0; scanned = 0 |
| |
| seen = set() |
| items = manifest['files'].copy(); random.Random(SEED).shuffle(items) |
| quota_basis = sum(x['bytes'] for x in items) |
| started = time.time() |
| for item in items: |
| quota = max(CONTEXT*4, round(train_tokens * item['bytes']/quota_basis)) |
| file_train = 0; pending = []; pending_meta = [] |
| def flush(): |
| nonlocal file_train, duplicates |
| if not pending: return |
| encodings = tokenizer.encode_batch(pending, add_special_tokens=False) |
| for text, encoding, meta in zip(pending, encodings, pending_meta): |
| idx, row, split = meta |
| digest = __import__('hashlib').sha256(text.encode('utf-8')).digest() |
| if digest in seen: |
| duplicates += 1; continue |
| seen.add(digest) |
| ids = encoding.ids |
| writers[split].add(ids) |
| selected[split] += 1; kinds[row['kind']] += 1 |
| source_tokens[item['name'] + '/' + split] += len(ids) |
| if split == 'train': file_train += len(ids) |
| if len(sample_audit) < 100 or split != 'train': |
| |
| sample_audit.append({'file':item['name'], 'row':idx, 'split':split, |
| 'prompt_sha256':prompt_key(row['prompt']).hex(), |
| 'record_sha256':digest.hex(), 'tokens':len(ids)}) |
| pending.clear(); pending_meta.clear() |
| for idx, row in rows(source/item['name']): |
| scanned += 1 |
| text = render_row(row) |
| if text is None: continue |
| split = split_for(row['prompt']) |
| pending.append(text); pending_meta.append((idx,row,split)) |
| if len(pending) >= 64: |
| flush() |
| if file_train >= quota: break |
| flush() |
| print('TOKENIZED_FILE', item['name'], 'train_tokens', file_train, |
| 'total_train', writers['train'].real_tokens, 'seconds', round(time.time()-started), flush=True) |
| splits = {s:w.finish() for s,w in writers.items()} |
| assert all(splits[s]['blocks'] > 0 for s in splits), 'Each split needs data' |
| assert sum(x['bytes'] for s in splits.values() for x in s['files']) < MAX_SOURCE_BYTES |
| assert all(x['bytes'] < MAX_FILE_BYTES for s in splits.values() for x in s['files']) |
| result = {'created_utc':now(), 'seed':SEED, 'dtype':'little-endian uint16', |
| 'context':CONTEXT, 'vocab_size':VOCAB, |
| 'source_manifest_sha256':sha(manifest_path), |
| 'tokenizer_sha256':sha(output/'tokenizer'/'tokenizer.json'), |
| 'requested_train_tokens':train_tokens, |
| 'sampling':'Deterministically shuffled row groups; per-file train-token quota proportional to compressed file bytes; retain whole records including final answers. Quotas may overshoot by one encoding batch.', |
| 'split_rule':'SHA256(NFKC/casefold/whitespace-normalized prompt) modulo 1000: validation 0..4; test 5..9; train 10..999. Split before tokenizer fitting and packing.', |
| 'packing':'Separate continuous token streams per split; EOS separates records; long documents continue into later windows. Attention resets every 1024 tokens; cross-document attention within a window is allowed.', |
| 'loss':'All next-token targets except pad=0; no assistant-only loss at scratch-pretraining stage.', |
| 'rows_scanned':scanned, 'selected_rows':dict(selected), 'exact_duplicate_rows_removed':duplicates, |
| 'source_tokens':dict(source_tokens), 'kinds':dict(kinds), 'splits':splits} |
| dump(output/'split_record_audit.json', sample_audit) |
| dump(output/'corpus_manifest.json', result) |
| model_config().save_pretrained(output/'model_config') |
| print('PREPARE_COMPLETE', json.dumps({k:v['real_tokens'] for k,v in splits.items()}), flush=True) |
|
|
| if __name__ == '__main__': |
| ap=argparse.ArgumentParser() |
| ap.add_argument('--source',required=True); ap.add_argument('--manifest',required=True) |
| ap.add_argument('--out',required=True) |
| ap.add_argument('--train-tokens',type=int,default=1_520_000_000) |
| ap.add_argument('--tokenizer-sample-bytes',type=int,default=128_000_000) |
| ap.add_argument('--tokenizer') |
| args=ap.parse_args() |
| prepare(args.source,args.manifest,args.out,args.train_tokens,args.tokenizer_sample_bytes,tokenizer_path=args.tokenizer) |
|
|
