初始化项目，由ModelHub XC社区提供模型
Model: pathcosmos/frankenstallm Source: Original Platform
This commit is contained in:
ModelHub XC
348
data/prepare.py
Normal file
348
data/prepare.py
Normal file
@@ -0,0 +1,348 @@
|
||||
"""
|
||||
Prepare raw text files (or a HuggingFace dataset) for LLM training.
|
||||
|
||||
Tokenizes all input text, concatenates all token IDs into a single flat
|
||||
sequence, splits into train / validation sets, and saves each as a uint16
|
||||
numpy binary file (.bin) ready for TextDataset / PackedDataset.
|
||||
|
||||
Usage — glob of local text files:
|
||||
python data/prepare.py \
|
||||
--input "data/raw/*.txt" \
|
||||
--output data/train.bin \
|
||||
--val_output data/val.bin \
|
||||
--tokenizer tokenizer/tokenizer.json \
|
||||
--val_split 0.005 \
|
||||
--seed 42
|
||||
|
||||
Usage — HuggingFace dataset (streaming):
|
||||
python data/prepare.py \
|
||||
--hf_dataset allenai/c4 \
|
||||
--hf_subset en \
|
||||
--hf_split train \
|
||||
--hf_text_col text \
|
||||
--output data/train.bin \
|
||||
--val_output data/val.bin \
|
||||
--tokenizer tokenizer/tokenizer.json \
|
||||
--val_split 0.005
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import glob
|
||||
import os
|
||||
import random
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
from tokenizers import Tokenizer
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def load_tokenizer(tokenizer_path: str) -> Tokenizer:
|
||||
path = Path(tokenizer_path)
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"Tokenizer not found: {path}")
|
||||
return Tokenizer.from_file(str(path))
|
||||
|
||||
|
||||
def find_input_files(pattern: str) -> list[str]:
|
||||
"""Resolve a glob pattern or a plain file path to a list of files."""
|
||||
if any(c in pattern for c in ("*", "?", "[")):
|
||||
files = sorted(glob.glob(pattern, recursive=True))
|
||||
else:
|
||||
files = [pattern] if Path(pattern).exists() else []
|
||||
if not files:
|
||||
raise FileNotFoundError(f"No files matched pattern: {pattern!r}")
|
||||
return files
|
||||
|
||||
|
||||
def tokenize_file(path: str, tokenizer: Tokenizer) -> list[int]:
|
||||
"""Read a single text file and return its token IDs."""
|
||||
with open(path, "r", encoding="utf-8", errors="replace") as fh:
|
||||
text = fh.read()
|
||||
return tokenizer.encode(text).ids
|
||||
|
||||
|
||||
def derive_val_path(output_path: Path, val_output_arg: str | None) -> Path:
|
||||
"""Return the val .bin path, either explicitly provided or auto-derived."""
|
||||
if val_output_arg:
|
||||
return Path(val_output_arg)
|
||||
# If the stem contains "train", swap it for "val".
|
||||
if "train" in output_path.name:
|
||||
candidate = output_path.parent / output_path.name.replace("train", "val")
|
||||
if candidate != output_path:
|
||||
return candidate
|
||||
# Generic fallback: append _val before the suffix.
|
||||
return output_path.with_name(output_path.stem + "_val" + output_path.suffix)
|
||||
|
||||
|
||||
def save_bin(tokens: list[int], path: Path) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
np.array(tokens, dtype=np.uint16).tofile(str(path))
|
||||
|
||||
|
||||
def _fmt_bytes(n_tokens: int) -> str:
|
||||
"""Return a human-readable size string for a uint16 token array."""
|
||||
nbytes = n_tokens * 2 # uint16 = 2 bytes per token
|
||||
for unit in ("B", "KB", "MB", "GB", "TB"):
|
||||
if nbytes < 1024:
|
||||
return f"{nbytes:.1f} {unit}"
|
||||
nbytes /= 1024
|
||||
return f"{nbytes:.1f} PB"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Source iterators
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def iter_tokens_from_files(
|
||||
input_files: list[str],
|
||||
tokenizer: Tokenizer,
|
||||
seed: int,
|
||||
) -> tuple[list[int], int]:
|
||||
"""
|
||||
Tokenize every file, shuffle at file level, flatten, and return
|
||||
(all_tokens_shuffled, file_count).
|
||||
"""
|
||||
per_file_tokens: list[list[int]] = []
|
||||
for fpath in tqdm(input_files, desc="Tokenizing", unit="file"):
|
||||
per_file_tokens.append(tokenize_file(fpath, tokenizer))
|
||||
|
||||
rng = random.Random(seed)
|
||||
rng.shuffle(per_file_tokens)
|
||||
|
||||
all_tokens: list[int] = []
|
||||
for toks in per_file_tokens:
|
||||
all_tokens.extend(toks)
|
||||
|
||||
return all_tokens, len(input_files)
|
||||
|
||||
|
||||
def iter_tokens_from_hf(
|
||||
hf_dataset: str,
|
||||
hf_subset: str | None,
|
||||
hf_split: str,
|
||||
hf_text_col: str,
|
||||
tokenizer: Tokenizer,
|
||||
) -> tuple[list[int], int]:
|
||||
"""
|
||||
Stream a HuggingFace dataset row-by-row, tokenize each row's text column,
|
||||
and return (all_tokens, row_count).
|
||||
|
||||
Rows are appended in streaming order; no shuffle is performed here because
|
||||
the stream may be very large. A seed-based split by position is used later.
|
||||
"""
|
||||
try:
|
||||
from datasets import load_dataset
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"The 'datasets' package is required for --hf_dataset. "
|
||||
"Install it with: pip install datasets"
|
||||
)
|
||||
|
||||
print(f"Streaming HuggingFace dataset: {hf_dataset}"
|
||||
+ (f" / {hf_subset}" if hf_subset else "")
|
||||
+ f" split={hf_split}")
|
||||
|
||||
ds = load_dataset(
|
||||
hf_dataset,
|
||||
hf_subset,
|
||||
split=hf_split,
|
||||
streaming=True,
|
||||
trust_remote_code=True,
|
||||
)
|
||||
|
||||
all_tokens: list[int] = []
|
||||
row_count = 0
|
||||
pbar = tqdm(desc="Tokenizing rows", unit="row")
|
||||
for row in ds:
|
||||
text = row.get(hf_text_col, "")
|
||||
if text:
|
||||
all_tokens.extend(tokenizer.encode(text).ids)
|
||||
row_count += 1
|
||||
pbar.update(1)
|
||||
pbar.close()
|
||||
|
||||
return all_tokens, row_count
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Argument parsing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description=(
|
||||
"Tokenize text sources and save as uint16 binary files for LLM training. "
|
||||
"Accepts either a glob of local text files (--input) or a HuggingFace "
|
||||
"dataset (--hf_dataset)."
|
||||
)
|
||||
)
|
||||
|
||||
# --- Input source (mutually exclusive) ---
|
||||
source = parser.add_mutually_exclusive_group()
|
||||
source.add_argument(
|
||||
"--input",
|
||||
default=None,
|
||||
help='Glob pattern or path to a single text file, e.g. "data/raw/*.txt"',
|
||||
)
|
||||
source.add_argument(
|
||||
"--hf_dataset",
|
||||
default=None,
|
||||
metavar="DATASET",
|
||||
help="HuggingFace dataset name, e.g. allenai/c4 (alternative to --input)",
|
||||
)
|
||||
|
||||
# --- HuggingFace-specific options ---
|
||||
parser.add_argument(
|
||||
"--hf_subset",
|
||||
default=None,
|
||||
metavar="SUBSET",
|
||||
help="Dataset subset / config name, e.g. 'en' for allenai/c4",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hf_split",
|
||||
default="train",
|
||||
metavar="SPLIT",
|
||||
help="Dataset split to use (default: train)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hf_text_col",
|
||||
default="text",
|
||||
metavar="COLUMN",
|
||||
help="Name of the text column in the dataset (default: text)",
|
||||
)
|
||||
|
||||
# --- Output paths ---
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
required=True,
|
||||
help="Output path for the training binary, e.g. data/train.bin",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--val_output",
|
||||
default=None,
|
||||
metavar="PATH",
|
||||
help=(
|
||||
"Explicit output path for the validation binary "
|
||||
"(default: auto-derived from --output, e.g. train.bin → val.bin)"
|
||||
),
|
||||
)
|
||||
|
||||
# --- Tokenizer ---
|
||||
parser.add_argument(
|
||||
"--tokenizer",
|
||||
default="tokenizer/tokenizer.json",
|
||||
help="Path to a trained tokenizer JSON file (default: tokenizer/tokenizer.json)",
|
||||
)
|
||||
|
||||
# --- Split / reproducibility ---
|
||||
parser.add_argument(
|
||||
"--val_split",
|
||||
type=float,
|
||||
default=0.005,
|
||||
help="Fraction of tokens reserved for validation (default: 0.005)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--seed",
|
||||
type=int,
|
||||
default=42,
|
||||
help="Random seed for reproducible train/val split (default: 42)",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Require at least one input source.
|
||||
if args.input is None and args.hf_dataset is None:
|
||||
parser.error("One of --input or --hf_dataset is required.")
|
||||
|
||||
return args
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
|
||||
# ---- Load tokenizer ----
|
||||
tokenizer = load_tokenizer(args.tokenizer)
|
||||
vocab_size = tokenizer.get_vocab_size()
|
||||
|
||||
# Warn early if IDs could overflow uint16.
|
||||
if vocab_size > 65535:
|
||||
print(
|
||||
"WARNING: vocab_size > 65535; token IDs above 65535 will be "
|
||||
"truncated when cast to uint16.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
# ---- Collect tokens from the chosen source ----
|
||||
if args.hf_dataset:
|
||||
all_tokens, source_count = iter_tokens_from_hf(
|
||||
hf_dataset=args.hf_dataset,
|
||||
hf_subset=args.hf_subset,
|
||||
hf_split=args.hf_split,
|
||||
hf_text_col=args.hf_text_col,
|
||||
tokenizer=tokenizer,
|
||||
)
|
||||
source_label = f"{source_count:,} rows"
|
||||
else:
|
||||
input_files = find_input_files(args.input)
|
||||
print(f"Found {len(input_files)} input file(s).")
|
||||
all_tokens, source_count = iter_tokens_from_files(
|
||||
input_files=input_files,
|
||||
tokenizer=tokenizer,
|
||||
seed=args.seed,
|
||||
)
|
||||
source_label = f"{source_count:,} files"
|
||||
|
||||
total_tokens = len(all_tokens)
|
||||
|
||||
# ---- Split into train / val ----
|
||||
val_size = max(1, int(total_tokens * args.val_split))
|
||||
train_size = total_tokens - val_size
|
||||
|
||||
train_tokens = all_tokens[:train_size]
|
||||
val_tokens = all_tokens[train_size:]
|
||||
|
||||
# ---- Resolve output paths ----
|
||||
train_path = Path(args.output)
|
||||
val_path = derive_val_path(train_path, args.val_output)
|
||||
|
||||
# ---- Save ----
|
||||
print(f"\nSaving train data -> {train_path}")
|
||||
save_bin(train_tokens, train_path)
|
||||
|
||||
print(f"Saving val data -> {val_path}")
|
||||
save_bin(val_tokens, val_path)
|
||||
|
||||
# ---- Final stats ----
|
||||
tokens_per_step = 8 * 2048 * 4 * 8 # bs=8, seq=2048, accum=4, 8 GPUs
|
||||
estimated_steps = train_size // tokens_per_step
|
||||
|
||||
print()
|
||||
print(f"Tokenizer: {args.tokenizer} (vocab_size={vocab_size:,})")
|
||||
print(f"Total tokens: {total_tokens:,}")
|
||||
print(
|
||||
f"Train tokens: {train_size:,}"
|
||||
f" (stored in {train_path}, {_fmt_bytes(train_size)})"
|
||||
)
|
||||
print(
|
||||
f"Val tokens: {val_size:,}"
|
||||
f" (stored in {val_path}, {_fmt_bytes(val_size)})"
|
||||
)
|
||||
print(
|
||||
f"Estimated steps (bs=8, seq=2048, 8 GPUs, accum=4): {estimated_steps:,}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user