From 1b6b428aaa836e4fd45f0058d44c292f6c9337f8 Mon Sep 17 00:00:00 2001 From: Alpamys Date: Fri, 3 Apr 2026 16:35:23 +0500 Subject: [PATCH] =?UTF-8?q?feat:=20v0.24.0=20=E2=80=94=20Dataset=20Hub,=20?= =?UTF-8?q?Freeze=20Training,=20Loss=20Watchdog,=20Dataset=20Registry?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Part A: HuggingFace Dataset browser - soup data search: search HF Hub for datasets (sort by downloads/likes) - soup data preview: preview remote dataset metadata, splits, features - soup data download: stream HF dataset to local JSONL (with format conversion) - Security: trust_remote_code=False, path traversal protection, samples cap at 1M Part B: Freeze training (like LLaMA-Factory finetuning_type: freeze) - freeze_layers / freeze_ratio config fields - soup_cli/utils/freeze.py: detect layers, freeze bottom N - Wired into SFT trainer before LoRA application - Supports LLaMA (layers.N) and GPT-2 (h.N) naming Part C: Loss watchdog (like Axolotl loss_watchdog_threshold) - loss_watchdog, loss_watchdog_threshold, loss_watchdog_patience config - Implemented in SoupTrainerCallback with patience counter - Rich warning panel (stops Live display first), fires only once - Wired into all 11 trainers via callback kwargs Part D: Dataset info registry - soup data register/unregister/registry commands - ~/.soup/datasets.json local name→path+format mapping - Name validation, path traversal protection, Rich markup escaping 82 new tests (2061 total), 74 test files. --- .claude/CLAUDE.md | 33 ++- README.md | 32 +++ SECURITY.md | 6 +- pyproject.toml | 2 +- soup_cli/__init__.py | 2 +- soup_cli/commands/data.py | 451 +++++++++++++++++++++++++++++++ soup_cli/config/schema.py | 30 ++ soup_cli/monitoring/callback.py | 36 +++ soup_cli/monitoring/display.py | 3 +- soup_cli/trainer/dpo.py | 7 +- soup_cli/trainer/embedding.py | 7 +- soup_cli/trainer/grpo.py | 7 +- soup_cli/trainer/ipo.py | 7 +- soup_cli/trainer/kto.py | 7 +- soup_cli/trainer/orpo.py | 7 +- soup_cli/trainer/ppo.py | 7 +- soup_cli/trainer/pretrain.py | 7 +- soup_cli/trainer/reward_model.py | 7 +- soup_cli/trainer/sft.py | 20 +- soup_cli/trainer/simpo.py | 7 +- soup_cli/utils/freeze.py | 73 +++++ soup_cli/utils/registry.py | 100 +++++++ tests/test_dataset_hub.py | 357 ++++++++++++++++++++++++ tests/test_dataset_registry.py | 254 +++++++++++++++++ tests/test_freeze_training.py | 272 +++++++++++++++++++ tests/test_loss_watchdog.py | 218 +++++++++++++++ 26 files changed, 1936 insertions(+), 23 deletions(-) create mode 100644 soup_cli/utils/freeze.py create mode 100644 soup_cli/utils/registry.py create mode 100644 tests/test_dataset_hub.py create mode 100644 tests/test_dataset_registry.py create mode 100644 tests/test_freeze_training.py create mode 100644 tests/test_loss_watchdog.py diff --git a/.claude/CLAUDE.md b/.claude/CLAUDE.md index 9c1569a..19ef0f4 100644 --- a/.claude/CLAUDE.md +++ b/.claude/CLAUDE.md @@ -1,12 +1,12 @@ # Soup CLI — Project CLAUDE.md -Soup is a CLI-first LLM fine-tuning tool (v0.23.1). Python 3.9+, MIT license. +Soup is a CLI-first LLM fine-tuning tool (v0.24.0). Python 3.9+, MIT license. ## Build & Development ```bash pip install -e ".[dev]" # Install editable + test deps -pytest tests/ -v --tb=short # Run all tests (1979 tests) +pytest tests/ -v --tb=short # Run all tests (2061 tests) ruff check soup_cli/ tests/ # Lint (must pass before commit) ruff check --fix soup_cli/ tests/ # Auto-fix lint issues ``` @@ -49,7 +49,7 @@ soup_cli/ rewards.py # Built-in reward fns (accuracy, format) + custom .py loader commands/ ... - data.py # soup data (inspect/validate/convert/merge/dedup/stats/filter/split) + data.py # soup data (inspect/validate/convert/merge/dedup/stats/filter/split/search/preview/download/register) monitoring/ callback.py # HF TrainerCallback -> Rich display + SQLite tracker display.py # Rich Live terminal dashboard (2Hz refresh) @@ -112,8 +112,10 @@ soup_cli/ ollama.py # Ollama integration (detect, deploy, list, remove, Modelfile gen) profiler.py # Training memory/speed estimator (GPU lookup, model arch) curriculum.py # Curriculum learning: sort by difficulty, create buckets + freeze.py # Freeze training: freeze bottom N layers + registry.py # Dataset registry: name → path + format mapping constants.py # APP_NAME, paths, default chat template -tests/ # 70 test files, 1979 tests +tests/ # 74 test files, 2061 tests examples/ configs/ # 7 production-ready YAML examples data/ # Sample datasets @@ -156,6 +158,12 @@ soup recipes search # Search recipes by keyword, task, or model size soup data filter # Quality filter (perplexity + coherence scoring) soup data sample # Sample subset: random, diverse (TF-IDF + clusters), hard (by length) soup data split # Split dataset into train/val/test files (random or stratified) +soup data search # Search HuggingFace Hub for datasets +soup data preview # Preview remote HF dataset metadata, splits, features +soup data download # Download HF dataset to local JSONL (streaming) +soup data register # Register local dataset by name for use in soup.yaml +soup data unregister # Remove dataset from local registry +soup data registry # List all registered datasets soup profile # Estimate memory, speed, GPU requirements before training soup adapters list # Scan directory for LoRA adapters soup adapters info # Show adapter metadata (base model, rank, size) @@ -179,7 +187,7 @@ soup version # Show version (--full for details) - **SoupConfig**: base (required), task (sft/dpo/kto/orpo/simpo/ipo/grpo/ppo/reward_model/pretrain/embedding), modality (text/vision/audio), backend (transformers/unsloth), data, training, output, eval - **EvalConfig**: auto_eval, benchmarks, custom_tasks, judge - **DataConfig**: train, format (alpaca/sharegpt/chatml/dpo/kto/llava/sharegpt4v/plaintext/embedding/audio/auto), val_split, max_length, image_dir, audio_dir -- **TrainingConfig**: epochs, lr, batch_size (int or "auto"), quantization (4bit/8bit/none), quantization_aware, optimizer, scheduler, dpo_beta, kto_beta, orpo_beta, simpo_gamma, cpo_alpha, ipo_tau, grpo_beta, num_generations, reward_fn, ppo_epochs, ppo_clip_ratio, ppo_kl_penalty, reward_model, loraplus_lr_ratio, use_galore, galore_rank, galore_update_proj_gap, galore_scale, moe_lora, moe_aux_loss_coeff, use_liger, use_flash_attn, use_ring_attention, rope_scaling_type, gradient_checkpointing, embedding_loss, embedding_margin, embedding_pooling, embedding_temperature, neftune_alpha, packing, curriculum, curriculum_metric, curriculum_buckets +- **TrainingConfig**: epochs, lr, batch_size (int or "auto"), quantization (4bit/8bit/none), quantization_aware, optimizer, scheduler, dpo_beta, kto_beta, orpo_beta, simpo_gamma, cpo_alpha, ipo_tau, grpo_beta, num_generations, reward_fn, ppo_epochs, ppo_clip_ratio, ppo_kl_penalty, reward_model, loraplus_lr_ratio, use_galore, galore_rank, galore_update_proj_gap, galore_scale, moe_lora, moe_aux_loss_coeff, use_liger, use_flash_attn, use_ring_attention, rope_scaling_type, gradient_checkpointing, embedding_loss, embedding_margin, embedding_pooling, embedding_temperature, neftune_alpha, packing, curriculum, curriculum_metric, curriculum_buckets, loss_watchdog, loss_watchdog_threshold, loss_watchdog_patience, freeze_layers, freeze_ratio - **LoraConfig**: r, alpha, dropout, target_modules, use_dora, use_rslora 15 built-in templates: chat, code, medical, reasoning, vision, audio, kto, orpo, simpo, ipo, embedding, rlhf, pretrain, moe, longcontext. @@ -277,6 +285,15 @@ soup version # Show version (--full for details) - **Curriculum config**: curriculum_buckets bounded ge=1, le=20 (v0.23.0) - **AWQ/GPTQ export**: trust_remote_code warning panel before model loading (v0.23.0) - **Data split**: output files written to input's parent dir (consistent with sample_data) (v0.23.0) +- **HF download**: trust_remote_code=False + warning panel before download (v0.24.0) +- **HF download**: default output path sanitized via Path.name (v0.24.0) +- **HF download**: --samples capped at 1,000,000 (v0.24.0) +- **Dataset registry**: name validation — no path separators or null bytes (v0.24.0) +- **Dataset registry**: register path traversal protection — resolve + relative_to(cwd) (v0.24.0) +- **Dataset registry**: Rich markup escaped in table output (v0.24.0) +- **Dataset registry**: JSON validation on load — catches corruption + type mismatch (v0.24.0) +- **Loss watchdog**: threshold bounded le=100.0, patience bounded le=1000 (v0.24.0) +- **Freeze training**: freeze_layers bounded le=1000 (v0.24.0) ## Code Conventions @@ -350,7 +367,7 @@ soup version # Show version (--full for details) 15. **Tag**: `git tag v0.X.Y && git push origin v0.X.Y` 16. **Release**: `gh release create v0.X.Y` with changelog (What's New, Install/Upgrade) -## Tests (70 test files, 1979 tests) +## Tests (74 test files, 2061 tests) | File | Covers | |------|--------| @@ -423,3 +440,7 @@ soup version # Show version (--full for details) | test_packing.py | Sample packing: config, YAML, trainer integration, sweep | | test_data_split.py | Data split: ratio/absolute/stratified splits, seed, edge cases | | test_curriculum.py | Curriculum learning: config, length sort, buckets, sweep | +| test_dataset_hub.py | HF dataset search, preview, download, format conversion, security | +| test_freeze_training.py | Freeze training: config, layer freezing, GPT-2 naming, sweep | +| test_loss_watchdog.py | Loss watchdog: config, callback behavior, patience, sweep | +| test_dataset_registry.py | Dataset registry: CRUD, CLI, name validation, error handling | diff --git a/README.md b/README.md index 5faf081..c22fd6b 100644 --- a/README.md +++ b/README.md @@ -636,6 +636,30 @@ training: curriculum_buckets: 4 # Number of difficulty stages ``` +## Freeze Training + +Freeze bottom layers of the model — train only the top layers (like LLaMA-Factory's `finetuning_type: freeze`): + +```yaml +training: + freeze_layers: 24 # Freeze first 24 layers, train the rest + # OR + freeze_ratio: 0.75 # Freeze 75% of layers from the bottom +``` + +Works with and without LoRA. When used with LoRA, LoRA is applied only to unfrozen layers. + +## Loss Watchdog + +Auto-stop training when loss spikes above a threshold (like Axolotl's `loss_watchdog_threshold`): + +```yaml +training: + loss_watchdog: true # Enable loss spike detection + loss_watchdog_threshold: 3.0 # Stop if loss exceeds this value + loss_watchdog_patience: 5 # Consecutive steps above threshold before stopping +``` + ## GaLore (Memory-Efficient Full-Parameter Training) Train without LoRA using gradient low-rank projection — saves optimizer memory: @@ -1386,6 +1410,14 @@ soup data sample --pct 10 Sample by percentage soup data split --val 10 --test 10 Split into train/val/test soup data split --val 500 --absolute Split with absolute counts soup data split --val 10 --stratify category Stratified by field +soup data search "code instructions" Search HuggingFace Hub for datasets +soup data search --sort likes --limit 10 Sort and paginate search results +soup data preview teknium/OpenHermes-2.5 Preview remote dataset metadata +soup data download user/dataset -o data.jsonl Download HF dataset as JSONL +soup data download user/ds --samples 1000 Stream first 1000 samples +soup data register --name my-ds --path d.jsonl --format alpaca Register dataset +soup data unregister --name my-ds Remove from registry +soup data registry List all registered datasets soup profile --config soup.yaml Estimate memory/speed before training soup profile --config soup.yaml --gpu a100 Estimate for specific GPU soup profile --config soup.yaml --json Machine-readable output diff --git a/SECURITY.md b/SECURITY.md index 9123934..3f1ffaa 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -9,9 +9,9 @@ We provide security updates for the following versions: - **Versions older than 3 minor versions:** No support Example: -- v0.23.0-0.23.x → Full support (latest) -- v0.22.0-0.22.x → Bug-fix support only -- v0.21.x and below → No support +- v0.24.0-0.24.x → Full support (latest) +- v0.23.0-0.23.x → Bug-fix support only +- v0.22.x and below → No support ## Reporting a Vulnerability diff --git a/pyproject.toml b/pyproject.toml index bcba346..b40c265 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "soup-cli" -version = "0.23.1" +version = "0.24.0" description = "Fine-tune LLMs in one command. No SSH, no config hell." readme = "README.md" license = "MIT" diff --git a/soup_cli/__init__.py b/soup_cli/__init__.py index af37867..00f8ff5 100644 --- a/soup_cli/__init__.py +++ b/soup_cli/__init__.py @@ -1,3 +1,3 @@ """Soup CLI — Fine-tune LLMs in one command.""" -__version__ = "0.23.1" +__version__ = "0.24.0" diff --git a/soup_cli/commands/data.py b/soup_cli/commands/data.py index 5390a1c..445eddc 100644 --- a/soup_cli/commands/data.py +++ b/soup_cli/commands/data.py @@ -925,3 +925,454 @@ def _stratified_split( test_data = [data[idx] for idx in test_indices] return train_data, val_data, test_data + + +# --------------------------------------------------------------------------- +# HuggingFace Dataset Hub helpers +# --------------------------------------------------------------------------- + + +def list_datasets(search: str, sort: str = "downloads", limit: int = 20) -> list: + """Search HuggingFace Hub for datasets. Returns list of DatasetInfo objects.""" + from huggingface_hub import HfApi + + api = HfApi() + return list(api.list_datasets(search=search, sort=sort, limit=limit)) + + +def _hf_dataset_info(dataset_id: str) -> dict: + """Fetch metadata about a HuggingFace dataset.""" + from huggingface_hub import HfApi + + api = HfApi() + try: + info = api.dataset_info(dataset_id) + except Exception as exc: + raise ValueError(f"Dataset not found: {dataset_id} — {exc}") from exc + + # Extract split sizes + splits: dict[str, int] = {} + if hasattr(info, "card_data") and info.card_data: + ds_info = getattr(info.card_data, "dataset_info", None) + if ds_info and isinstance(ds_info, dict): + for config_data in ds_info.values(): + if isinstance(config_data, dict) and "splits" in config_data: + for split_name, split_data in config_data["splits"].items(): + if isinstance(split_data, dict): + splits[split_name] = split_data.get("num_examples", 0) + + # Extract feature names + features: list[str] = [] + if hasattr(info, "card_data") and info.card_data: + ds_info = getattr(info.card_data, "dataset_info", None) + if ds_info and isinstance(ds_info, dict): + for config_data in ds_info.values(): + if isinstance(config_data, dict) and "features" in config_data: + feat_list = config_data["features"] + if isinstance(feat_list, list): + for feat in feat_list: + if isinstance(feat, dict) and "name" in feat: + features.append(feat["name"]) + break + + return { + "id": info.id, + "description": getattr(info, "description", "") or "", + "downloads": getattr(info, "downloads", 0) or 0, + "likes": getattr(info, "likes", 0) or 0, + "size_bytes": getattr(info, "size", None), + "splits": splits, + "features": features, + "tags": list(info.tags) if info.tags else [], + } + + +def _hf_download_dataset( + dataset_id: str, + split: str = "train", + samples: int | None = None, +) -> list[dict]: + """Download a dataset from HuggingFace Hub and return as list of dicts.""" + from datasets import load_dataset + + try: + ds = load_dataset( + dataset_id, split=split, streaming=True, trust_remote_code=False, + ) + except Exception as exc: + raise ValueError(f"Failed to load dataset {dataset_id}: {exc}") from exc + + rows: list[dict] = [] + for idx, row in enumerate(ds): + if samples is not None and idx >= samples: + break + rows.append(dict(row)) + + return rows + + +def _format_size_bytes(size_bytes: int | None) -> str: + """Format byte count as human-readable string.""" + if size_bytes is None: + return "unknown" + if size_bytes == 0: + return "0 B" + units = ["B", "KB", "MB", "GB", "TB"] + unit_idx = 0 + size = float(size_bytes) + while size >= 1024 and unit_idx < len(units) - 1: + size /= 1024 + unit_idx += 1 + if unit_idx == 0: + return f"{int(size)} {units[unit_idx]}" + return f"{size:.1f} {units[unit_idx]}" + + +def _format_count(count: int) -> str: + """Format large numbers with K/M suffix.""" + if count >= 1_000_000: + return f"{count / 1_000_000:.1f}M" + if count >= 1_000: + return f"{count / 1_000:.1f}K" + return str(count) + + +# --------------------------------------------------------------------------- +# HuggingFace Dataset Hub CLI commands +# --------------------------------------------------------------------------- + + +@app.command(name="search") +def search_datasets( + query: str = typer.Argument(..., help="Search query for HuggingFace datasets"), + limit: int = typer.Option(20, "--limit", "-l", help="Maximum results to show"), + sort: str = typer.Option( + "downloads", "--sort", "-s", + help="Sort by: downloads, likes, lastModified, trending, createdAt", + ), +): + """Search HuggingFace Hub for datasets.""" + valid_sorts = {"downloads", "likes", "lastModified", "trending", "createdAt"} + if sort not in valid_sorts: + console.print( + f"[red]Invalid sort: {sort}[/]\n" + f"Valid options: {', '.join(sorted(valid_sorts))}" + ) + raise typer.Exit(1) + + try: + datasets = list_datasets(search=query, sort=sort, limit=limit) + except ImportError: + console.print( + "[red]huggingface_hub not available.[/]\n" + "Install with: [bold]pip install huggingface-hub[/]" + ) + raise typer.Exit(1) + except Exception as exc: + console.print(f"[red]Search failed: {exc}[/]") + raise typer.Exit(1) + + if not datasets: + console.print(f"[yellow]No datasets found for '{query}'.[/]") + return + + table = Table(title=f"HuggingFace Datasets: '{query}'") + table.add_column("Dataset", style="bold cyan", max_width=45) + table.add_column("Downloads", justify="right") + table.add_column("Likes", justify="right") + table.add_column("Tags", max_width=30) + + for ds_item in datasets[:limit]: + ds_tags = getattr(ds_item, "tags", []) or [] + tag_str = ", ".join(ds_tags[:5]) + if len(ds_tags) > 5: + tag_str += "..." + table.add_row( + ds_item.id, + _format_count(getattr(ds_item, "downloads", 0) or 0), + _format_count(getattr(ds_item, "likes", 0) or 0), + tag_str, + ) + + console.print(table) + console.print(f"[dim]Showing {min(limit, len(datasets))} results.[/]") + + +@app.command(name="preview") +def preview_dataset( + dataset_id: str = typer.Argument( + ..., help="HuggingFace dataset ID (e.g. teknium/OpenHermes-2.5)" + ), +): + """Preview a remote HuggingFace dataset: metadata, splits, features.""" + try: + info = _hf_dataset_info(dataset_id) + except ValueError as exc: + console.print(f"[red]{exc}[/]") + raise typer.Exit(1) + except ImportError: + console.print( + "[red]huggingface_hub not available.[/]\n" + "Install with: [bold]pip install huggingface-hub[/]" + ) + raise typer.Exit(1) + + table = Table(title=f"Dataset: {info['id']}") + table.add_column("Field", style="bold") + table.add_column("Value", max_width=80) + + table.add_row("ID", info["id"]) + desc = info["description"] + if len(desc) > 200: + desc = desc[:200] + "..." + table.add_row("Description", desc or "[dim]No description[/]") + table.add_row("Downloads", _format_count(info["downloads"])) + table.add_row("Likes", _format_count(info["likes"])) + table.add_row("Size", _format_size_bytes(info["size_bytes"])) + + if info["splits"]: + splits_str = ", ".join( + f"{name} ({_format_count(count)})" + for name, count in info["splits"].items() + ) + table.add_row("Splits", splits_str) + else: + table.add_row("Splits", "[dim]Not available (use streaming to explore)[/]") + + if info["features"]: + table.add_row("Features", ", ".join(info["features"])) + + if info["tags"]: + table.add_row("Tags", ", ".join(info["tags"][:10])) + + console.print(table) + + +@app.command(name="download") +def download_dataset( + dataset_id: str = typer.Argument( + ..., help="HuggingFace dataset ID (e.g. teknium/OpenHermes-2.5)" + ), + output: str = typer.Option( + None, "--output", "-o", + help="Output file path (default: .jsonl)", + ), + split: str = typer.Option( + "train", "--split", + help="Dataset split to download (e.g. train, test, train[:1000])", + ), + samples: int = typer.Option( + None, "--samples", "-n", + help="Max number of samples to download (streams, no full download)", + ), + fmt: str = typer.Option( + None, "--format", "-f", + help="Convert to Soup format after download: alpaca, sharegpt, chatml", + ), +): + """Download a HuggingFace dataset and save as JSONL.""" + max_download_samples = 1_000_000 + if samples is not None and samples > max_download_samples: + console.print( + f"[red]--samples cannot exceed {max_download_samples:,}.[/]" + ) + raise typer.Exit(1) + + # Resolve output path + if output is None: + ds_name = dataset_id.split("/")[-1] if "/" in dataset_id else dataset_id + # Strip embedded path separators to prevent traversal + ds_name = Path(ds_name).name + out_path = (Path.cwd() / f"{ds_name}.jsonl").resolve() + cwd = Path.cwd().resolve() + try: + out_path.relative_to(cwd) + except ValueError: + console.print( + "[red]Derived output path escapes working directory.[/]" + ) + raise typer.Exit(1) + else: + out_path = Path(output).resolve() + cwd = Path.cwd().resolve() + try: + out_path.relative_to(cwd) + except ValueError: + console.print( + "[red]Output path must be under the current working directory.[/]" + ) + raise typer.Exit(1) + + from rich.panel import Panel + + console.print(Panel( + "[bold yellow]Warning:[/] Downloading this dataset may execute a " + "remote dataset loading script from HuggingFace Hub.\n\n" + "Only download datasets from sources you trust.", + title="Remote Code Warning", + border_style="yellow", + )) + console.print(f"[dim]Downloading {dataset_id} (split={split})...[/]") + + try: + data = _hf_download_dataset( + dataset_id, split=split, samples=samples, + ) + except ValueError as exc: + console.print(f"[red]{exc}[/]") + raise typer.Exit(1) + except ImportError: + console.print( + "[red]datasets library not available.[/]\n" + "Install with: [bold]pip install datasets[/]" + ) + raise typer.Exit(1) + + if not data: + console.print("[red]No data downloaded (dataset may be empty).[/]") + raise typer.Exit(1) + + # Optional format conversion + if fmt: + from soup_cli.data.formats import ( + CONVERTIBLE_FORMATS, + detect_format, + format_to_messages, + messages_to_format, + ) + + if fmt not in CONVERTIBLE_FORMATS: + console.print( + f"[red]Invalid format: {fmt}[/]\n" + f"Supported: {', '.join(CONVERTIBLE_FORMATS)}" + ) + raise typer.Exit(1) + + try: + src_fmt = detect_format(data) + except ValueError: + src_fmt = None + + if src_fmt and src_fmt != fmt: + converted = [] + for row in data: + messages = format_to_messages(row, src_fmt) + if messages is not None: + result = messages_to_format(messages, fmt) + if result is not None: + converted.append(result) + if converted: + data = converted + console.print( + f"[dim]Converted {len(data)} rows to {fmt} format.[/]" + ) + + # Apply samples limit if data came from non-streaming path + if samples is not None and len(data) > samples: + data = data[:samples] + + _write_jsonl(out_path, data) + console.print( + f"[green]Downloaded {len(data)} rows.[/]\n" + f"Output: [bold]{out_path}[/]" + ) + + +# --------------------------------------------------------------------------- +# Dataset registry CLI commands +# --------------------------------------------------------------------------- + + +def _get_registry_path() -> Path: + """Get the default registry path (~/.soup/datasets.json).""" + from soup_cli.utils.registry import _default_registry_path + + return _default_registry_path() + + +@app.command(name="register") +def register_data( + name: str = typer.Option(..., "--name", "-n", help="Dataset name"), + path: str = typer.Option(..., "--path", "-p", help="Path to dataset file"), + fmt: str = typer.Option( + "auto", "--format", "-f", + help="Dataset format: alpaca, sharegpt, chatml, dpo, kto, auto", + ), +): + """Register a local dataset by name for use in soup.yaml.""" + from soup_cli.utils.registry import register_dataset + + # Path traversal protection + resolved = Path(path).resolve() + cwd = Path.cwd().resolve() + try: + resolved.relative_to(cwd) + except ValueError: + console.print( + "[red]Dataset path must be under the current working directory.[/]" + ) + raise typer.Exit(1) + + registry_path = _get_registry_path() + + try: + register_dataset(name, str(resolved), fmt, registry_path=registry_path) + except ValueError as exc: + console.print(f"[red]{exc}[/]") + raise typer.Exit(1) + + console.print( + f"[green]Registered dataset '[bold]{name}[/bold]'[/]\n" + f" Path: {path}\n" + f" Format: {fmt}" + ) + + +@app.command(name="unregister") +def unregister_data( + name: str = typer.Option(..., "--name", "-n", help="Dataset name to remove"), +): + """Remove a dataset from the local registry.""" + from soup_cli.utils.registry import unregister_dataset + + registry_path = _get_registry_path() + removed = unregister_dataset(name, registry_path=registry_path) + + if removed: + console.print(f"[green]Removed dataset '{name}' from registry.[/]") + else: + console.print(f"[red]Dataset '{name}' not found in registry.[/]") + raise typer.Exit(1) + + +@app.command(name="registry") +def list_registry(): + """List all registered datasets.""" + from soup_cli.utils.registry import load_registry + + registry_path = _get_registry_path() + registry = load_registry(registry_path) + + if not registry: + console.print("[yellow]No datasets registered.[/]") + console.print( + "[dim]Register with: " + "soup data register --name my-data --path data.jsonl --format alpaca[/]" + ) + return + + table = Table(title="Registered Datasets") + table.add_column("Name", style="bold cyan") + table.add_column("Path") + table.add_column("Format") + + from rich.markup import escape + + for ds_name, ds_info in sorted(registry.items()): + table.add_row( + escape(ds_name), + escape(ds_info.get("path", "")), + escape(ds_info.get("format", "")), + ) + + console.print(table) diff --git a/soup_cli/config/schema.py b/soup_cli/config/schema.py index 78ee574..0ba816f 100644 --- a/soup_cli/config/schema.py +++ b/soup_cli/config/schema.py @@ -204,6 +204,36 @@ class TrainingConfig(BaseModel): default=4, ge=1, le=20, description="Number of difficulty stages for curriculum learning", ) + # Loss watchdog — auto-stop on loss spikes + loss_watchdog: bool = Field( + default=False, + description="Enable loss spike detection (auto-stop if loss exceeds threshold)", + ) + loss_watchdog_threshold: float = Field( + default=3.0, + gt=0, + le=100.0, + description="Stop training if loss exceeds this threshold", + ) + loss_watchdog_patience: int = Field( + default=5, + ge=1, + le=1000, + description="Consecutive high-loss steps before stopping", + ) + # Freeze training — freeze bottom layers for parameter-efficient training + freeze_layers: Optional[int] = Field( + default=None, + ge=1, + le=1000, + description="Freeze first N layers (from bottom). Train only remaining layers.", + ) + freeze_ratio: Optional[float] = Field( + default=None, + gt=0.0, + lt=1.0, + description="Freeze this fraction of layers (0.75 = freeze 75% from bottom).", + ) # Sample packing — pack multiple short samples into one sequence packing: bool = Field( default=False, diff --git a/soup_cli/monitoring/callback.py b/soup_cli/monitoring/callback.py index 4a2fda1..28a7b71 100644 --- a/soup_cli/monitoring/callback.py +++ b/soup_cli/monitoring/callback.py @@ -27,12 +27,21 @@ class SoupTrainerCallback(TrainerCallback): run_id: str = "", eval_config: Optional[object] = None, output_dir: str = "", + loss_watchdog: bool = False, + loss_watchdog_threshold: float = 3.0, + loss_watchdog_patience: int = 5, ): self.display = display self.tracker = tracker self.run_id = run_id self.eval_config = eval_config self.output_dir = output_dir + # Loss watchdog state + self._watchdog_enabled = loss_watchdog + self._watchdog_threshold = loss_watchdog_threshold + self._watchdog_patience = loss_watchdog_patience + self._watchdog_counter = 0 + self._watchdog_fired = False def on_train_begin( self, args: TrainingArguments, state: TrainerState, @@ -76,6 +85,33 @@ class SoupTrainerCallback(TrainerCallback): gpu_mem=gpu_mem, ) + # Loss watchdog — detect loss spikes and auto-stop + if self._watchdog_enabled and not self._watchdog_fired and "loss" in logs: + if loss > self._watchdog_threshold: + self._watchdog_counter += 1 + if self._watchdog_counter >= self._watchdog_patience: + self._watchdog_fired = True + # Stop Live display before printing panel + self.display.stop() + + from rich.console import Console as WatchdogConsole + from rich.panel import Panel + + wc = WatchdogConsole() + wc.print(Panel( + f"[bold red]Loss watchdog triggered![/]\n\n" + f"Loss {loss:.4f} exceeded threshold " + f"{self._watchdog_threshold} for " + f"{self._watchdog_counter} consecutive steps " + f"(patience={self._watchdog_patience}).\n\n" + f"Training will stop.", + title="Loss Watchdog", + border_style="red", + )) + control.should_training_stop = True + else: + self._watchdog_counter = 0 + # Log to experiment tracker if self.tracker and self.run_id: self.tracker.log_metrics( diff --git a/soup_cli/monitoring/display.py b/soup_cli/monitoring/display.py index c914886..231bb1f 100644 --- a/soup_cli/monitoring/display.py +++ b/soup_cli/monitoring/display.py @@ -47,9 +47,10 @@ class TrainingDisplay: self._live.update(self._render()) def stop(self): - """Stop the live display.""" + """Stop the live display. Safe to call multiple times.""" if self._live: self._live.stop() + self._live = None def _render(self) -> Panel: """Render the dashboard panel.""" diff --git a/soup_cli/trainer/dpo.py b/soup_cli/trainer/dpo.py index 99e6dcc..ec7f22b 100644 --- a/soup_cli/trainer/dpo.py +++ b/soup_cli/trainer/dpo.py @@ -232,7 +232,12 @@ class DPOTrainerWrapper: from soup_cli.monitoring.callback import SoupTrainerCallback self.trainer.add_callback( - SoupTrainerCallback(display, tracker=tracker, run_id=run_id) + SoupTrainerCallback( + display, tracker=tracker, run_id=run_id, + loss_watchdog=self.config.training.loss_watchdog, + loss_watchdog_threshold=self.config.training.loss_watchdog_threshold, + loss_watchdog_patience=self.config.training.loss_watchdog_patience, + ) ) self.trainer.train(resume_from_checkpoint=resume_from_checkpoint) diff --git a/soup_cli/trainer/embedding.py b/soup_cli/trainer/embedding.py index 9ee6d1b..a091f7a 100644 --- a/soup_cli/trainer/embedding.py +++ b/soup_cli/trainer/embedding.py @@ -254,7 +254,12 @@ class EmbeddingTrainerWrapper: from soup_cli.monitoring.callback import SoupTrainerCallback self.trainer.add_callback( - SoupTrainerCallback(display, tracker=tracker, run_id=run_id) + SoupTrainerCallback( + display, tracker=tracker, run_id=run_id, + loss_watchdog=self.config.training.loss_watchdog, + loss_watchdog_threshold=self.config.training.loss_watchdog_threshold, + loss_watchdog_patience=self.config.training.loss_watchdog_patience, + ) ) self.trainer.train(resume_from_checkpoint=resume_from_checkpoint) diff --git a/soup_cli/trainer/grpo.py b/soup_cli/trainer/grpo.py index 1e90dfa..6651e75 100644 --- a/soup_cli/trainer/grpo.py +++ b/soup_cli/trainer/grpo.py @@ -284,7 +284,12 @@ class GRPOTrainerWrapper: from soup_cli.monitoring.callback import SoupTrainerCallback self.trainer.add_callback( - SoupTrainerCallback(display, tracker=tracker, run_id=run_id) + SoupTrainerCallback( + display, tracker=tracker, run_id=run_id, + loss_watchdog=self.config.training.loss_watchdog, + loss_watchdog_threshold=self.config.training.loss_watchdog_threshold, + loss_watchdog_patience=self.config.training.loss_watchdog_patience, + ) ) self.trainer.train(resume_from_checkpoint=resume_from_checkpoint) diff --git a/soup_cli/trainer/ipo.py b/soup_cli/trainer/ipo.py index 548969b..c194bca 100644 --- a/soup_cli/trainer/ipo.py +++ b/soup_cli/trainer/ipo.py @@ -237,7 +237,12 @@ class IPOTrainerWrapper: from soup_cli.monitoring.callback import SoupTrainerCallback self.trainer.add_callback( - SoupTrainerCallback(display, tracker=tracker, run_id=run_id) + SoupTrainerCallback( + display, tracker=tracker, run_id=run_id, + loss_watchdog=self.config.training.loss_watchdog, + loss_watchdog_threshold=self.config.training.loss_watchdog_threshold, + loss_watchdog_patience=self.config.training.loss_watchdog_patience, + ) ) self.trainer.train(resume_from_checkpoint=resume_from_checkpoint) diff --git a/soup_cli/trainer/kto.py b/soup_cli/trainer/kto.py index a90e182..a9d577a 100644 --- a/soup_cli/trainer/kto.py +++ b/soup_cli/trainer/kto.py @@ -236,7 +236,12 @@ class KTOTrainerWrapper: from soup_cli.monitoring.callback import SoupTrainerCallback self.trainer.add_callback( - SoupTrainerCallback(display, tracker=tracker, run_id=run_id) + SoupTrainerCallback( + display, tracker=tracker, run_id=run_id, + loss_watchdog=self.config.training.loss_watchdog, + loss_watchdog_threshold=self.config.training.loss_watchdog_threshold, + loss_watchdog_patience=self.config.training.loss_watchdog_patience, + ) ) self.trainer.train(resume_from_checkpoint=resume_from_checkpoint) diff --git a/soup_cli/trainer/orpo.py b/soup_cli/trainer/orpo.py index 19353cf..191bea1 100644 --- a/soup_cli/trainer/orpo.py +++ b/soup_cli/trainer/orpo.py @@ -234,7 +234,12 @@ class ORPOTrainerWrapper: from soup_cli.monitoring.callback import SoupTrainerCallback self.trainer.add_callback( - SoupTrainerCallback(display, tracker=tracker, run_id=run_id) + SoupTrainerCallback( + display, tracker=tracker, run_id=run_id, + loss_watchdog=self.config.training.loss_watchdog, + loss_watchdog_threshold=self.config.training.loss_watchdog_threshold, + loss_watchdog_patience=self.config.training.loss_watchdog_patience, + ) ) self.trainer.train(resume_from_checkpoint=resume_from_checkpoint) diff --git a/soup_cli/trainer/ppo.py b/soup_cli/trainer/ppo.py index fc70237..d5a2525 100644 --- a/soup_cli/trainer/ppo.py +++ b/soup_cli/trainer/ppo.py @@ -430,7 +430,12 @@ class PPOTrainerWrapper: from soup_cli.monitoring.callback import SoupTrainerCallback self.trainer.add_callback( - SoupTrainerCallback(display, tracker=tracker, run_id=run_id) + SoupTrainerCallback( + display, tracker=tracker, run_id=run_id, + loss_watchdog=self.config.training.loss_watchdog, + loss_watchdog_threshold=self.config.training.loss_watchdog_threshold, + loss_watchdog_patience=self.config.training.loss_watchdog_patience, + ) ) # trl experimental PPOTrainer.train() does not accept resume_from_checkpoint diff --git a/soup_cli/trainer/pretrain.py b/soup_cli/trainer/pretrain.py index 8072af6..4e702ac 100644 --- a/soup_cli/trainer/pretrain.py +++ b/soup_cli/trainer/pretrain.py @@ -289,7 +289,12 @@ class PretrainTrainerWrapper: from soup_cli.monitoring.callback import SoupTrainerCallback self.trainer.add_callback( - SoupTrainerCallback(display, tracker=tracker, run_id=run_id) + SoupTrainerCallback( + display, tracker=tracker, run_id=run_id, + loss_watchdog=self.config.training.loss_watchdog, + loss_watchdog_threshold=self.config.training.loss_watchdog_threshold, + loss_watchdog_patience=self.config.training.loss_watchdog_patience, + ) ) self.trainer.train(resume_from_checkpoint=resume_from_checkpoint) diff --git a/soup_cli/trainer/reward_model.py b/soup_cli/trainer/reward_model.py index 52acb22..b6095f0 100644 --- a/soup_cli/trainer/reward_model.py +++ b/soup_cli/trainer/reward_model.py @@ -215,7 +215,12 @@ class RewardModelTrainerWrapper: from soup_cli.monitoring.callback import SoupTrainerCallback self.trainer.add_callback( - SoupTrainerCallback(display, tracker=tracker, run_id=run_id) + SoupTrainerCallback( + display, tracker=tracker, run_id=run_id, + loss_watchdog=self.config.training.loss_watchdog, + loss_watchdog_threshold=self.config.training.loss_watchdog_threshold, + loss_watchdog_patience=self.config.training.loss_watchdog_patience, + ) ) self.trainer.train(resume_from_checkpoint=resume_from_checkpoint) diff --git a/soup_cli/trainer/sft.py b/soup_cli/trainer/sft.py index b84b93a..80a0b8c 100644 --- a/soup_cli/trainer/sft.py +++ b/soup_cli/trainer/sft.py @@ -314,6 +314,19 @@ class SFTTrainerWrapper: if tcfg.quantization in ("4bit", "8bit"): self.model = prepare_model_for_kbit_training(self.model) + # Freeze training — freeze bottom layers before LoRA + if tcfg.freeze_layers is not None or tcfg.freeze_ratio is not None: + from soup_cli.utils.freeze import freeze_model_layers + + frozen = freeze_model_layers( + self.model, + freeze_layers=tcfg.freeze_layers, + freeze_ratio=tcfg.freeze_ratio, + ) + console.print( + f"[green]Freeze training:[/] {frozen} parameters frozen" + ) + # LoRA — with MoE-aware target modules if moe_lora is enabled target_modules = tcfg.lora.target_modules if target_modules == "auto": @@ -600,7 +613,12 @@ class SFTTrainerWrapper: from soup_cli.monitoring.callback import SoupTrainerCallback self.trainer.add_callback( - SoupTrainerCallback(display, tracker=tracker, run_id=run_id) + SoupTrainerCallback( + display, tracker=tracker, run_id=run_id, + loss_watchdog=self.config.training.loss_watchdog, + loss_watchdog_threshold=self.config.training.loss_watchdog_threshold, + loss_watchdog_patience=self.config.training.loss_watchdog_patience, + ) ) self.trainer.train(resume_from_checkpoint=resume_from_checkpoint) diff --git a/soup_cli/trainer/simpo.py b/soup_cli/trainer/simpo.py index cfce9ec..16f9375 100644 --- a/soup_cli/trainer/simpo.py +++ b/soup_cli/trainer/simpo.py @@ -237,7 +237,12 @@ class SimPOTrainerWrapper: from soup_cli.monitoring.callback import SoupTrainerCallback self.trainer.add_callback( - SoupTrainerCallback(display, tracker=tracker, run_id=run_id) + SoupTrainerCallback( + display, tracker=tracker, run_id=run_id, + loss_watchdog=self.config.training.loss_watchdog, + loss_watchdog_threshold=self.config.training.loss_watchdog_threshold, + loss_watchdog_patience=self.config.training.loss_watchdog_patience, + ) ) self.trainer.train(resume_from_checkpoint=resume_from_checkpoint) diff --git a/soup_cli/utils/freeze.py b/soup_cli/utils/freeze.py new file mode 100644 index 0000000..a7db407 --- /dev/null +++ b/soup_cli/utils/freeze.py @@ -0,0 +1,73 @@ +"""Freeze training: freeze bottom N layers of a model for parameter-efficient training.""" + +from __future__ import annotations + +import logging +import re +from typing import Any, Optional + +logger = logging.getLogger(__name__) + + +def _detect_num_layers(model: Any) -> int: + """Detect total number of transformer layers from model parameter names. + + Looks for patterns like 'model.layers.N.' or 'transformer.h.N.' and + returns max(N) + 1. + """ + max_layer = -1 + pattern = re.compile(r"(?:layers|h)\.(\d+)\.") + for name, _ in model.named_parameters(): + match = pattern.search(name) + if match: + layer_idx = int(match.group(1)) + if layer_idx > max_layer: + max_layer = layer_idx + return max_layer + 1 if max_layer >= 0 else 0 + + +def freeze_model_layers( + model: Any, + freeze_layers: Optional[int] = None, + freeze_ratio: Optional[float] = None, +) -> int: + """Freeze the bottom layers of a model. + + Args: + model: A PyTorch model with named_parameters(). + freeze_layers: Freeze the first N layers. Takes priority over freeze_ratio. + freeze_ratio: Freeze this fraction of layers (e.g. 0.75 = 75% from bottom). + + Returns: + Number of parameters frozen. + """ + if freeze_layers is None and freeze_ratio is None: + return 0 + + total_layers = _detect_num_layers(model) + if total_layers == 0: + logger.warning( + "freeze_model_layers: could not detect numbered layers in model " + "parameter names. Freezing has no effect. Check that your model " + "uses 'layers.N.' or 'h.N.' naming." + ) + return 0 + + # Determine cutoff + if freeze_layers is not None: + cutoff = min(freeze_layers, total_layers) + else: + cutoff = int(total_layers * freeze_ratio) + + # Freeze parameters in layers below cutoff + frozen_count = 0 + pattern = re.compile(r"(?:layers|h)\.(\d+)\.") + for name, param in model.named_parameters(): + match = pattern.search(name) + if match: + layer_idx = int(match.group(1)) + if layer_idx < cutoff: + param.requires_grad = False + frozen_count += 1 + + return frozen_count diff --git a/soup_cli/utils/registry.py b/soup_cli/utils/registry.py new file mode 100644 index 0000000..7f9841f --- /dev/null +++ b/soup_cli/utils/registry.py @@ -0,0 +1,100 @@ +"""Dataset info registry — local name → path + format mapping. + +Stores registry at ~/.soup/datasets.json so datasets can be referenced +by name in soup.yaml instead of by path. +""" + +from __future__ import annotations + +import json +import re +from pathlib import Path +from typing import Optional + + +def _default_registry_path() -> Path: + """Default registry path: ~/.soup/datasets.json.""" + return Path.home() / ".soup" / "datasets.json" + + +def _validate_name(name: str) -> None: + """Validate dataset name — no path separators, null bytes, or empty.""" + if not name: + raise ValueError("Dataset name must not be empty") + if re.search(r'[/\\:\x00]', name): + raise ValueError( + f"Dataset name '{name}' must not contain " + "path separators (/ \\ :) or null bytes" + ) + + +def load_registry(registry_path: Optional[Path] = None) -> dict[str, dict]: + """Load the dataset registry. Returns empty dict if file doesn't exist.""" + path = registry_path or _default_registry_path() + if not path.exists(): + return {} + try: + with open(path, encoding="utf-8") as fh: + data = json.load(fh) + except json.JSONDecodeError as exc: + raise ValueError(f"Registry file is corrupted ({path}): {exc}") from exc + if not isinstance(data, dict): + raise ValueError(f"Registry file has unexpected format in {path}") + return data + + +def _save_registry(registry: dict, registry_path: Optional[Path] = None) -> None: + """Save the dataset registry to disk.""" + path = registry_path or _default_registry_path() + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, "w", encoding="utf-8") as fh: + json.dump(registry, fh, indent=2, ensure_ascii=False) + + +def register_dataset( + name: str, + data_path: str, + data_format: str, + registry_path: Optional[Path] = None, +) -> None: + """Register a dataset by name.""" + _validate_name(name) + registry = load_registry(registry_path) + registry[name] = { + "path": data_path, + "format": data_format, + } + _save_registry(registry, registry_path) + + +def unregister_dataset( + name: str, + registry_path: Optional[Path] = None, +) -> bool: + """Unregister a dataset. Returns True if removed, False if not found.""" + registry = load_registry(registry_path) + if name not in registry: + return False + del registry[name] + _save_registry(registry, registry_path) + return True + + +def resolve_dataset( + name: str, + registry_path: Optional[Path] = None, +) -> Optional[dict]: + """Resolve a dataset name to its entry (path + format). None if not found. + + Callers must apply resolve + relative_to(cwd) before trusting the path. + The returned path is validated for null bytes. + """ + registry = load_registry(registry_path) + entry = registry.get(name) + if entry is not None: + stored_path = entry.get("path", "") + if "\x00" in stored_path: + raise ValueError( + f"Registry entry '{name}' contains null bytes in path" + ) + return entry diff --git a/tests/test_dataset_hub.py b/tests/test_dataset_hub.py new file mode 100644 index 0000000..549f0e2 --- /dev/null +++ b/tests/test_dataset_hub.py @@ -0,0 +1,357 @@ +"""Tests for HuggingFace Dataset Hub browser: search, preview, download.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +from typer.testing import CliRunner + +from soup_cli.commands.data import app + +runner = CliRunner() + + +# --------------------------------------------------------------------------- +# soup data search +# --------------------------------------------------------------------------- + + +class TestDataSearch: + """Tests for `soup data search` subcommand.""" + + @patch("soup_cli.commands.data.list_datasets") + def test_search_basic(self, mock_list): + """Basic search returns results table.""" + mock_ds = MagicMock() + mock_ds.id = "teknium/OpenHermes-2.5" + mock_ds.downloads = 50000 + mock_ds.likes = 200 + mock_ds.tags = ["en", "sft"] + mock_list.return_value = [mock_ds] + + result = runner.invoke(app, ["search", "openhermes"]) + assert result.exit_code == 0 + assert "teknium/OpenHermes-2.5" in result.output + + @patch("soup_cli.commands.data.list_datasets") + def test_search_no_results(self, mock_list): + """Search with no results shows message.""" + mock_list.return_value = [] + + result = runner.invoke(app, ["search", "nonexistent_dataset_xyz_123"]) + assert result.exit_code == 0 + assert "No datasets found" in result.output + + @patch("soup_cli.commands.data.list_datasets") + def test_search_limit(self, mock_list): + """--limit controls number of results.""" + datasets = [] + for idx in range(5): + mock_ds = MagicMock() + mock_ds.id = f"user/dataset-{idx}" + mock_ds.downloads = 100 + mock_ds.likes = 10 + mock_ds.tags = [] + datasets.append(mock_ds) + mock_list.return_value = datasets + + result = runner.invoke(app, ["search", "dataset", "--limit", "3"]) + assert result.exit_code == 0 + + @patch("soup_cli.commands.data.list_datasets") + def test_search_sort_downloads(self, mock_list): + """--sort downloads sorts by download count.""" + mock_list.return_value = [] + + result = runner.invoke( + app, ["search", "code", "--sort", "downloads"] + ) + assert result.exit_code == 0 + mock_list.assert_called_once() + call_kwargs = mock_list.call_args + assert call_kwargs[1].get("sort") == "downloads" + + @patch("soup_cli.commands.data.list_datasets") + def test_search_sort_likes(self, mock_list): + """--sort likes sorts by likes.""" + mock_list.return_value = [] + + result = runner.invoke(app, ["search", "code", "--sort", "likes"]) + assert result.exit_code == 0 + call_kwargs = mock_list.call_args + assert call_kwargs[1].get("sort") == "likes" + + @patch("soup_cli.commands.data.list_datasets") + def test_search_sort_invalid(self, mock_list): + """Invalid --sort value is rejected.""" + result = runner.invoke(app, ["search", "code", "--sort", "invalid"]) + assert result.exit_code != 0 + + @patch("soup_cli.commands.data.list_datasets") + def test_search_huggingface_hub_not_installed(self, mock_list): + """Graceful error when huggingface_hub is missing.""" + mock_list.side_effect = ImportError("No module named 'huggingface_hub'") + + result = runner.invoke(app, ["search", "code"]) + assert result.exit_code == 1 + assert "huggingface" in result.output.lower() + + +# --------------------------------------------------------------------------- +# soup data preview +# --------------------------------------------------------------------------- + + +class TestDataPreview: + """Tests for `soup data preview` subcommand.""" + + @patch("soup_cli.commands.data._hf_dataset_info") + def test_preview_basic(self, mock_info): + """Preview shows dataset info table.""" + mock_info.return_value = { + "id": "teknium/OpenHermes-2.5", + "description": "A large collection of instruction pairs.", + "downloads": 50000, + "likes": 200, + "size_bytes": 2_000_000_000, + "splits": {"train": 1_000_000, "test": 10_000}, + "features": ["conversations"], + "tags": ["en", "sft"], + } + + result = runner.invoke(app, ["preview", "teknium/OpenHermes-2.5"]) + assert result.exit_code == 0 + assert "teknium/OpenHermes-2.5" in result.output + + @patch("soup_cli.commands.data._hf_dataset_info") + def test_preview_not_found(self, mock_info): + """Preview of nonexistent dataset shows error.""" + mock_info.side_effect = ValueError("Dataset not found") + + result = runner.invoke(app, ["preview", "nonexistent/dataset"]) + assert result.exit_code == 1 + assert "not found" in result.output.lower() or "error" in result.output.lower() + + @patch("soup_cli.commands.data._hf_dataset_info") + def test_preview_shows_splits(self, mock_info): + """Preview shows split information.""" + mock_info.return_value = { + "id": "test/ds", + "description": "Test", + "downloads": 100, + "likes": 5, + "size_bytes": 1000, + "splits": {"train": 500, "validation": 100}, + "features": ["text"], + "tags": [], + } + + result = runner.invoke(app, ["preview", "test/ds"]) + assert result.exit_code == 0 + assert "train" in result.output + + @patch("soup_cli.commands.data._hf_dataset_info") + def test_preview_shows_features(self, mock_info): + """Preview shows feature columns.""" + mock_info.return_value = { + "id": "test/ds", + "description": "Test", + "downloads": 100, + "likes": 5, + "size_bytes": 1000, + "splits": {"train": 500}, + "features": ["instruction", "output", "input"], + "tags": [], + } + + result = runner.invoke(app, ["preview", "test/ds"]) + assert result.exit_code == 0 + + +# --------------------------------------------------------------------------- +# soup data download +# --------------------------------------------------------------------------- + + +class TestDataDownload: + """Tests for `soup data download` subcommand.""" + + @patch("soup_cli.commands.data._hf_download_dataset") + def test_download_basic(self, mock_download, tmp_path, monkeypatch): + """Basic download writes JSONL output.""" + monkeypatch.chdir(tmp_path) + output_file = tmp_path / "data.jsonl" + mock_download.return_value = [ + {"instruction": "What is 2+2?", "output": "4"}, + {"instruction": "Hello", "output": "Hi there"}, + ] + + result = runner.invoke( + app, ["download", "test/dataset", "-o", str(output_file)] + ) + assert result.exit_code == 0 + assert output_file.exists() + lines = output_file.read_text(encoding="utf-8").strip().split("\n") + assert len(lines) == 2 + + @patch("soup_cli.commands.data._hf_download_dataset") + def test_download_with_split(self, mock_download, tmp_path, monkeypatch): + """--split flag is passed to download function.""" + monkeypatch.chdir(tmp_path) + output_file = tmp_path / "data.jsonl" + mock_download.return_value = [{"text": "hello"}] + + result = runner.invoke( + app, + ["download", "test/dataset", "--split", "train[:100]", "-o", str(output_file)], + ) + assert result.exit_code == 0 + mock_download.assert_called_once() + call_kwargs = mock_download.call_args + assert call_kwargs[1].get("split") == "train[:100]" + + @patch("soup_cli.commands.data._hf_download_dataset") + def test_download_with_format_conversion(self, mock_download, tmp_path, monkeypatch): + """--format converts downloaded data.""" + monkeypatch.chdir(tmp_path) + output_file = tmp_path / "data.jsonl" + # Alpaca format input + mock_download.return_value = [ + {"instruction": "Q1", "input": "", "output": "A1"}, + ] + + result = runner.invoke( + app, + [ + "download", "test/dataset", + "--format", "sharegpt", + "-o", str(output_file), + ], + ) + assert result.exit_code == 0 + assert output_file.exists() + + @patch("soup_cli.commands.data._hf_download_dataset") + def test_download_empty_dataset(self, mock_download, tmp_path, monkeypatch): + """Empty download result shows error.""" + monkeypatch.chdir(tmp_path) + mock_download.return_value = [] + output_file = tmp_path / "data.jsonl" + + result = runner.invoke( + app, ["download", "test/empty", "-o", str(output_file)] + ) + assert result.exit_code == 1 + assert "empty" in result.output.lower() or "no data" in result.output.lower() + + @patch("soup_cli.commands.data._hf_download_dataset") + def test_download_error_handling(self, mock_download, tmp_path): + """Download failure shows friendly error.""" + mock_download.side_effect = ValueError("Dataset not found on HuggingFace Hub") + output_file = tmp_path / "data.jsonl" + + result = runner.invoke( + app, ["download", "nonexistent/dataset", "-o", str(output_file)] + ) + assert result.exit_code == 1 + + @patch("soup_cli.commands.data._hf_download_dataset") + def test_download_output_path_traversal(self, mock_download, tmp_path): + """Output path must stay under cwd.""" + mock_download.return_value = [{"text": "hello"}] + + result = runner.invoke( + app, + ["download", "test/dataset", "-o", "/etc/passwd"], + ) + assert result.exit_code == 1 + assert "current working directory" in result.output.lower() + + @patch("soup_cli.commands.data._hf_download_dataset") + def test_download_default_output(self, mock_download, tmp_path, monkeypatch): + """Default output file uses dataset name.""" + monkeypatch.chdir(tmp_path) + mock_download.return_value = [{"text": "hello"}] + + result = runner.invoke( + app, ["download", "user/my-dataset"] + ) + assert result.exit_code == 0 + # Default should be my-dataset.jsonl in cwd + expected = tmp_path / "my-dataset.jsonl" + assert expected.exists() + + @patch("soup_cli.commands.data._hf_download_dataset") + def test_download_samples_limit(self, mock_download, tmp_path, monkeypatch): + """--samples limits number of downloaded rows.""" + monkeypatch.chdir(tmp_path) + output_file = tmp_path / "data.jsonl" + mock_download.return_value = [ + {"text": f"row {idx}"} for idx in range(10) + ] + + result = runner.invoke( + app, ["download", "test/dataset", "--samples", "5", "-o", str(output_file)] + ) + assert result.exit_code == 0 + lines = output_file.read_text(encoding="utf-8").strip().split("\n") + assert len(lines) == 5 + + +# --------------------------------------------------------------------------- +# Helper function tests +# --------------------------------------------------------------------------- + + +class TestHelperFunctions: + """Tests for internal helper functions.""" + + def test_format_size_bytes(self): + """_format_size_bytes formats various sizes correctly.""" + from soup_cli.commands.data import _format_size_bytes + + assert _format_size_bytes(0) == "0 B" + assert _format_size_bytes(500) == "500 B" + assert "KB" in _format_size_bytes(1024) + assert "MB" in _format_size_bytes(1024 * 1024) + assert "GB" in _format_size_bytes(1024 * 1024 * 1024) + + def test_format_size_bytes_none(self): + """_format_size_bytes handles None.""" + from soup_cli.commands.data import _format_size_bytes + + assert _format_size_bytes(None) == "unknown" + + def test_format_count(self): + """_format_count formats large numbers.""" + from soup_cli.commands.data import _format_count + + assert _format_count(500) == "500" + assert "K" in _format_count(1500) + assert "M" in _format_count(1_500_000) + + +# --------------------------------------------------------------------------- +# Security edge cases +# --------------------------------------------------------------------------- + + +class TestSecurityEdgeCases: + """Security-relevant edge case tests.""" + + def test_download_samples_over_limit(self): + """--samples above 1M is rejected.""" + result = runner.invoke( + app, + ["download", "test/dataset", "--samples", "2000000", "-o", "out.jsonl"], + ) + assert result.exit_code == 1 + assert "1,000,000" in result.output + + @patch("soup_cli.commands.data._hf_dataset_info") + def test_preview_huggingface_hub_not_installed(self, mock_info): + """Preview gracefully handles missing huggingface_hub.""" + mock_info.side_effect = ImportError("No module named 'huggingface_hub'") + result = runner.invoke(app, ["preview", "test/ds"]) + assert result.exit_code == 1 + assert "huggingface" in result.output.lower() diff --git a/tests/test_dataset_registry.py b/tests/test_dataset_registry.py new file mode 100644 index 0000000..d82dcb2 --- /dev/null +++ b/tests/test_dataset_registry.py @@ -0,0 +1,254 @@ +"""Tests for dataset info registry: register, unregister, list, resolve.""" + +from __future__ import annotations + +from unittest.mock import patch + +from typer.testing import CliRunner + +from soup_cli.commands.data import app + +runner = CliRunner() + + +# --------------------------------------------------------------------------- +# Registry CRUD +# --------------------------------------------------------------------------- + + +class TestDatasetRegistry: + """Tests for dataset registry utility functions.""" + + def test_register_dataset(self, tmp_path): + """register_dataset adds entry to registry.""" + from soup_cli.utils.registry import ( + load_registry, + register_dataset, + ) + + registry_path = tmp_path / "datasets.json" + register_dataset( + "my-instruct", + str(tmp_path / "data.jsonl"), + "alpaca", + registry_path=registry_path, + ) + reg = load_registry(registry_path) + assert "my-instruct" in reg + assert reg["my-instruct"]["format"] == "alpaca" + + def test_register_overwrites(self, tmp_path): + """Registering same name overwrites previous entry.""" + from soup_cli.utils.registry import ( + load_registry, + register_dataset, + ) + + registry_path = tmp_path / "datasets.json" + register_dataset("ds", "/a.jsonl", "alpaca", registry_path=registry_path) + register_dataset("ds", "/b.jsonl", "sharegpt", registry_path=registry_path) + reg = load_registry(registry_path) + assert reg["ds"]["path"].endswith("b.jsonl") + assert reg["ds"]["format"] == "sharegpt" + + def test_unregister_dataset(self, tmp_path): + """unregister_dataset removes entry.""" + from soup_cli.utils.registry import ( + load_registry, + register_dataset, + unregister_dataset, + ) + + registry_path = tmp_path / "datasets.json" + register_dataset("ds", "/a.jsonl", "alpaca", registry_path=registry_path) + result = unregister_dataset("ds", registry_path=registry_path) + assert result is True + reg = load_registry(registry_path) + assert "ds" not in reg + + def test_unregister_nonexistent(self, tmp_path): + """unregister_dataset returns False for missing name.""" + from soup_cli.utils.registry import unregister_dataset + + registry_path = tmp_path / "datasets.json" + result = unregister_dataset("nope", registry_path=registry_path) + assert result is False + + def test_load_empty_registry(self, tmp_path): + """load_registry returns empty dict for missing file.""" + from soup_cli.utils.registry import load_registry + + registry_path = tmp_path / "datasets.json" + assert load_registry(registry_path) == {} + + def test_resolve_registered_name(self, tmp_path): + """resolve_dataset returns path for registered name.""" + from soup_cli.utils.registry import register_dataset, resolve_dataset + + registry_path = tmp_path / "datasets.json" + data_path = str(tmp_path / "data.jsonl") + register_dataset("my-ds", data_path, "alpaca", registry_path=registry_path) + resolved = resolve_dataset("my-ds", registry_path=registry_path) + assert resolved is not None + assert resolved["path"] == data_path + + def test_resolve_unregistered_name(self, tmp_path): + """resolve_dataset returns None for unknown name.""" + from soup_cli.utils.registry import resolve_dataset + + registry_path = tmp_path / "datasets.json" + assert resolve_dataset("unknown", registry_path=registry_path) is None + + def test_list_registry(self, tmp_path): + """load_registry returns all entries.""" + from soup_cli.utils.registry import ( + load_registry, + register_dataset, + ) + + registry_path = tmp_path / "datasets.json" + register_dataset("a", "/a.jsonl", "alpaca", registry_path=registry_path) + register_dataset("b", "/b.jsonl", "dpo", registry_path=registry_path) + reg = load_registry(registry_path) + assert len(reg) == 2 + assert "a" in reg + assert "b" in reg + + +# --------------------------------------------------------------------------- +# CLI commands +# --------------------------------------------------------------------------- + + +class TestRegistryCLI: + """Tests for `soup data register/unregister/list` CLI commands.""" + + @patch("soup_cli.commands.data._get_registry_path") + def test_register_cli(self, mock_path, tmp_path, monkeypatch): + """CLI register creates registry entry.""" + monkeypatch.chdir(tmp_path) + registry_path = tmp_path / "datasets.json" + mock_path.return_value = registry_path + + # Create a dummy data file + data_file = tmp_path / "train.jsonl" + data_file.write_text('{"instruction": "hi", "output": "hello"}\n') + + result = runner.invoke( + app, + ["register", "--name", "my-ds", "--path", str(data_file), "--format", "alpaca"], + ) + assert result.exit_code == 0 + assert "registered" in result.output.lower() + + @patch("soup_cli.commands.data._get_registry_path") + def test_unregister_cli(self, mock_path, tmp_path): + """CLI unregister removes entry.""" + from soup_cli.utils.registry import register_dataset + + registry_path = tmp_path / "datasets.json" + mock_path.return_value = registry_path + register_dataset("my-ds", "/data.jsonl", "alpaca", registry_path=registry_path) + + result = runner.invoke(app, ["unregister", "--name", "my-ds"]) + assert result.exit_code == 0 + assert "removed" in result.output.lower() + + @patch("soup_cli.commands.data._get_registry_path") + def test_unregister_missing_cli(self, mock_path, tmp_path): + """CLI unregister for missing name shows error.""" + registry_path = tmp_path / "datasets.json" + mock_path.return_value = registry_path + + result = runner.invoke(app, ["unregister", "--name", "nonexistent"]) + assert result.exit_code == 1 + + @patch("soup_cli.commands.data._get_registry_path") + def test_list_registry_cli(self, mock_path, tmp_path): + """CLI list shows registered datasets.""" + from soup_cli.utils.registry import register_dataset + + registry_path = tmp_path / "datasets.json" + mock_path.return_value = registry_path + register_dataset("ds1", "/a.jsonl", "alpaca", registry_path=registry_path) + register_dataset("ds2", "/b.jsonl", "dpo", registry_path=registry_path) + + # list is a Python builtin, command is "registry" or "list" + result = runner.invoke(app, ["registry"]) + assert result.exit_code == 0 + assert "ds1" in result.output + assert "ds2" in result.output + + @patch("soup_cli.commands.data._get_registry_path") + def test_list_empty_registry_cli(self, mock_path, tmp_path): + """CLI list on empty registry shows message.""" + registry_path = tmp_path / "datasets.json" + mock_path.return_value = registry_path + + result = runner.invoke(app, ["registry"]) + assert result.exit_code == 0 + assert "no datasets" in result.output.lower() or "empty" in result.output.lower() + + +# --------------------------------------------------------------------------- +# Name validation +# --------------------------------------------------------------------------- + + +class TestRegistryValidation: + """Tests for dataset name validation.""" + + def test_name_with_path_separator_rejected(self, tmp_path): + """Names with / or \\ are rejected.""" + from soup_cli.utils.registry import register_dataset + + registry_path = tmp_path / "datasets.json" + with __import__("pytest").raises(ValueError): + register_dataset( + "bad/name", "/data.jsonl", "alpaca", registry_path=registry_path + ) + + def test_name_with_null_byte_rejected(self, tmp_path): + """Names with null bytes are rejected.""" + from soup_cli.utils.registry import register_dataset + + registry_path = tmp_path / "datasets.json" + with __import__("pytest").raises(ValueError): + register_dataset( + "bad\x00name", "/data.jsonl", "alpaca", registry_path=registry_path + ) + + def test_empty_name_rejected(self, tmp_path): + """Empty names are rejected.""" + from soup_cli.utils.registry import register_dataset + + registry_path = tmp_path / "datasets.json" + with __import__("pytest").raises(ValueError): + register_dataset("", "/data.jsonl", "alpaca", registry_path=registry_path) + + +# --------------------------------------------------------------------------- +# Error handling +# --------------------------------------------------------------------------- + + +class TestRegistryErrorHandling: + """Tests for registry error handling.""" + + def test_load_corrupted_registry(self, tmp_path): + """Corrupted JSON file raises ValueError.""" + from soup_cli.utils.registry import load_registry + + path = tmp_path / "datasets.json" + path.write_text("{not valid json", encoding="utf-8") + with __import__("pytest").raises(ValueError, match="corrupted"): + load_registry(path) + + def test_load_non_dict_registry(self, tmp_path): + """Non-dict JSON content raises ValueError.""" + from soup_cli.utils.registry import load_registry + + path = tmp_path / "datasets.json" + path.write_text("[1, 2, 3]", encoding="utf-8") + with __import__("pytest").raises(ValueError, match="unexpected format"): + load_registry(path) diff --git a/tests/test_freeze_training.py b/tests/test_freeze_training.py new file mode 100644 index 0000000..2fc12e3 --- /dev/null +++ b/tests/test_freeze_training.py @@ -0,0 +1,272 @@ +"""Tests for freeze training: freeze_layers, freeze_ratio config and layer freezing.""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import pytest + +from soup_cli.config.schema import SoupConfig, TrainingConfig + +# --------------------------------------------------------------------------- +# Config validation +# --------------------------------------------------------------------------- + + +class TestFreezeConfig: + """Tests for freeze training config fields.""" + + def test_freeze_layers_default_none(self): + """freeze_layers defaults to None.""" + cfg = TrainingConfig() + assert cfg.freeze_layers is None + + def test_freeze_ratio_default_none(self): + """freeze_ratio defaults to None.""" + cfg = TrainingConfig() + assert cfg.freeze_ratio is None + + def test_freeze_layers_valid(self): + """freeze_layers accepts positive int.""" + cfg = TrainingConfig(freeze_layers=24) + assert cfg.freeze_layers == 24 + + def test_freeze_layers_zero_rejected(self): + """freeze_layers must be >= 1.""" + with pytest.raises(Exception): + TrainingConfig(freeze_layers=0) + + def test_freeze_layers_negative_rejected(self): + """freeze_layers must be positive.""" + with pytest.raises(Exception): + TrainingConfig(freeze_layers=-5) + + def test_freeze_ratio_valid(self): + """freeze_ratio accepts float in (0, 1).""" + cfg = TrainingConfig(freeze_ratio=0.75) + assert cfg.freeze_ratio == 0.75 + + def test_freeze_ratio_zero_rejected(self): + """freeze_ratio must be > 0.""" + with pytest.raises(Exception): + TrainingConfig(freeze_ratio=0.0) + + def test_freeze_ratio_one_rejected(self): + """freeze_ratio must be < 1 (can't freeze everything).""" + with pytest.raises(Exception): + TrainingConfig(freeze_ratio=1.0) + + def test_freeze_ratio_over_one_rejected(self): + """freeze_ratio must be < 1.""" + with pytest.raises(Exception): + TrainingConfig(freeze_ratio=1.5) + + def test_freeze_layers_and_ratio_both_set(self): + """Both freeze_layers and freeze_ratio can be set (layers takes priority).""" + cfg = TrainingConfig(freeze_layers=10, freeze_ratio=0.5) + assert cfg.freeze_layers == 10 + assert cfg.freeze_ratio == 0.5 + + def test_freeze_in_yaml_roundtrip(self): + """freeze fields survive YAML round-trip via SoupConfig.""" + cfg = SoupConfig( + base="test/model", + data={"train": "data.jsonl"}, + training={"freeze_layers": 16}, + ) + assert cfg.training.freeze_layers == 16 + + def test_freeze_ratio_in_yaml_roundtrip(self): + """freeze_ratio survives YAML round-trip.""" + cfg = SoupConfig( + base="test/model", + data={"train": "data.jsonl"}, + training={"freeze_ratio": 0.5}, + ) + assert cfg.training.freeze_ratio == 0.5 + + +# --------------------------------------------------------------------------- +# Layer freezing logic +# --------------------------------------------------------------------------- + + +class TestFreezeModelLayers: + """Tests for freeze_model_layers utility function.""" + + def _make_mock_model(self, num_layers: int = 32): + """Create a mock model with named_parameters.""" + model = MagicMock() + params = [] + for layer_idx in range(num_layers): + parts = [ + "self_attn.q_proj.weight", + "self_attn.v_proj.weight", + "mlp.up_proj.weight", + ] + for part in parts: + param = MagicMock() + param.requires_grad = True + name = f"model.layers.{layer_idx}.{part}" + params.append((name, param)) + # Add non-layer params (embed, lm_head) + embed_param = MagicMock() + embed_param.requires_grad = True + params.append(("model.embed_tokens.weight", embed_param)) + + head_param = MagicMock() + head_param.requires_grad = True + params.append(("lm_head.weight", head_param)) + + model.named_parameters.return_value = params + return model, params + + def test_freeze_by_layer_count(self): + """freeze_model_layers freezes first N layers.""" + from soup_cli.utils.freeze import freeze_model_layers + + model, params = self._make_mock_model(32) + frozen_count = freeze_model_layers(model, freeze_layers=24) + + # First 24 layers' params should have requires_grad = False + for name, param in params: + if "layers." in name: + layer_idx = int(name.split("layers.")[1].split(".")[0]) + if layer_idx < 24: + assert param.requires_grad is False + else: + assert param.requires_grad is True + + assert frozen_count > 0 + + def test_freeze_by_ratio(self): + """freeze_model_layers freezes by ratio.""" + from soup_cli.utils.freeze import freeze_model_layers + + model, params = self._make_mock_model(32) + frozen_count = freeze_model_layers(model, freeze_ratio=0.75) + + # 75% of 32 = 24 layers frozen + for name, param in params: + if "layers." in name: + layer_idx = int(name.split("layers.")[1].split(".")[0]) + if layer_idx < 24: + assert param.requires_grad is False + + assert frozen_count > 0 + + def test_freeze_layers_priority_over_ratio(self): + """freeze_layers takes priority when both specified.""" + from soup_cli.utils.freeze import freeze_model_layers + + model, params = self._make_mock_model(32) + freeze_model_layers(model, freeze_layers=10, freeze_ratio=0.75) + + # Should freeze 10 layers, not 24 + for name, param in params: + if "layers." in name: + layer_idx = int(name.split("layers.")[1].split(".")[0]) + if layer_idx < 10: + assert param.requires_grad is False + else: + assert param.requires_grad is True + + def test_freeze_does_not_freeze_embeddings(self): + """Embeddings are not frozen (they're not layer params).""" + from soup_cli.utils.freeze import freeze_model_layers + + model, params = self._make_mock_model(32) + freeze_model_layers(model, freeze_layers=24) + + # embed_tokens and lm_head should remain trainable + for name, param in params: + if "embed_tokens" in name or "lm_head" in name: + assert param.requires_grad is True + + def test_freeze_more_than_total_layers(self): + """Freezing more layers than model has freezes all layers.""" + from soup_cli.utils.freeze import freeze_model_layers + + model, params = self._make_mock_model(8) + freeze_model_layers(model, freeze_layers=100) + + # All 8 layers frozen + for name, param in params: + if "layers." in name: + assert param.requires_grad is False + + def test_freeze_returns_count(self): + """freeze_model_layers returns number of frozen parameters.""" + from soup_cli.utils.freeze import freeze_model_layers + + model, params = self._make_mock_model(32) + frozen = freeze_model_layers(model, freeze_layers=16) + # 16 layers × 3 params each = 48 + assert frozen == 48 + + def test_no_freeze_when_none(self): + """No freezing when both are None.""" + from soup_cli.utils.freeze import freeze_model_layers + + model, params = self._make_mock_model(8) + frozen = freeze_model_layers(model, freeze_layers=None, freeze_ratio=None) + assert frozen == 0 + + def test_detect_num_layers(self): + """_detect_num_layers extracts layer count from model params.""" + from soup_cli.utils.freeze import _detect_num_layers + + model = MagicMock() + params = [ + (f"model.layers.{idx}.self_attn.weight", MagicMock()) + for idx in range(32) + ] + model.named_parameters.return_value = params + assert _detect_num_layers(model) == 32 + + def test_detect_num_layers_no_layers(self): + """_detect_num_layers returns 0 for models without numbered layers.""" + from soup_cli.utils.freeze import _detect_num_layers + + model = MagicMock() + model.named_parameters.return_value = [ + ("embed.weight", MagicMock()), + ] + assert _detect_num_layers(model) == 0 + + def test_detect_num_layers_gpt2_style(self): + """_detect_num_layers handles GPT-2 style 'transformer.h.N.' naming.""" + from soup_cli.utils.freeze import _detect_num_layers + + model = MagicMock() + params = [ + (f"transformer.h.{idx}.attn.weight", MagicMock()) + for idx in range(12) + ] + model.named_parameters.return_value = params + assert _detect_num_layers(model) == 12 + + +# --------------------------------------------------------------------------- +# Sweep integration +# --------------------------------------------------------------------------- + + +class TestFreezeSweep: + """Tests for freeze fields in sweep param support.""" + + def test_freeze_layers_in_sweep(self): + """freeze_layers is a valid sweep param.""" + from soup_cli.commands.sweep import _parse_sweep_params + + params = _parse_sweep_params(["training.freeze_layers=8,16,24"]) + assert "training.freeze_layers" in params + assert len(params["training.freeze_layers"]) == 3 + + def test_freeze_ratio_in_sweep(self): + """freeze_ratio is a valid sweep param.""" + from soup_cli.commands.sweep import _parse_sweep_params + + params = _parse_sweep_params(["training.freeze_ratio=0.25,0.5,0.75"]) + assert "training.freeze_ratio" in params + assert len(params["training.freeze_ratio"]) == 3 diff --git a/tests/test_loss_watchdog.py b/tests/test_loss_watchdog.py new file mode 100644 index 0000000..7b1fbe7 --- /dev/null +++ b/tests/test_loss_watchdog.py @@ -0,0 +1,218 @@ +"""Tests for loss watchdog: auto-stop on loss spikes.""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import pytest + +from soup_cli.config.schema import TrainingConfig + +# --------------------------------------------------------------------------- +# Config validation +# --------------------------------------------------------------------------- + + +class TestWatchdogConfig: + """Tests for loss_watchdog config fields.""" + + def test_loss_watchdog_default_false(self): + """loss_watchdog defaults to False.""" + cfg = TrainingConfig() + assert cfg.loss_watchdog is False + + def test_loss_watchdog_threshold_default(self): + """loss_watchdog_threshold defaults to 3.0.""" + cfg = TrainingConfig() + assert cfg.loss_watchdog_threshold == 3.0 + + def test_loss_watchdog_patience_default(self): + """loss_watchdog_patience defaults to 5.""" + cfg = TrainingConfig() + assert cfg.loss_watchdog_patience == 5 + + def test_loss_watchdog_threshold_positive(self): + """loss_watchdog_threshold must be positive.""" + with pytest.raises(Exception): + TrainingConfig(loss_watchdog_threshold=0.0) + + def test_loss_watchdog_threshold_negative_rejected(self): + """loss_watchdog_threshold rejects negative values.""" + with pytest.raises(Exception): + TrainingConfig(loss_watchdog_threshold=-1.0) + + def test_loss_watchdog_patience_positive(self): + """loss_watchdog_patience must be >= 1.""" + with pytest.raises(Exception): + TrainingConfig(loss_watchdog_patience=0) + + def test_loss_watchdog_enable(self): + """loss_watchdog can be enabled.""" + cfg = TrainingConfig(loss_watchdog=True) + assert cfg.loss_watchdog is True + + def test_loss_watchdog_custom_values(self): + """Custom threshold and patience values are accepted.""" + cfg = TrainingConfig( + loss_watchdog=True, + loss_watchdog_threshold=5.0, + loss_watchdog_patience=10, + ) + assert cfg.loss_watchdog_threshold == 5.0 + assert cfg.loss_watchdog_patience == 10 + + +# --------------------------------------------------------------------------- +# Watchdog logic in callback +# --------------------------------------------------------------------------- + + +class TestWatchdogCallback: + """Tests for loss watchdog logic in SoupTrainerCallback.""" + + def _make_callback(self, threshold: float = 3.0, patience: int = 5): + """Create a callback with watchdog enabled.""" + from soup_cli.monitoring.callback import SoupTrainerCallback + from soup_cli.monitoring.display import TrainingDisplay + + display = MagicMock(spec=TrainingDisplay) + cb = SoupTrainerCallback( + display=display, + loss_watchdog=True, + loss_watchdog_threshold=threshold, + loss_watchdog_patience=patience, + ) + return cb + + def test_watchdog_no_stop_normal_loss(self): + """Normal loss values do not trigger watchdog.""" + cb = self._make_callback(threshold=3.0, patience=3) + + state = MagicMock() + state.global_step = 1 + state.epoch = 1.0 + args = MagicMock() + control = MagicMock() + control.should_training_stop = False + + # Simulate 5 normal loss values + for step in range(5): + state.global_step = step + 1 + cb.on_log(args, state, control, logs={"loss": 1.5}) + + assert control.should_training_stop is False + + def test_watchdog_triggers_after_patience(self): + """Watchdog stops training after patience consecutive high losses.""" + cb = self._make_callback(threshold=3.0, patience=3) + + state = MagicMock() + args = MagicMock() + control = MagicMock() + control.should_training_stop = False + + # Simulate patience+1 high loss values + for step in range(4): + state.global_step = step + 1 + state.epoch = 1.0 + cb.on_log(args, state, control, logs={"loss": 5.0}) + + assert control.should_training_stop is True + + def test_watchdog_resets_on_good_loss(self): + """Counter resets when a normal loss appears.""" + cb = self._make_callback(threshold=3.0, patience=5) + + state = MagicMock() + args = MagicMock() + control = MagicMock() + control.should_training_stop = False + + # 3 high, then 1 normal, then 3 high — should NOT trigger (patience=5) + for loss_val in [5.0, 5.0, 5.0, 1.0, 5.0, 5.0, 5.0]: + state.global_step = 1 + state.epoch = 1.0 + cb.on_log(args, state, control, logs={"loss": loss_val}) + + assert control.should_training_stop is False + + def test_watchdog_disabled_by_default(self): + """Watchdog does not interfere when disabled.""" + from soup_cli.monitoring.callback import SoupTrainerCallback + from soup_cli.monitoring.display import TrainingDisplay + + display = MagicMock(spec=TrainingDisplay) + cb = SoupTrainerCallback(display=display) + + state = MagicMock() + args = MagicMock() + control = MagicMock() + control.should_training_stop = False + + # Very high loss — should NOT trigger since watchdog is off + for step in range(20): + state.global_step = step + state.epoch = 1.0 + cb.on_log(args, state, control, logs={"loss": 100.0}) + + assert control.should_training_stop is False + + def test_watchdog_exact_patience_boundary(self): + """Exactly patience high-loss steps triggers on the next one.""" + cb = self._make_callback(threshold=3.0, patience=3) + + state = MagicMock() + args = MagicMock() + control = MagicMock() + control.should_training_stop = False + + # Exactly 3 high losses (patience=3) — should trigger on 4th + for step in range(3): + state.global_step = step + 1 + state.epoch = 1.0 + cb.on_log(args, state, control, logs={"loss": 5.0}) + + # After patience steps of high loss, next one triggers + assert control.should_training_stop is True + + def test_watchdog_no_loss_in_logs(self): + """Watchdog handles logs without loss key gracefully.""" + cb = self._make_callback(threshold=3.0, patience=3) + + state = MagicMock() + args = MagicMock() + control = MagicMock() + control.should_training_stop = False + + # Logs without loss key + cb.on_log(args, state, control, logs={"learning_rate": 1e-5}) + assert control.should_training_stop is False + + +# --------------------------------------------------------------------------- +# Sweep integration +# --------------------------------------------------------------------------- + + +class TestWatchdogSweep: + """Tests for watchdog fields in sweep param support.""" + + def test_watchdog_threshold_in_sweep(self): + """loss_watchdog_threshold is a valid sweep param.""" + from soup_cli.commands.sweep import _parse_sweep_params + + params = _parse_sweep_params( + ["training.loss_watchdog_threshold=2.0,3.0,5.0"] + ) + assert "training.loss_watchdog_threshold" in params + assert len(params["training.loss_watchdog_threshold"]) == 3 + + def test_watchdog_patience_in_sweep(self): + """loss_watchdog_patience is a valid sweep param.""" + from soup_cli.commands.sweep import _parse_sweep_params + + params = _parse_sweep_params( + ["training.loss_watchdog_patience=3,5,10"] + ) + assert "training.loss_watchdog_patience" in params + assert len(params["training.loss_watchdog_patience"]) == 3