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()