| 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): |
| |
| |
| |
| |
| 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() |
|
|
