File size: 7,933 Bytes
187a491 | 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 | """
BoccaccioAI - Corpus Pre-Tokenization Script (Streaming)
Pre-tokenizes filtered JSONL text shards into binary memmap format
for efficient training data loading.
Streaming architecture: reads and tokenizes one shard at a time,
writing token IDs directly to a flat binary file. Memory usage is
O(1 shard) ~500MB instead of O(full corpus) ~30GB+.
De Lauretis Tech
"""
import argparse
import json
import logging
import struct
from pathlib import Path
import numpy as np
from tokenizers import Tokenizer
from tqdm import tqdm
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s [%(levelname)s] %(name)s - %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger("boccaccio.tokenize_corpus")
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Pre-tokenize filtered corpus into binary format for BoccaccioAI training."
)
parser.add_argument(
"--input-dir",
type=str,
default="data/filtered",
help="Directory containing filtered JSONL shards.",
)
parser.add_argument(
"--output-dir",
type=str,
default="data/tokenized/pretrain",
help="Directory for tokenized binary output.",
)
parser.add_argument(
"--tokenizer-path",
type=str,
default="tokenizer/boccaccio-32k.json",
help="Path to trained tokenizer JSON file.",
)
parser.add_argument(
"--sequence-length",
type=int,
default=2048,
help="Context window length for packed sequences.",
)
parser.add_argument(
"--val-split",
type=float,
default=0.005,
help="Fraction of data to reserve for validation.",
)
return parser.parse_args()
def main() -> None:
args = parse_args()
input_dir = Path(args.input_dir)
output_dir = Path(args.output_dir)
tokenizer_path = Path(args.tokenizer_path)
sequence_length: int = args.sequence_length
val_split: float = args.val_split
logger.info("=== BoccaccioAI Corpus Pre-Tokenization ===")
logger.info("Input dir: %s", input_dir)
logger.info("Output dir: %s", output_dir)
logger.info("Tokenizer: %s", tokenizer_path)
logger.info("Sequence length: %d", sequence_length)
logger.info("Val split: %.4f", val_split)
# Load tokenizer.
if not tokenizer_path.exists():
raise FileNotFoundError(f"Tokenizer not found: {tokenizer_path}")
tokenizer = Tokenizer.from_file(str(tokenizer_path))
vocab_size = tokenizer.get_vocab_size()
logger.info("Loaded tokenizer with vocab size %d", vocab_size)
eos_token_id = tokenizer.token_to_id("</s>")
if eos_token_id is None:
raise ValueError("EOS token '</s>' not found in tokenizer vocabulary")
logger.info("EOS token id: %d", eos_token_id)
# Verify token IDs fit in uint16.
if vocab_size > np.iinfo(np.uint16).max + 1:
raise ValueError(
f"Vocab size {vocab_size} exceeds uint16 max ({np.iinfo(np.uint16).max + 1}). "
"Cannot store token IDs as uint16."
)
# Find shards.
shard_files = sorted(input_dir.glob("*.jsonl"))
if not shard_files:
raise FileNotFoundError(f"No JSONL files found in {input_dir}")
logger.info("Found %d JSONL shard(s) in %s", len(shard_files), input_dir)
# ─── Streaming: tokenize shard-by-shard, write to flat binary ──
output_dir.mkdir(parents=True, exist_ok=True)
all_tokens_path = output_dir / "all_tokens.bin"
logger.info("Streaming tokenization to %s ...", all_tokens_path)
total_tokens = 0
total_docs = 0
# Write token IDs as uint16 directly to binary file
with open(all_tokens_path, "wb") as bin_file:
# Use a buffer for efficient writes
buffer = np.empty(100_000, dtype=np.uint16)
buf_idx = 0
for shard_path in tqdm(shard_files, desc="Tokenizing shards"):
with open(shard_path, "r", encoding="utf-8") as f:
for line in f:
line = line.strip()
if not line:
continue
record = json.loads(line)
text = record.get("text", "")
if not text:
continue
encoding = tokenizer.encode(text)
ids = encoding.ids
total_docs += 1
# Write token IDs + EOS
for tid in ids:
buffer[buf_idx] = tid
buf_idx += 1
if buf_idx >= len(buffer):
bin_file.write(buffer[:buf_idx].tobytes())
total_tokens += buf_idx
buf_idx = 0
# EOS token
buffer[buf_idx] = eos_token_id
buf_idx += 1
if buf_idx >= len(buffer):
bin_file.write(buffer[:buf_idx].tobytes())
total_tokens += buf_idx
buf_idx = 0
# Flush remaining buffer
if buf_idx > 0:
bin_file.write(buffer[:buf_idx].tobytes())
total_tokens += buf_idx
logger.info("Tokenized %s documents into %s tokens", f"{total_docs:,}", f"{total_tokens:,}")
# ─── Pack into sequences and split train/val ──────────
num_sequences = total_tokens // sequence_length
usable_tokens = num_sequences * sequence_length
if usable_tokens < total_tokens:
logger.info("Discarding last %d tokens (partial chunk)", total_tokens - usable_tokens)
logger.info("Packing into %d sequences of length %d", num_sequences, sequence_length)
# Memory-map the flat file for efficient splitting
memmap = np.memmap(all_tokens_path, dtype=np.uint16, mode="r", shape=(total_tokens,))
num_val = max(1, int(num_sequences * val_split))
num_train = num_sequences - num_val
train_tokens = num_train * sequence_length
val_tokens = num_val * sequence_length
logger.info("Train: %d sequences (%s tokens)", num_train, f"{train_tokens:,}")
logger.info("Val: %d sequences (%s tokens)", num_val, f"{val_tokens:,}")
# Write train split
train_path = output_dir / "train.bin"
logger.info("Saving train split to %s", train_path)
train_memmap = np.memmap(train_path, dtype=np.uint16, mode="w+", shape=(train_tokens,))
train_memmap[:] = memmap[:train_tokens]
train_memmap.flush()
del train_memmap
# Write val split
val_path = output_dir / "val.bin"
logger.info("Saving val split to %s", val_path)
val_memmap = np.memmap(val_path, dtype=np.uint16, mode="w+", shape=(val_tokens,))
val_memmap[:] = memmap[train_tokens:train_tokens + val_tokens]
val_memmap.flush()
del val_memmap
del memmap
# Save metadata
meta_path = output_dir / "meta.json"
meta = {
"vocab_size": vocab_size,
"train_tokens": train_tokens,
"val_tokens": val_tokens,
"sequence_length": sequence_length,
"total_documents": total_docs,
}
with open(meta_path, "w", encoding="utf-8") as f:
json.dump(meta, f, indent=2)
logger.info("Saved metadata to %s", meta_path)
# Clean up intermediate file
all_tokens_path.unlink()
logger.info("Cleaned up intermediate file: %s", all_tokens_path)
# Summary.
logger.info("=== Pre-tokenization complete ===")
logger.info("Total tokens: %s", f"{total_tokens:,}")
logger.info("Train tokens: %s", f"{train_tokens:,}")
logger.info("Val tokens: %s", f"{val_tokens:,}")
logger.info("Train sequences: %d", num_train)
logger.info("Val sequences: %d", num_val)
logger.info("Sequence length: %d", sequence_length)
if __name__ == "__main__":
main()
|