sway Public

Watch 0 Fork 0 Star 0
Python · 8835 bytes Raw Blame History
1 """Tests for :mod:`dlm_sway.mining.paraphrase_miner`.
2
3 We stub the embedder + the nlpaug generator so tests run without the
4 80-MB MiniLM load or the nlpaug wheel. What's under test is the ranker
5 and the diversity filter — the infrastructure around them (the
6 embedder, the generator) is exercised by the integration test.
7 """
8
9 from __future__ import annotations
10
11 import numpy as np
12 import pytest
13
14 from dlm_sway.backends.dummy import DummyDifferentialBackend, DummyResponses
15 from dlm_sway.mining.paraphrase_miner import (
16 MiningResult,
17 ParaphraseCandidate,
18 mine_paraphrases,
19 )
20
21
22 def _stub_embedder(text_to_vec: dict[str, np.ndarray]): # type: ignore[no-untyped-def]
23 def _encode(texts: list[str]): # type: ignore[no-untyped-def]
24 return np.stack([text_to_vec[t] for t in texts])
25
26 return _encode
27
28
29 @pytest.fixture
30 def monkeyed_embed(monkeypatch: pytest.MonkeyPatch) -> dict[str, np.ndarray]:
31 table: dict[str, np.ndarray] = {}
32 monkeypatch.setattr(
33 "dlm_sway.mining.paraphrase_miner._load_embedder",
34 lambda _model_id: _stub_embedder(table), # type: ignore[arg-type]
35 )
36 return table
37
38
39 def _canned_candidates(candidates: list[str]): # type: ignore[no-untyped-def]
40 """Build a generator closure that returns ``candidates`` verbatim."""
41
42 def _gen(_prompt: str, *, n: int, seed: int) -> list[str]:
43 del n, seed
44 return list(candidates)
45
46 return _gen
47
48
49 class TestRanker:
50 def test_gap_sort_ranks_hardest_first(self, monkeyed_embed: dict[str, np.ndarray]) -> None:
51 """The paraphrase with the largest verbatim-vs-paraphrase lift
52 gap ranks first. Planted: three candidates with known
53 ``(base, ft)`` logprob pairs; the ranker must surface them in
54 gap-descending order regardless of the generator's input
55 order."""
56 # Seed + candidates
57 prompt = "Q: what is the capital?"
58 gold = " Paris"
59 candidates = ["Q1", "Q2", "Q3"]
60 # Embed all four distinctly so the diversity filter keeps
61 # every candidate.
62 monkeyed_embed[prompt] = np.array([1.0, 0.0, 0.0], dtype=np.float32)
63 monkeyed_embed["Q1"] = np.array([0.0, 1.0, 0.0], dtype=np.float32)
64 monkeyed_embed["Q2"] = np.array([0.0, 0.0, 1.0], dtype=np.float32)
65 monkeyed_embed["Q3"] = np.array([0.5, 0.5, 0.0], dtype=np.float32)
66
67 # Logprob table — verbatim lift is large; paraphrase lifts
68 # vary by design: Q1 lift ≈ 0 (big gap), Q2 lift ≈ verbatim
69 # (no gap), Q3 lift somewhere in between.
70 base_logprobs = {
71 (prompt, gold): -4.0,
72 ("Q1", gold): -4.0, # base unchanged
73 ("Q2", gold): -4.0,
74 ("Q3", gold): -4.0,
75 }
76 ft_logprobs = {
77 (prompt, gold): -1.0, # verbatim lift: +3.0
78 ("Q1", gold): -4.0, # paraphrase lift: 0 (big gap of 3)
79 ("Q2", gold): -1.0, # paraphrase lift: +3 (gap of 0)
80 ("Q3", gold): -2.5, # paraphrase lift: +1.5 (gap of 1.5)
81 }
82 backend = DummyDifferentialBackend(
83 base=DummyResponses(logprobs=base_logprobs),
84 ft=DummyResponses(logprobs=ft_logprobs),
85 )
86
87 result = mine_paraphrases(
88 prompt=prompt,
89 gold=gold,
90 backend=backend,
91 generate_candidates=_canned_candidates(candidates),
92 n_candidates=3,
93 top_k=3,
94 seed=0,
95 )
96
97 assert isinstance(result, MiningResult)
98 assert [c.prompt for c in result.candidates] == ["Q1", "Q3", "Q2"]
99 # Gaps are in descending order.
100 gaps = [c.gap for c in result.candidates]
101 assert gaps == sorted(gaps, reverse=True)
102 # Top candidate's gap is meaningfully larger than 0.
103 assert result.candidates[0].gap > 0.1
104
105 def test_dedup_drops_verbatim_seed(self, monkeyed_embed: dict[str, np.ndarray]) -> None:
106 """If the generator echoes the seed prompt back in its output,
107 it's dropped before the ranker sees it — otherwise the seed
108 would bubble to the top with gap = 0 and pollute the list."""
109 prompt = "prompt"
110 gold = " gold"
111 # Generator returns the seed + two real candidates; the seed
112 # must not appear in the final result.
113 candidates = ["prompt", "C1", "C2"]
114 for p in candidates + [prompt]:
115 monkeyed_embed[p] = (
116 np.random.RandomState(hash(p) & 0xFFFFFFFF).randn(4).astype(np.float32)
117 )
118
119 base_logprobs = {(p, gold): -3.0 for p in [prompt, "C1", "C2"]}
120 ft_logprobs = {(prompt, gold): -1.0, ("C1", gold): -2.5, ("C2", gold): -2.0}
121 backend = DummyDifferentialBackend(
122 base=DummyResponses(logprobs=base_logprobs),
123 ft=DummyResponses(logprobs=ft_logprobs),
124 )
125
126 result = mine_paraphrases(
127 prompt=prompt,
128 gold=gold,
129 backend=backend,
130 generate_candidates=_canned_candidates(candidates),
131 n_candidates=3,
132 top_k=3,
133 seed=0,
134 )
135 assert prompt not in {c.prompt for c in result.candidates}
136
137 def test_empty_candidate_list_returns_empty_result(
138 self, monkeyed_embed: dict[str, np.ndarray]
139 ) -> None:
140 del monkeyed_embed
141 backend = DummyDifferentialBackend(base=DummyResponses(), ft=DummyResponses())
142 result = mine_paraphrases(
143 prompt="x",
144 gold=" y",
145 backend=backend,
146 generate_candidates=_canned_candidates([]),
147 n_candidates=3,
148 top_k=3,
149 )
150 assert isinstance(result, MiningResult)
151 assert result.candidates == []
152
153
154 class TestDiversityFilter:
155 def test_keeps_farthest_when_candidates_cluster(
156 self, monkeyed_embed: dict[str, np.ndarray]
157 ) -> None:
158 """Two near-duplicate candidates + one distant one → the
159 distant candidate must survive the k=2 diversity filter."""
160 prompt = "seed"
161 gold = " gold"
162 candidates = ["near1", "near2", "far"]
163 monkeyed_embed[prompt] = np.array([1.0, 0.0, 0.0], dtype=np.float32)
164 # near1 / near2 collapse onto the same embedding; far lives
165 # orthogonally. With k=2, the filter must pick one near + far
166 # (not both nears).
167 monkeyed_embed["near1"] = np.array([0.9, 0.1, 0.0], dtype=np.float32)
168 monkeyed_embed["near2"] = np.array([0.9, 0.1, 0.0], dtype=np.float32)
169 monkeyed_embed["far"] = np.array([0.0, 0.0, 1.0], dtype=np.float32)
170
171 base_logprobs = {("near1", gold): -3.0, ("near2", gold): -3.0, ("far", gold): -3.0}
172 ft_logprobs = {("near1", gold): -2.0, ("near2", gold): -2.0, ("far", gold): -1.0}
173 base_logprobs[(prompt, gold)] = -3.0
174 ft_logprobs[(prompt, gold)] = -1.0
175 backend = DummyDifferentialBackend(
176 base=DummyResponses(logprobs=base_logprobs),
177 ft=DummyResponses(logprobs=ft_logprobs),
178 )
179
180 result = mine_paraphrases(
181 prompt=prompt,
182 gold=gold,
183 backend=backend,
184 generate_candidates=_canned_candidates(candidates),
185 n_candidates=3,
186 top_k=2,
187 seed=0,
188 )
189 chosen = {c.prompt for c in result.candidates}
190 assert "far" in chosen
191 assert len(chosen & {"near1", "near2"}) == 1
192
193
194 class TestInputValidation:
195 def test_rejects_top_k_zero(self, monkeyed_embed: dict[str, np.ndarray]) -> None:
196 del monkeyed_embed
197 backend = DummyDifferentialBackend(base=DummyResponses(), ft=DummyResponses())
198 with pytest.raises(ValueError, match="top_k must be positive"):
199 mine_paraphrases(
200 prompt="x",
201 gold=" y",
202 backend=backend,
203 generate_candidates=_canned_candidates(["a"]),
204 n_candidates=5,
205 top_k=0,
206 )
207
208 def test_rejects_n_candidates_below_top_k(self, monkeyed_embed: dict[str, np.ndarray]) -> None:
209 del monkeyed_embed
210 backend = DummyDifferentialBackend(base=DummyResponses(), ft=DummyResponses())
211 with pytest.raises(ValueError, match="must be ≥ top_k"):
212 mine_paraphrases(
213 prompt="x",
214 gold=" y",
215 backend=backend,
216 generate_candidates=_canned_candidates(["a"]),
217 n_candidates=2,
218 top_k=5,
219 )
220
221
222 class TestParaphraseCandidate:
223 def test_is_frozen_dataclass(self) -> None:
224 c = ParaphraseCandidate(
225 prompt="p", gap=0.5, verbatim_lift=1.0, paraphrase_lift=0.5, diversity_rank=0
226 )
227 with pytest.raises(Exception): # noqa: B017, PT011 — FrozenInstanceError varies
228 c.gap = 0.0 # type: ignore[misc]