初始化项目，由ModelHub XC社区提供模型
Model: iamrahulreddy/Quintus Source: Original Platform
This commit is contained in:
ModelHub XC
574
src/download.py
Normal file
574
src/download.py
Normal file
@@ -0,0 +1,574 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user