mirror of https://github.com/razor-ai/soup.git
371 lines
12 KiB
Python
371 lines
12 KiB
Python
"""v0.68.0 Part B — ``soup distill-prompt``.
|
|
|
|
Distill prompt-heavy traces (large-prompt GPT-5 / Claude calls) into a small
|
|
FT plan. Bridge between prompt-engineering and FT worlds. Composes with the
|
|
v0.70 Part B cross-tokenizer KD when that ships.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import dataclasses
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from typer.testing import CliRunner
|
|
|
|
|
|
class TestPublicSurface:
|
|
def test_module_importable(self) -> None:
|
|
from soup_cli.utils import prompt_distill
|
|
|
|
assert hasattr(prompt_distill, "SUPPORTED_DISTILL_STRATEGIES")
|
|
assert hasattr(prompt_distill, "validate_distill_strategy")
|
|
assert hasattr(prompt_distill, "validate_teacher_id")
|
|
assert hasattr(prompt_distill, "validate_student_id")
|
|
assert hasattr(prompt_distill, "validate_traces_path")
|
|
assert hasattr(prompt_distill, "DistillPromptPlan")
|
|
assert hasattr(prompt_distill, "build_distill_prompt_plan")
|
|
assert hasattr(prompt_distill, "prepare_distill_dataset")
|
|
|
|
|
|
class TestAllowlist:
|
|
def test_frozenset(self) -> None:
|
|
from soup_cli.utils.prompt_distill import SUPPORTED_DISTILL_STRATEGIES
|
|
|
|
assert isinstance(SUPPORTED_DISTILL_STRATEGIES, frozenset)
|
|
assert "sft" in SUPPORTED_DISTILL_STRATEGIES
|
|
assert "preference" in SUPPORTED_DISTILL_STRATEGIES
|
|
assert "kl" in SUPPORTED_DISTILL_STRATEGIES
|
|
|
|
def test_immutable(self) -> None:
|
|
from soup_cli.utils.prompt_distill import SUPPORTED_DISTILL_STRATEGIES
|
|
|
|
with pytest.raises(AttributeError):
|
|
SUPPORTED_DISTILL_STRATEGIES.add("evil") # type: ignore[attr-defined]
|
|
|
|
|
|
class TestValidateStrategy:
|
|
def test_happy(self) -> None:
|
|
from soup_cli.utils.prompt_distill import validate_distill_strategy
|
|
|
|
assert validate_distill_strategy("sft") == "sft"
|
|
|
|
def test_case_insensitive(self) -> None:
|
|
from soup_cli.utils.prompt_distill import validate_distill_strategy
|
|
|
|
assert validate_distill_strategy("SFT") == "sft"
|
|
|
|
def test_bool_rejected(self) -> None:
|
|
from soup_cli.utils.prompt_distill import validate_distill_strategy
|
|
|
|
with pytest.raises(TypeError):
|
|
validate_distill_strategy(True) # type: ignore[arg-type]
|
|
|
|
def test_unknown_rejected(self) -> None:
|
|
from soup_cli.utils.prompt_distill import validate_distill_strategy
|
|
|
|
with pytest.raises(ValueError, match="unknown"):
|
|
validate_distill_strategy("evil")
|
|
|
|
def test_empty_rejected(self) -> None:
|
|
from soup_cli.utils.prompt_distill import validate_distill_strategy
|
|
|
|
with pytest.raises(ValueError):
|
|
validate_distill_strategy("")
|
|
|
|
def test_null_byte_rejected(self) -> None:
|
|
from soup_cli.utils.prompt_distill import validate_distill_strategy
|
|
|
|
with pytest.raises(ValueError):
|
|
validate_distill_strategy("sft\x00")
|
|
|
|
|
|
class TestValidateModelId:
|
|
def test_teacher_happy(self) -> None:
|
|
from soup_cli.utils.prompt_distill import validate_teacher_id
|
|
|
|
assert validate_teacher_id("anthropic/claude-3-5-sonnet") == "anthropic/claude-3-5-sonnet"
|
|
|
|
def test_teacher_oversize_rejected(self) -> None:
|
|
from soup_cli.utils.prompt_distill import validate_teacher_id
|
|
|
|
with pytest.raises(ValueError):
|
|
validate_teacher_id("a" * 513)
|
|
|
|
def test_teacher_null_byte_rejected(self) -> None:
|
|
from soup_cli.utils.prompt_distill import validate_teacher_id
|
|
|
|
with pytest.raises(ValueError):
|
|
validate_teacher_id("model\x00")
|
|
|
|
def test_teacher_bool_rejected(self) -> None:
|
|
from soup_cli.utils.prompt_distill import validate_teacher_id
|
|
|
|
with pytest.raises(TypeError):
|
|
validate_teacher_id(True) # type: ignore[arg-type]
|
|
|
|
def test_teacher_empty_rejected(self) -> None:
|
|
from soup_cli.utils.prompt_distill import validate_teacher_id
|
|
|
|
with pytest.raises(ValueError):
|
|
validate_teacher_id("")
|
|
|
|
def test_student_happy(self) -> None:
|
|
from soup_cli.utils.prompt_distill import validate_student_id
|
|
|
|
assert validate_student_id("meta-llama/Llama-3.2-1B") == "meta-llama/Llama-3.2-1B"
|
|
|
|
|
|
class TestValidateTracesPath:
|
|
def test_happy(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from soup_cli.utils.prompt_distill import validate_traces_path
|
|
|
|
monkeypatch.chdir(tmp_path)
|
|
p = tmp_path / "traces.jsonl"
|
|
p.write_text('{"prompt":"x","response":"y"}\n', encoding="utf-8")
|
|
assert validate_traces_path(str(p)).endswith("traces.jsonl")
|
|
|
|
def test_outside_cwd_rejected(
|
|
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
from soup_cli.utils.prompt_distill import validate_traces_path
|
|
|
|
outside = tmp_path / "outside"
|
|
outside.mkdir()
|
|
p = outside / "t.jsonl"
|
|
p.write_text("[]", encoding="utf-8")
|
|
sub = tmp_path / "sub"
|
|
sub.mkdir()
|
|
monkeypatch.chdir(sub)
|
|
with pytest.raises(ValueError):
|
|
validate_traces_path(str(p))
|
|
|
|
def test_null_byte_rejected(self) -> None:
|
|
from soup_cli.utils.prompt_distill import validate_traces_path
|
|
|
|
with pytest.raises(ValueError):
|
|
validate_traces_path("t\x00.jsonl")
|
|
|
|
@pytest.mark.skipif(
|
|
not hasattr(__import__("os"), "symlink"), reason="POSIX symlink only"
|
|
)
|
|
def test_symlink_rejected(
|
|
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
import os
|
|
|
|
from soup_cli.utils.prompt_distill import validate_traces_path
|
|
|
|
monkeypatch.chdir(tmp_path)
|
|
real = tmp_path / "r.jsonl"
|
|
real.write_text("[]", encoding="utf-8")
|
|
link = tmp_path / "link.jsonl"
|
|
try:
|
|
os.symlink(real, link)
|
|
except (OSError, NotImplementedError):
|
|
pytest.skip("symlinks not creatable")
|
|
with pytest.raises(ValueError, match="symlink"):
|
|
validate_traces_path(str(link))
|
|
|
|
|
|
class TestDistillPromptPlan:
|
|
def test_frozen(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from soup_cli.utils.prompt_distill import DistillPromptPlan
|
|
|
|
monkeypatch.chdir(tmp_path)
|
|
traces = tmp_path / "traces.jsonl"
|
|
traces.write_text("[]", encoding="utf-8")
|
|
|
|
plan = DistillPromptPlan(
|
|
traces_path=str(traces),
|
|
teacher="anthropic/claude-3-5-sonnet",
|
|
student="meta-llama/Llama-3.2-1B",
|
|
strategy="sft",
|
|
output_path="distilled.jsonl",
|
|
)
|
|
with pytest.raises(dataclasses.FrozenInstanceError):
|
|
plan.strategy = "preference" # type: ignore[misc]
|
|
|
|
def test_invalid_strategy_propagates(
|
|
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
from soup_cli.utils.prompt_distill import DistillPromptPlan
|
|
|
|
monkeypatch.chdir(tmp_path)
|
|
traces = tmp_path / "traces.jsonl"
|
|
traces.write_text("[]", encoding="utf-8")
|
|
|
|
with pytest.raises(ValueError):
|
|
DistillPromptPlan(
|
|
traces_path=str(traces),
|
|
teacher="t",
|
|
student="s",
|
|
strategy="evil",
|
|
output_path="o.jsonl",
|
|
)
|
|
|
|
|
|
class TestBuildPlan:
|
|
def test_happy(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from soup_cli.utils.prompt_distill import build_distill_prompt_plan
|
|
|
|
monkeypatch.chdir(tmp_path)
|
|
traces = tmp_path / "traces.jsonl"
|
|
traces.write_text("[]", encoding="utf-8")
|
|
plan = build_distill_prompt_plan(
|
|
traces_path=str(traces),
|
|
teacher="t/x",
|
|
student="s/y",
|
|
strategy="sft",
|
|
output_path="o.jsonl",
|
|
)
|
|
assert plan.strategy == "sft"
|
|
|
|
|
|
class TestPrepareDataset:
|
|
def test_live_writes_rows(
|
|
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
# v0.71.13 #226: live dataset preparation with an injected teacher.
|
|
from soup_cli.utils.prompt_distill import (
|
|
build_distill_prompt_plan,
|
|
prepare_distill_dataset,
|
|
)
|
|
|
|
monkeypatch.chdir(tmp_path)
|
|
traces = tmp_path / "traces.jsonl"
|
|
traces.write_text(
|
|
json.dumps({"prompt": "q"}) + "\n", encoding="utf-8"
|
|
)
|
|
plan = build_distill_prompt_plan(
|
|
traces_path=str(traces),
|
|
teacher="t/x",
|
|
student="s/y",
|
|
strategy="sft",
|
|
output_path="o.jsonl",
|
|
)
|
|
n = prepare_distill_dataset(plan, teacher_fn=lambda p: {"text": "T"})
|
|
assert n == 1
|
|
assert (tmp_path / "o.jsonl").is_file()
|
|
|
|
def test_non_plan_rejected(self) -> None:
|
|
from soup_cli.utils.prompt_distill import prepare_distill_dataset
|
|
|
|
with pytest.raises(TypeError):
|
|
prepare_distill_dataset({}) # type: ignore[arg-type]
|
|
|
|
|
|
class TestCli:
|
|
def test_help(self) -> None:
|
|
from soup_cli.cli import app
|
|
|
|
runner = CliRunner()
|
|
result = runner.invoke(app, ["distill-prompt", "--help"])
|
|
assert result.exit_code == 0, (result.output, repr(result.exception))
|
|
|
|
def test_plan_only(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from soup_cli.cli import app
|
|
|
|
monkeypatch.chdir(tmp_path)
|
|
traces = tmp_path / "traces.jsonl"
|
|
traces.write_text(json.dumps({"prompt": "x"}) + "\n", encoding="utf-8")
|
|
runner = CliRunner()
|
|
result = runner.invoke(
|
|
app,
|
|
[
|
|
"distill-prompt",
|
|
"--traces",
|
|
str(traces),
|
|
"--teacher",
|
|
"anthropic/claude-3-5-sonnet",
|
|
"--student",
|
|
"meta-llama/Llama-3.2-1B",
|
|
"--strategy",
|
|
"sft",
|
|
"--plan-only",
|
|
],
|
|
)
|
|
assert result.exit_code == 0, (result.output, repr(result.exception))
|
|
|
|
def test_unknown_strategy_exits_2(
|
|
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
from soup_cli.cli import app
|
|
|
|
monkeypatch.chdir(tmp_path)
|
|
traces = tmp_path / "traces.jsonl"
|
|
traces.write_text("[]", encoding="utf-8")
|
|
runner = CliRunner()
|
|
result = runner.invoke(
|
|
app,
|
|
[
|
|
"distill-prompt",
|
|
"--traces",
|
|
str(traces),
|
|
"--teacher",
|
|
"t",
|
|
"--student",
|
|
"s",
|
|
"--strategy",
|
|
"evil",
|
|
"--plan-only",
|
|
],
|
|
)
|
|
assert result.exit_code == 2
|
|
|
|
def test_live_writes_output(
|
|
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
# v0.71.13 #226: live runner writes a distilled dataset (provider
|
|
# mocked so the test never touches the network).
|
|
import soup_cli.utils.prompt_distill as pd
|
|
from soup_cli.cli import app
|
|
|
|
monkeypatch.chdir(tmp_path)
|
|
monkeypatch.setattr(
|
|
pd, "_build_provider_fn", lambda *a, **k: (lambda p: {"text": "R"})
|
|
)
|
|
traces = tmp_path / "traces.jsonl"
|
|
traces.write_text(
|
|
json.dumps({"prompt": "x"}) + "\n", encoding="utf-8"
|
|
)
|
|
runner = CliRunner()
|
|
result = runner.invoke(
|
|
app,
|
|
[
|
|
"distill-prompt",
|
|
"--traces",
|
|
str(traces),
|
|
"--teacher",
|
|
"t",
|
|
"--student",
|
|
"s",
|
|
"--strategy",
|
|
"sft",
|
|
"--output",
|
|
"o.jsonl",
|
|
],
|
|
)
|
|
assert result.exit_code == 0, (result.output, repr(result.exception))
|
|
assert (tmp_path / "o.jsonl").is_file()
|
|
|
|
|
|
class TestSourceWiring:
|
|
def test_no_top_level_heavy_imports(self) -> None:
|
|
path = (
|
|
Path(__file__).resolve().parent.parent
|
|
/ "src" / "soup_cli"
|
|
/ "utils"
|
|
/ "prompt_distill.py"
|
|
)
|
|
text = path.read_text(encoding="utf-8")
|
|
for token in ("\nimport torch", "\nimport transformers", "\nimport anthropic"):
|
|
assert token not in text
|
|
|
|
def test_cli_registered(self) -> None:
|
|
from soup_cli.cli import app
|
|
|
|
names = [c.name for c in app.registered_commands]
|
|
assert "distill-prompt" in names
|