File size: 5,669 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 | """
BoccaccioAI - Merge tokenized binary data
Concatenates new tokenized data (e.g. Wikipedia) to the existing
pre-training binary files (train.bin, val.bin) and updates metadata.
Usage:
python -m src.data.merge_tokenized \
--existing-dir data/tokenized/pretrain \
--new-dir data/tokenized/wiki \
--output-dir data/tokenized/pretrain_v2
De Lauretis Tech
"""
from __future__ import annotations
import argparse
import json
import logging
import shutil
from pathlib import Path
import numpy as np
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.data.merge_tokenized")
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Merge tokenized binary datasets for BoccaccioAI pre-training."
)
parser.add_argument(
"--existing-dir",
type=str,
required=True,
help="Directory with existing tokenized data (train.bin, val.bin, meta.json).",
)
parser.add_argument(
"--new-dir",
type=str,
required=True,
help="Directory with new tokenized data to append.",
)
parser.add_argument(
"--output-dir",
type=str,
required=True,
help="Output directory for merged data.",
)
parser.add_argument(
"--sequence-length",
type=int,
default=2048,
help="Sequence length for alignment (default: 2048).",
)
return parser.parse_args()
def load_bin(path: Path) -> np.ndarray:
"""Load a binary file as uint16 array."""
return np.memmap(path, dtype=np.uint16, mode="r")
def main() -> None:
args = parse_args()
existing_dir = Path(args.existing_dir)
new_dir = Path(args.new_dir)
output_dir = Path(args.output_dir)
seq_len = args.sequence_length
logger.info("=== BoccaccioAI - Merge Tokenized Data ===")
logger.info("Existing: %s", existing_dir)
logger.info("New: %s", new_dir)
logger.info("Output: %s", output_dir)
# Validate inputs
existing_train = existing_dir / "train.bin"
existing_val = existing_dir / "val.bin"
new_train = new_dir / "train.bin"
new_val = new_dir / "val.bin"
for f in [existing_train, existing_val, new_train, new_val]:
if not f.exists():
raise FileNotFoundError(f"Required file not found: {f}")
# Load existing data
logger.info("Loading existing train.bin ...")
old_train = load_bin(existing_train)
logger.info(" %s tokens", f"{len(old_train):,}")
logger.info("Loading existing val.bin ...")
old_val = load_bin(existing_val)
logger.info(" %s tokens", f"{len(old_val):,}")
# Load new data
logger.info("Loading new train.bin ...")
add_train = load_bin(new_train)
logger.info(" %s tokens", f"{len(add_train):,}")
logger.info("Loading new val.bin ...")
add_val = load_bin(new_val)
logger.info(" %s tokens", f"{len(add_val):,}")
# Concatenate
merged_train_tokens = len(old_train) + len(add_train)
merged_val_tokens = len(old_val) + len(add_val)
# Align to sequence length
merged_train_tokens = (merged_train_tokens // seq_len) * seq_len
merged_val_tokens = (merged_val_tokens // seq_len) * seq_len
logger.info("Merged train: %s tokens (%d sequences)",
f"{merged_train_tokens:,}", merged_train_tokens // seq_len)
logger.info("Merged val: %s tokens (%d sequences)",
f"{merged_val_tokens:,}", merged_val_tokens // seq_len)
# Write output
output_dir.mkdir(parents=True, exist_ok=True)
# Train
logger.info("Writing merged train.bin ...")
out_train = np.memmap(
output_dir / "train.bin", dtype=np.uint16, mode="w+",
shape=(merged_train_tokens,)
)
old_n = min(len(old_train), merged_train_tokens)
out_train[:old_n] = old_train[:old_n]
remaining = merged_train_tokens - old_n
if remaining > 0:
out_train[old_n:old_n + remaining] = add_train[:remaining]
out_train.flush()
del out_train
# Val
logger.info("Writing merged val.bin ...")
out_val = np.memmap(
output_dir / "val.bin", dtype=np.uint16, mode="w+",
shape=(merged_val_tokens,)
)
old_vn = min(len(old_val), merged_val_tokens)
out_val[:old_vn] = old_val[:old_vn]
remaining_v = merged_val_tokens - old_vn
if remaining_v > 0:
out_val[old_vn:old_vn + remaining_v] = add_val[:remaining_v]
out_val.flush()
del out_val
# Load existing metadata if available
total_docs = 0
existing_meta = existing_dir / "meta.json"
new_meta = new_dir / "meta.json"
if existing_meta.exists():
with open(existing_meta) as f:
total_docs += json.load(f).get("total_documents", 0)
if new_meta.exists():
with open(new_meta) as f:
total_docs += json.load(f).get("total_documents", 0)
# Write metadata
meta = {
"vocab_size": 32000,
"train_tokens": merged_train_tokens,
"val_tokens": merged_val_tokens,
"sequence_length": seq_len,
"total_documents": total_docs,
"sources": [str(existing_dir), str(new_dir)],
}
with open(output_dir / "meta.json", "w", encoding="utf-8") as f:
json.dump(meta, f, indent=2)
logger.info("=== Merge complete ===")
logger.info("Total train tokens: %s", f"{merged_train_tokens:,}")
logger.info("Total val tokens: %s", f"{merged_val_tokens:,}")
logger.info("Output: %s", output_dir)
if __name__ == "__main__":
main()
|