| """ |
| 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) |
|
|
| |
| 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}") |
|
|
| |
| 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):,}") |
|
|
| |
| 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):,}") |
|
|
| |
| merged_train_tokens = len(old_train) + len(add_train) |
| merged_val_tokens = len(old_val) + len(add_val) |
|
|
| |
| 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) |
|
|
| |
| output_dir.mkdir(parents=True, exist_ok=True) |
|
|
| |
| 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 |
|
|
| |
| 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 |
|
|
| |
| 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) |
|
|
| |
| 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() |
|
|
