File size: 6,488 Bytes
d74d56c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
import json
import sys
import types
import unittest
from unittest.mock import MagicMock, patch

from owmi.analysis import calibration_scores
from owmi.backends import HFBackend
from owmi.benchmarks.base import BenchmarkExample
from owmi.benchmarks.evaluate import (
    characterization_scores,
    evaluate_probe_output,
    reconstruction_scores,
    semantic_similarity,
)
from owmi.benchmarks.runner import BenchmarkRunner
from owmi.benchmarks.suite import (
    SuiteConfig,
    build_job_manifest,
    build_manifest,
)


class RevisionRoundTripTests(unittest.TestCase):
    def _config(self):
        return SuiteConfig.from_dict({
            'suite_name': 'revision-test',
            'models': [{
                'name': 'org/model',
                'revision': '0123456789abcdef0123456789abcdef01234567',
                'load_in_4bit': False,
            }],
            'benchmarks': [{
                'name': 'toy', 'schema': 'short_answer', 'dataset_name': 'unused',
            }],
            'objects': [{'kind': 'residual_stream', 'layer_index': 1}],
            'interventions': [{'kind': 'zero', 'mode': 'zero'}],
            'probes': [{'task': 'detection'}],
        })

    def test_revision_round_trips_from_suite_through_hf_loading(self):
        example = BenchmarkExample('toy', 'one', 'short_answer', 'q', {'answer': 'a'})
        revision = self._config().models[0].revision
        with patch('owmi.benchmarks.suite._load_examples_for_spec', return_value=[example]):
            expanded = build_manifest(self._config()).rows[0].to_dict()
            compact = build_job_manifest(self._config())[0].to_dict()
        expanded = json.loads(json.dumps(expanded))
        compact = json.loads(json.dumps(compact))
        self.assertEqual(expanded['model']['revision'], revision)
        self.assertEqual(compact['model']['revision'], revision)

        config = BenchmarkRunner(expanded)._build_experiment_config()
        self.assertEqual(config.revision, revision)
        tokenizer = MagicMock()
        tokenizer.pad_token = '<pad>'
        model = MagicMock()
        auto_tokenizer = MagicMock()
        auto_tokenizer.from_pretrained.return_value = tokenizer
        auto_model = MagicMock()
        auto_model.from_pretrained.return_value = model
        transformers = types.ModuleType('transformers')
        transformers.AutoTokenizer = auto_tokenizer
        transformers.AutoModelForCausalLM = auto_model
        with patch.dict(sys.modules, {'transformers': transformers}):
            HFBackend(config.model_name, config).load()
        self.assertEqual(auto_tokenizer.from_pretrained.call_args.kwargs['revision'], revision)
        self.assertEqual(auto_model.from_pretrained.call_args.kwargs['revision'], revision)


class ConfidenceFamilyTests(unittest.TestCase):
    def test_selective_prediction_auroc_uses_confidence_for_correctness(self):
        scores = calibration_scores(
            confidence=[0.1, 0.9, 0.2, 0.8],
            correctness=[0.0, 1.0, 0.0, 1.0],
        )
        self.assertEqual(scores['selective_prediction_auroc'], 1.0)
        self.assertEqual(scores['selective_auroc'], 1.0)
        tied = calibration_scores([0.5, 0.5], [0.0, 1.0])
        self.assertEqual(tied['selective_prediction_auroc'], 0.5)


class ProbeScoringTests(unittest.TestCase):
    def test_characterization_is_closed_set_and_reports_macro_f1(self):
        metrics = characterization_scores(
            ['zeroing', 'noise', 'masked', 'noise happened'],
            ['zero', 'noise', 'zero', 'noise'],
            ('zero', 'noise'),
        )
        self.assertEqual(metrics['category_accuracy'], 0.75)
        self.assertAlmostEqual(metrics['macro_f1'], 5.0 / 6.0)
        example = BenchmarkExample('toy', 'one', 'short_answer', 'q', {'answer': 'a'})
        self.assertEqual(
            evaluate_probe_output(
                '{"change_type": "noise happened"}', 'characterization', example,
                {'mode': 'noise', 'categories': ['zero', 'noise']},
            ),
            0.0,
        )

    @patch('owmi.benchmarks.evaluate.semantic_similarity', return_value=(0.95, 'test-embedding'))
    def test_reconstruction_combines_exact_category_and_embedding(self, _similarity):
        categories = ('animal', 'vehicle')
        gamed = reconstruction_scores(
            'A domesticated canine', 'A household dog',
            predicted_category='vehicle', target_category='animal', categories=categories,
        )
        self.assertGreaterEqual(gamed['embedding_similarity'], gamed['similarity_threshold'])
        self.assertEqual(gamed['category_match'], 0.0)
        self.assertEqual(gamed['combined_match'], 0.0)

        valid_paraphrase = reconstruction_scores(
            'A domesticated canine', 'A household dog',
            predicted_category='animal', target_category='animal', categories=categories,
        )
        self.assertEqual(valid_paraphrase['combined_match'], 1.0)
        exact = reconstruction_scores(
            'A household dog', 'the household dog',
            predicted_category=None, target_category='animal', categories=categories,
        )
        self.assertEqual(exact['exact_match'], 1.0)
        self.assertEqual(exact['combined_match'], 1.0)

    def test_reconstruction_similarity_has_dependency_free_tfidf_fallback(self):
        with patch('owmi.benchmarks.evaluate._SENTENCE_TRANSFORMER_UNAVAILABLE', True):
            similarity, backend = semantic_similarity('masked household dog', 'household dog')
        self.assertGreater(similarity, 0.0)
        self.assertEqual(backend, 'tfidf-word-1-2gram-cosine')

    def test_embedding_similarity_ranks_known_similar_above_known_dissimilar(self):
        # Item #19: the embedding-similarity metric must rank a known-similar
        # pair above a known-dissimilar pair, under whichever backend is
        # actually available in this environment (sentence-transformers if
        # cached locally, otherwise the dependency-free TF-IDF fallback).
        anchor = 'a domesticated canine that barks and wags its tail'
        similar, _ = semantic_similarity(anchor, 'a household dog that barks and wags its tail')
        dissimilar, _ = semantic_similarity(anchor, 'quarterly tax filing deadlines for small businesses')
        self.assertGreater(similar, dissimilar)
        identical, _ = semantic_similarity(anchor, anchor)
        self.assertGreaterEqual(identical, similar)


if __name__ == '__main__':
    unittest.main()