575 lines
22 KiB
Python
575 lines
22 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import platform
|
||
|
|
import sys
|
||
|
|
import time
|
||
|
|
import warnings
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
_REPO_ROOT = Path(__file__).resolve().parents[1]
|
||
|
|
if str(_REPO_ROOT) not in sys.path:
|
||
|
|
sys.path.insert(0, str(_REPO_ROOT))
|
||
|
|
|
||
|
|
os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "0")
|
||
|
|
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
|
||
|
|
warnings.filterwarnings("ignore", category=FutureWarning)
|
||
|
|
warnings.filterwarnings("ignore", category=UserWarning)
|
||
|
|
|
||
|
|
import torch
|
||
|
|
from datasets import load_dataset
|
||
|
|
from huggingface_hub import snapshot_download
|
||
|
|
from transformers import AutoTokenizer
|
||
|
|
|
||
|
|
from configs import cfg, emit_log_spacing, setup_logger
|
||
|
|
from src.transformers_compat import format_model_load_error
|
||
|
|
|
||
|
|
_IGNORE_PATTERNS = ["*.msgpack", "*.h5", "*.bin", "optimizer.pt", "optimizer.safetensors"]
|
||
|
|
_TOKENIZER_ALLOW_PATTERNS = [
|
||
|
|
"tokenizer.json",
|
||
|
|
"tokenizer.model",
|
||
|
|
"tokenizer_config.json",
|
||
|
|
"special_tokens_map.json",
|
||
|
|
"added_tokens.json",
|
||
|
|
"vocab.json",
|
||
|
|
"merges.txt",
|
||
|
|
"generation_config.json",
|
||
|
|
]
|
||
|
|
_MIN_TOKEN_LENGTH = 10
|
||
|
|
_DATA_STATS_FILENAME = "_data_stats.json"
|
||
|
|
_ASSISTANT_MASK_KEYS = ("assistant_masks", "assistant_mask", "assistant_tokens_mask")
|
||
|
|
|
||
|
|
|
||
|
|
def _build_chat_template_error_types() -> tuple[type[BaseException], ...]:
|
||
|
|
error_types: list[type[BaseException]] = [
|
||
|
|
AttributeError,
|
||
|
|
IndexError,
|
||
|
|
KeyError,
|
||
|
|
RuntimeError,
|
||
|
|
TypeError,
|
||
|
|
ValueError,
|
||
|
|
]
|
||
|
|
try:
|
||
|
|
from jinja2 import TemplateError
|
||
|
|
error_types.append(TemplateError)
|
||
|
|
except ImportError:
|
||
|
|
pass
|
||
|
|
return tuple(error_types)
|
||
|
|
|
||
|
|
|
||
|
|
_CHAT_TEMPLATE_ERRORS = _build_chat_template_error_types()
|
||
|
|
|
||
|
|
|
||
|
|
def _config_revision(value: str | None) -> str | None:
|
||
|
|
if value is None:
|
||
|
|
return None
|
||
|
|
stripped = value.strip()
|
||
|
|
return stripped or None
|
||
|
|
|
||
|
|
|
||
|
|
def _download_tokenizer_artifacts(tokenizer_model: str, tokenizer_revision: str | None, tokenizer_dir: str, log) -> None:
|
||
|
|
log.info(f"Downloading tokenizer -> ./{tokenizer_dir}/")
|
||
|
|
t0 = time.time()
|
||
|
|
try:
|
||
|
|
snapshot_download(
|
||
|
|
repo_id=tokenizer_model,
|
||
|
|
local_dir=tokenizer_dir,
|
||
|
|
revision=tokenizer_revision,
|
||
|
|
allow_patterns=_TOKENIZER_ALLOW_PATTERNS,
|
||
|
|
)
|
||
|
|
size_mb = sum(f.stat().st_size for f in Path(tokenizer_dir).rglob("*") if f.is_file()) / 1e6
|
||
|
|
log.info(f"Tokenizer downloaded: {size_mb:.1f} MB in {time.time() - t0:.0f}s")
|
||
|
|
except Exception as exc:
|
||
|
|
log.error(f"Failed to download tokenizer: {exc}")
|
||
|
|
sys.exit(1)
|
||
|
|
|
||
|
|
|
||
|
|
def write_system_info(output_path: str, logger) -> None:
|
||
|
|
output_dir = os.path.dirname(output_path)
|
||
|
|
if output_dir:
|
||
|
|
os.makedirs(output_dir, exist_ok=True)
|
||
|
|
info = {
|
||
|
|
"timestamp": time.strftime("%Y-%m-%d %H:%M:%S %Z"),
|
||
|
|
"platform": platform.platform(),
|
||
|
|
"python": sys.version.split()[0],
|
||
|
|
"torch": torch.__version__,
|
||
|
|
"cuda": torch.version.cuda if torch.cuda.is_available() else None,
|
||
|
|
"gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,
|
||
|
|
"cpu_count": os.cpu_count(),
|
||
|
|
}
|
||
|
|
with open(output_path, "w", encoding="utf-8") as f:
|
||
|
|
json.dump(info, f, indent=2)
|
||
|
|
logger.info(f"System info -> {output_path}")
|
||
|
|
|
||
|
|
|
||
|
|
def write_data_stats(output_path: str, stats: dict, dataset_id: str, config_name: str, target_samples: int, max_seq_len: int, logger) -> None:
|
||
|
|
output_dir = os.path.dirname(output_path)
|
||
|
|
if output_dir:
|
||
|
|
os.makedirs(output_dir, exist_ok=True)
|
||
|
|
meta = {
|
||
|
|
"dataset": dataset_id,
|
||
|
|
"config": config_name,
|
||
|
|
"target_samples": target_samples,
|
||
|
|
"max_seq_len": max_seq_len,
|
||
|
|
"stats": stats,
|
||
|
|
}
|
||
|
|
with open(output_path, "w", encoding="utf-8") as f:
|
||
|
|
json.dump(meta, f, indent=2)
|
||
|
|
logger.info(f"Dataset stats -> {output_path}")
|
||
|
|
|
||
|
|
|
||
|
|
def _coerce_content_str(content) -> str | None:
|
||
|
|
if isinstance(content, str):
|
||
|
|
return content.strip() or None
|
||
|
|
if isinstance(content, dict):
|
||
|
|
for key in ("answer_content", "text", "value", "content", "think_content"):
|
||
|
|
value = content.get(key)
|
||
|
|
if isinstance(value, str) and value.strip():
|
||
|
|
return value.strip()
|
||
|
|
return None
|
||
|
|
if isinstance(content, list):
|
||
|
|
parts = []
|
||
|
|
for item in content:
|
||
|
|
if isinstance(item, str) and item.strip():
|
||
|
|
parts.append(item.strip())
|
||
|
|
elif isinstance(item, dict):
|
||
|
|
value = item.get("text", item.get("value", ""))
|
||
|
|
if isinstance(value, str) and value.strip():
|
||
|
|
parts.append(value.strip())
|
||
|
|
joined = " ".join(parts).strip()
|
||
|
|
return joined or None
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def _extract_clean_sample(row: dict) -> dict | None:
|
||
|
|
messages = row.get("messages", row.get("conversations", []))
|
||
|
|
if not messages and "instruction" in row and "output" in row:
|
||
|
|
messages = [
|
||
|
|
{"role": "user", "content": row["instruction"]},
|
||
|
|
{"role": "assistant", "content": row["output"]},
|
||
|
|
]
|
||
|
|
if isinstance(messages, str):
|
||
|
|
try:
|
||
|
|
messages = json.loads(messages)
|
||
|
|
except (json.JSONDecodeError, TypeError):
|
||
|
|
return None
|
||
|
|
if not isinstance(messages, list) or len(messages) < 2:
|
||
|
|
return None
|
||
|
|
|
||
|
|
clean_messages: list[dict] = []
|
||
|
|
for message in messages:
|
||
|
|
if not isinstance(message, dict):
|
||
|
|
return None
|
||
|
|
role = message.get("role", message.get("from", ""))
|
||
|
|
content = _coerce_content_str(message.get("content", message.get("value", "")))
|
||
|
|
if not isinstance(role, str) or not role.strip() or content is None:
|
||
|
|
return None
|
||
|
|
role = role.strip().lower()
|
||
|
|
if role == "human":
|
||
|
|
role = "user"
|
||
|
|
elif role == "gpt":
|
||
|
|
role = "assistant"
|
||
|
|
clean_messages.append({"role": role, "content": content})
|
||
|
|
|
||
|
|
source = str(row.get("source", "unknown"))
|
||
|
|
return {"messages": clean_messages, "source": source}
|
||
|
|
|
||
|
|
|
||
|
|
def _coerce_token_ids(token_ids) -> list[int]:
|
||
|
|
if hasattr(token_ids, "input_ids"):
|
||
|
|
token_ids = token_ids.input_ids
|
||
|
|
if isinstance(token_ids, dict):
|
||
|
|
token_ids = token_ids.get("input_ids", token_ids)
|
||
|
|
if hasattr(token_ids, "tolist"):
|
||
|
|
token_ids = token_ids.tolist()
|
||
|
|
if isinstance(token_ids, tuple):
|
||
|
|
token_ids = list(token_ids)
|
||
|
|
if not isinstance(token_ids, list):
|
||
|
|
raise TypeError(f"Unexpected token id payload: {type(token_ids).__name__}")
|
||
|
|
if token_ids and isinstance(token_ids[0], list):
|
||
|
|
raise TypeError("Expected a single token sequence, not a batched payload")
|
||
|
|
return [int(token_id) for token_id in token_ids]
|
||
|
|
|
||
|
|
|
||
|
|
def _coerce_binary_mask(mask_values, expected_len: int) -> list[int]:
|
||
|
|
if hasattr(mask_values, "tolist"):
|
||
|
|
mask_values = mask_values.tolist()
|
||
|
|
if isinstance(mask_values, tuple):
|
||
|
|
mask_values = list(mask_values)
|
||
|
|
if not isinstance(mask_values, list):
|
||
|
|
raise TypeError(f"Unexpected assistant mask payload: {type(mask_values).__name__}")
|
||
|
|
if mask_values and isinstance(mask_values[0], list):
|
||
|
|
raise TypeError("Expected a single assistant mask, not a batched payload")
|
||
|
|
mask = [1 if int(value) != 0 else 0 for value in mask_values]
|
||
|
|
if len(mask) != expected_len:
|
||
|
|
raise ValueError(f"Assistant mask length mismatch: got {len(mask)}, expected {expected_len}")
|
||
|
|
return mask
|
||
|
|
|
||
|
|
|
||
|
|
def _try_builtin_assistant_mask(sample: dict, tokenizer: AutoTokenizer, input_ids: list[int]) -> list[int] | None:
|
||
|
|
try:
|
||
|
|
with warnings.catch_warnings():
|
||
|
|
warnings.filterwarnings("ignore", message="return_assistant_tokens_mask")
|
||
|
|
encoded = tokenizer.apply_chat_template(
|
||
|
|
sample["messages"],
|
||
|
|
tokenize=True,
|
||
|
|
add_generation_prompt=False,
|
||
|
|
return_dict=True,
|
||
|
|
return_assistant_tokens_mask=True,
|
||
|
|
)
|
||
|
|
except TypeError:
|
||
|
|
return None
|
||
|
|
except _CHAT_TEMPLATE_ERRORS:
|
||
|
|
return None
|
||
|
|
if not hasattr(encoded, "get"):
|
||
|
|
return None
|
||
|
|
try:
|
||
|
|
encoded_ids = _coerce_token_ids(encoded)
|
||
|
|
except _CHAT_TEMPLATE_ERRORS:
|
||
|
|
return None
|
||
|
|
if encoded_ids != input_ids:
|
||
|
|
return None
|
||
|
|
for key in _ASSISTANT_MASK_KEYS:
|
||
|
|
if key not in encoded:
|
||
|
|
continue
|
||
|
|
try:
|
||
|
|
mask = _coerce_binary_mask(encoded[key], expected_len=len(input_ids))
|
||
|
|
except (TypeError, ValueError):
|
||
|
|
return None
|
||
|
|
if any(mask):
|
||
|
|
return mask
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def _build_assistant_mask_from_prefixes(sample: dict, tokenizer: AutoTokenizer, input_ids: list[int]) -> list[int]:
|
||
|
|
loss_mask = [0] * len(input_ids)
|
||
|
|
prefix_ids: list[int] = []
|
||
|
|
for turn_index, message in enumerate(sample["messages"], start=1):
|
||
|
|
role = str(message.get("role", "")).strip().lower()
|
||
|
|
if role == "assistant":
|
||
|
|
# prompt_ids contains everything up to user message + assistant header (<|im_start|>assistant\n)
|
||
|
|
prompt_ids = _coerce_token_ids(
|
||
|
|
tokenizer.apply_chat_template(
|
||
|
|
sample["messages"][:turn_index-1],
|
||
|
|
tokenize=True,
|
||
|
|
add_generation_prompt=True,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
# full_ids contains prompt + assistant response content + eos
|
||
|
|
full_ids = _coerce_token_ids(
|
||
|
|
tokenizer.apply_chat_template(
|
||
|
|
sample["messages"][:turn_index],
|
||
|
|
tokenize=True,
|
||
|
|
add_generation_prompt=False,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
if len(full_ids) < len(prompt_ids) or full_ids[:len(prompt_ids)] != prompt_ids:
|
||
|
|
raise ValueError("Chat template is not prefix-stable enough to derive assistant-only targets")
|
||
|
|
|
||
|
|
# Loss mask is 1 only for assistant's content tokens (after prompt_ids)
|
||
|
|
for j in range(len(prompt_ids), len(full_ids)):
|
||
|
|
loss_mask[j] = 1
|
||
|
|
prefix_ids = _coerce_token_ids(
|
||
|
|
tokenizer.apply_chat_template(
|
||
|
|
sample["messages"][:turn_index],
|
||
|
|
tokenize=True,
|
||
|
|
add_generation_prompt=False,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
if prefix_ids != input_ids:
|
||
|
|
raise ValueError("Prefix tokenization mismatch while deriving assistant-only targets")
|
||
|
|
if len(loss_mask) != len(input_ids):
|
||
|
|
raise ValueError(f"Assistant mask length mismatch: got {len(loss_mask)}, expected {len(input_ids)}")
|
||
|
|
return loss_mask
|
||
|
|
|
||
|
|
|
||
|
|
def _build_assistant_loss_mask(sample: dict, tokenizer: AutoTokenizer, input_ids: list[int]) -> list[int]:
|
||
|
|
builtin_mask = _try_builtin_assistant_mask(sample, tokenizer, input_ids)
|
||
|
|
if builtin_mask is not None:
|
||
|
|
return builtin_mask
|
||
|
|
return _build_assistant_mask_from_prefixes(sample, tokenizer, input_ids)
|
||
|
|
|
||
|
|
|
||
|
|
def write_tokenized_dataset(tokenizer, num_samples: int, out_file: str, log) -> dict:
|
||
|
|
if num_samples <= 0:
|
||
|
|
raise RuntimeError("num_samples must be positive when tokenization is enabled")
|
||
|
|
output_dir = os.path.dirname(out_file)
|
||
|
|
if output_dir:
|
||
|
|
os.makedirs(output_dir, exist_ok=True)
|
||
|
|
|
||
|
|
log.info(f"Streaming {cfg.data.dataset_path}...")
|
||
|
|
ds = load_dataset(cfg.data.dataset_path, split="train", streaming=True, token=cfg.hub.token)
|
||
|
|
buffer_size = int(getattr(cfg.data, "stream_shuffle_buffer_size", 0) or 0)
|
||
|
|
if buffer_size > 0 and hasattr(ds, "shuffle"):
|
||
|
|
seed = int(getattr(cfg.data, "stream_shuffle_seed", 42))
|
||
|
|
log.info(f"Shuffling stream with buffer_size={buffer_size:,} seed={seed}")
|
||
|
|
ds = ds.shuffle(seed=seed, buffer_size=buffer_size)
|
||
|
|
|
||
|
|
# Check if dataset is pre-tokenized
|
||
|
|
try:
|
||
|
|
first_sample = next(iter(ds))
|
||
|
|
is_pre_tokenized = "input_ids" in first_sample and "loss_mask" in first_sample
|
||
|
|
except StopIteration:
|
||
|
|
raise RuntimeError("Loaded dataset is empty")
|
||
|
|
|
||
|
|
stats = {
|
||
|
|
"scanned": 0,
|
||
|
|
"written": 0,
|
||
|
|
"too_long_tokens": 0,
|
||
|
|
"too_short_tokens": 0,
|
||
|
|
"template_errors": 0,
|
||
|
|
"no_target_tokens": 0,
|
||
|
|
"invalid_messages": 0,
|
||
|
|
"total_tokens_written": 0,
|
||
|
|
"total_target_tokens_written": 0,
|
||
|
|
"min_tokens_written": 0,
|
||
|
|
"max_tokens_written": 0,
|
||
|
|
"min_target_tokens_written": 0,
|
||
|
|
"max_target_tokens_written": 0,
|
||
|
|
}
|
||
|
|
|
||
|
|
if is_pre_tokenized:
|
||
|
|
log.info("Auto-detected pre-tokenized dataset on HF Hub. Writing directly to train.jsonl...")
|
||
|
|
# Re-initialize to avoid losing the first element consumed by next(iter())
|
||
|
|
ds = load_dataset(cfg.data.dataset_path, split="train", streaming=True, token=cfg.hub.token)
|
||
|
|
if buffer_size > 0 and hasattr(ds, "shuffle"):
|
||
|
|
ds = ds.shuffle(seed=seed, buffer_size=buffer_size)
|
||
|
|
|
||
|
|
with open(out_file, "w", encoding="utf-8") as f:
|
||
|
|
for row in ds:
|
||
|
|
stats["scanned"] += 1
|
||
|
|
input_ids = row["input_ids"]
|
||
|
|
loss_mask = row["loss_mask"]
|
||
|
|
token_len = len(input_ids)
|
||
|
|
target_tokens = sum(loss_mask)
|
||
|
|
|
||
|
|
out_row = {
|
||
|
|
"input_ids": input_ids,
|
||
|
|
"loss_mask": loss_mask,
|
||
|
|
"length": token_len,
|
||
|
|
"target_tokens": target_tokens,
|
||
|
|
"source": row.get("source", "unknown"),
|
||
|
|
}
|
||
|
|
f.write(json.dumps(out_row) + "\n")
|
||
|
|
|
||
|
|
stats["written"] += 1
|
||
|
|
stats["total_tokens_written"] += token_len
|
||
|
|
stats["total_target_tokens_written"] += target_tokens
|
||
|
|
if stats["written"] == 1:
|
||
|
|
stats["min_tokens_written"] = token_len
|
||
|
|
stats["max_tokens_written"] = token_len
|
||
|
|
stats["min_target_tokens_written"] = target_tokens
|
||
|
|
stats["max_target_tokens_written"] = target_tokens
|
||
|
|
else:
|
||
|
|
stats["min_tokens_written"] = min(stats["min_tokens_written"], token_len)
|
||
|
|
stats["max_tokens_written"] = max(stats["max_tokens_written"], token_len)
|
||
|
|
stats["min_target_tokens_written"] = min(stats["min_target_tokens_written"], target_tokens)
|
||
|
|
stats["max_target_tokens_written"] = max(stats["max_target_tokens_written"], target_tokens)
|
||
|
|
|
||
|
|
if stats["written"] >= num_samples:
|
||
|
|
break
|
||
|
|
return stats
|
||
|
|
|
||
|
|
log.info("Standard raw text dataset detected. Running tokenization locally...")
|
||
|
|
with open(out_file, "w", encoding="utf-8") as f:
|
||
|
|
for row in ds:
|
||
|
|
stats["scanned"] += 1
|
||
|
|
sample = _extract_clean_sample(row)
|
||
|
|
if sample is None:
|
||
|
|
stats["invalid_messages"] += 1
|
||
|
|
continue
|
||
|
|
try:
|
||
|
|
input_ids = _coerce_token_ids(
|
||
|
|
tokenizer.apply_chat_template(
|
||
|
|
sample["messages"],
|
||
|
|
tokenize=True,
|
||
|
|
add_generation_prompt=False,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
token_len = len(input_ids)
|
||
|
|
if token_len < _MIN_TOKEN_LENGTH:
|
||
|
|
stats["too_short_tokens"] += 1
|
||
|
|
continue
|
||
|
|
if token_len > cfg.data.max_seq_len:
|
||
|
|
stats["too_long_tokens"] += 1
|
||
|
|
continue
|
||
|
|
loss_mask = _build_assistant_loss_mask(sample, tokenizer, input_ids)
|
||
|
|
except _CHAT_TEMPLATE_ERRORS:
|
||
|
|
stats["template_errors"] += 1
|
||
|
|
continue
|
||
|
|
|
||
|
|
target_tokens = sum(loss_mask)
|
||
|
|
if target_tokens == 0:
|
||
|
|
stats["no_target_tokens"] += 1
|
||
|
|
continue
|
||
|
|
|
||
|
|
out_row = {
|
||
|
|
"input_ids": input_ids,
|
||
|
|
"loss_mask": loss_mask,
|
||
|
|
"length": token_len,
|
||
|
|
"target_tokens": target_tokens,
|
||
|
|
"source": sample.get("source", "unknown"),
|
||
|
|
}
|
||
|
|
f.write(json.dumps(out_row) + "\n")
|
||
|
|
|
||
|
|
stats["written"] += 1
|
||
|
|
stats["total_tokens_written"] += token_len
|
||
|
|
stats["total_target_tokens_written"] += target_tokens
|
||
|
|
if stats["written"] == 1:
|
||
|
|
stats["min_tokens_written"] = token_len
|
||
|
|
stats["max_tokens_written"] = token_len
|
||
|
|
stats["min_target_tokens_written"] = target_tokens
|
||
|
|
stats["max_target_tokens_written"] = target_tokens
|
||
|
|
else:
|
||
|
|
stats["min_tokens_written"] = min(stats["min_tokens_written"], token_len)
|
||
|
|
stats["max_tokens_written"] = max(stats["max_tokens_written"], token_len)
|
||
|
|
stats["min_target_tokens_written"] = min(stats["min_target_tokens_written"], target_tokens)
|
||
|
|
stats["max_target_tokens_written"] = max(stats["max_target_tokens_written"], target_tokens)
|
||
|
|
|
||
|
|
if stats["written"] >= num_samples:
|
||
|
|
break
|
||
|
|
|
||
|
|
return stats
|
||
|
|
|
||
|
|
def main() -> None:
|
||
|
|
parser = argparse.ArgumentParser(description="Download models and tokenize data")
|
||
|
|
parser.add_argument("--num_samples", type=int, default=cfg.data.num_samples)
|
||
|
|
parser.add_argument("--skip_teacher", action="store_true", help="Skip teacher download.")
|
||
|
|
parser.add_argument("--skip_tokenization", action="store_true", help="Skip data tokenization.")
|
||
|
|
parser.add_argument("--tokenizer_only", action="store_true", help="Download tokenizer artifacts only")
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
log = setup_logger("DOWNLOAD")
|
||
|
|
log.info("=" * 70)
|
||
|
|
log.info("Download and tokenize")
|
||
|
|
log.info("=" * 70)
|
||
|
|
|
||
|
|
write_system_info(cfg.paths.system_info, log)
|
||
|
|
|
||
|
|
log.info(f" Teacher: {cfg.model.teacher}")
|
||
|
|
log.info(f" Teacher rev: {_config_revision(cfg.model.teacher_revision) or 'unversioned'}")
|
||
|
|
tokenizer_model = getattr(cfg.model, "tokenizer", cfg.model.student)
|
||
|
|
tokenizer_revision = _config_revision(getattr(cfg.model, "tokenizer_revision", cfg.model.student_revision))
|
||
|
|
tokenizer_dir = getattr(cfg.paths, "tokenizer_dir", cfg.paths.student_dir)
|
||
|
|
log.info(f" Student: {cfg.model.student}")
|
||
|
|
log.info(f" Student rev: {_config_revision(cfg.model.student_revision) or 'unversioned'}")
|
||
|
|
log.info(f" Tokenizer: {tokenizer_model}")
|
||
|
|
log.info(f" Tokenizer rev:{tokenizer_revision or 'unversioned'}")
|
||
|
|
log.info(f" Student dir: {cfg.paths.student_dir}")
|
||
|
|
log.info(f" Tokenizer dir:{tokenizer_dir}")
|
||
|
|
log.info(f" Remote code: {cfg.model.allow_remote_code}")
|
||
|
|
log.info(f" Dataset: {cfg.data.dataset_path}")
|
||
|
|
log.info(f" Num samples: {args.num_samples:,}")
|
||
|
|
log.info(f" Max seq len: {cfg.data.max_seq_len}")
|
||
|
|
if torch.cuda.is_available():
|
||
|
|
log.info(f" GPU: {torch.cuda.get_device_name(0)}")
|
||
|
|
|
||
|
|
if not args.tokenizer_only:
|
||
|
|
emit_log_spacing(log)
|
||
|
|
log.info("-" * 70)
|
||
|
|
log.info(f"Downloading student -> ./{cfg.paths.student_dir}/")
|
||
|
|
t0 = time.time()
|
||
|
|
try:
|
||
|
|
snapshot_download(
|
||
|
|
repo_id=cfg.model.student,
|
||
|
|
local_dir=cfg.paths.student_dir,
|
||
|
|
revision=_config_revision(cfg.model.student_revision),
|
||
|
|
ignore_patterns=_IGNORE_PATTERNS,
|
||
|
|
)
|
||
|
|
size_gb = sum(f.stat().st_size for f in Path(cfg.paths.student_dir).rglob("*") if f.is_file()) / 1e9
|
||
|
|
log.info(f"Student downloaded: {size_gb:.1f} GB in {time.time() - t0:.0f}s")
|
||
|
|
except Exception as exc:
|
||
|
|
log.error(f"Failed to download student: {exc}")
|
||
|
|
sys.exit(1)
|
||
|
|
|
||
|
|
if args.tokenizer_only or Path(tokenizer_dir).resolve() != Path(cfg.paths.student_dir).resolve():
|
||
|
|
emit_log_spacing(log)
|
||
|
|
log.info("-" * 70)
|
||
|
|
_download_tokenizer_artifacts(tokenizer_model, tokenizer_revision, tokenizer_dir, log)
|
||
|
|
|
||
|
|
if not args.skip_tokenization:
|
||
|
|
emit_log_spacing(log)
|
||
|
|
log.info("-" * 70)
|
||
|
|
log.info(f"Preparing dataset: {cfg.data.dataset_path}")
|
||
|
|
|
||
|
|
try:
|
||
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
||
|
|
tokenizer_dir,
|
||
|
|
trust_remote_code=cfg.model.allow_remote_code,
|
||
|
|
)
|
||
|
|
except Exception as exc:
|
||
|
|
log.error(format_model_load_error("Student tokenizer load", exc))
|
||
|
|
sys.exit(1)
|
||
|
|
if tokenizer.pad_token is None:
|
||
|
|
tokenizer.pad_token = tokenizer.eos_token
|
||
|
|
|
||
|
|
os.makedirs(cfg.paths.tokenized_dir, exist_ok=True)
|
||
|
|
out_file = os.path.join(cfg.paths.tokenized_dir, "train.jsonl")
|
||
|
|
stats_file = os.path.join(cfg.paths.tokenized_dir, _DATA_STATS_FILENAME)
|
||
|
|
|
||
|
|
log.info(
|
||
|
|
f"Streaming + tokenizing up to {args.num_samples:,} samples "
|
||
|
|
f"(max_seq_len={cfg.data.max_seq_len}, strict token limit, no truncation)"
|
||
|
|
)
|
||
|
|
t0 = time.time()
|
||
|
|
try:
|
||
|
|
token_stats = write_tokenized_dataset(
|
||
|
|
tokenizer=tokenizer,
|
||
|
|
num_samples=args.num_samples,
|
||
|
|
out_file=out_file,
|
||
|
|
log=log,
|
||
|
|
)
|
||
|
|
except RuntimeError as exc:
|
||
|
|
log.error(str(exc))
|
||
|
|
sys.exit(1)
|
||
|
|
|
||
|
|
write_data_stats(
|
||
|
|
output_path=stats_file,
|
||
|
|
stats=token_stats,
|
||
|
|
dataset_id=cfg.data.dataset_path,
|
||
|
|
config_name="default",
|
||
|
|
target_samples=args.num_samples,
|
||
|
|
max_seq_len=cfg.data.max_seq_len,
|
||
|
|
logger=log,
|
||
|
|
)
|
||
|
|
n_written = token_stats["written"]
|
||
|
|
|
||
|
|
if n_written == 0:
|
||
|
|
log.error("Tokenization produced 0 usable rows - aborting.")
|
||
|
|
sys.exit(1)
|
||
|
|
|
||
|
|
log.info(f"Pretokenization complete: {n_written:,} samples -> {out_file}")
|
||
|
|
else:
|
||
|
|
log.info("Skipping dataset tokenization (--skip_tokenization)")
|
||
|
|
|
||
|
|
if not args.skip_teacher:
|
||
|
|
emit_log_spacing(log)
|
||
|
|
log.info("-" * 70)
|
||
|
|
log.info(f"Downloading teacher -> ./{cfg.paths.teacher_dir}/")
|
||
|
|
t0 = time.time()
|
||
|
|
try:
|
||
|
|
snapshot_download(
|
||
|
|
repo_id=cfg.model.teacher,
|
||
|
|
local_dir=cfg.paths.teacher_dir,
|
||
|
|
revision=_config_revision(cfg.model.teacher_revision),
|
||
|
|
ignore_patterns=_IGNORE_PATTERNS,
|
||
|
|
)
|
||
|
|
size_gb = sum(f.stat().st_size for f in Path(cfg.paths.teacher_dir).rglob("*") if f.is_file()) / 1e9
|
||
|
|
log.info(f"Teacher downloaded: {size_gb:.1f} GB in {time.time() - t0:.0f}s")
|
||
|
|
except Exception as exc:
|
||
|
|
log.error(f"Failed to download teacher: {exc}")
|
||
|
|
sys.exit(1)
|
||
|
|
else:
|
||
|
|
log.info("Skipping teacher download (--skip_teacher)")
|
||
|
|
|
||
|
|
emit_log_spacing(log)
|
||
|
|
log.info("-" * 70)
|
||
|
|
log.info("Download complete")
|
||
|
|
log.info("-" * 70)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|