`ed303dd`

sway(probes): A2 adapter_revert via sentence embeddings

Authored by

espadonne 3 weeks ago

SHA: ed303dd9f10fceae4dab13c6c7e56637717ad1fa
Parents: c5cfd2e
Tree: fe04c0d

2 changed files

Status	File	+	-
A	`src/dlm_sway/probes/adapter_revert.py`	178	0
A	`tests/unit/test_probe_adapter_revert.py`	170	0

src/dlm_sway/probes/adapter_revert.pyadded

 +"""A2 AdapterRevert — does the fine-tuned model drift back to base under pressure?
++
 +For each test case the user provides a prompt, a "gold" answer (the
 +adapter's intended response), and one or more adversarial paraphrases of
 +the prompt. We generate base-model and ft-model completions on every
 +paraphrase and ask: does the ft output cluster semantically with the
 +base's output (revert) or with the gold (adhere)?
++
 +Signal: ``revert_rate`` = fraction of (case, paraphrase) pairs where
 +``cos(ft, base) > cos(ft, gold)``. A healthy fine-tune holds below 25%.
++
 +Needs sentence embeddings. Without the ``semsim`` extra installed the
 +probe returns :attr:`Verdict.SKIP` with a pip hint — deterministic
 +n-gram fallbacks don't carry semantic equivalence reliably enough to
 +drive a revert decision, and we'd rather be honest than lossy.
 +"""
++
 +from __future__ import annotations
++
 +from typing import Any, Literal
++
 +from pydantic import BaseModel, ConfigDict, Field
++
 +from dlm_sway.core.errors import BackendNotAvailableError
 +from dlm_sway.core.result import ProbeResult, Verdict
 +from dlm_sway.probes.base import Probe, ProbeSpec, RunContext
++
++
 +class AdapterRevertCase(BaseModel):
 +    """One revert test case."""
++
 +    model_config = ConfigDict(extra="forbid", frozen=True)
++
 +    prompt: str
 +    gold: str
 +    """What the adapter is supposed to produce."""
 +    paraphrases: list[str] = Field(default_factory=list, min_length=1)
 +    """At least one paraphrase is required — revert is observed under
 +    reframing, not on the original prompt."""
++
++
 +class AdapterRevertSpec(ProbeSpec):
 +    kind: Literal["adapter_revert"] = "adapter_revert"
 +    cases: list[AdapterRevertCase] = Field(default_factory=list)
 +    max_new_tokens: int = 64
 +    embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2"
 +    """HF id of the embedder. Default is ~80 MB, CPU-friendly."""
 +    base_gold_similarity_cap: float = 0.75
 +    """Skip pairs where base and gold are trivially similar — those
 +    can't distinguish revert from adherence, and including them would
 +    inflate the revert rate with noise."""
 +    assert_revert_rate_lt: float = 0.25
++
++
 +class AdapterRevertProbe(Probe):
 +    kind = "adapter_revert"
 +    spec_cls = AdapterRevertSpec
 +    category = "adherence"
++
 +    def run(self, spec: ProbeSpec, ctx: RunContext) -> ProbeResult:
 +        assert isinstance(spec, AdapterRevertSpec)
 +        if not spec.cases:
 +            return ProbeResult(
 +                name=spec.name,
 +                kind=spec.kind,
 +                verdict=Verdict.ERROR,
 +                score=None,
 +                message="no cases provided",
 +            )
++
 +        try:
 +            embed = _load_embedder(spec.embedding_model)
 +        except BackendNotAvailableError as exc:
 +            return ProbeResult(
 +                name=spec.name,
 +                kind=spec.kind,
 +                verdict=Verdict.SKIP,
 +                score=None,
 +                message=str(exc),
 +            )
++
 +        import numpy as np
++
 +        total = 0
 +        reverts = 0
 +        dropped_trivial = 0
 +        per_case: list[dict[str, Any]] = []
 +        for case in spec.cases:
 +            gold_vec = embed([case.gold])[0]
 +            for pp in case.paraphrases:
 +                with ctx.backend.as_base() as bv:
 +                    base_gen = bv.generate(pp, max_new_tokens=spec.max_new_tokens, seed=ctx.seed)
 +                with ctx.backend.as_finetuned() as fv:
 +                    ft_gen = fv.generate(pp, max_new_tokens=spec.max_new_tokens, seed=ctx.seed)
 +                vecs = embed([base_gen, ft_gen])
 +                base_vec, ft_vec = vecs[0], vecs[1]
 +                base_gold = _cosine(base_vec, gold_vec)
 +                if base_gold > spec.base_gold_similarity_cap:
 +                    dropped_trivial += 1
 +                    continue
 +                cos_ft_base = _cosine(ft_vec, base_vec)
 +                cos_ft_gold = _cosine(ft_vec, gold_vec)
 +                total += 1
 +                if cos_ft_base > cos_ft_gold:
 +                    reverts += 1
 +                per_case.append(
 +                    {
 +                        "prompt": pp[:80],
 +                        "cos_ft_base": cos_ft_base,
 +                        "cos_ft_gold": cos_ft_gold,
 +                        "reverted": cos_ft_base > cos_ft_gold,
 +                    }
 +                )
