diff --git a/CHANGELOG.md b/CHANGELOG.md index 6270219..5fad6b7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,61 @@ reproducing 70+ versions of notes. ## [Unreleased] +## [0.71.22] - 2026-06-10 + +### Added +- **Perf & measure polish** — a 4-issue patch tightening four live paths from + the recent BETA lifts. Pure code, validated on Windows + RTX 3050. + - **MiniLLM on-policy KV-cache (closes #263).** The on-policy distillation + rollout (`soup train` with `training.minillm_on_policy: true`) now threads + `past_key_values` so each step forwards only the new token instead of + re-feeding the whole prefix — resolving the O(L²) per-step cost from + v0.71.18. A LoRA student (the common distill case) activates the cache + too: the new `_supports_kv_cache` probe unwraps the PEFT model via + `get_base_model()` before deciding. The teacher is always cached; the + student cache respects the retained autograd graph and degrades gracefully + if a model returns no cache mid-loop. + - **`soup serve --mole` KV-cache (closes #262).** Each of the N task adapters + in a served MoLE now keeps its own KV cache in lockstep, created fresh per + `generate()` call (never stored on the instance, so there is no + cross-request leak). Top-k zero-weight adapters are still skipped, and the + output is byte-identical to the no-cache path on a real MoLE. + - **Deploy-autopilot live measure factories (closes #143).** `soup deploy + autopilot --measure` ships a first-party transformers loader factory (lazy + import, per-candidate quant config via the Quant Menu loader; `before` = + base, `after` = quantised) replacing the inject-only test hooks. The + baseline is now scored **once** and the whole candidate list is + **pre-validated up front**, so a typo in `--measure-candidates` raises + before any model load instead of burning N live loads or doubling peak + VRAM. + - **Live-codec TTS via SNAC, partial (#265-partial).** The live-codec + encode path (`data.format='audio'`) is validated for **Orpheus**: + `load_audio_mono` now probes `soundfile.info` (duration + byte cap) + *before* `soundfile.read` (no multi-GB decode into RAM) and reads through + an `O_NOFOLLOW` file descriptor; a real SNAC-backed encode of a 24 kHz wav + produced 42 Orpheus codec tokens. + +### Fixed +- MiniLLM on-policy KV-cache was silently disabled for LoRA students (the + PEFT wrapper hid the base model's `past_key_values` support) — now probed + via `get_base_model()`. +- Deploy-measure no longer re-scores the baseline once per candidate or burns + live model loads on a bad candidate (per-candidate validation moved up front). +- `load_audio_mono` capped audio duration only *after* decoding into RAM — + the cap is now checked from `soundfile.info` before reading. + +### Known limitations +- KV-cache correctness is validated (cache == no-cache equality on real tiny + artifacts) but large-model throughput gains were not measured on the 4 GB + dev box. +- **#265 stays open** — the live-codec `data.format='audio'` SNAC encode path + is validated for Orpheus only; the other four TTS families keep their + per-family codec dependency gate. +- The deploy-measure first-party factory's real quantized (bitsandbytes 4-bit) + load is CUDA + bitsandbytes-gated; on Windows / no-bnb the injected test + seams are the validated path. +- The MoLE serve KV-cache assumes single-sequence (`B == 1`) decode. + ## [0.71.21] - 2026-06-10 ### Added diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 7101e85..e48fb89 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -120,7 +120,7 @@ src/soup_cli/ templates/ - 17 built-in soup.yaml templates (YAML + manifest.json) with load_template loader (v0.39.0, +bco v0.40.0) ui/ - Web UI (FastAPI + HTML/JS SPA) -tests/ - Test suite (292 files, 14084 tests) +tests/ - Test suite (293 files, 14184 tests) examples/ - Real-world config examples and datasets ``` @@ -150,7 +150,7 @@ pytest tests/test_data.py::test_detect_alpaca_format -v pytest tests/ --cov=soup_cli --cov-report=html ``` -### Test Files (292 files) +### Test Files (293 files) > A representative sample of the suite below. The full table lives in > [`.claude/CLAUDE.md`](.claude/CLAUDE.md); run `pytest tests/ -v` for the complete list. diff --git a/README.md b/README.md index 5c53bbd..8dd1b2c 100644 --- a/README.md +++ b/README.md @@ -49,23 +49,22 @@ infrastructure instead of improving models. Soup fixes that. ## What's New -**v0.71.21 — Precision & rollout lift (BETA, hardware-gated).** Five long-deferred stubs go live: +**v0.71.22 — Perf & measure polish.** Four live paths from the recent BETA lifts get tighter: -- **Multi-turn agent rollouts for GRPO** — `training.rollout_backend: openenv` + - `training.rollout_func: my_module:my_fn` runs a live multi-turn rollout at the start of GRPO - training; the rows your function returns replace the prompt dataset. ART / RULER / NeMo-Gym - adapters ship behind honest dependency gates. -- **FP8 attention + NVFP4 training** — `training.fp8_attention: true` converts the attention - projections to torchao float8 (Hopper-gated) and `training.nvfp4: true` applies torchao NVFP4 - quantization (Blackwell-gated). Unsupported hardware degrades to a clear advisory, never a crash. -- **vLLM sleep mode for RL** — `training.vllm_sleep_mode: true` puts the vLLM engine on standby - between GRPO rollouts (vLLM ≥ 0.7), freeing VRAM for the training step. -- **Apple-adapter conversion is live** — `soup apple-adapter --direction hf-to-mlx | mlx-to-hf` - converts PEFT LoRA safetensors ↔ mlx-lm adapters with a numerically-equal round trip - (rank / scale / dropout / num_layers carried). -- **Llama-4 expert delinearization is live** — `soup delinearize-llama4 --target ` - reshapes fused `[E*din, dout]` expert weights to `[E, din, dout]` shard-by-shard and copies the - JSON sidecars so the target stays loadable. +- **MiniLLM on-policy distillation is fast** — the on-policy rollout + (`training.minillm_on_policy: true`) now threads a KV cache so each step forwards only the new + token instead of re-feeding the whole prefix, resolving the earlier O(L²) cost. A LoRA student + activates the cache too (the PEFT wrapper is unwrapped before the cache check). +- **`soup serve --mole` uses a KV cache** — each task adapter in a served MoLE keeps its own cache + in lockstep, created fresh per request (no cross-request leak). Output is byte-identical to the + previous no-cache path. +- **`soup deploy autopilot --measure` is real and frugal** — a first-party transformers loader + factory loads each candidate live (per-candidate quant config), scores the baseline **once**, and + pre-validates the whole `--measure-candidates` list up front, so a typo fails before any model + load instead of burning N loads or doubling peak VRAM. +- **Live-codec TTS via SNAC (Orpheus)** — encoding raw audio at train time (`data.format: audio`) + is validated for Orpheus: audio is duration- and byte-capped from `soundfile.info` **before** + decoding into RAM, and read through a symlink-safe file descriptor. Full history: [CHANGELOG.md](CHANGELOG.md) · [GitHub Releases](https://github.com/MakazhanAlpamys/Soup/releases). diff --git a/pyproject.toml b/pyproject.toml index 3dc78b8..9610c00 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "soup-cli" -version = "0.71.21" +version = "0.71.22" description = "Fine-tune and post-train LLMs in one command. No SSH, no config hell." readme = "README.md" license = "Apache-2.0" @@ -22,7 +22,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] # Core install (`pip install soup-cli`) is intentionally light: the CLI, config -# system, and data tools — no PyTorch. v0.71.0 split the heavy training stack +# system, and data tools — no PyTorch. v0.71.0 split the heavy training stack # (torch / transformers / peft / trl / datasets / bitsandbytes / accelerate) # into the `[train]` extra below. dependencies = [ @@ -35,7 +35,7 @@ dependencies = [ ] [project.optional-dependencies] -# v0.71.0 — heavy training stack. `pip install 'soup-cli[train]'` to fine-tune. +# v0.71.0 — heavy training stack. `pip install 'soup-cli[train]'` to fine-tune. # These were core dependencies through v0.70.0; pins are unchanged. train = [ "torch>=2.0.0", @@ -46,7 +46,7 @@ train = [ "bitsandbytes>=0.41.0", "accelerate>=0.25.0", ] -# v0.71.0 — convenience meta-extra pulling the main optional stacks. +# v0.71.0 — convenience meta-extra pulling the main optional stacks. all = ["soup-cli[train,serve,ui,data]"] eval = ["lm-eval>=0.4.0"] data = ["datasketch>=1.6.0"] @@ -74,31 +74,31 @@ sglang = ["sglang>=0.2.0", "fastapi>=0.104.0", "uvicorn>=0.24.0"] mlx = ["mlx>=0.20.0", "mlx-lm>=0.20.0"] cce = ["cut-cross-entropy>=24.10.0"] tui = ["textual>=0.50.0"] -# v0.53.8 #89 — bundle MLflow / SwanLab / Trackio for `--tracker` users. +# v0.53.8 #89 — bundle MLflow / SwanLab / Trackio for `--tracker` users. trackers = ["mlflow>=2.0.0", "swanlab>=0.3.0", "trackio>=0.0.1"] -# v0.53.8 #85 — fsspec backends for remote dataset loading (s3 / gs / az / oci). +# v0.53.8 #85 — fsspec backends for remote dataset loading (s3 / gs / az / oci). remote = ["fsspec>=2024.1.0", "s3fs>=2024.1.0", "gcsfs>=2024.1.0", "adlfs>=2024.1.0"] -# v0.53.10 #150 — bundle scikit-optimize so `soup data mix --optimize` runs the +# v0.53.10 #150 — bundle scikit-optimize so `soup data mix --optimize` runs the # Bayesian-style loop instead of falling back to the v0.48.0 Dirichlet sampler. mix = ["scikit-optimize>=0.9.0"] -# v0.53.10 #113 — production-grade data quality: langdetect (language) + +# v0.53.10 #113 — production-grade data quality: langdetect (language) + # presidio-analyzer (PII). Llama-Guard-3-1B is documented as a manual recipe # (license + ~600 MB weight blob too large to bundle by default). data-pro = ["langdetect>=1.0.9", "presidio-analyzer>=2.2.0"] -# v0.71.2 #179/#185 — ed25519 detached signing for `soup attest` / `soup +# v0.71.2 #179/#185 — ed25519 detached signing for `soup attest` / `soup # adapters sign`. Pure-offline; Sigstore keyless (OIDC + Fulcio/Rekor network) # stays infra-blocked and is NOT bundled here. sign = ["cryptography>=41.0.0"] -# v0.71.3 #181 — reportlab PDF rendering for `soup train --annex-xi *.pdf`. +# v0.71.3 #181 — reportlab PDF rendering for `soup train --annex-xi *.pdf`. pdf = ["reportlab>=4.0.0"] -# v0.71.3 #180 — codecarbon offline energy/CO2 measurement for +# v0.71.3 #180 — codecarbon offline energy/CO2 measurement for # `soup train --track-energy` (offline; no IP-geolocation network call). carbon = ["codecarbon>=2.0.0"] -# v0.71.13 #225/#227 — prompt-program / tool-schema compilers +# v0.71.13 #225/#227 — prompt-program / tool-schema compilers # (`soup compile` / `soup compile-tools`). Lazy-imported with a friendly # ImportError; not installed by default (heavy dependency trees). compile = ["dspy-ai>=2.5.0", "textgrad>=0.1.0", "gepa>=0.0.1"] -# v0.71.18 #16 — serverless cloud GPU training (`soup train --cloud modal`). +# v0.71.18 #16 — serverless cloud GPU training (`soup train --cloud modal`). # Lazy-imported; only needed for `--cloud-submit` (plan-only render needs no # dependency). Modal auth is via `modal setup`. modal = ["modal>=0.60.0"] @@ -113,7 +113,7 @@ Issues = "https://github.com/MakazhanAlpamys/Soup/issues" [tool.hatch.build.targets.wheel] packages = ["src/soup_cli"] -# v0.53.8 #93 — include bundled fixture JSONLs as package data so +# v0.53.8 #93 — include bundled fixture JSONLs as package data so # `soup data demo` works in zipapp / namespace-package installs. # Hatchling's ``packages = ["src/soup_cli"]`` already recurses into the # package directory, so we use the artifacts directive (NOT @@ -146,7 +146,7 @@ follow_imports = "silent" testpaths = ["tests"] markers = [ "smoke: slow smoke tests that download models and run training (run with: pytest -m smoke)", - "unit: fast isolated tests — no subprocess, network, filesystem, or real model load", + "unit: fast isolated tests — no subprocess, network, filesystem, or real model load", "integration: tests that touch real subprocess, SQLite, filesystem, or HTTP", ] addopts = "-m 'not smoke' --cov=soup_cli --cov-fail-under=77 --cov-report=term-missing:skip-covered" diff --git a/src/soup_cli/__init__.py b/src/soup_cli/__init__.py index b9f617f..81495ae 100644 --- a/src/soup_cli/__init__.py +++ b/src/soup_cli/__init__.py @@ -1,3 +1,3 @@ -"""Soup CLI — Fine-tune and post-train LLMs in one command.""" +"""Soup CLI — Fine-tune and post-train LLMs in one command.""" -__version__ = "0.71.21" +__version__ = "0.71.22" diff --git a/src/soup_cli/commands/deploy.py b/src/soup_cli/commands/deploy.py index 19f0ad4..a59491b 100644 --- a/src/soup_cli/commands/deploy.py +++ b/src/soup_cli/commands/deploy.py @@ -4,7 +4,7 @@ from __future__ import annotations import os from pathlib import Path -from typing import TYPE_CHECKING, Callable, List, Optional +from typing import TYPE_CHECKING, List, Optional import typer @@ -792,25 +792,21 @@ def _run_deploy_autopilot_measure( f"against {Path(tasks_file).name}...[/]" ) - def _placeholder_before(prompt: str) -> str: - # The full v0.46.1 live measurement plumbs in real - # transformers / vllm generators. v0.53.1 ships the orchestrator - # surface; callers / smoke runs can monkeypatch this in. - return "" - - def _placeholder_after_factory(candidate: str) -> Callable[[str], str]: - def _gen(prompt: str) -> str: - return "" - return _gen - - # Pull injected generators if the caller registered them via env (escape - # hatch for tests + advanced operator workflows) + # Injected generators (test seam / advanced operator workflows) still win; + # otherwise v0.71.22 #143 first-party transformers factories load the real + # base (before) + per-candidate Quant-Menu-quantized model (after), lazily + # on first prompt so a cache hit never loads a model. from soup_cli.utils import deploy_measure as _dm - before_gen = getattr(_dm, "_DEPLOY_MEASURE_BEFORE_GEN", None) or _placeholder_before - after_factory = ( - getattr(_dm, "_DEPLOY_MEASURE_AFTER_FACTORY", None) - or _placeholder_after_factory - ) + + injected_before = getattr(_dm, "_DEPLOY_MEASURE_BEFORE_GEN", None) + injected_after = getattr(_dm, "_DEPLOY_MEASURE_AFTER_FACTORY", None) + if injected_before is None and injected_after is None: + console.print( + "[dim]Using live transformers generators (greedy decode; models " + "load lazily per candidate)...[/]" + ) + before_gen = injected_before or _dm.build_before_generator(base) + after_factory = injected_after or _dm.build_after_generator_factory(base) try: results, cache_hit = run_measure( @@ -824,10 +820,18 @@ def _run_deploy_autopilot_measure( except (TypeError, ValueError, FileNotFoundError) as exc: console.print(f"[red]Measure failed:[/] {escape(str(exc))}") raise typer.Exit(1) from exc + except (RuntimeError, ImportError, OSError) as exc: + # Model-load failures from the live factories (missing quant kernel, + # OOM, network) — friendly exit, no traceback dump. + console.print(f"[red]Live measure failed:[/] {escape(str(exc))}") + raise typer.Exit(1) from exc console.print(render_measure_table(results)) if cache_hit: - console.print("[dim](cache hit — re-run with --no-cache to refresh)[/]") + console.print( + "[dim](cache hit — delete ~/.soup/deploy_autopilot_cache.json or " + "change --tasks to refresh)[/]" + ) best = pick_best(results) if best is not None: console.print( diff --git a/src/soup_cli/trainer/tts.py b/src/soup_cli/trainer/tts.py index 9d6d8c8..798a0e8 100644 --- a/src/soup_cli/trainer/tts.py +++ b/src/soup_cli/trainer/tts.py @@ -64,20 +64,27 @@ class TTSTrainerWrapper(SFTTrainerWrapper): data_format = getattr(cfg.data, "format", None) if data_format in _LIVE_CODEC_FORMATS: - # Live-codec mode: encoding raw audio needs the family's codec. + # Live-codec mode (v0.71.22 #265-partial): encode raw audio into + # codec-token strings at setup time, then fall through to the + # validated pre-encoded CE path. The per-family codec dep gate + # fires first (friendly pip hint); families without a validated + # encoder raise inside encode_tts_dataset (tracked in #265 — + # Orpheus/SNAC is the v0.71.22 live family). self._require_tts_codec(family) - # If the codec IS importable (never validated on the maintainer's - # box — no codecs/large models installed), the encode-then-CE path - # would run here. It is intentionally surfaced as an explicit - # not-yet-validated error so we never silently ship an unrun path. - pkg = tts_codec_package(family) - raise RuntimeError( - f"TTS live-codec mode (data.format='audio') for family " - f"'{family}' requires encoding audio with the {pkg!r} codec at " - "train time. This path is hardware/dependency-gated and not " - "yet validated. Use the pre-encoded chat workflow (encode " - "audio to codec tokens offline, then train with " - "data.format=chat) for a runnable TTS fine-tune." + from soup_cli.utils.tts_codec import encode_tts_dataset + + console.print( + f"[green]TTS live-codec:[/] encoding audio with the " + f"{tts_codec_package(family)!r} codec (family={family})" + ) + # device is set by the SFTTrainerWrapper.__init__; getattr keeps a + # bare object.__new__ test fixture (no device) working — default + # None falls through to CPU encoding in encode_tts_dataset. + dataset = encode_tts_dataset( + dataset, + family, + device=getattr(self, "device", None), + console=console, ) # Pre-encoded chat mode: plain SFT cross-entropy over the codec tokens. diff --git a/src/soup_cli/utils/deploy_measure.py b/src/soup_cli/utils/deploy_measure.py index 19b1893..ed22674 100644 --- a/src/soup_cli/utils/deploy_measure.py +++ b/src/soup_cli/utils/deploy_measure.py @@ -16,6 +16,7 @@ from __future__ import annotations import hashlib import json import os +import re import stat import tempfile from dataclasses import asdict, dataclass @@ -31,12 +32,21 @@ DEFAULT_MAJOR_THRESHOLD: float = 0.05 _MAX_CACHE_BYTES: int = 1 * 1024 * 1024 # 1 MB cap on cache file _MAX_CANDIDATES: int = 32 +_MAX_BASE_LEN: int = 512 # mirrors v0.40.5 reward_model / v0.62.0 --base policy + +# v0.71.22 #143 — closed candidate allowlist for the first-party generator +# factories. Mirrors the TrainingConfig.quantization Quant Menu surface +# ("none" = the unquantized baseline; HQQ uses the hqq:Nbit shape). +MEASURABLE_QUANT_CANDIDATES: frozenset[str] = frozenset( + {"none", "4bit", "8bit", "gptq", "awq", "aqlm", "eetq", "mxfp4", "fp8"} +) +_HQQ_CANDIDATE_RE = re.compile(r"^hqq:[12348]bit$") # Test + advanced-operator escape hatch: when set on this module, the deploy -# CLI uses these callables instead of the v0.46.1 model-loading generators. -# Documented as v0.53.1 deferral — see the v0.53.1 Known Limitations entry -# in plan.md. NOT a public API; the live transformers / vLLM generators land -# alongside v0.53.2. +# CLI uses these callables INSTEAD of the first-party transformers factories +# below (v0.71.22 #143 lifted the v0.53.1 placeholder deferral — the live +# loaders are ``build_before_generator`` / ``build_after_generator_factory``). +# NOT a public API; kept as the test seam. _DEPLOY_MEASURE_BEFORE_GEN: Optional[Callable[[str], str]] = None _DEPLOY_MEASURE_AFTER_FACTORY: Optional[ Callable[[str], Callable[[str], str]] @@ -227,12 +237,21 @@ def measure_candidate( *, candidate: str, tasks_file: str, - before_gen: Callable[[str], str], + before_gen: Optional[Callable[[str], str]] = None, after_gen: Callable[[str], str], + before_score: Optional[float] = None, ) -> MeasureResult: """Score one quant candidate against the baseline. Returns a :class:`MeasureResult` with classify_delta verdict. + + The baseline is deterministic across candidates, so the orchestrator + (:func:`run_measure`) computes it ONCE and threads the float in via + ``before_score`` — this avoids re-running full greedy generation over + every task once per candidate (N× redundant GPU time) and keeps the + baseline model out of VRAM while a quantized candidate is loaded. + Direct / back-compat callers may instead pass ``before_gen`` and the + baseline is scored here. Exactly one of the two must be supplied. """ if isinstance(candidate, bool) or not isinstance(candidate, str): raise TypeError("candidate must be a non-bool str") @@ -241,7 +260,18 @@ def measure_candidate( if "\x00" in candidate: raise ValueError("candidate must not contain null bytes") - before = _score_tasks(tasks_file, before_gen) + if before_score is not None: + if isinstance(before_score, bool) or not isinstance( + before_score, (int, float) + ): + raise TypeError("before_score must be a non-bool number") + before = float(before_score) + elif before_gen is not None: + before = _score_tasks(tasks_file, before_gen) + else: + raise ValueError( + "measure_candidate requires either before_gen or before_score" + ) after = _score_tasks(tasks_file, after_gen) delta = after - before if delta >= 0: @@ -276,6 +306,221 @@ def pick_best(results: Sequence[MeasureResult]) -> Optional[MeasureResult]: return max(results, key=lambda r: r.delta) +# --- v0.71.22 #143 — first-party transformers generator factories ----------- + + +def validate_measure_candidate(candidate: object) -> str: + """Validate + canonicalise a quant candidate for the live measure loop. + + Accepts the closed :data:`MEASURABLE_QUANT_CANDIDATES` allowlist plus the + ``hqq:{1,2,3,4,8}bit`` shape. Case-insensitive; returns the lowercase + canonical form. Bool rejected before the str check (project policy). + """ + if isinstance(candidate, bool): + raise TypeError("candidate must not be bool") + if not isinstance(candidate, str): + raise TypeError(f"candidate must be str, got {type(candidate).__name__}") + if not candidate: + raise ValueError("candidate must be non-empty") + if "\x00" in candidate: + raise ValueError("candidate must not contain null bytes") + canonical = candidate.lower() + if canonical in MEASURABLE_QUANT_CANDIDATES or _HQQ_CANDIDATE_RE.match( + canonical + ): + return canonical + supported = ", ".join(sorted(MEASURABLE_QUANT_CANDIDATES)) + # Truncate the echoed candidate (mirrors longlora._truncate_for_message + # policy) so a pathologically long value can't bloat the error message. + shown = candidate if len(candidate) <= 64 else candidate[:64] + "…" + raise ValueError( + f"candidate {shown!r} is not a measurable quant. Supported: " + f"{supported}, hqq:Nbit (N in 1/2/3/4/8)" + ) + + +def _check_base_id(base: object) -> str: + """Shape-validate a base model id / path for the factory builders.""" + # NOTE: this raises ValueError (not TypeError) for bool/non-str wrong-type + # input — DELIBERATELY inconsistent with validate_measure_candidate (which + # raises TypeError for the same class). The ValueError behaviour is pinned + # by tests (test_v07122.test_builders_validate_max_new_tokens); changing it + # would break them, so we keep it and document the seam here. + if isinstance(base, bool) or not isinstance(base, str): + raise ValueError("base must be a non-bool string") + if not base: + raise ValueError("base must be non-empty") + if "\x00" in base: + raise ValueError("base must not contain null bytes") + if len(base) > _MAX_BASE_LEN: + raise ValueError(f"base exceeds {_MAX_BASE_LEN} chars") + return base + + +def _check_max_new_tokens(value: object) -> int: + # NOTE: raises ValueError (not TypeError) for bool/non-int wrong-type input + # — deliberately inconsistent with validate_measure_candidate's TypeError; + # the ValueError behaviour is pinned by tests (see _check_base_id note). + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError("max_new_tokens must be a non-bool int") + if value < 1: + raise ValueError(f"max_new_tokens must be >= 1, got {value}") + return value + + +def _free_accelerator_memory() -> None: + """Best-effort GC + CUDA cache flush between candidate model loads.""" + import gc + + gc.collect() + try: + import torch + + if torch.cuda.is_available(): + torch.cuda.empty_cache() + except Exception: # noqa: BLE001 — torch missing / CUDA error: best-effort + pass + + +def _import_transformers(): + """Lazy transformers import — module-level seam for tests.""" + import transformers + + return transformers + + +def _load_measure_model( + base: str, + *, + quantization: str, + device: Optional[str] = None, + trust_remote_code: bool = False, +): + """Load ``base`` with the candidate quant via the Quant Menu loader. + + Returns a ``(model, tokenizer, device)`` triple shaped for + :func:`soup_cli.utils.live_eval.make_generator`'s ``loaded=`` kwarg. + ``quantization='none'`` loads the plain baseline. Heavy imports are lazy. + """ + from soup_cli.config.schema import TrainingConfig + from soup_cli.utils.live_eval import resolve_device + from soup_cli.utils.quant_menu import build_quantization_config_for_loader + + transformers = _import_transformers() + dev = resolve_device(device) + tcfg = TrainingConfig(quantization=quantization) + quant_config = build_quantization_config_for_loader( + tcfg=tcfg, base=base, console=None + ) + tokenizer = transformers.AutoTokenizer.from_pretrained( + base, trust_remote_code=trust_remote_code + ) + if getattr(tokenizer, "pad_token", None) is None: + tokenizer.pad_token = tokenizer.eos_token + if quant_config is not None: + # Quantized loads place weights themselves (bnb/gptq/awq need + # device_map on CUDA; a post-hoc .to() would break bnb layouts). + model = transformers.AutoModelForCausalLM.from_pretrained( + base, + quantization_config=quant_config, + device_map="auto" if dev == "cuda" else None, + trust_remote_code=trust_remote_code, + ) + else: + model = transformers.AutoModelForCausalLM.from_pretrained( + base, trust_remote_code=trust_remote_code + ).to(dev) + model.eval() + return model, tokenizer, dev + + +def _lazy_generator( + base: str, + *, + quantization: str, + device: Optional[str], + max_new_tokens: int, + trust_remote_code: bool, +) -> Callable[[str], str]: + """Closure that loads the model on FIRST call (cache hits never load).""" + holder: dict[str, Callable[[str], str]] = {} + + def _gen(prompt: str) -> str: + if "fn" not in holder: + _free_accelerator_memory() + loaded = _load_measure_model( + base, + quantization=quantization, + device=device, + trust_remote_code=trust_remote_code, + ) + from soup_cli.utils import live_eval + + holder["fn"] = live_eval.make_generator( + base, loaded=loaded, max_new_tokens=max_new_tokens + ) + return holder["fn"](prompt) + + return _gen + + +def build_before_generator( + base: str, + *, + device: Optional[str] = None, + max_new_tokens: int = 64, + trust_remote_code: bool = False, +) -> Callable[[str], str]: + """First-party baseline generator for ``deploy autopilot --measure``. + + Lazy-loads the UNQUANTIZED ``base`` (greedy decode via + :func:`soup_cli.utils.live_eval.make_generator`) on the first prompt — + a cache-hit measure run never touches the model. (v0.71.22 #143) + """ + _check_base_id(base) + _check_max_new_tokens(max_new_tokens) + return _lazy_generator( + base, + quantization="none", + device=device, + max_new_tokens=max_new_tokens, + trust_remote_code=trust_remote_code, + ) + + +def build_after_generator_factory( + base: str, + *, + device: Optional[str] = None, + max_new_tokens: int = 64, + trust_remote_code: bool = False, +) -> Callable[[str], Callable[[str], str]]: + """First-party per-candidate quantized generator factory (#143). + + ``factory(candidate)`` validates the candidate eagerly (fail-fast at the + measure loop boundary) and returns a generator that lazy-loads ``base`` + with that candidate's quant config from the Quant Menu loader on first + use. CUDA memory from the previous candidate is freed before each load. + """ + _check_base_id(base) + _check_max_new_tokens(max_new_tokens) + + def factory(candidate: str) -> Callable[[str], str]: + canonical = validate_measure_candidate(candidate) + return _lazy_generator( + base, + quantization=canonical, + device=device, + max_new_tokens=max_new_tokens, + trust_remote_code=trust_remote_code, + ) + + # Marker so run_measure can pre-validate the whole candidate list up front + # for the first-party path (M2) without touching injected test seams. + factory._soup_first_party_measure_factory = True # type: ignore[attr-defined] + return factory + + def run_measure( *, profile_name: str, @@ -303,6 +548,15 @@ def run_measure( f"too many candidates ({len(candidates)}; cap {_MAX_CANDIDATES})" ) + # M2 — pre-validate the ENTIRE candidate list up front so a typo'd Nth + # candidate fails BEFORE any (expensive) model load, rather than burning a + # live load + eval on every preceding candidate. Only enforced when the + # first-party factory is in play; an injected test seam factory is left + # untouched (it may legitimately use non-Quant-Menu candidate names). + if getattr(after_gen_factory, "_soup_first_party_measure_factory", False): + for candidate in candidates: + validate_measure_candidate(candidate) + tasks_sha = sha_of_file(tasks_file) key = compute_cache_key( base_sha=base_sha, profile_name=profile_name, tasks_sha=tasks_sha, @@ -320,13 +574,16 @@ def run_measure( except (TypeError, ValueError): pass - # Cache miss — run the eval loop + # Cache miss — compute the deterministic baseline ONCE (M3), then run each + # candidate against it. This keeps the baseline model out of VRAM while a + # quantized candidate is loaded and avoids N× redundant baseline scoring. + before_score = _score_tasks(tasks_file, before_gen) results: list[MeasureResult] = [] for candidate in candidates: after_gen = after_gen_factory(candidate) result = measure_candidate( candidate=candidate, tasks_file=tasks_file, - before_gen=before_gen, after_gen=after_gen, + after_gen=after_gen, before_score=before_score, ) results.append(result) diff --git a/src/soup_cli/utils/minillm.py b/src/soup_cli/utils/minillm.py index de87c52..b2f62b3 100644 --- a/src/soup_cli/utils/minillm.py +++ b/src/soup_cli/utils/minillm.py @@ -235,6 +235,39 @@ def _multinomial_sample(probs): return torch.multinomial(probs, num_samples=1) +def _supports_kv_cache(model: object) -> bool: + """Capability probe: does ``model.forward`` explicitly take a KV cache? + + Requires BOTH ``past_key_values`` and ``use_cache`` to be declared as + explicit parameters (every HF ``*ForCausalLM`` does). A bare ``**kwargs`` + does NOT count — a model that merely swallows the kwarg would silently + ignore the cache and the rollout would feed it a single token per step + against an empty context. PEFT wrappers + (``PeftModel.forward(*args, **kwargs)``) — the common live distill case — + are probed through ``get_base_model()`` so a LoRA student still gets the + cached path (mirrors :func:`mole_routing._model_supports_cache`). The real + forwards already go through the wrapper; only the probe is unwrapped. Never + raises. + """ + import inspect + + target = model + get_base = getattr(model, "get_base_model", None) + if callable(get_base): + try: + target = get_base() + except Exception: # noqa: BLE001 — probe must never raise + target = model + forward = getattr(target, "forward", None) + if forward is None: + return False + try: + params = inspect.signature(forward).parameters + except (TypeError, ValueError): + return False + return "past_key_values" in params and "use_cache" in params + + def minillm_on_policy_rollout( student_model, teacher_model, @@ -245,6 +278,7 @@ def minillm_on_policy_rollout( max_new_tokens: int, temperature: float = 1.0, sample_fn=None, + use_cache: bool = True, ): """True on-policy MiniLLM teacher-mixed rollout (v0.71.18 #257). @@ -263,6 +297,28 @@ def minillm_on_policy_rollout( reverse-KL is averaged over the rollout length so longer rollouts don't dominate the gradient. + KV-cache (v0.71.22 #263): when ``use_cache`` is True (default) and BOTH + models declare ``past_key_values``/``use_cache`` on their forward + (:func:`_supports_kv_cache`), each step forwards only the new token + against the cached context — O(L) compute instead of the O(L²) full + re-forward. Forward logits are identical either way (modulo float + noise), so the loss trajectory matches the no-cache path. Gradients are + mathematically equivalent too — both paths are full + backprop-through-rollout — so only the retained-graph *shape* differs: + + - **Teacher** cache lives under ``torch.no_grad()`` — trivially safe. + - **Student** cache is grad-carrying: cached KVs from earlier steps stay + in the autograd graph, so the cached path keeps ONE shared graph + spanning the whole rollout, whereas the no-cache path builds L separate + per-step graphs (each its own fresh full-prefix re-forward). Both + backprop through every step; the cached shared-graph form has retained + activations of ~one full-sequence forward — at or below the no-cache + peak. + + Models whose forward does not declare the cache params (or that ignore + ``use_cache`` and return no ``past_key_values``) transparently keep the + legacy full re-forward path. + Returns ``(loss, num_steps)`` where ``loss`` is a scalar differentiable w.r.t. the student parameters and scaled by ``temperature**2`` (the Hinton convention shared with the offline term). @@ -279,6 +335,8 @@ def minillm_on_policy_rollout( raise ValueError(f"max_new_tokens must be >= 1, got {max_new_tokens}") if isinstance(temperature, bool) or not isinstance(temperature, (int, float)): raise TypeError("temperature must be a non-bool number") + if not isinstance(use_cache, bool): + raise TypeError("use_cache must be bool") t = float(temperature) if not math.isfinite(t) or t <= 0.0: raise ValueError("temperature must be finite and positive") @@ -290,13 +348,51 @@ def minillm_on_policy_rollout( cur_mask = torch.ones_like(input_ids) else: cur_mask = attention_mask + cache_ok = ( + use_cache + and _supports_kv_cache(student_model) + and _supports_kv_cache(teacher_model) + ) + s_past = None + t_past = None + step_ids = cur_ids # full prompt on the first step; the delta afterwards kl_terms = [] for _ in range(int(max_new_tokens)): - s_out = student_model(input_ids=cur_ids, attention_mask=cur_mask) + if cache_ok: + s_out = student_model( + input_ids=step_ids, + attention_mask=cur_mask, + past_key_values=s_past, + use_cache=True, + ) + else: + s_out = student_model(input_ids=cur_ids, attention_mask=cur_mask) s_logits = s_out.logits[:, -1, :] / t # [B, Vs] (grad ON) with torch.no_grad(): - te_out = teacher_model(input_ids=cur_ids, attention_mask=cur_mask) + if cache_ok: + te_out = teacher_model( + input_ids=step_ids, + attention_mask=cur_mask, + past_key_values=t_past, + use_cache=True, + ) + else: + te_out = teacher_model( + input_ids=cur_ids, attention_mask=cur_mask + ) te_logits = te_out.logits[:, -1, :] / t # [B, Vt] (grad OFF) + if cache_ok: + new_s_past = getattr(s_out, "past_key_values", None) + new_t_past = getattr(te_out, "past_key_values", None) + if new_s_past is None or new_t_past is None: + # Model ignored use_cache — degrade to full re-forwards for + # the remaining steps (cur_ids is maintained either way). + cache_ok = False + s_past = None + t_past = None + else: + s_past = new_s_past + t_past = new_t_past # MiniLLM assumes a shared vocab; clamp to the common min so a # mismatched teacher never index-errors (cross-tokenizer = ULD path). common = min(s_logits.shape[-1], te_logits.shape[-1]) @@ -315,6 +411,7 @@ def minillm_on_policy_rollout( next_tok = next_tok.to(cur_ids.dtype) cur_ids = torch.cat([cur_ids, next_tok], dim=1) cur_mask = torch.cat([cur_mask, torch.ones_like(next_tok)], dim=1) + step_ids = next_tok if cache_ok else cur_ids stacked = torch.stack(kl_terms, dim=1) # [B, L] if config.length_normalize: @@ -348,11 +445,14 @@ class MiniLLMCallback(_TrainerCallbackBase): # type: ignore[misc, valid-type] lazily from ``pretrain_anchor_path``, scaled by ``pretrain_anchor_weight`` (prevents the student drifting away from coherent language). + - :meth:`on_policy_term` — true on-policy teacher-mixed autoregressive + rollout (sample → teacher-score) via + :func:`minillm_on_policy_rollout` (shipped v0.71.18 #257; KV-cache + accelerated v0.71.22 #263). - Honest scope: the teacher-mix is the offline distribution-blend - analog of MiniLLM's on-policy teacher-mixed *sampling*; a full - autoregressive rollout loop (sample → teacher-score) is a larger - follow-up documented as a known limitation. + The offline :meth:`distill_term` teacher-mix and the live on-policy + rollout are both available; pick per the ``training.minillm_on_policy`` + flag. """ def __init__( @@ -394,6 +494,7 @@ class MiniLLMCallback(_TrainerCallbackBase): # type: ignore[misc, valid-type] attention_mask=None, *, sample_fn=None, + use_cache: bool = True, ): """True on-policy teacher-mixed rollout reverse-KL (v0.71.18 #257). @@ -401,6 +502,11 @@ class MiniLLMCallback(_TrainerCallbackBase): # type: ignore[misc, valid-type] rollout from the teacher/student mixture and returns the length-normalised reverse-KL along that path. Returns the scalar loss only (the step count is internal to the rollout). + + ``use_cache`` (v0.71.22 #263) threads through to + :func:`minillm_on_policy_rollout` — cache-capable models forward only + the new token per step (see the rollout docstring for the gradient + trade-off). """ loss, _ = minillm_on_policy_rollout( student_model, @@ -411,6 +517,7 @@ class MiniLLMCallback(_TrainerCallbackBase): # type: ignore[misc, valid-type] max_new_tokens=self.config.rollout_length, temperature=self.temperature, sample_fn=sample_fn, + use_cache=use_cache, ) return loss diff --git a/src/soup_cli/utils/mole_routing.py b/src/soup_cli/utils/mole_routing.py index d3156e2..924cfd7 100644 --- a/src/soup_cli/utils/mole_routing.py +++ b/src/soup_cli/utils/mole_routing.py @@ -33,7 +33,7 @@ import math import os import stat import tempfile -from dataclasses import dataclass +from dataclasses import dataclass, field from functools import lru_cache from typing import Any, Optional, Tuple @@ -61,42 +61,42 @@ _MAX_ADAPTER_PATH_LEN = 4096 # --------------------------------------------------------------------------- -def _check_int(value: object, field: str, lo: int, hi: int) -> int: +def _check_int(value: object, field_name: str, lo: int, hi: int) -> int: if isinstance(value, bool) or not isinstance(value, int): - raise TypeError(f"{field} must be int") + raise TypeError(f"{field_name} must be int") if value < lo: - raise ValueError(f"{field} {value} below floor {lo}") + raise ValueError(f"{field_name} {value} below floor {lo}") if value > hi: - raise ValueError(f"{field} {value} above cap {hi}") + raise ValueError(f"{field_name} {value} above cap {hi}") return value def _check_finite_positive( - value: object, field: str, lo: float, hi: float + value: object, field_name: str, lo: float, hi: float ) -> float: if isinstance(value, bool): - raise TypeError(f"{field} must not be bool") + raise TypeError(f"{field_name} must not be bool") if not isinstance(value, (int, float)): - raise TypeError(f"{field} must be numeric") + raise TypeError(f"{field_name} must be numeric") val = float(value) if not math.isfinite(val): - raise ValueError(f"{field} must be finite") + raise ValueError(f"{field_name} must be finite") if val < lo: - raise ValueError(f"{field} {val} below floor {lo}") + raise ValueError(f"{field_name} {val} below floor {lo}") if val > hi: - raise ValueError(f"{field} {val} above cap {hi}") + raise ValueError(f"{field_name} {val} above cap {hi}") return val -def _check_str_field(value: object, field: str) -> str: +def _check_str_field(value: object, field_name: str) -> str: if isinstance(value, bool): - raise TypeError(f"{field} must not be bool") + raise TypeError(f"{field_name} must not be bool") if not isinstance(value, str): - raise TypeError(f"{field} must be str") + raise TypeError(f"{field_name} must be str") if not value: - raise ValueError(f"{field} must be non-empty") + raise ValueError(f"{field_name} must be non-empty") if "\x00" in value: - raise ValueError(f"{field} must not contain null bytes") + raise ValueError(f"{field_name} must not contain null bytes") return value @@ -441,6 +441,56 @@ def load_mole_manifest(directory: str) -> MoleServeManifest: return _manifest_from_dict(data) +def _model_supports_cache(model: object) -> bool: + """Capability probe: can this model take a per-call KV cache? (#262) + + Requires BOTH ``past_key_values`` and ``use_cache`` as EXPLICIT params on + the forward signature — a bare ``**kwargs`` does not count (a model that + merely swallows the kwarg would silently ignore the cache). PEFT wrappers + (``PeftModel.forward(*args, **kwargs)``) are probed through + ``get_base_model()`` so real HF causal LMs get the cached path. Never + raises. + """ + import inspect + + target = model + get_base = getattr(model, "get_base_model", None) + if callable(get_base): + try: + target = get_base() + except Exception: # noqa: BLE001 — probe must never raise + target = model + forward = getattr(target, "forward", None) + if forward is None: + return False + try: + params = inspect.signature(forward).parameters + except (TypeError, ValueError): + return False + return "past_key_values" in params and "use_cache" in params + + +@dataclass +class _MoleKvState: + """Per-``generate()`` KV-cache state — one cache per adapter + the base. + + ``caches[i]`` / ``lens[i]`` track adapter ``i``'s cache object and how + many tokens of the current sequence it has seen. An adapter skipped by + top-k masking simply falls behind; when it becomes active again it is fed + the missed tokens in one catch-up forward (caches stay in lockstep with + the sequence without paying for inactive adapters every step). + ``enabled`` flips off when the model ignores ``use_cache`` (no + ``past_key_values`` on the output) — the call degrades to the legacy + full-re-forward path. + """ + + enabled: bool = True + base_cache: Any = None + base_len: int = 0 + caches: dict = field(default_factory=dict) + lens: dict = field(default_factory=dict) + + class LoadedMole: """Runtime serve-time MoLE: base + N frozen task LoRAs + the trained gate. @@ -450,10 +500,14 @@ class LoadedMole: blended by the per-token gate weights. At decode we only need the LAST token's blend per step. - Generation recomputes the full sequence each step (no shared KV cache — - each task adapter would need its own cache, so recompute keeps the blend - correct). This is fine for short demo generations on a tiny model; large - serving deployments should expect linear-per-step cost. + KV-cache (v0.71.22 #262): cache-capable models (probe: + :func:`_model_supports_cache`) keep one KV cache per adapter plus one for + the base/router forward, all scoped to a single ``generate()`` call. Each + step then feeds only the new token (or, for an adapter that was skipped + by top-k masking, the tokens it missed — a catch-up delta) instead of + re-forwarding the whole sequence through every adapter. Models without + explicit cache support keep the legacy full-re-forward path, which is + fine for short demo generations on a tiny model. """ def __init__( @@ -507,6 +561,68 @@ class LoadedMole: model.set_adapter(self.adapter_names[0]) return blended + def _blended_last_logits_cached(self, input_ids, attention_mask, state): + """KV-cached per-token blend (#262) — returns ``[B, V]``. + + Feeds each forward only the tokens its cache has not seen yet + (``input_ids[:, cached_len:]``) with the FULL attention mask, so the + base/router and every active adapter pay one-token cost per step + instead of re-processing the whole sequence. Falls back to the legacy + path (and disables the state) when the model ignores ``use_cache``. + """ + import torch + + model = self.model + seq_len = int(input_ids.shape[1]) + base_delta = input_ids[:, state.base_len:] + with torch.no_grad(), model.disable_adapter(): + base_out = model( + input_ids=base_delta, + attention_mask=attention_mask, + past_key_values=state.base_cache, + use_cache=True, + output_hidden_states=True, + ) + new_base_cache = getattr(base_out, "past_key_values", None) + if new_base_cache is None: + state.enabled = False + return self._blended_last_logits(input_ids, attention_mask) + state.base_cache = new_base_cache + state.base_len = seq_len + # hidden_states cover only the fed tokens; -1 is the newest token. + router_hidden = base_out.hidden_states[-1][:, -1, :] # [B, H] + weights = self.gate(router_hidden.to(self.gate.gate.weight.dtype)) # [B, N] + col_max = weights.abs().amax(dim=tuple(range(weights.dim() - 1))) # [N] + active = [i for i in range(len(self.adapter_names)) if float(col_max[i]) > 0.0] + if not active: # degenerate (softmax should always keep >=1) — run all + active = list(range(len(self.adapter_names))) + blended = None + for i in active: + model.set_adapter(self.adapter_names[i]) + # Catch-up delta: an adapter skipped by top-k on earlier steps is + # behind the sequence — feed it everything it missed at once. + delta_i = input_ids[:, state.lens.get(i, 0):] + with torch.no_grad(): + out_i = model( + input_ids=delta_i, + attention_mask=attention_mask, + past_key_values=state.caches.get(i), + use_cache=True, + ) + cache_i = getattr(out_i, "past_key_values", None) + if cache_i is None: + state.enabled = False + model.set_adapter(self.adapter_names[0]) + return self._blended_last_logits(input_ids, attention_mask) + state.caches[i] = cache_i + state.lens[i] = seq_len + logits_i = out_i.logits[:, -1, :].detach() # [B, V] + w_i = weights[..., i : i + 1].to(logits_i.dtype) + term = w_i * logits_i + blended = term if blended is None else blended + term + model.set_adapter(self.adapter_names[0]) + return blended + def generate( self, input_ids, @@ -522,6 +638,9 @@ class LoadedMole: Single-sequence (``B == 1``): the EOS stop inspects row 0, so a batched caller would stop the whole batch on the first row's EOS. The serve handler always sends one sequence; ``B > 1`` is rejected. + + Cache state (#262) is local to this call — two concurrent or + sequential ``generate()`` calls never share KV caches. """ import torch @@ -538,8 +657,12 @@ class LoadedMole: do_sample = isinstance(temperature, (int, float)) and not isinstance( temperature, bool ) and float(temperature) > 0.0 + state = _MoleKvState() if _model_supports_cache(self.model) else None for _ in range(int(max_new_tokens)): - logits = self._blended_last_logits(seq, attn) # [B, V] + if state is not None and state.enabled: + logits = self._blended_last_logits_cached(seq, attn, state) + else: + logits = self._blended_last_logits(seq, attn) # [B, V] if do_sample: next_id = _sample_next(logits, float(temperature), float(top_p)) else: diff --git a/src/soup_cli/utils/tts_codec.py b/src/soup_cli/utils/tts_codec.py new file mode 100644 index 0000000..a3f9c21 --- /dev/null +++ b/src/soup_cli/utils/tts_codec.py @@ -0,0 +1,373 @@ +"""v0.71.22 #265-partial — live-codec TTS audio encoding (Orpheus via SNAC). + +Lifts the v0.71.20 ``data.format='audio'`` hardware gate into a real +encode-at-train-time path. Rows shaped ``{"audio": , "messages": [...]}`` +(the v0.17.0 audio format — paths are resolved + containment-checked by the +data loader) are encoded into discrete codec-token strings that become the +assistant turn, after which training proceeds through the validated +pre-encoded SFT cross-entropy path. + +The generic pipeline (``encode_tts_dataset`` → ``encode_tts_row`` → +per-family encoder) is family-agnostic; v0.71.22 ships ONE validated encoder: + +* **Orpheus / SNAC** (``pip install snac``) — 24 kHz SNAC produces 3 + codebooks at a 1/2/4 frame ratio; Orpheus interleaves them 7 tokens per + frame with per-slot codebook offsets and renders each code as + ```` where ``N = code + 10 + slot*4096`` (the official + Orpheus id layout: ``id = 128256 + N = 128266 + code + slot*4096`` on the + Orpheus tokenizer, whose ```` maps to ``128256 + i``). + +The other four families (sesame_csm / llasa / spark / oute) keep their +per-family codec dep gate; their encoders are tracked in #265 and raise a +friendly ``RuntimeError`` pointing at the offline pre-encode workflow. + +Security / robustness: +- Heavy imports (numpy / soundfile / torch / snac) are lazy — module import + stays light for the CLI hot path. +- Audio paths: null-byte rejection, symlink rejection (``os.lstat``, + defence-in-depth — the loader already containment-checks under + ``data.audio_dir``), duration cap (``_MAX_AUDIO_SECONDS``). +- Pure interleave kernel validates the 1/2/4 codebook ratio, per-code bounds + ``[0, 4096)``, and rejects bool codes (project bool-as-int policy). +- Rows are deep-copied — the caller's dataset is never mutated (v0.33.0 #47 + immutability policy). +""" + +from __future__ import annotations + +import os +import stat +from typing import Any, Callable, Optional + +from soup_cli.utils.tts import tts_codec_package, validate_tts_family + +# SNAC checkpoint for the Orpheus family (24 kHz, 3 codebooks). +SNAC_MODEL_ID = "hubertsiuzdak/snac_24khz" +SNAC_SAMPLE_RATE = 24_000 + +# Orpheus token layout: with N = code + OFFSET + slot*4096. +# Slots 0..9 are Orpheus control tokens; audio codes start at 10. +ORPHEUS_CODE_OFFSET = 10 +ORPHEUS_CODEBOOK_SIZE = 4096 +_ORPHEUS_FRAME_SLOTS = 7 + +# Families whose live-codec encoder is implemented AND validated. +LIVE_CODEC_FAMILIES: frozenset[str] = frozenset({"orpheus"}) + +_MAX_AUDIO_SECONDS = 600.0 # 10 min cap — defends against runaway encodes +# Byte cap (~600 s of 24 kHz stereo float32 + headroom) — rejected before any +# read so a runaway file never materialises into RAM. +_MAX_AUDIO_BYTES = 512 * 1024 * 1024 + +# Process-level SNAC model cache (one load per device). NOTE: single-threaded +# assumption — the live-codec trainer encode path runs on one thread, so this +# unbounded, unguarded cache is intentional (no eviction / no lock needed). +_SNAC_CACHE: dict[str, Any] = {} + + +# --------------------------------------------------------------------------- +# Pure interleave kernel (no torch — unit-testable on plain lists) +# --------------------------------------------------------------------------- + + +def _check_codes(codes, name: str, expected_len: int) -> list[int]: + out = [] + for code in codes: + if isinstance(code, bool) or not isinstance(code, int): + raise TypeError(f"{name} codes must be non-bool ints") + if not (0 <= code < ORPHEUS_CODEBOOK_SIZE): + raise ValueError( + f"{name} code {code} out of range [0, {ORPHEUS_CODEBOOK_SIZE})" + ) + out.append(code) + if len(out) != expected_len: + raise ValueError( + f"{name} codebook length {len(out)} != expected {expected_len} " + "(SNAC codebooks must have the 1/2/4 frame ratio)" + ) + return out + + +def interleave_orpheus_codes(coarse, medium, fine) -> list[int]: + """Interleave the 3 SNAC codebooks into Orpheus token indices. + + ``coarse``/``medium``/``fine`` are the per-frame code sequences at the + SNAC 1/2/4 ratio (lengths T / 2T / 4T). Returns 7 token indices per + frame in the official Orpheus slot order, each offset by + ``ORPHEUS_CODE_OFFSET + slot * ORPHEUS_CODEBOOK_SIZE`` so the rendered + ```` strings map onto the Orpheus tokenizer's audio ids. + """ + coarse = list(coarse) + if not coarse: + raise ValueError("coarse codebook must be non-empty") + frames = len(coarse) + coarse = _check_codes(coarse, "coarse", frames) + medium = _check_codes(list(medium), "medium", 2 * frames) + fine = _check_codes(list(fine), "fine", 4 * frames) + + size = ORPHEUS_CODEBOOK_SIZE + base = ORPHEUS_CODE_OFFSET + out: list[int] = [] + for i in range(frames): + out.append(coarse[i] + base) + out.append(medium[2 * i] + base + size) + out.append(fine[4 * i] + base + 2 * size) + out.append(fine[4 * i + 1] + base + 3 * size) + out.append(medium[2 * i + 1] + base + 4 * size) + out.append(fine[4 * i + 2] + base + 5 * size) + out.append(fine[4 * i + 3] + base + 6 * size) + return out + + +def orpheus_tokens_to_string(indices) -> str: + """Render Orpheus token indices as a ```` string. + + Indices must be non-bool ints (mirrors :func:`_check_codes` — project + bool-as-int policy; no float/bool coercion). + """ + indices = list(indices) + if not indices: + raise ValueError("token index list must be non-empty") + for n in indices: + if isinstance(n, bool) or not isinstance(n, int): + raise TypeError("token indices must be non-bool ints") + return "".join(f"" for n in indices) + + +# --------------------------------------------------------------------------- +# Audio loading (lazy soundfile / numpy) +# --------------------------------------------------------------------------- + + +def load_audio_mono(path: str, *, target_sr: int = SNAC_SAMPLE_RATE): + """Load an audio file as mono float32 at ``target_sr``. + + Lazy-imports soundfile (friendly ImportError naming the install). + Off-rate audio is resampled with basic linear interpolation — adequate + for codec encoding of speech; operators wanting studio-grade resampling + should resample offline. Returns a 1-D ``np.float32`` array. + """ + if not isinstance(path, str) or not path: + raise ValueError("audio path must be a non-empty string") + if "\x00" in path: + raise ValueError("audio path must not contain null bytes") + if not os.path.exists(path): + raise FileNotFoundError( + f"audio file not found: {os.path.basename(path)!r}" + ) + # Symlink rejection — defence-in-depth; the loader already containment- + # checks rows under data.audio_dir (v0.17.0 policy). + if stat.S_ISLNK(os.lstat(path).st_mode): + raise ValueError("audio path must not be a symlink") + # Byte-size cap before any read — defends against a runaway file. + if os.path.getsize(path) > _MAX_AUDIO_BYTES: + raise ValueError( + f"audio file exceeds the {_MAX_AUDIO_BYTES} byte live-codec cap; " + "split the clip or pre-encode offline" + ) + try: + import soundfile + except ImportError as exc: # pragma: no cover - dep present in CI + raise ImportError( + "Live-codec TTS needs the 'soundfile' package to read audio. " + "Install it with `pip install soundfile` (or `pip install " + "'soup-cli[audio]'`)." + ) from exc + import numpy as np + + # Probe metadata FIRST — reject an over-long clip before materialising the + # whole file into RAM (the post-read check below stays as defence-in-depth). + info = soundfile.info(path) + if info.samplerate and info.frames / float(info.samplerate) > _MAX_AUDIO_SECONDS: + raise ValueError( + f"audio is {info.frames / float(info.samplerate):.1f} seconds — " + f"exceeds the {_MAX_AUDIO_SECONDS:.0f} seconds live-codec cap; " + "split the clip or pre-encode offline" + ) + # Open via O_NOFOLLOW and hand the fd to soundfile — closes the TOCTOU + # window between the lstat symlink check above and the read (project + # v0.65/v0.67 O_NOFOLLOW reader policy). soundfile.read accepts an open + # file object. + fd = os.open(path, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)) + try: + with os.fdopen(fd, "rb") as handle: + data, sr = soundfile.read(handle, dtype="float32", always_2d=True) + except OSError as exc: + # O_NOFOLLOW raises ELOOP if the path became a symlink after the lstat. + raise ValueError(f"audio path could not be read: {exc}") from exc + duration = data.shape[0] / float(sr) + if duration > _MAX_AUDIO_SECONDS: + raise ValueError( + f"audio is {duration:.1f} seconds — exceeds the " + f"{_MAX_AUDIO_SECONDS:.0f} seconds live-codec cap; split the " + "clip or pre-encode offline" + ) + mono = data.mean(axis=1) + if int(sr) != int(target_sr): + n_out = max(1, int(round(len(mono) * target_sr / float(sr)))) + x_old = np.linspace(0.0, 1.0, num=len(mono), endpoint=False) + x_new = np.linspace(0.0, 1.0, num=n_out, endpoint=False) + mono = np.interp(x_new, x_old, mono) + return mono.astype(np.float32) + + +# --------------------------------------------------------------------------- +# Orpheus / SNAC encoder +# --------------------------------------------------------------------------- + + +def _get_snac_model(device: Optional[str] = None): + """Load (once per device) the SNAC codec model. Lazy snac import.""" + dev = device or "cpu" + cached = _SNAC_CACHE.get(dev) + if cached is not None: + return cached + try: + from snac import SNAC + except ImportError as exc: + raise ImportError( + "TTS family 'orpheus' live-codec training requires the 'snac' " + "package (audio codec). Install it with `pip install snac`, or " + "pre-encode your audio to codec tokens offline and train with " + "data.format=chat." + ) from exc + model = SNAC.from_pretrained(SNAC_MODEL_ID).eval().to(dev) + _SNAC_CACHE[dev] = model + return model + + +def encode_audio_orpheus( + path: str, + *, + snac_model: Optional[Any] = None, + device: Optional[str] = None, +) -> str: + """Encode one audio file into the Orpheus ```` string. + + ``snac_model`` is injectable (tests / pre-loaded models); when ``None`` + the 24 kHz SNAC checkpoint is loaded once per process per device. + """ + audio = load_audio_mono(path, target_sr=SNAC_SAMPLE_RATE) + import torch + + model = snac_model if snac_model is not None else _get_snac_model(device) + wav = torch.from_numpy(audio).reshape(1, 1, -1) + if device: + wav = wav.to(device) + with torch.no_grad(): + codes = model.encode(wav) + coarse = [int(c) for c in codes[0][0].tolist()] + medium = [int(c) for c in codes[1][0].tolist()] + fine = [int(c) for c in codes[2][0].tolist()] + return orpheus_tokens_to_string( + interleave_orpheus_codes(coarse, medium, fine) + ) + + +def tts_encoder_for_family( + family: str, *, device: Optional[str] = None +) -> Callable[[str], str]: + """Return the live audio→codec-string encoder for ``family``. + + Only the families in :data:`LIVE_CODEC_FAMILIES` have a validated + encoder (v0.71.22 ships Orpheus/SNAC). The remaining families raise a + friendly ``RuntimeError`` naming the offline workflow — their encoders + are tracked in #265 (the per-family codec dep gate fires earlier, at + trainer setup). + """ + canonical = validate_tts_family(family) + if canonical == "orpheus": + + def _encode(path: str) -> str: + return encode_audio_orpheus(path, device=device) + + return _encode + pkg = tts_codec_package(canonical) + raise RuntimeError( + f"Live-codec encoding for TTS family '{canonical}' is not yet " + f"implemented (v0.71.22 ships the Orpheus/SNAC encoder; the " + f"'{canonical}' encoder via {pkg!r} is tracked in #265). Pre-encode " + "your audio to codec tokens offline and train with data.format=chat." + ) + + +# --------------------------------------------------------------------------- +# Dataset mapping +# --------------------------------------------------------------------------- + + +def encode_tts_row(row: dict, encoder: Callable[[str], str]) -> dict: + """Encode one ``{"audio", "messages"}`` row into a chat row. + + The codec-token string becomes the FINAL assistant turn: an existing + trailing assistant message has its content replaced (the audio is the + training target), otherwise an assistant turn is appended. Returns a NEW + deep-copied row without the ``audio`` key — the caller's row is never + mutated. + """ + import copy + + if not isinstance(row, dict): + raise TypeError(f"row must be a dict, got {type(row).__name__}") + audio = row.get("audio") + if not isinstance(audio, str) or not audio: + raise ValueError("live-codec TTS row must have a non-empty 'audio' path") + messages = row.get("messages") + if not isinstance(messages, (list, tuple)): + raise ValueError("live-codec TTS row must have a 'messages' list") + + codec_string = encoder(audio) + new_messages = [copy.deepcopy(m) for m in messages] + if new_messages and isinstance(new_messages[-1], dict) and ( + new_messages[-1].get("role") == "assistant" + ): + new_messages[-1]["content"] = codec_string + else: + new_messages.append({"role": "assistant", "content": codec_string}) + out = { + key: copy.deepcopy(value) + for key, value in row.items() + if key not in ("audio", "messages") + } + out["messages"] = new_messages + return out + + +def encode_tts_dataset( + dataset: dict, + family: str, + *, + encoder: Optional[Callable[[str], str]] = None, + device: Optional[str] = None, + console: Optional[Any] = None, +) -> dict: + """Encode every ``train`` / ``val`` row's audio into codec-token chat rows. + + Returns a NEW dataset dict (input never mutated). ``encoder`` is + injectable for tests; by default it is resolved per-family via + :func:`tts_encoder_for_family` (which raises for not-yet-implemented + families — #265). + """ + if not isinstance(dataset, dict): + raise TypeError( + f"dataset must be a dict, got {type(dataset).__name__}" + ) + canonical = validate_tts_family(family) + encode = encoder if encoder is not None else tts_encoder_for_family( + canonical, device=device + ) + + new_dataset = dict(dataset) + total = 0 + for split in ("train", "val"): + rows = new_dataset.get(split) + if not isinstance(rows, (list, tuple)): + continue + new_dataset[split] = [encode_tts_row(row, encode) for row in rows] + total += len(rows) + if console is not None: + console.print( + f"[green]TTS live-codec:[/] encoded {total} row(s) " + f"(family={canonical})" + ) + return new_dataset diff --git a/tests/test_v07120.py b/tests/test_v07120.py index 7de6dea..704c409 100644 --- a/tests/test_v07120.py +++ b/tests/test_v07120.py @@ -200,7 +200,20 @@ class TestTtsTrainerSetup: with pytest.raises(RuntimeError, match="sparktts"): w.setup({"train": []}) - def test_live_codec_mode_orpheus_names_snac(self): + def test_live_codec_mode_orpheus_names_snac(self, monkeypatch): + # v0.71.22 #265 lifted the unconditional not-yet-validated gate — + # orpheus is live when snac IS installed, so force the missing-dep + # branch to keep exercising the friendly pip hint. + import importlib.util as ilu + + real_find_spec = ilu.find_spec + monkeypatch.setattr( + ilu, + "find_spec", + lambda name, *a, **k: None + if name == "snac" + else real_find_spec(name, *a, **k), + ) w = self._wrapper(data_format="audio", family="orpheus") with pytest.raises(RuntimeError, match="snac"): w.setup({"train": []}) diff --git a/tests/test_v07122.py b/tests/test_v07122.py new file mode 100644 index 0000000..96ee56a --- /dev/null +++ b/tests/test_v07122.py @@ -0,0 +1,1343 @@ +"""v0.71.22 "Perf & measure polish" tests. + +Closes #263 (MiniLLM on-policy KV-cache), #262 (`serve --mole` KV-cache), +#143 (first-party generator factories for `deploy autopilot --measure`), +#265-partial (live-codec TTS — Orpheus via SNAC; the other 4 families stay +dep-gated with the encoder tracked in #265). + +All four are pure code, fully validatable on the RTX 3050 / CPU budget. The +KV-cache paths are exercised with cache-capable fake LMs (prefix-sum caches +so cached logits provably equal full-re-forward logits); models without an +explicit ``past_key_values`` forward param transparently keep the legacy +full-re-forward path (capability probe — old fakes / exotic models stay +correct). +""" + +from __future__ import annotations + +import math +import types +import wave +from pathlib import Path +from types import SimpleNamespace + +import pytest +from typer.testing import CliRunner + +_SRC = Path(__file__).resolve().parent.parent / "src" / "soup_cli" + +runner = CliRunner() + + +def _module_head(rel_path: str) -> str: + """Return a module's source text up to the first ``def``/``class``.""" + src = (_SRC / rel_path).read_text(encoding="utf-8").replace("\r\n", "\n") + # Split on the first top-level def or class so we only inspect the head. + for marker in ("\ndef ", "\nclass "): + src = src.split(marker, 1)[0] + return src + + +def _assert_no_top_level_import(rel_path: str, mod: str) -> None: + """Reject BOTH ``import {mod}...`` and ``from {mod} ...`` at module top. + + Stronger than the ``\\nimport {mod}\\n`` idiom — catches ``import numpy as + np``, ``import torch, snac``, and ``from numpy import ...`` (L11). + """ + head = _module_head(rel_path) + assert f"\nimport {mod}" not in head, ( + f"top-level `import {mod}` in {rel_path}" + ) + assert f"\nfrom {mod} " not in head, ( + f"top-level `from {mod} ` in {rel_path}" + ) + + +def _torch_or_skip(): + return pytest.importorskip("torch") + + +# --------------------------------------------------------------------------- +# #263 — MiniLLM on-policy rollout KV-cache +# --------------------------------------------------------------------------- + + +def _make_cache_fake_lm(vocab: int = 11, hidden: int = 6, seed: int = 0): + """Prefix-sum LM with an explicit past_key_values cache. + + The last-token logits depend on the FULL prefix (cumulative embedding + sum), so a correct cache implementation produces bit-identical logits to + a full re-forward — the property the equality tests rely on. The cache is + the running prefix-sum ``[B, 1, H]`` (grad-carrying, like real KV caches). + """ + torch = _torch_or_skip() + from torch import nn + + torch.manual_seed(seed) + + class _CacheFakeLM(nn.Module): + def __init__(self) -> None: + super().__init__() + self.emb = nn.Embedding(vocab, hidden) + self.head = nn.Linear(hidden, vocab) + self.config = types.SimpleNamespace(vocab_size=vocab) + self.calls: list[int] = [] # fed token counts per forward + + def forward( + self, + input_ids=None, + attention_mask=None, + past_key_values=None, + use_cache=False, + ): + self.calls.append(int(input_ids.shape[1])) + h = self.emb(input_ids) + csum = torch.cumsum(h, dim=1) + if past_key_values is not None: + csum = csum + past_key_values[0] + logits = self.head(csum) + out = SimpleNamespace(logits=logits) + if use_cache: + out.past_key_values = (csum[:, -1:, :],) + return out + + return _CacheFakeLM() + + +def _make_legacy_fake_lm(vocab: int = 11, hidden: int = 6, seed: int = 0): + """The pre-#263 fake — no explicit past_key_values param (``**kw`` only).""" + torch = _torch_or_skip() + from torch import nn + + torch.manual_seed(seed) + + class _LegacyFakeLM(nn.Module): + def __init__(self) -> None: + super().__init__() + self.emb = nn.Embedding(vocab, hidden) + self.head = nn.Linear(hidden, vocab) + self.calls: list[int] = [] + + def forward(self, input_ids=None, attention_mask=None, **kw): + self.calls.append(int(input_ids.shape[1])) + h = self.emb(input_ids) + logits = self.head(torch.cumsum(h, dim=1)) + return SimpleNamespace(logits=logits) + + return _LegacyFakeLM() + + +def _make_no_past_fake_lm(vocab: int = 11, hidden: int = 6, seed: int = 0): + """Declares past_key_values/use_cache but never returns a cache.""" + torch = _torch_or_skip() + from torch import nn + + torch.manual_seed(seed) + + class _NoPastFakeLM(nn.Module): + def __init__(self) -> None: + super().__init__() + self.emb = nn.Embedding(vocab, hidden) + self.head = nn.Linear(hidden, vocab) + self.calls: list[int] = [] + + def forward( + self, + input_ids=None, + attention_mask=None, + past_key_values=None, + use_cache=False, + ): + self.calls.append(int(input_ids.shape[1])) + h = self.emb(input_ids) + logits = self.head(torch.cumsum(h, dim=1)) + return SimpleNamespace(logits=logits) # no past_key_values attr + + return _NoPastFakeLM() + + +def _greedy_sampler(probs): + return probs.argmax(dim=-1, keepdim=True) + + +class TestSupportsKvCache: + def test_explicit_params_true(self): + from soup_cli.utils.minillm import _supports_kv_cache + + assert _supports_kv_cache(_make_cache_fake_lm()) + + def test_var_keyword_only_false(self): + """``**kw`` does NOT count — a swallowed kwarg is not cache support.""" + from soup_cli.utils.minillm import _supports_kv_cache + + assert not _supports_kv_cache(_make_legacy_fake_lm()) + + def test_no_params_false(self): + from soup_cli.utils.minillm import _supports_kv_cache + + def plain(input_ids=None, attention_mask=None): + return None + + assert not _supports_kv_cache(SimpleNamespace(forward=plain)) + + def test_non_model_false(self): + from soup_cli.utils.minillm import _supports_kv_cache + + assert not _supports_kv_cache(object()) + + def test_peft_wrapper_probes_base_model(self): + """A ``*args, **kwargs`` wrapper (PeftModel-style LoRA student — the + common live distill case) is probed through ``get_base_model()`` so the + base's explicit cache params are seen and KV-cache activates (#263).""" + from soup_cli.utils.minillm import _supports_kv_cache + + inner = _make_cache_fake_lm() + + class _Wrapper: + def forward(self, *args, **kwargs): + return inner(*args, **kwargs) + + def get_base_model(self): + return inner + + assert _supports_kv_cache(_Wrapper()) + + def test_peft_wrapper_get_base_model_raising_is_safe(self): + """If ``get_base_model()`` raises, the probe swallows it and falls back + to inspecting the wrapper itself (never raises).""" + from soup_cli.utils.minillm import _supports_kv_cache + + class _Bad: + def forward(self, *args, **kwargs): + return None + + def get_base_model(self): + raise RuntimeError("boom") + + # Wrapper forward is *args/**kwargs only — falls back to it → False. + assert not _supports_kv_cache(_Bad()) + + +class TestOnPolicyKvCache: + def _rollout(self, student, teacher, *, use_cache, steps=4): + torch = _torch_or_skip() + from soup_cli.utils.minillm import ( + MiniLLMConfig, + minillm_on_policy_rollout, + ) + + cfg = MiniLLMConfig(teacher_mix_ratio=0.5, on_policy=True) + ids = torch.tensor([[1, 2, 3]]) + loss, n = minillm_on_policy_rollout( + student, + teacher, + ids, + None, + config=cfg, + max_new_tokens=steps, + sample_fn=_greedy_sampler, + use_cache=use_cache, + ) + return loss, n + + def test_cached_loss_matches_uncached(self): + torch = _torch_or_skip() + s1, t1 = _make_cache_fake_lm(seed=0), _make_cache_fake_lm(seed=1) + s2, t2 = _make_cache_fake_lm(seed=0), _make_cache_fake_lm(seed=1) + loss_cached, n1 = self._rollout(s1, t1, use_cache=True) + loss_full, n2 = self._rollout(s2, t2, use_cache=False) + assert n1 == n2 == 4 + assert torch.allclose(loss_cached, loss_full, rtol=1e-5, atol=1e-6) + + def test_cached_path_feeds_single_tokens(self): + student, teacher = _make_cache_fake_lm(seed=0), _make_cache_fake_lm(seed=1) + self._rollout(student, teacher, use_cache=True, steps=3) + # prompt (3 tokens) once, then 1 token per remaining step. + assert student.calls == [3, 1, 1] + assert teacher.calls == [3, 1, 1] + + def test_uncached_path_refeeds_full_prefix(self): + student, teacher = _make_cache_fake_lm(seed=0), _make_cache_fake_lm(seed=1) + self._rollout(student, teacher, use_cache=False, steps=3) + assert student.calls == [3, 4, 5] + + def test_legacy_kwargs_model_routes_uncached(self): + """A model without explicit cache params keeps the legacy path even + with use_cache=True (capability probe).""" + student, teacher = _make_legacy_fake_lm(seed=0), _make_legacy_fake_lm(seed=1) + loss, _ = self._rollout(student, teacher, use_cache=True, steps=3) + assert student.calls == [3, 4, 5] + assert math.isfinite(float(loss)) + + def test_missing_past_in_output_falls_back(self): + """A model that declares but ignores use_cache degrades to full + re-forwards after the first step (no crash, finite loss).""" + student, teacher = _make_no_past_fake_lm(seed=0), _make_no_past_fake_lm(seed=1) + loss, n = self._rollout(student, teacher, use_cache=True, steps=3) + assert n == 3 + assert math.isfinite(float(loss)) + # step 0 fed the prompt; the fallback re-feeds full prefixes after. + assert student.calls == [3, 4, 5] + + def test_grad_flows_through_cached_rollout(self): + student, teacher = _make_cache_fake_lm(seed=0), _make_cache_fake_lm(seed=1) + loss, _ = self._rollout(student, teacher, use_cache=True) + loss.backward() + assert student.emb.weight.grad is not None + assert float(student.emb.weight.grad.abs().sum()) > 0.0 + # Teacher is no_grad throughout. + assert teacher.emb.weight.grad is None + + def test_use_cache_bool_rejected(self): + student, teacher = _make_cache_fake_lm(), _make_cache_fake_lm(seed=1) + with pytest.raises(TypeError, match="use_cache"): + self._rollout(student, teacher, use_cache="yes") + + def test_mixed_capability_routes_uncached(self): + """Cache requires BOTH models capable — a legacy teacher disables it.""" + student = _make_cache_fake_lm(seed=0) + teacher = _make_legacy_fake_lm(seed=1) + self._rollout(student, teacher, use_cache=True, steps=3) + assert student.calls == [3, 4, 5] + + def test_on_policy_term_threads_use_cache(self, monkeypatch): + _torch_or_skip() + from soup_cli.utils import minillm as m + + captured = {} + + def fake_rollout(*args, **kwargs): + captured.update(kwargs) + import torch + + return torch.tensor(0.0), 1 + + monkeypatch.setattr(m, "minillm_on_policy_rollout", fake_rollout) + cb = m.build_minillm_callback(m.MiniLLMConfig(on_policy=True)) + import torch + + cb.on_policy_term(object(), object(), torch.tensor([[1]]), use_cache=False) + assert captured["use_cache"] is False + cb.on_policy_term(object(), object(), torch.tensor([[1]])) + assert captured["use_cache"] is True + + +# --------------------------------------------------------------------------- +# #262 — serve --mole per-adapter KV-cache +# --------------------------------------------------------------------------- + + +def _make_cached_mole_model(hidden: int = 4, vocab: int = 6): + """Adapter-switching fake with an explicit per-call KV cache. + + PREFIX-DEPENDENT logits (mirrors the minillm ``_CacheFakeLM`` prefix-sum + design): the active row's logit encodes the TOTAL prefix length the cache + claims to have seen (``past + t``) plus the adapter id. A correct cache + therefore yields bit-identical logits to a full re-forward, while a broken + cache (wrong ``state.lens`` slice, corrupted ``past_key_values``, off-by-one + catch-up bookkeeping) would corrupt the encoded prefix length and the + equality assertions would FAIL — the property M4 needs. ``calls`` records + ``(mode, fed_tokens)`` so the delta-feeding / catch-up tests can assert + exact feed lengths. + """ + torch = _torch_or_skip() + from torch import nn + + class _CachedMoleModel(nn.Module): + def __init__(self): + super().__init__() + self.hidden = hidden + self.vocab = vocab + self._adapter = "task_0" + self._disabled = False + self.calls: list[tuple[str, int]] = [] + + @property + def device(self): + return torch.device("cpu") + + def disable_adapter(self): + import contextlib + + @contextlib.contextmanager + def _cm(): + self._disabled = True + try: + yield + finally: + self._disabled = False + + return _cm() + + def set_adapter(self, name): + self._adapter = name + + def forward( + self, + input_ids=None, + attention_mask=None, + output_hidden_states=False, + past_key_values=None, + use_cache=False, + ): + mode = "base" if self._disabled else self._adapter + self.calls.append((mode, int(input_ids.shape[1]))) + b, t = input_ids.shape + hs = torch.ones(b, t, self.hidden) + idx = 0 if self._disabled else int(self._adapter.split("_")[1]) + past = int(past_key_values[0]) if past_key_values is not None else 0 + # Total prefix length this stream has now seen — bakes the cache + # bookkeeping into the output so a wrong cache produces wrong logits. + seen = past + t + logits = torch.zeros(b, t, self.vocab) + logits[..., idx % self.vocab] = float(idx + 1) + float(seen) + out = SimpleNamespace(logits=logits) + if output_hidden_states: + out.hidden_states = (hs,) + if use_cache: + out.past_key_values = (seen,) + return out + + return _CachedMoleModel() + + +def _make_legacy_mole_model(hidden: int = 4, vocab: int = 6): + """The pre-#262 fake — no cache params on forward.""" + torch = _torch_or_skip() + from torch import nn + + class _LegacyMoleModel(nn.Module): + def __init__(self): + super().__init__() + self.hidden = hidden + self.vocab = vocab + self._adapter = "task_0" + self._disabled = False + self.calls: list[tuple[str, int]] = [] + + @property + def device(self): + return torch.device("cpu") + + def disable_adapter(self): + import contextlib + + @contextlib.contextmanager + def _cm(): + self._disabled = True + try: + yield + finally: + self._disabled = False + + return _cm() + + def set_adapter(self, name): + self._adapter = name + + def forward( + self, input_ids=None, attention_mask=None, output_hidden_states=False + ): + mode = "base" if self._disabled else self._adapter + self.calls.append((mode, int(input_ids.shape[1]))) + b, t = input_ids.shape + hs = torch.ones(b, t, self.hidden) + idx = 0 if self._disabled else int(self._adapter.split("_")[1]) + # Re-feeds the full prefix each step, so t == sequence length so far + # — encode it identically to the cached model's (past + t) so the + # legacy and cached blends match exactly (M4 equality property). + logits = torch.zeros(b, t, self.vocab) + logits[..., idx % self.vocab] = float(idx + 1) + float(t) + out = SimpleNamespace(logits=logits) + if output_hidden_states: + out.hidden_states = (hs,) + return out + + return _LegacyMoleModel() + + +def _make_uniform_gate(hidden: int = 4, n: int = 2): + torch = _torch_or_skip() + from torch import nn + + class _UniformGate(nn.Module): + def __init__(self): + super().__init__() + self.gate = nn.Linear(hidden, n, bias=False) + + def forward(self, h): + b = h.shape[0] + return torch.full((b, n), 1.0 / n) + + return _UniformGate() + + +def _make_scripted_gate(script, hidden: int = 4, n: int = 2): + """Gate returning a scripted weight row per call (cache-catch-up test).""" + torch = _torch_or_skip() + from torch import nn + + class _ScriptedGate(nn.Module): + def __init__(self): + super().__init__() + self.gate = nn.Linear(hidden, n, bias=False) + self._i = 0 + + def forward(self, h): + row = script[min(self._i, len(script) - 1)] + self._i += 1 + return torch.tensor([row]) + + return _ScriptedGate() + + +class TestModelSupportsCache: + def test_explicit_params_true(self): + from soup_cli.utils.mole_routing import _model_supports_cache + + assert _model_supports_cache(_make_cached_mole_model()) + + def test_legacy_fake_false(self): + from soup_cli.utils.mole_routing import _model_supports_cache + + assert not _model_supports_cache(_make_legacy_mole_model()) + + def test_peft_wrapper_probes_base_model(self): + """A ``*args, **kwargs`` wrapper (PeftModel-style) is probed through + ``get_base_model()`` so real PEFT models get the cached path.""" + from soup_cli.utils.mole_routing import _model_supports_cache + + inner = _make_cached_mole_model() + + class _Wrapper: + def forward(self, *args, **kwargs): + return inner(*args, **kwargs) + + def get_base_model(self): + return inner + + assert _model_supports_cache(_Wrapper()) + + def test_non_model_false(self): + from soup_cli.utils.mole_routing import _model_supports_cache + + assert not _model_supports_cache(object()) + + +class TestMoleKvCache: + def _generate(self, model, gate, *, steps=3): + torch = _torch_or_skip() + from soup_cli.utils.mole_routing import LoadedMole + + mole = LoadedMole(model, object(), gate, ["task_0", "task_1"]) + ids = torch.tensor([[1, 2]]) + attn = torch.ones_like(ids) + return mole.generate(ids, attn, max_new_tokens=steps) + + def test_cached_output_matches_legacy(self): + torch = _torch_or_skip() + out_cached = self._generate(_make_cached_mole_model(), _make_uniform_gate()) + out_legacy = self._generate(_make_legacy_mole_model(), _make_uniform_gate()) + assert torch.equal(out_cached, out_legacy) + + def test_cached_path_feeds_deltas(self): + model = _make_cached_mole_model() + self._generate(model, _make_uniform_gate(), steps=3) + base_feeds = [t for mode, t in model.calls if mode == "base"] + a0_feeds = [t for mode, t in model.calls if mode == "task_0"] + # prompt (2 tokens) once, then 1 token per subsequent step. + assert base_feeds == [2, 1, 1] + assert a0_feeds == [2, 1, 1] + + def test_legacy_model_refeeds_full_prefix(self): + model = _make_legacy_mole_model() + self._generate(model, _make_uniform_gate(), steps=3) + base_feeds = [t for mode, t in model.calls if mode == "base"] + assert base_feeds == [2, 3, 4] + + def test_topk_skip_catches_up_cache(self): + """An adapter skipped by top-k catches up its cache (fed the missed + tokens) when it becomes active again — caches stay in lockstep.""" + model = _make_cached_mole_model() + # step 1: adapter 0 only; step 2: adapter 1 only; step 3: adapter 0. + gate = _make_scripted_gate([[1.0, 0.0], [0.0, 1.0], [1.0, 0.0]]) + self._generate(model, gate, steps=3) + a0_feeds = [t for mode, t in model.calls if mode == "task_0"] + a1_feeds = [t for mode, t in model.calls if mode == "task_1"] + # adapter 0: prompt (2) at step 1, skipped at step 2, catch-up of the + # 2 missed tokens at step 3. + assert a0_feeds == [2, 2] + # adapter 1: skipped at step 1, catch-up (prompt + 1 new = 3) at step 2. + assert a1_feeds == [3] + + def test_topk_skip_catchup_output_matches_legacy(self): + """The cached top-k skip/catch-up path produces bit-identical output to + a full-re-forward legacy run driven by the SAME scripted gate. The fake + bakes the per-adapter prefix length into its logits, so an off-by-one in + ``state.lens`` bookkeeping or a corrupted catch-up cache would diverge + here (M4 — the equality must hold THROUGH the skip/catch-up sequence).""" + torch = _torch_or_skip() + script = [[1.0, 0.0], [0.0, 1.0], [1.0, 0.0], [0.0, 1.0]] + out_cached = self._generate( + _make_cached_mole_model(), _make_scripted_gate(script), steps=4 + ) + out_legacy = self._generate( + _make_legacy_mole_model(), _make_scripted_gate(script), steps=4 + ) + assert torch.equal(out_cached, out_legacy) + + def test_no_past_in_output_falls_back(self): + """A capable-signature model that ignores use_cache degrades to the + legacy full-re-forward path mid-call without crashing.""" + _torch_or_skip() + from torch import nn + + base = _make_cached_mole_model() + + class _NoPast(nn.Module): + def __init__(self): + super().__init__() + self.inner = base + self.calls = base.calls + + @property + def device(self): + return base.device + + def disable_adapter(self): + return base.disable_adapter() + + def set_adapter(self, name): + base.set_adapter(name) + + def forward( + self, + input_ids=None, + attention_mask=None, + output_hidden_states=False, + past_key_values=None, + use_cache=False, + ): + out = base( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=output_hidden_states, + ) + return out # never returns past_key_values + + out = self._generate(_NoPast(), _make_uniform_gate(), steps=2) + assert out.shape[1] == 4 # 2 prompt + 2 generated + + def test_state_is_per_call(self): + """Two sequential generate() calls don't share cache state.""" + torch = _torch_or_skip() + from soup_cli.utils.mole_routing import LoadedMole + + model = _make_cached_mole_model() + mole = LoadedMole(model, object(), _make_uniform_gate(), ["task_0", "task_1"]) + ids = torch.tensor([[1, 2]]) + attn = torch.ones_like(ids) + mole.generate(ids, attn, max_new_tokens=2) + model.calls.clear() + mole.generate(ids, attn, max_new_tokens=2) + base_feeds = [t for mode, t in model.calls if mode == "base"] + # Second call starts fresh: full prompt again, then a delta. + assert base_feeds == [2, 1] + + +# --------------------------------------------------------------------------- +# #143 — first-party deploy-measure generator factories +# --------------------------------------------------------------------------- + + +class TestValidateMeasureCandidate: + @pytest.mark.parametrize( + "candidate", + ["none", "4bit", "8bit", "gptq", "awq", "aqlm", "eetq", "mxfp4", "fp8", + "hqq:4bit", "hqq:8bit", "hqq:1bit"], + ) + def test_known_candidates_accepted(self, candidate): + from soup_cli.utils.deploy_measure import validate_measure_candidate + + assert validate_measure_candidate(candidate) == candidate + + def test_case_insensitive(self): + from soup_cli.utils.deploy_measure import validate_measure_candidate + + assert validate_measure_candidate("GPTQ") == "gptq" + + @pytest.mark.parametrize("candidate", ["evil", "", "hqq:5bit", "4 bit"]) + def test_unknown_rejected(self, candidate): + from soup_cli.utils.deploy_measure import validate_measure_candidate + + with pytest.raises(ValueError): + validate_measure_candidate(candidate) + + def test_bool_rejected(self): + from soup_cli.utils.deploy_measure import validate_measure_candidate + + with pytest.raises(TypeError): + validate_measure_candidate(True) + + def test_null_byte_rejected(self): + from soup_cli.utils.deploy_measure import validate_measure_candidate + + with pytest.raises(ValueError): + validate_measure_candidate("4bit\x00") + + +class TestMeasureGeneratorFactories: + def _patch_loader(self, monkeypatch): + """Stub the model loader + generator builder; record invocations.""" + from soup_cli.utils import deploy_measure as dm + + loads: list[dict] = [] + + def fake_load(base, *, quantization, device=None, trust_remote_code=False): + loads.append({"base": base, "quantization": quantization}) + return ("model", "tok", "cpu") + + def fake_make_generator(model_id, *, loaded=None, max_new_tokens=64, **kw): + assert loaded == ("model", "tok", "cpu") + return lambda prompt: f"gen:{prompt}" + + monkeypatch.setattr(dm, "_load_measure_model", fake_load) + import soup_cli.utils.live_eval as live_eval + + monkeypatch.setattr(live_eval, "make_generator", fake_make_generator) + return loads + + def test_before_generator_is_lazy(self, monkeypatch): + from soup_cli.utils.deploy_measure import build_before_generator + + loads = self._patch_loader(monkeypatch) + gen = build_before_generator("org/tiny") + assert loads == [] # nothing loaded at build time (cache-hit safety) + assert gen("hello") == "gen:hello" + assert loads == [{"base": "org/tiny", "quantization": "none"}] + gen("again") + assert len(loads) == 1 # loaded once, reused + + def test_after_factory_threads_candidate_quant(self, monkeypatch): + from soup_cli.utils.deploy_measure import build_after_generator_factory + + loads = self._patch_loader(monkeypatch) + factory = build_after_generator_factory("org/tiny") + gen = factory("4bit") + assert loads == [] + assert gen("p") == "gen:p" + assert loads == [{"base": "org/tiny", "quantization": "4bit"}] + + def test_after_factory_rejects_unknown_candidate_eagerly(self, monkeypatch): + from soup_cli.utils.deploy_measure import build_after_generator_factory + + self._patch_loader(monkeypatch) + factory = build_after_generator_factory("org/tiny") + with pytest.raises(ValueError, match="evil"): + factory("evil") + + def test_builders_validate_base(self): + from soup_cli.utils.deploy_measure import ( + build_after_generator_factory, + build_before_generator, + ) + + for builder in (build_before_generator, build_after_generator_factory): + with pytest.raises(ValueError): + builder("") + with pytest.raises(ValueError): + builder("a\x00b") + + def test_builders_validate_max_new_tokens(self): + from soup_cli.utils.deploy_measure import build_before_generator + + with pytest.raises(ValueError): + build_before_generator("org/tiny", max_new_tokens=0) + with pytest.raises(ValueError): + build_before_generator("org/tiny", max_new_tokens=True) + + def test_builders_base_len_boundary(self, monkeypatch): + """``_MAX_BASE_LEN`` is 512: exactly 512 accepted, 513 rejected.""" + from soup_cli.utils.deploy_measure import ( + build_after_generator_factory, + build_before_generator, + ) + + # Build is lazy (no load) — stub the loader so the 512-char case can't + # accidentally touch the network even if a later .call happened. + self._patch_loader(monkeypatch) + for builder in (build_before_generator, build_after_generator_factory): + # 512 chars: accepted (returns a callable / factory, no raise). + assert builder("a" * 512) is not None + with pytest.raises(ValueError, match="512"): + builder("a" * 513) + + def test_load_measure_model_builds_quant_config(self, monkeypatch): + """_load_measure_model routes the candidate through the Quant Menu + loader (capture tcfg.quantization) without loading a real model.""" + import soup_cli.utils.deploy_measure as dm + + captured = {} + + import soup_cli.utils.quant_menu as quant_menu + + def fake_build(*, tcfg, base, console=None): + captured["quantization"] = tcfg.quantization + captured["base"] = base + return None # behave like quantization='none' + + monkeypatch.setattr( + quant_menu, "build_quantization_config_for_loader", fake_build + ) + + class _FakeTok: + pad_token = "x" + eos_token = "x" + + class _FakeModel: + def to(self, dev): + return self + + def eval(self): + return self + + fake_tf = types.SimpleNamespace( + AutoModelForCausalLM=types.SimpleNamespace( + from_pretrained=lambda *a, **k: _FakeModel() + ), + AutoTokenizer=types.SimpleNamespace( + from_pretrained=lambda *a, **k: _FakeTok() + ), + ) + monkeypatch.setattr(dm, "_import_transformers", lambda: fake_tf) + model, tok, dev = dm._load_measure_model( + "org/tiny", quantization="4bit", device="cpu" + ) + assert captured == {"quantization": "4bit", "base": "org/tiny"} + assert dev == "cpu" + + def test_deploy_cli_wires_first_party_factories(self): + """Source-grep: the CLI uses the first-party builders; the empty-string + placeholders are gone. The injected seams still win when set.""" + src = (_SRC / "commands" / "deploy.py").read_text(encoding="utf-8") + assert "build_before_generator" in src + assert "build_after_generator_factory" in src + assert "_placeholder_before" not in src + assert "_DEPLOY_MEASURE_BEFORE_GEN" in src # seam retained + + +def _write_measure_tasks(tmp_path: Path) -> Path: + """Write a tiny 2-row JSONL eval task file (mirrors test_v0531_109).""" + f = tmp_path / "tasks.jsonl" + f.write_text( + '{"prompt": "say hello", "expected": "hello", "scoring": "exact"}\n' + '{"prompt": "say world", "expected": "world", "scoring": "exact"}\n', + encoding="utf-8", + ) + return f + + +class TestDeployMeasureLiveFailure: + def test_live_measure_failure_exits_1_no_traceback(self, tmp_path, monkeypatch): + """#143 — a model-load failure from the live factory (missing quant + kernel, OOM) surfaces as a friendly exit 1 with no traceback leak.""" + import typer + + from soup_cli.commands.deploy import autopilot + from soup_cli.utils import deploy_measure as _dm + + monkeypatch.chdir(tmp_path) + tasks = _write_measure_tasks(tmp_path) + + def boom(prompt): + raise RuntimeError("no awq kernel") + + # Injected before-gen raises mid-eval; surfaces through run_measure to + # the CLI's (RuntimeError, ImportError, OSError) -> exit 1 branch. + monkeypatch.setattr( + _dm, "_DEPLOY_MEASURE_BEFORE_GEN", boom, raising=False + ) + monkeypatch.setattr( + _dm, + "_DEPLOY_MEASURE_AFTER_FACTORY", + lambda candidate: (lambda p: "x"), + raising=False, + ) + monkeypatch.setenv( + "SOUP_DEPLOY_AUTOPILOT_CACHE", str(tmp_path / "cache.json") + ) + + app = typer.Typer() + app.command()(autopilot) + result = runner.invoke( + app, + [ + "--target", "rtx-4090-24gb", + "--base", "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "--recipe-out", str(tmp_path / "recipe.yaml"), + "--script-out", str(tmp_path / "deploy.sh"), + "--measure", + "--tasks", str(tasks), + "--measure-candidates", "awq", + ], + ) + assert result.exit_code == 1, (result.output, repr(result.exception)) + assert "Live measure failed" in result.output + # No raw traceback text leaks into the output. + assert "Traceback" not in result.output + assert "no awq kernel" in result.output + + +# --------------------------------------------------------------------------- +# #265-partial — live-codec TTS (Orpheus via SNAC) +# --------------------------------------------------------------------------- + + +def _write_wav(path: Path, *, sr: int = 24_000, seconds: float = 0.05, + channels: int = 1) -> None: + """Write a tiny PCM16 sine WAV using only the stdlib.""" + import struct + + n = int(sr * seconds) + frames = bytearray() + for i in range(n): + val = int(12_000 * math.sin(2 * math.pi * 220.0 * i / sr)) + for _ in range(channels): + frames += struct.pack("" + ) + + def test_empty_rejected(self): + from soup_cli.utils.tts_codec import orpheus_tokens_to_string + + with pytest.raises(ValueError): + orpheus_tokens_to_string([]) + + def test_bool_index_rejected(self): + from soup_cli.utils.tts_codec import orpheus_tokens_to_string + + with pytest.raises(TypeError): + orpheus_tokens_to_string([True]) + + def test_float_index_rejected(self): + from soup_cli.utils.tts_codec import orpheus_tokens_to_string + + with pytest.raises(TypeError): + orpheus_tokens_to_string([1.5]) + + +class TestLoadAudioMono: + def test_loads_mono_24k(self, tmp_path): + np = pytest.importorskip("numpy") + pytest.importorskip("soundfile") + from soup_cli.utils.tts_codec import load_audio_mono + + wav = tmp_path / "a.wav" + _write_wav(wav, sr=24_000, seconds=0.05) + audio = load_audio_mono(str(wav)) + assert audio.dtype == np.float32 + assert audio.ndim == 1 + assert abs(len(audio) - 1200) <= 2 + + def test_resamples_other_rates(self, tmp_path): + pytest.importorskip("soundfile") + from soup_cli.utils.tts_codec import load_audio_mono + + wav = tmp_path / "b.wav" + _write_wav(wav, sr=8_000, seconds=0.05) + audio = load_audio_mono(str(wav)) + # 0.05s at 24k target ≈ 1200 samples after resample. + assert abs(len(audio) - 1200) <= 8 + + def test_stereo_mixdown(self, tmp_path): + pytest.importorskip("soundfile") + from soup_cli.utils.tts_codec import load_audio_mono + + wav = tmp_path / "c.wav" + _write_wav(wav, sr=24_000, seconds=0.05, channels=2) + audio = load_audio_mono(str(wav)) + assert audio.ndim == 1 + + def test_missing_file_rejected(self, tmp_path): + pytest.importorskip("soundfile") + from soup_cli.utils.tts_codec import load_audio_mono + + with pytest.raises(FileNotFoundError): + load_audio_mono(str(tmp_path / "nope.wav")) + + def test_null_byte_rejected(self): + from soup_cli.utils.tts_codec import load_audio_mono + + with pytest.raises(ValueError): + load_audio_mono("a\x00b.wav") + + def test_non_string_path_rejected(self): + from soup_cli.utils.tts_codec import load_audio_mono + + with pytest.raises((ValueError, TypeError)): + load_audio_mono(123) + + def test_empty_path_rejected(self): + from soup_cli.utils.tts_codec import load_audio_mono + + with pytest.raises(ValueError): + load_audio_mono("") + + @pytest.mark.skipif( + not hasattr(__import__("os"), "symlink"), + reason="symlink rejection needs os.symlink", + ) + def test_symlink_rejected(self, tmp_path): + import os + + pytest.importorskip("soundfile") + from soup_cli.utils.tts_codec import load_audio_mono + + real = tmp_path / "real.wav" + _write_wav(real, sr=24_000, seconds=0.05) + link = tmp_path / "link.wav" + try: + os.symlink(real, link) + except (OSError, NotImplementedError): + pytest.skip("symlink creation unavailable (needs privilege)") + with pytest.raises(ValueError, match="symlink"): + load_audio_mono(str(link)) + + def test_byte_size_cap(self, tmp_path, monkeypatch): + pytest.importorskip("soundfile") + import soup_cli.utils.tts_codec as tc + + monkeypatch.setattr(tc, "_MAX_AUDIO_BYTES", 16) + wav = tmp_path / "big.wav" + _write_wav(wav, sr=24_000, seconds=0.05) + with pytest.raises(ValueError, match="byte"): + tc.load_audio_mono(str(wav)) + + def test_duration_cap(self, tmp_path, monkeypatch): + pytest.importorskip("soundfile") + import soup_cli.utils.tts_codec as tc + + monkeypatch.setattr(tc, "_MAX_AUDIO_SECONDS", 0.01) + wav = tmp_path / "long.wav" + _write_wav(wav, sr=24_000, seconds=0.05) + with pytest.raises(ValueError, match="seconds"): + tc.load_audio_mono(str(wav)) + + +class _FakeSnac: + """SNAC stand-in: .encode(wav) -> 3 code tensors of the 1/2/4 ratio.""" + + def __init__(self, frames: int = 2): + self.frames = frames + + def encode(self, wav): + import torch + + t = self.frames + return [ + torch.arange(t).unsqueeze(0), + torch.arange(2 * t).unsqueeze(0), + torch.arange(4 * t).unsqueeze(0), + ] + + +class TestEncodeAudioOrpheus: + def test_encodes_with_injected_model(self, tmp_path): + _torch_or_skip() + pytest.importorskip("soundfile") + from soup_cli.utils.tts_codec import encode_audio_orpheus + + wav = tmp_path / "a.wav" + _write_wav(wav) + out = encode_audio_orpheus(str(wav), snac_model=_FakeSnac(frames=2)) + assert out.startswith("" + + def test_appends_assistant_turn(self): + from soup_cli.utils.tts_codec import encode_tts_row + + row = { + "audio": "a.wav", + "messages": [{"role": "user", "content": "Say hi"}], + } + out = encode_tts_row(row, self._encoder) + assert "audio" not in out + assert out["messages"][-1] == { + "role": "assistant", + "content": "", + } + + def test_replaces_existing_assistant_turn(self): + from soup_cli.utils.tts_codec import encode_tts_row + + original_assistant = {"role": "assistant", "content": "placeholder"} + row = { + "audio": "a.wav", + "messages": [ + {"role": "user", "content": "Say hi"}, + original_assistant, + ], + } + out = encode_tts_row(row, self._encoder) + assert out["messages"][-1]["content"] == "" + assert len(out["messages"]) == 2 + # L8 — the deep-copy guarantee: the ORIGINAL assistant dict is NOT + # mutated (no nested aliasing between the input row and the output). + assert original_assistant["content"] == "placeholder" + assert out["messages"][-1] is not original_assistant + + def test_caller_row_not_mutated(self): + from soup_cli.utils.tts_codec import encode_tts_row + + messages = [{"role": "user", "content": "Say hi"}] + row = {"audio": "a.wav", "messages": messages} + encode_tts_row(row, self._encoder) + assert row["messages"] is messages + assert len(messages) == 1 + assert "audio" in row + + def test_missing_audio_rejected(self): + from soup_cli.utils.tts_codec import encode_tts_row + + with pytest.raises(ValueError, match="audio"): + encode_tts_row({"messages": []}, self._encoder) + + def test_missing_messages_rejected(self): + from soup_cli.utils.tts_codec import encode_tts_row + + with pytest.raises(ValueError, match="messages"): + encode_tts_row({"audio": "a.wav"}, self._encoder) + + +class TestEncodeTtsDataset: + def test_maps_train_and_val(self): + from soup_cli.utils.tts_codec import encode_tts_dataset + + rows = [ + {"audio": "a.wav", "messages": [{"role": "user", "content": "x"}]}, + ] + ds = {"train": list(rows), "val": list(rows)} + out = encode_tts_dataset( + ds, "orpheus", encoder=lambda p: "" + ) + assert out["train"][0]["messages"][-1]["role"] == "assistant" + assert out["val"][0]["messages"][-1]["role"] == "assistant" + # input dataset untouched + assert "audio" in ds["train"][0] + + def test_family_validated(self): + from soup_cli.utils.tts_codec import encode_tts_dataset + + with pytest.raises(ValueError): + encode_tts_dataset({"train": []}, "klingon", encoder=lambda p: "x") + + def test_non_dict_rejected(self): + from soup_cli.utils.tts_codec import encode_tts_dataset + + with pytest.raises(TypeError): + encode_tts_dataset([], "orpheus", encoder=lambda p: "x") + + +class TestTtsTrainerLiveCodecWiring: + def test_live_codec_setup_encodes_then_delegates(self, monkeypatch): + """Orpheus + data.format='audio' + snac present: setup() encodes the + rows then falls through to the pre-encoded SFT path (no gate raise).""" + _torch_or_skip() + pytest.importorskip("snac") + from soup_cli.trainer.sft import SFTTrainerWrapper + from soup_cli.trainer.tts import TTSTrainerWrapper + from soup_cli.utils import tts_codec + + monkeypatch.setattr( + tts_codec, + "tts_encoder_for_family", + lambda family, device=None: (lambda path: ""), + ) + captured = {} + monkeypatch.setattr( + SFTTrainerWrapper, + "setup", + lambda self, dataset: captured.update(dataset=dataset), + ) + w = object.__new__(TTSTrainerWrapper) + w.config = SimpleNamespace( + training=SimpleNamespace(tts_family="orpheus", tts_emotion=None), + data=SimpleNamespace(format="audio", new_special_tokens=None), + ) + w._tts_family = None + w.setup( + { + "train": [ + { + "audio": "a.wav", + "messages": [{"role": "user", "content": "Say hi"}], + } + ] + } + ) + row = captured["dataset"]["train"][0] + assert "audio" not in row + assert row["messages"][-1] == { + "role": "assistant", + "content": "", + } + + def test_trainer_calls_encode_dataset(self): + src = (_SRC / "trainer" / "tts.py").read_text(encoding="utf-8") + assert "encode_tts_dataset" in src + # The unconditional not-yet-validated raise is gone. + assert "not yet validated" not in src + + def test_dep_gate_still_present(self): + src = (_SRC / "trainer" / "tts.py").read_text(encoding="utf-8") + assert "_require_tts_codec" in src + + def test_no_top_level_heavy_imports_in_tts_codec(self): + for mod in ("torch", "transformers", "snac", "soundfile", "numpy"): + _assert_no_top_level_import("utils/tts_codec.py", mod) + + +# --------------------------------------------------------------------------- +# Cross-cutting patch invariants +# --------------------------------------------------------------------------- + + +class TestPatchInvariants: + def test_version_bumped(self): + import soup_cli + + parts = tuple(int(p) for p in soup_cli.__version__.split(".")[:3]) + assert parts >= (0, 71, 22) + + def test_no_top_level_heavy_imports_in_minillm(self): + for mod in ("torch", "numpy", "transformers", "peft"): + _assert_no_top_level_import("utils/minillm.py", mod) + + def test_no_top_level_heavy_imports_in_mole_routing(self): + for mod in ("torch", "numpy", "safetensors", "transformers", "peft"): + _assert_no_top_level_import("utils/mole_routing.py", mod) + + def test_no_top_level_heavy_imports_in_deploy_measure(self): + for mod in ("torch", "numpy", "transformers", "peft", "safetensors"): + _assert_no_top_level_import("utils/deploy_measure.py", mod)