feat: v0.24.0 — Dataset Hub, Freeze Training, Loss Watchdog, Dataset Registry

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.
This commit is contained in:
Alpamys 2026-04-03 16:35:23 +05:00
parent ada4a078b6
commit 1b6b428aaa
26 changed files with 1936 additions and 23 deletions

View File

@ -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 |

View File

@ -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 <path> --pct 10 Sample by percentage
soup data split <path> --val 10 --test 10 Split into train/val/test
soup data split <path> --val 500 --absolute Split with absolute counts
soup data split <path> --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

View File

@ -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

View File

@ -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"

View File

@ -1,3 +1,3 @@
"""Soup CLI — Fine-tune LLMs in one command."""
__version__ = "0.23.1"
__version__ = "0.24.0"

View File

@ -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: <dataset-name>.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)

View File

@ -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,

View File

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

View File

@ -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."""

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

@ -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

View File

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

View File

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

View File

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

View File

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

73
soup_cli/utils/freeze.py Normal file
View File

@ -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

100
soup_cli/utils/registry.py Normal file
View File

@ -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

357
tests/test_dataset_hub.py Normal file
View File

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

View File

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

View File

@ -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

218
tests/test_loss_watchdog.py Normal file
View File

@ -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