++
 +        if total == 0:
 +            return ProbeResult(
 +                name=spec.name,
 +                kind=spec.kind,
 +                verdict=Verdict.WARN,
 +                score=0.5,
 +                message=(
 +                    f"all {dropped_trivial} cases had base≈gold (> "
 +                    f"{spec.base_gold_similarity_cap}) — no separable signal"
 +                ),
 +                evidence={"dropped_trivial": dropped_trivial, "weight": spec.weight},
 +            )
++
 +        rate = reverts / total
 +        verdict = Verdict.PASS if rate < spec.assert_revert_rate_lt else Verdict.FAIL
 +        score = max(0.0, 1.0 - rate / max(spec.assert_revert_rate_lt, 1e-6))
 +        score = float(np.clip(score, 0.0, 1.0))
++
 +        return ProbeResult(
 +            name=spec.name,
 +            kind=spec.kind,
 +            verdict=verdict,
 +            score=score,
 +            raw=rate,
 +            evidence={
 +                "revert_rate": rate,
 +                "reverts": reverts,
 +                "total": total,
 +                "dropped_trivial": dropped_trivial,
 +                "per_case": per_case[:8],  # cap to keep JSON bounded
 +                "weight": spec.weight,
 +            },
 +            message=f"revert_rate={rate:.2%} (reverts={reverts}/{total}, dropped_trivial={dropped_trivial})",
 +        )
++
++
 +def _load_embedder(model_id: str):  # type: ignore[no-untyped-def]
 +    """Return a callable ``list[str] -> np.ndarray`` over encoded vectors."""
 +    try:
 +        from sentence_transformers import SentenceTransformer
 +    except ImportError as exc:
 +        raise BackendNotAvailableError(
 +            "adapter_revert",
 +            extra="semsim",
 +            hint="adapter_revert relies on sentence embeddings.",
 +        ) from exc
 +    st = SentenceTransformer(model_id)
++
 +    def _embed(texts: list[str]):  # type: ignore[no-untyped-def]
 +        return st.encode(texts, convert_to_numpy=True, normalize_embeddings=True)
++
 +    return _embed
++
++
 +def _cosine(a: Any, b: Any) -> float:
 +    import numpy as np
++
 +    av = np.asarray(a, dtype=np.float64)
 +    bv = np.asarray(b, dtype=np.float64)
 +    na = float(np.linalg.norm(av))
 +    nb = float(np.linalg.norm(bv))
 +    if na == 0.0 or nb == 0.0:
 +        return 0.0
 +    return float(np.dot(av, bv) / (na * nb))

tests/unit/test_probe_adapter_revert.pyadded

 +"""Tests for :mod:`dlm_sway.probes.adapter_revert`.
++
 +We stub out the embedder so these tests don't need sentence-transformers
 +installed. The ``probe.py`` SKIP path for the missing-extra case is
 +covered separately by monkeypatching the importer.
 +"""
++
 +from __future__ import annotations
++
 +from typing import Any
++
 +import numpy as np
 +import pytest
++
 +from dlm_sway.backends.dummy import DummyDifferentialBackend, DummyResponses
 +from dlm_sway.core.result import Verdict
 +from dlm_sway.probes.adapter_revert import AdapterRevertProbe
 +from dlm_sway.probes.base import RunContext, build_probe
++
++
 +def _backend(*, ft_like_base: bool = False) -> DummyDifferentialBackend:
 +    base = DummyResponses(
 +        generations={
 +            "pp1": "cats are mammals",
 +            "pp2": "cats have fur",
 +        }
 +    )
 +    if ft_like_base:
 +        ft_gens = dict(base.generations)
 +    else:
 +        ft_gens = {
 +            "pp1": "dolphins are mammals",
 +            "pp2": "dolphins are smart",
 +        }
 +    ft = DummyResponses(generations=ft_gens)
 +    return DummyDifferentialBackend(base=base, ft=ft)
++
++
 +def _stub_embedder(text_to_vec: dict[str, np.ndarray]):  # type: ignore[no-untyped-def]
 +    def _encode(texts: list[str]):  # type: ignore[no-untyped-def]
 +        return np.stack([text_to_vec[t] for t in texts])
++
 +    return _encode
++
++
 +@pytest.fixture
 +def monkeyed_embed(monkeypatch: pytest.MonkeyPatch) -> dict[str, np.ndarray]:
 +    """Install a stub embedder with a controllable text→vec mapping.
++
 +    Tests populate the dict before calling ``probe.run()``.
 +    """
 +    table: dict[str, np.ndarray] = {}
 +    monkeypatch.setattr(
 +        "dlm_sway.probes.adapter_revert._load_embedder",
 +        lambda _model_id: _stub_embedder(table),  # type: ignore[arg-type]
 +    )
 +    return table
