mirror of https://github.com/razor-ai/soup.git
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:
parent
ada4a078b6
commit
1b6b428aaa
|
|
@ -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 |
|
||||
|
|
|
|||
32
README.md
32
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 <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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -1,3 +1,3 @@
|
|||
"""Soup CLI — Fine-tune LLMs in one command."""
|
||||
|
||||
__version__ = "0.23.1"
|
||||
__version__ = "0.24.0"
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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()
|
||||
|
|
@ -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)
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
Loading…
Reference in New Issue