"""Download and prepare the pinned FineWeb-Edu 10B sample. The heavy ``huggingface_hub`` and ``pyarrow`` dependencies are imported only by the operations that need them. Raw Parquet files are cached locally, then converted into independently replaceable uint16/uint32 shards with deterministic document-level train/validation assignment. """ from __future__ import annotations import argparse import hashlib import json import os from dataclasses import dataclass from pathlib import Path from typing import Any, Iterable, Iterator, Literal, Sequence import numpy as np from diffusion_lm.data import PACKED_MANIFEST_FORMAT from diffusion_lm.tokenizer import ( load_tokenizer, special_token_ids, train_tokenizer_from_iterator, ) FINEWEB_EDU_REPO_ID = "HuggingFaceFW/fineweb-edu" FINEWEB_EDU_CONFIG = "sample-10BT" FINEWEB_EDU_REVISION = "87f09149ef4734204d70ed1d046ddc9ca3f2b8f9" FINEWEB_EDU_PATH_PREFIX = "sample/10BT/" SOURCE_STATE_FORMAT = "mini-diffusion-lm-corpus-source-v1" SPLIT_HASH_PERSON = b"mini-mdlm-split" @dataclass(frozen=True) class CorpusSource: """One pinned Hub dataset: where its Parquets live and how rows are read. ``id_column=None`` derives the split id from a sha256 of the text, which keeps the train/validation assignment order-independent for datasets without a stable row id. """ name: str repo_id: str revision: str path_prefix: str text_column: str = "text" id_column: str | None = "id" config: str | None = None CORPUS_SOURCES: dict[str, CorpusSource] = { source.name: source for source in ( CorpusSource( name="fineweb-edu", repo_id=FINEWEB_EDU_REPO_ID, revision=FINEWEB_EDU_REVISION, path_prefix=FINEWEB_EDU_PATH_PREFIX, config=FINEWEB_EDU_CONFIG, ), CorpusSource( name="ultra-fineweb-en", repo_id="openbmb/Ultra-FineWeb", revision="7ddd4170ce03e0afbd7d9b80d4bc0b8eebf877e4", path_prefix="data/ultrafineweb_en/", text_column="content", id_column=None, ), CorpusSource( name="cosmopedia-v2", repo_id="HuggingFaceTB/smollm-corpus", revision="3ba9d605774198c5868892d7a8deda78031a781f", path_prefix="cosmopedia-v2/", id_column=None, config="cosmopedia-v2", ), CorpusSource( name="finemath-4plus", repo_id="HuggingFaceTB/finemath", revision="e92b25a616738fe95dc186b64dfb19f9c8525594", path_prefix="finemath-4plus/", id_column=None, config="finemath-4plus", ), ) } def _require_huggingface_hub(): try: from huggingface_hub import HfApi, snapshot_download except ImportError as exc: # pragma: no cover - exercised in minimal installations. raise RuntimeError( "FineWeb-Edu download requires huggingface_hub; install the corpus dependencies" ) from exc return HfApi, snapshot_download def _require_parquet(): try: import pyarrow.parquet as parquet except ImportError as exc: # pragma: no cover - exercised in minimal installations. raise RuntimeError( "FineWeb-Edu preparation requires pyarrow; install the corpus dependencies" ) from exc return parquet def _atomic_json(path: Path, value: dict[str, Any]) -> None: path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_name(f".{path.name}.tmp") temporary.unlink(missing_ok=True) with temporary.open("w", encoding="utf-8") as handle: json.dump(value, handle, indent=2, sort_keys=True) handle.write("\n") handle.flush() os.fsync(handle.fileno()) temporary.replace(path) def _read_json(path: Path) -> dict[str, Any]: try: with path.open("r", encoding="utf-8") as handle: value = json.load(handle) except (OSError, json.JSONDecodeError) as exc: raise ValueError(f"could not read corpus state {path}: {exc}") from exc if not isinstance(value, dict): raise ValueError(f"corpus state must be a JSON object: {path}") return value def _sha256_file(path: Path, *, chunk_size: int = 1024 * 1024) -> str: digest = hashlib.sha256() with path.open("rb") as handle: for chunk in iter(lambda: handle.read(chunk_size), b""): digest.update(chunk) return digest.hexdigest() def download_source( source: CorpusSource, raw_dir: str | Path, *, max_files: int | None = None, max_workers: int = 1, ) -> list[Path]: """Download a pinned source's Parquets into ``raw_dir`` using the Hub cache.""" if max_files is not None and max_files <= 0: raise ValueError("max_files must be positive") if max_workers <= 0: raise ValueError("max_workers must be positive") HfApi, snapshot_download = _require_huggingface_hub() api = HfApi() repository_files = api.list_repo_files( repo_id=source.repo_id, repo_type="dataset", revision=source.revision, ) parquet_names = sorted( name for name in repository_files if name.startswith(source.path_prefix) and name.endswith(".parquet") ) if max_files is not None: parquet_names = parquet_names[:max_files] if not parquet_names: raise RuntimeError( f"no Parquet files found for {source.repo_id}@{source.revision} " f"under {source.path_prefix}" ) destination = Path(raw_dir) destination.mkdir(parents=True, exist_ok=True) snapshot_root = Path( snapshot_download( repo_id=source.repo_id, repo_type="dataset", revision=source.revision, allow_patterns=parquet_names, local_dir=str(destination), max_workers=max_workers, ) ) paths: list[Path] = [] for name in parquet_names: path = snapshot_root / name if not path.is_file(): raise FileNotFoundError(f"Hub download did not produce {path}") paths.append(path) return paths def download_fineweb_edu( raw_dir: str | Path, *, revision: str = FINEWEB_EDU_REVISION, max_files: int | None = None, max_workers: int = 1, ) -> list[Path]: """Download the pinned FineWeb-Edu sample (compatibility entry point).""" from dataclasses import replace source = replace(CORPUS_SOURCES["fineweb-edu"], revision=revision) return download_source(source, raw_dir, max_files=max_files, max_workers=max_workers) def find_local_parquets( raw_dir: str | Path, path_prefix: str = FINEWEB_EDU_PATH_PREFIX ) -> list[Path]: """Find already downloaded source Parquets in stable source order.""" root = Path(raw_dir) preferred = sorted((root / path_prefix).glob("*.parquet")) paths = preferred or sorted(root.rglob("*.parquet")) if not paths: raise FileNotFoundError(f"no Parquet files found below {root}") return paths def document_split( document_id: str, *, validation_modulus: int = 1024, validation_bucket: int = 0, ) -> Literal["train", "validation"]: """Assign a stable document ID to train or validation without depending on row order.""" if not document_id: raise ValueError("document_id must be non-empty") if validation_modulus <= 1: raise ValueError("validation_modulus must be greater than one") if not 0 <= validation_bucket < validation_modulus: raise ValueError("validation_bucket must be inside validation_modulus") digest = hashlib.blake2b( document_id.encode("utf-8"), digest_size=8, person=SPLIT_HASH_PERSON ).digest() bucket = int.from_bytes(digest, "little") % validation_modulus return "validation" if bucket == validation_bucket else "train" def _iter_parquet_batches( paths: Sequence[Path], *, batch_size: int, text_column: str = "text", id_column: str | None = "id", ) -> Iterator[tuple[Path, list[str], list[str]]]: if batch_size <= 0: raise ValueError("batch_size must be positive") parquet = _require_parquet() read_columns = [text_column] if id_column is None else [id_column, text_column] for path in paths: source = parquet.ParquetFile(path) try: batches = source.iter_batches(batch_size=batch_size, columns=read_columns) for batch in batches: columns = batch.to_pydict() texts = columns[text_column] if id_column is not None: ids = columns[id_column] else: # Content-derived ids keep the split assignment order-independent. ids = [ hashlib.sha256(text.encode("utf-8")).hexdigest() if isinstance(text, str) else "" for text in texts ] if len(ids) != len(texts): raise ValueError(f"mismatched id/text columns in {path}") yield path, ids, texts except (KeyError, ValueError) as exc: raise ValueError( f"expected columns {read_columns} in {path}: {exc}" ) from exc def iter_tokenizer_text( paths: Iterable[str | Path], *, max_utf8_bytes: int = 1 << 29, batch_size: int = 512, validation_modulus: int = 1024, validation_bucket: int = 0, text_column: str = "text", id_column: str | None = "id", ) -> Iterator[list[str]]: """Yield bounded train-only text batches for iterator-based tokenizer training.""" if max_utf8_bytes <= 0: raise ValueError("max_utf8_bytes must be positive") sources = sorted(Path(path) for path in paths) used_bytes = 0 output: list[str] = [] for _path, ids, texts in _iter_parquet_batches( sources, batch_size=batch_size, text_column=text_column, id_column=id_column ): for document_id, text in zip(ids, texts, strict=True): if not isinstance(document_id, str) or not isinstance(text, str) or not text: continue if document_split( document_id, validation_modulus=validation_modulus, validation_bucket=validation_bucket, ) != "train": continue encoded_bytes = len(text.encode("utf-8")) if used_bytes and used_bytes + encoded_bytes > max_utf8_bytes: if output: yield output return output.append(text) used_bytes += encoded_bytes if len(output) >= batch_size: yield output output = [] if used_bytes >= max_utf8_bytes: if output: yield output return if output: yield output @dataclass(frozen=True) class CorpusPreparationResult: train_manifest: Path validation_manifest: Path processed_sources: int resumed_sources: int class _SplitWriter: def __init__(self, final_path: Path, dtype: np.dtype[Any]) -> None: self.final_path = final_path self.dtype = dtype self.temporary_path = final_path.with_name(f".{final_path.name}.tmp") final_path.parent.mkdir(parents=True, exist_ok=True) self.temporary_path.unlink(missing_ok=True) self.handle = self.temporary_path.open("wb") self.digest = hashlib.sha256() self.token_count = 0 self.document_count = 0 def append(self, token_ids: list[int]) -> None: payload = np.asarray(token_ids, dtype=self.dtype).tobytes() self.handle.write(payload) self.digest.update(payload) self.token_count += len(token_ids) self.document_count += 1 def finish(self) -> dict[str, Any]: self.handle.flush() os.fsync(self.handle.fileno()) self.handle.close() self.temporary_path.replace(self.final_path) return { "token_count": self.token_count, "document_count": self.document_count, "sha256": self.digest.hexdigest(), } def abort(self) -> None: if not self.handle.closed: self.handle.close() self.temporary_path.unlink(missing_ok=True) def _relative_path(path: Path, root: Path) -> str: try: return path.relative_to(root).as_posix() except ValueError: return str(path) def _completed_source_state( state_path: Path, *, source: Path, corpus_source: CorpusSource, output_dir: Path, tokenizer_sha256: str, validation_modulus: int, validation_bucket: int, ) -> dict[str, Any] | None: if not state_path.is_file(): return None try: state = _read_json(state_path) except ValueError: return None if ( state.get("format") != SOURCE_STATE_FORMAT or state.get("source_name") != source.name or state.get("source_size") != source.stat().st_size # States written before multi-source support carry no dataset name. or state.get("source_dataset", "fineweb-edu") != corpus_source.name or state.get("dataset_revision") != corpus_source.revision or state.get("tokenizer_sha256") != tokenizer_sha256 or state.get("split_hash_person") != SPLIT_HASH_PERSON.decode("ascii") or state.get("validation_modulus") != validation_modulus or state.get("validation_bucket") != validation_bucket ): return None splits = state.get("splits") if not isinstance(splits, dict): return None try: dtype = np.dtype(state.get("dtype")) except TypeError: return None for split in ("train", "validation"): shard = splits.get(split) if not isinstance(shard, dict) or not isinstance(shard.get("path"), str): return None path = output_dir / shard["path"] expected_bytes = int(shard.get("token_count", -1)) * dtype.itemsize expected_sha256 = shard.get("sha256") if ( not path.is_file() or path.stat().st_size != expected_bytes or not isinstance(expected_sha256, str) or _sha256_file(path) != expected_sha256 ): return None return state def _encode_source( source: Path, *, corpus_source: CorpusSource, tokenizer_path: Path, output_dir: Path, batch_size: int, validation_modulus: int, validation_bucket: int, ) -> tuple[dict[str, Any], bool]: tokenizer_sha256 = hashlib.sha256(tokenizer_path.read_bytes()).hexdigest() # Hub Parquet names are unique and stable. Avoid list indexes so a partial download can be # expanded later without invalidating already completed source shards. source_key = source.stem state_path = output_dir / "state" / f"{source_key}.json" resumed = _completed_source_state( state_path, source=source, corpus_source=corpus_source, output_dir=output_dir, tokenizer_sha256=tokenizer_sha256, validation_modulus=validation_modulus, validation_bucket=validation_bucket, ) if resumed is not None: return resumed, True tokenizer = load_tokenizer(tokenizer_path) role_ids = special_token_ids(tokenizer) reserved_ids = set(role_ids.values()) vocab_size = tokenizer.get_vocab_size(with_added_tokens=True) dtype = np.dtype("uint16" if vocab_size <= np.iinfo(np.uint16).max else "uint32") final_paths = { split: output_dir / "shards" / split / f"{source_key}.bin" for split in ("train", "validation") } writers = {split: _SplitWriter(path, dtype) for split, path in final_paths.items()} skipped_empty = 0 skipped_invalid = 0 skipped_special = 0 rows_seen = 0 try: for _path, ids, texts in _iter_parquet_batches( [source], batch_size=batch_size, text_column=corpus_source.text_column, id_column=corpus_source.id_column, ): valid_rows: list[tuple[str, str]] = [] for document_id, text in zip(ids, texts, strict=True): rows_seen += 1 if not isinstance(document_id, str) or not document_id: skipped_invalid += 1 elif not isinstance(text, str): skipped_invalid += 1 elif not text: skipped_empty += 1 else: valid_rows.append((document_id, text)) if not valid_rows: continue encodings = tokenizer.encode_batch( [text for _document_id, text in valid_rows], add_special_tokens=False ) for (document_id, _text), encoding in zip(valid_rows, encodings, strict=True): token_ids = encoding.ids if reserved_ids.intersection(token_ids): # The new sentinel strings make this practically impossible, but skipping is # preferable to losing hours of preprocessing if an exact literal is present. skipped_special += 1 continue split = document_split( document_id, validation_modulus=validation_modulus, validation_bucket=validation_bucket, ) writers[split].append([*token_ids, role_ids["eos"]]) split_metadata = {split: writer.finish() for split, writer in writers.items()} except BaseException: for writer in writers.values(): writer.abort() raise for split, metadata in split_metadata.items(): metadata["path"] = _relative_path(final_paths[split], output_dir) metadata["source"] = source.name state: dict[str, Any] = { "format": SOURCE_STATE_FORMAT, "source_name": source.name, "source_size": source.stat().st_size, "source_rows": rows_seen, "source_dataset": corpus_source.name, "dataset_revision": corpus_source.revision, "dtype": dtype.name, "vocab_size": vocab_size, "mask_token_id": role_ids["mask"], "eos_token_id": role_ids["eos"], "special_token_ids": role_ids, "tokenizer_sha256": tokenizer_sha256, "split_hash_person": SPLIT_HASH_PERSON.decode("ascii"), "validation_modulus": validation_modulus, "validation_bucket": validation_bucket, "skipped_empty_documents": skipped_empty, "skipped_invalid_documents": skipped_invalid, "skipped_special_documents": skipped_special, "splits": split_metadata, } # The marker is written last: its presence commits both split files as one source unit. _atomic_json(state_path, state) return state, False def _build_manifest( split: Literal["train", "validation"], *, corpus_source: CorpusSource, source_paths: Sequence[Path], source_states: Sequence[dict[str, Any]], validation_modulus: int, validation_bucket: int, ) -> dict[str, Any]: first = source_states[0] compatible_keys = ( "dtype", "vocab_size", "mask_token_id", "eos_token_id", "special_token_ids", "tokenizer_sha256", ) for state in source_states[1:]: if any(state.get(key) != first.get(key) for key in compatible_keys): raise ValueError("source states use incompatible tokenizer or token formats") shards = [dict(state["splits"][split]) for state in source_states] return { "format": PACKED_MANIFEST_FORMAT, "split": split, "dtype": first["dtype"], "token_count": sum(int(shard["token_count"]) for shard in shards), "document_count": sum(int(shard["document_count"]) for shard in shards), "vocab_size": first["vocab_size"], "mask_token_id": first["mask_token_id"], "eos_token_id": first["eos_token_id"], "special_token_ids": first["special_token_ids"], "tokenizer_sha256": first["tokenizer_sha256"], "dataset": { "repo_id": corpus_source.repo_id, "config": corpus_source.config, "revision": corpus_source.revision, "path_prefix": corpus_source.path_prefix, }, "split_rule": { "algorithm": "blake2b-64", "person": SPLIT_HASH_PERSON.decode("ascii"), "validation_modulus": validation_modulus, "validation_bucket": validation_bucket, }, "source_files": [path.name for path in source_paths], "skipped_documents": { reason: sum(int(state[reason]) for state in source_states) for reason in ( "skipped_empty_documents", "skipped_invalid_documents", "skipped_special_documents", ) }, "shards": shards, } def prepare_corpus( tokenizer_path: str | Path, output_dir: str | Path, *, corpus_source: CorpusSource, source_paths: Iterable[str | Path] | None = None, raw_dir: str | Path | None = None, batch_size: int = 256, validation_modulus: int = 1024, validation_bucket: int = 0, max_files: int | None = None, ) -> CorpusPreparationResult: """Convert one pinned source's Parquets into resumable train/validation manifests.""" # Validate split arguments before performing any download. document_split( "argument-validation", validation_modulus=validation_modulus, validation_bucket=validation_bucket, ) tokenizer = Path(tokenizer_path) load_tokenizer(tokenizer) destination = Path(output_dir) destination.mkdir(parents=True, exist_ok=True) if source_paths is None: raw = Path(raw_dir) if raw_dir is not None else destination / "raw" try: sources = find_local_parquets(raw, corpus_source.path_prefix) except FileNotFoundError: sources = download_source(corpus_source, raw, max_files=max_files) else: sources = sorted(Path(path) for path in source_paths) if not sources: raise ValueError("at least one source Parquet is required") missing = [str(path) for path in sources if not path.is_file()] if missing: raise FileNotFoundError(f"missing source Parquets: {missing}") source_keys = [source.stem for source in sources] if len(set(source_keys)) != len(source_keys): raise ValueError("source Parquet filenames must have unique stems") states: list[dict[str, Any]] = [] resumed_sources = 0 for source in sources: state, resumed = _encode_source( source, corpus_source=corpus_source, tokenizer_path=tokenizer, output_dir=destination, batch_size=batch_size, validation_modulus=validation_modulus, validation_bucket=validation_bucket, ) states.append(state) resumed_sources += int(resumed) manifest_paths = { "train": destination / "train.manifest.json", "validation": destination / "validation.manifest.json", } for split, manifest_path in manifest_paths.items(): manifest = _build_manifest( split, # type: ignore[arg-type] corpus_source=corpus_source, source_paths=sources, source_states=states, validation_modulus=validation_modulus, validation_bucket=validation_bucket, ) _atomic_json(manifest_path, manifest) return CorpusPreparationResult( train_manifest=manifest_paths["train"], validation_manifest=manifest_paths["validation"], processed_sources=len(sources) - resumed_sources, resumed_sources=resumed_sources, ) def prepare_fineweb_edu( tokenizer_path: str | Path, output_dir: str | Path, *, source_paths: Iterable[str | Path] | None = None, raw_dir: str | Path | None = None, batch_size: int = 256, validation_modulus: int = 1024, validation_bucket: int = 0, ) -> CorpusPreparationResult: """Convert pinned FineWeb-Edu Parquets into manifests (compatibility entry point).""" return prepare_corpus( tokenizer_path, output_dir, corpus_source=CORPUS_SOURCES["fineweb-edu"], source_paths=source_paths, raw_dir=raw_dir, batch_size=batch_size, validation_modulus=validation_modulus, validation_bucket=validation_bucket, ) def parse_token_budget(text: str) -> int: """Parse a token count with an optional K/M/B suffix (e.g. ``2.5B``, ``500M``).""" value = text.strip().upper() factor = 1 if value and value[-1] in "KMB": factor = {"K": 1_000, "M": 1_000_000, "B": 1_000_000_000}[value[-1]] value = value[:-1] try: tokens = int(float(value) * factor) except ValueError as exc: raise ValueError(f"invalid token budget {text!r}") from exc if tokens <= 0: raise ValueError(f"token budget must be positive: {text!r}") return tokens def mix_manifests( inputs: Sequence[tuple[Path, int | None]], output_dir: str | Path ) -> tuple[Path, Path]: """Combine prepared source directories into one mixture manifest pair. Each input contributes whole train shards in manifest order until its token budget is met (``None`` takes everything), so realized counts overshoot a budget by at most one shard; the overshoot is printed, never silent. Validation includes the validation shards of the sources whose train shards were selected. Shard paths in the mixed manifests are absolute so the inputs can live anywhere. """ from diffusion_lm.data import load_packed_manifest if not inputs: raise ValueError("at least one input directory is required") destination = Path(output_dir) destination.mkdir(parents=True, exist_ok=True) compatible_keys = ( "dtype", "vocab_size", "mask_token_id", "eos_token_id", "special_token_ids", "tokenizer_sha256", ) reference: dict[str, Any] | None = None shards: dict[str, list[dict[str, Any]]] = {"train": [], "validation": []} components: list[dict[str, Any]] = [] for input_dir, budget in inputs: manifests = { split: load_packed_manifest(input_dir / f"{split}.manifest.json") for split in ("train", "validation") } if reference is None: reference = manifests["train"] elif any( manifests["train"].get(key) != reference.get(key) for key in compatible_keys ): raise ValueError(f"{input_dir} uses an incompatible tokenizer or token format") taken = 0 selected_sources: set[str] = set() skipped = 0 for shard in manifests["train"]["shards"]: if budget is not None and taken >= budget: skipped += 1 continue entry = dict(shard) entry["path"] = str((input_dir / entry["path"]).resolve()) shards["train"].append(entry) taken += int(entry["token_count"]) selected_sources.add(str(entry.get("source"))) for shard in manifests["validation"]["shards"]: if str(shard.get("source")) not in selected_sources: continue entry = dict(shard) entry["path"] = str((input_dir / entry["path"]).resolve()) shards["validation"].append(entry) component = { "dataset": manifests["train"].get("dataset"), "directory": str(Path(input_dir).resolve()), "token_budget": budget, "train_tokens": taken, "skipped_shards": skipped, } components.append(component) print( f'{input_dir}: {taken:,} train tokens' + (f" (budget {budget:,}, {skipped} shards skipped)" if budget else "") ) assert reference is not None manifest_paths: dict[str, Path] = {} for split in ("train", "validation"): manifest = { "format": PACKED_MANIFEST_FORMAT, "split": split, "dtype": reference["dtype"], "token_count": sum(int(shard["token_count"]) for shard in shards[split]), "document_count": sum(int(shard["document_count"]) for shard in shards[split]), "vocab_size": reference["vocab_size"], "mask_token_id": reference["mask_token_id"], "eos_token_id": reference["eos_token_id"], "special_token_ids": reference["special_token_ids"], "tokenizer_sha256": reference["tokenizer_sha256"], "dataset": {"name": "mixture", "components": components}, "split_rule": reference.get("split_rule"), "source_files": [shard.get("source") for shard in shards[split]], "shards": shards[split], } manifest_paths[split] = destination / f"{split}.manifest.json" _atomic_json(manifest_paths[split], manifest) return manifest_paths["train"], manifest_paths["validation"] def _parse_mix_input(text: str) -> tuple[Path, int | None]: directory, separator, budget = text.partition("=") return Path(directory), parse_token_budget(budget) if separator else None def _build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description=__doc__) commands = parser.add_subparsers(dest="command", required=True) source_names = tuple(CORPUS_SOURCES) download = commands.add_parser("download", help="download a pinned source's Parquets") download.add_argument("--raw-dir", type=Path, required=True) download.add_argument("--source", choices=source_names, default="fineweb-edu") download.add_argument("--max-files", type=int) download.add_argument("--max-workers", type=int, default=1) tokenizer = commands.add_parser( "train-tokenizer", help="train a 32K tokenizer from cached Parquets" ) tokenizer.add_argument("--raw-dir", type=Path, required=True) tokenizer.add_argument("--source", choices=source_names, default="fineweb-edu") tokenizer.add_argument("--output", type=Path, required=True) tokenizer.add_argument("--sample-bytes", type=int, default=1 << 29) tokenizer.add_argument("--vocab-size", type=int, default=32_768) tokenizer.add_argument("--min-frequency", type=int, default=10) tokenizer.add_argument("--max-token-length", type=int, default=64) tokenizer.add_argument("--batch-size", type=int, default=512) prepare = commands.add_parser("prepare", help="encode cached Parquets into token shards") prepare.add_argument("--raw-dir", type=Path, required=True) prepare.add_argument("--source", choices=source_names, default="fineweb-edu") prepare.add_argument("--tokenizer", type=Path, required=True) prepare.add_argument("--output-dir", type=Path, required=True) prepare.add_argument("--batch-size", type=int, default=256) prepare.add_argument("--validation-modulus", type=int, default=1024) prepare.add_argument("--validation-bucket", type=int, default=0) mix = commands.add_parser( "mix", help="combine prepared source directories into one mixture manifest" ) mix.add_argument( "--input", action="append", required=True, metavar="DIR[=TOKENS]", help="prepared corpus directory with an optional train-token budget (e.g. 3.2B)", ) mix.add_argument("--output-dir", type=Path, required=True) return parser def main() -> None: args = _build_parser().parse_args() if args.command == "download": source = CORPUS_SOURCES[args.source] paths = download_source( source, args.raw_dir, max_files=args.max_files, max_workers=args.max_workers ) print(f"downloaded {len(paths)} Parquet shards below {args.raw_dir}") elif args.command == "train-tokenizer": source = CORPUS_SOURCES[args.source] paths = find_local_parquets(args.raw_dir, source.path_prefix) tokenizer = train_tokenizer_from_iterator( iter_tokenizer_text( paths, max_utf8_bytes=args.sample_bytes, batch_size=args.batch_size, text_column=source.text_column, id_column=source.id_column, ), args.output, vocab_size=args.vocab_size, min_frequency=args.min_frequency, max_token_length=args.max_token_length, ) print(f"saved {tokenizer.get_vocab_size():,}-token tokenizer to {args.output}") elif args.command == "prepare": source = CORPUS_SOURCES[args.source] paths = find_local_parquets(args.raw_dir, source.path_prefix) result = prepare_corpus( args.tokenizer, args.output_dir, corpus_source=source, source_paths=paths, batch_size=args.batch_size, validation_modulus=args.validation_modulus, validation_bucket=args.validation_bucket, ) print( f"prepared {result.processed_sources} sources " f"({result.resumed_sources} resumed); train={result.train_manifest}, " f"validation={result.validation_manifest}" ) elif args.command == "mix": inputs = [_parse_mix_input(item) for item in args.input] train_manifest, validation_manifest = mix_manifests(inputs, args.output_dir) print(f"mixed manifests: train={train_manifest}, validation={validation_manifest}") if __name__ == "__main__": main()