++
++
 +class TestAdapterRevert:
 +    def test_healthy_adapter_passes(self, monkeyed_embed: dict[str, np.ndarray]) -> None:
 +        # gold and ft-outputs cluster together, base outputs cluster elsewhere.
 +        monkeyed_embed["cats are mammals"] = np.array([1.0, 0.0])
 +        monkeyed_embed["cats have fur"] = np.array([1.0, 0.0])
 +        monkeyed_embed["dolphins are mammals"] = np.array([0.0, 1.0])
 +        monkeyed_embed["dolphins are smart"] = np.array([0.0, 1.0])
 +        monkeyed_embed["the answer is dolphins"] = np.array([0.0, 1.0])  # gold
++
 +        probe, spec = build_probe(
 +            {
 +                "name": "rev",
 +                "kind": "adapter_revert",
 +                "cases": [
 +                    {
 +                        "prompt": "anything",
 +                        "gold": "the answer is dolphins",
 +                        "paraphrases": ["pp1", "pp2"],
 +                    }
 +                ],
 +                "assert_revert_rate_lt": 0.25,
 +            }
 +        )
 +        ctx = RunContext(backend=_backend(ft_like_base=False))
 +        result = probe.run(spec, ctx)
 +        assert result.verdict == Verdict.PASS
 +        assert result.raw == 0.0
++
 +    def test_reverting_adapter_fails(self, monkeyed_embed: dict[str, np.ndarray]) -> None:
 +        # ft matches base (reverted), diverges from gold.
 +        monkeyed_embed["cats are mammals"] = np.array([1.0, 0.0])
 +        monkeyed_embed["cats have fur"] = np.array([1.0, 0.0])
 +        monkeyed_embed["the answer is dolphins"] = np.array([0.0, 1.0])  # gold
++
 +        probe, spec = build_probe(
 +            {
 +                "name": "rev",
 +                "kind": "adapter_revert",
 +                "cases": [
 +                    {
 +                        "prompt": "anything",
 +                        "gold": "the answer is dolphins",
 +                        "paraphrases": ["pp1", "pp2"],
 +                    }
 +                ],
 +            }
 +        )
 +        ctx = RunContext(backend=_backend(ft_like_base=True))
 +        result = probe.run(spec, ctx)
 +        assert result.verdict == Verdict.FAIL
 +        assert result.raw == 1.0  # 100% revert
++
 +    def test_trivially_similar_cases_dropped(self, monkeyed_embed: dict[str, np.ndarray]) -> None:
 +        # base and gold are identical → drop.
 +        v = np.array([1.0, 0.0])
 +        monkeyed_embed["cats are mammals"] = v
 +        monkeyed_embed["cats have fur"] = v
 +        monkeyed_embed["dolphins are mammals"] = np.array([0.0, 1.0])
 +        monkeyed_embed["dolphins are smart"] = np.array([0.0, 1.0])
 +        monkeyed_embed["cats are mammals too"] = v  # gold — matches base
++
 +        probe, spec = build_probe(
 +            {
 +                "name": "rev",
 +                "kind": "adapter_revert",
 +                "cases": [
 +                    {
 +                        "prompt": "anything",
 +                        "gold": "cats are mammals too",
 +                        "paraphrases": ["pp1", "pp2"],
 +                    }
 +                ],
 +            }
 +        )
 +        ctx = RunContext(backend=_backend(ft_like_base=False))
 +        result = probe.run(spec, ctx)
 +        # Both paraphrase pairs trivially similar → WARN (no separable signal).
 +        assert result.verdict == Verdict.WARN
 +        assert result.evidence["dropped_trivial"] == 2
++
 +    def test_no_cases_errors(self, monkeyed_embed: dict[str, np.ndarray]) -> None:
 +        probe, spec = build_probe({"name": "rev", "kind": "adapter_revert", "cases": []})
 +        ctx = RunContext(backend=_backend())
 +        result = probe.run(spec, ctx)
 +        assert result.verdict == Verdict.ERROR
++
++
 +class TestMissingSemsim:
 +    def test_skip_when_sentence_transformers_missing(self, monkeypatch: pytest.MonkeyPatch) -> None:
 +        from dlm_sway.core.errors import BackendNotAvailableError
++
 +        def raiser(_model_id: Any) -> Any:  # type: ignore[no-untyped-def]
 +            raise BackendNotAvailableError(
 +                "adapter_revert",
 +                extra="semsim",
 +                hint="adapter_revert relies on sentence embeddings.",
 +            )
++
 +        monkeypatch.setattr(
 +            "dlm_sway.probes.adapter_revert._load_embedder",
 +            raiser,  # type: ignore[arg-type]
 +        )
 +        probe = AdapterRevertProbe()
 +        spec = probe.spec_cls(
 +            name="rev",
 +            cases=[{"prompt": "x", "gold": "y", "paraphrases": ["pp1"]}],  # type: ignore[list-item]
 +        )
 +        ctx = RunContext(backend=_backend())
 +        result = probe.run(spec, ctx)
 +        assert result.verdict == Verdict.SKIP
 +        assert "semsim" in result.message