File size: 5,624 Bytes
196d16f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | """Candidate-route workflow tested with a real deterministic decoder."""
import math
from pathlib import Path
import sys
import unittest
import torch
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from demos.path_routes import aggregate_pair_effects, measure_receiver_paths, select_candidates
from test_path_patching import CircuitModel
from vlm_probing import CapabilityError, Prober
class CharacterTokenizer:
"""Exact token boundaries for the real-decoder integration test."""
chat_template = None
def __call__(self, text, **kwargs):
assert kwargs == {"return_tensors": "pt", "add_special_tokens": False}
ids = torch.tensor([[ord(character) - ord("a") + 1 for character in text]])
return {"input_ids": ids, "attention_mask": torch.ones_like(ids)}
class ReceiverPathTests(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.threads = torch.get_num_threads()
torch.set_num_threads(1)
@classmethod
def tearDownClass(cls):
torch.set_num_threads(cls.threads)
def test_candidate_selection_respects_decoder_order_and_physical_head_counts(self):
counts = {layer: 2 if layer % 2 else 3 for layer in range(20)}
ranking = [{"layer": 17, "head": 1, "mean_effect": 0.9},
{"layer": 19, "head": 1, "mean_effect": -0.8},
{"layer": 18, "head": 2, "mean_effect": 0.7},
{"layer": 2, "head": 2, "mean_effect": -0.6},
{"layer": 3, "head": 0, "mean_effect": 0.6}]
senders, receivers = select_candidates(list(reversed(ranking)), counts, 3, 2)
self.assertEqual(receivers, [(19, 1, "q"), (18, 2, "q")])
self.assertEqual(senders, [(17, 1), (2, 2), (3, 0)])
for sender in senders:
self.assertLess(sender[0], min(receiver[0] for receiver in receivers))
with self.assertRaisesRegex(ValueError, "unavailable"):
select_candidates([*ranking, {"layer": 19, "head": 2, "mean_effect": 5}], counts)
with self.assertRaisesRegex(ValueError, "duplicate"):
select_candidates([*ranking, ranking[0]], counts)
with self.assertRaisesRegex(ValueError, "not enough ranked sender"):
select_candidates(ranking, counts, 4, 2)
def test_aggregation_averages_per_pair_ratios_and_keeps_zero(self):
records = [{"raw_delta": 2, "normalizer": 4, "normalized_effect": 0.5},
{"raw_delta": 0, "normalizer": 1, "normalized_effect": 0.0},
{"raw_delta": 1, "normalizer": 0, "normalized_effect": None}]
self.assertEqual(aggregate_pair_effects(records), 0.25)
self.assertNotEqual(aggregate_pair_effects(records), (2 + 0) / (4 + 1))
self.assertIsNone(aggregate_pair_effects([records[-1]]))
with self.assertRaisesRegex(ValueError, "finite"):
aggregate_pair_effects([{"normalized_effect": math.inf}])
def test_workflow_records_actual_query_paths_and_zero_controls(self):
with torch.random.fork_rng(devices=[]):
torch.manual_seed(194)
model = CircuitModel()
probe = Prober(model)
pairs = [{"id": identifier, "clean_prompt": clean, "donor_prompt": donor,
"clean_answer": "c", "donor_answer": "k"}
for identifier, clean, donor in (("one", "abcd", "efgh"),
("two", "dcba", "hgfe"))]
artifact = measure_receiver_paths(probe, pairs, CharacterTokenizer(), "cpu",
[(0, 0), (1, 0)], [(2, 0, "q")])
self.assertEqual(len(artifact["edges"]), 2)
self.assertEqual(artifact["actual_forward_passes"], 28)
self.assertEqual(len(model.records), 28)
self.assertEqual(len(artifact["controls"]), 4)
self.assertTrue(all(control["passed"] for control in artifact["controls"]))
self.assertTrue(all(abs(control["raw_delta"]) <= 1e-8 for control in artifact["controls"]))
for edge in artifact["edges"]:
self.assertEqual(edge["receiver"], [2, 0, "q"])
self.assertEqual([pair["pair_id"] for pair in edge["per_pair"]], ["one", "two"])
for pair in edge["per_pair"]:
self.assertAlmostEqual(pair["raw_delta"], pair["patched_margin"] - pair["clean_margin"])
self.assertAlmostEqual(pair["normalizer"], pair["donor_margin"] - pair["clean_margin"])
self.assertAlmostEqual(pair["normalized_effect"], pair["raw_delta"] / pair["normalizer"], places=6)
self.assertEqual(edge["mean_normalized_effect"],
aggregate_pair_effects(edge["per_pair"]))
self.assertFalse(probe.adapter._active)
for module in model.modules():
self.assertFalse(module._forward_hooks)
self.assertFalse(module._forward_pre_hooks)
def test_invalid_query_capability_or_direction_fails_before_any_forward(self):
model = CircuitModel()
probe = Prober(model)
pair = {"id": "one", "clean_prompt": "abcd", "donor_prompt": "efgh",
"clean_answer": "c", "donor_answer": "k"}
with self.assertRaisesRegex(ValueError, "precede"):
measure_receiver_paths(probe, [pair], CharacterTokenizer(), "cpu", [(2, 0)], [(1, 0, "q")])
probe.spec.path_qkv.pop(2)
with self.assertRaises(CapabilityError):
measure_receiver_paths(probe, [pair], CharacterTokenizer(), "cpu", [(0, 0)], [(2, 0, "q")])
self.assertFalse(model.records)
if __name__ == "__main__":
unittest.main()
|