import unittest import numpy as np from baselines import ( BaselineDetectorName, BaselineExample, HallucinatedSpan, LabelCounts, LogisticConfig, MethodScoreAggregate, TokenEvalMetadata, TokenEvalRow, TokenPredictionRow, TokenScoreDetector, TokenScoreResult, TrainableTokenScoreDetector, ) from baselines.structs import BaselineDetectorName as StructBaselineDetectorName from baselines.structs import BaselineExample as StructBaselineExample from baselines.structs import HallucinatedSpan as StructHallucinatedSpan from baselines.structs import LabelCounts as StructLabelCounts from baselines.structs import LogisticConfig as StructLogisticConfig from baselines.structs import TokenEvalMetadata as StructTokenEvalMetadata from baselines.structs import TokenEvalRow as StructTokenEvalRow from baselines.structs import TokenPredictionRow as StructTokenPredictionRow from baselines.structs import TokenScoreResult as StructTokenScoreResult from baselines.helpers import prompt_tokenizer_kwargs, validate_attention_tensor class BaselineInterfaceTests(unittest.TestCase): def test_attention_validation_keeps_float32_precision(self): attention = np.array([[[[1.0, 0.0], [0.25, 0.75]]]], dtype=np.float16) validated = validate_attention_tensor(attention, [0, 1]) self.assertEqual(validated.dtype, np.float32) np.testing.assert_allclose(validated, attention.astype(np.float32)) def test_baseline_detector_name_values(self): self.assertEqual(BaselineDetectorName.EPR.value, "epr") self.assertEqual(BaselineDetectorName.LOOKBACK.value, "lookback") self.assertEqual(BaselineDetectorName.ATTENTION_ENTROPY.value, "attention_entropy") self.assertEqual(BaselineDetectorName.HIDDEN_STATE_PROBE.value, "hidden_state_probe") self.assertEqual(BaselineDetectorName.MVA_FEATURE_LOGISTIC.value, "mva_feature_logistic_baseline") self.assertEqual(BaselineDetectorName.MULTI_VIEW_ATTENTION_TRANSFORMER_CRF.value, "multi_view_attention_transformer_crf") def test_baseline_example_constructs_with_prompt_and_answer(self): example = BaselineExample( prompt="What is the capital of France?", answer="Paris.", ) self.assertEqual(example.prompt, "What is the capital of France?") self.assertEqual(example.answer, "Paris.") self.assertIsNone(example.labels) self.assertEqual(example.metadata, {}) def test_baseline_example_rejects_non_binary_labels(self): with self.assertRaises(ValueError): BaselineExample(prompt="q", answer="a", labels=[0, 2]) with self.assertRaises(ValueError): BaselineExample(prompt="q", answer="a", labels=[0.9, 1.1]) def test_token_score_result_accepts_matching_lengths(self): result = TokenScoreResult( tokens=["Paris", "."], scores=[0.1, 0.0], method="example", ) self.assertEqual(result.tokens, ["Paris", "."]) self.assertEqual(result.scores, [0.1, 0.0]) self.assertEqual(result.method, "example") self.assertEqual(result.metadata, {}) def test_token_score_result_rejects_mismatched_lengths(self): with self.assertRaises(ValueError): TokenScoreResult( tokens=["Paris", "."], scores=[0.1], method="example", ) def test_token_eval_row_converts_span_dicts_to_dataclasses(self): row = TokenEvalRow.from_json_dict( { "id": "ex1", "prompt": "What is the capital of France?", "answer": "The capital is Canberra.", "hallucinated_spans": [{"text": "Canberra."}], } ) self.assertEqual(row.id, "ex1") self.assertIsInstance(row.hallucinated_spans[0], HallucinatedSpan) self.assertEqual(row.hallucinated_spans[0].start, 15) self.assertEqual(row.hallucinated_spans[0].end, 24) def test_span_with_bounds_rejects_mismatched_text(self): with self.assertRaises(ValueError): HallucinatedSpan.from_json_dict( {"start": 15, "end": 24, "text": "Paris."}, "The capital is Canberra.", ) def test_empty_hallucinated_spans_still_mean_labels_provided(self): row = TokenEvalRow.from_json_dict( { "id": "clean", "prompt": "What is the capital of France?", "answer": "The capital is Paris.", "hallucinated_spans": [], } ) self.assertEqual(row.hallucinated_spans, []) self.assertTrue(row.span_labels_provided) self.assertTrue(row.has_label_signal()) def test_null_hallucinated_spans_do_not_mean_labels_provided(self): row = TokenEvalRow.from_json_dict( { "id": "unlabelled", "prompt": "q", "answer": "a", "hallucinated_spans": None, } ) self.assertEqual(row.hallucinated_spans, []) self.assertFalse(row.span_labels_provided) self.assertFalse(row.has_label_signal()) def test_unlabelled_eval_row_round_trip_stays_unlabelled(self): row = TokenEvalRow(prompt="q", answer="a") payload = row.to_json_dict() reloaded = TokenEvalRow.from_json_dict(payload) self.assertNotIn("hallucinated_spans", payload) self.assertFalse(reloaded.span_labels_provided) self.assertFalse(reloaded.has_label_signal()) def test_explicit_empty_span_labels_round_trip(self): row = TokenEvalRow(prompt="q", answer="a", hallucinated_spans=[]) payload = row.to_json_dict() reloaded = TokenEvalRow.from_json_dict(payload) self.assertEqual(payload["hallucinated_spans"], []) self.assertTrue(reloaded.span_labels_provided) self.assertTrue(reloaded.has_label_signal()) def test_token_eval_row_round_trips_to_json_dict(self): row = TokenEvalRow( id="ex1", prompt="q", answer="abc", labels=[0, 1], hallucinated_spans=[HallucinatedSpan(start=1, end=3, text="bc")], metadata=TokenEvalMetadata( source="unit", extra={"context": "ctx"}, generated_tokens=["a", "bc"], generated_token_ids=[10, 20], generated_logprobs=[-0.1, -0.2], token_label_reasons=["ok", "bad"], original_label_counts=LabelCounts(grounded=1, hallucinated=1), ), ) self.assertEqual( row.to_json_dict(), { "id": "ex1", "prompt": "q", "answer": "abc", "labels": [0, 1], "hallucinated_spans": [{"start": 1, "end": 3, "text": "bc"}], "metadata": { "source": "unit", "extra": {"context": "ctx"}, "generated_tokens": ["a", "bc"], "generated_token_ids": [10, 20], "generated_logprobs": [-0.1, -0.2], "token_label_reasons": ["ok", "bad"], "original_label_counts": {"G": 1, "H": 1}, }, }, ) def test_token_eval_row_rejects_non_binary_labels(self): with self.assertRaises(ValueError): TokenEvalRow(prompt="q", answer="a", labels=[0, -1]) with self.assertRaises(ValueError): TokenEvalRow(prompt="q", answer="a", labels=[0.9, 1.1]) def test_token_prediction_row_round_trips_to_json_dict(self): row = TokenPredictionRow.from_json_dict( { "index": 0, "id": "ex1", "method": "example_method", "tokens": ["Paris", "."], "scores": [0.1, 0.2], "labels": [0, 1], "seconds": 0.5, "labels_aligned": True, "metadata": {"detail": "kept"}, "prompt": "What is the capital of France?", "answer": "The capital is Canberra.", "hallucinated_spans": [{"text": "Canberra."}], "example_metadata": { "source": "truthfulqa", "reference_answer": "Paris.", }, } ) self.assertEqual(row.method, "example_method") self.assertEqual(row.tokens, ["Paris", "."]) self.assertEqual(row.scores, [0.1, 0.2]) self.assertEqual(row.labels, [0, 1]) self.assertEqual(row.metadata, {"detail": "kept"}) self.assertEqual(row.prompt, "What is the capital of France?") self.assertEqual(row.answer, "The capital is Canberra.") self.assertEqual(row.hallucinated_spans[0].text, "Canberra.") self.assertEqual(row.example_metadata.reference_answer, "Paris.") self.assertEqual( row.to_json_dict(), { "index": 0, "id": "ex1", "method": "example_method", "tokens": ["Paris", "."], "scores": [0.1, 0.2], "labels": [0, 1], "seconds": 0.5, "labels_aligned": True, "metadata": {"detail": "kept"}, "prompt": "What is the capital of France?", "answer": "The capital is Canberra.", "hallucinated_spans": [{"start": 15, "end": 24, "text": "Canberra."}], "example_metadata": { "source": "truthfulqa", "reference_answer": "Paris.", }, }, ) def test_token_prediction_row_rejects_mismatched_lengths(self): with self.assertRaises(ValueError): TokenPredictionRow( index=0, id=None, method="bad", tokens=["Paris"], scores=[0.1, 0.2], labels=None, seconds=0.0, labels_aligned=False, ) def test_token_prediction_row_rejects_false_aligned_metadata(self): with self.assertRaises(ValueError): TokenPredictionRow( index=0, id=None, method="bad", tokens=["Paris"], scores=[0.1], labels=None, seconds=0.0, labels_aligned=True, ) with self.assertRaises(ValueError): TokenPredictionRow( index=0, id=None, method="bad", tokens=["Paris", "."], scores=[0.1, 0.2], labels=[1], seconds=0.0, labels_aligned=True, ) def test_token_prediction_row_preserves_uncertain_labels(self): row = TokenPredictionRow( index=0, id=None, method="example", tokens=["Paris"], scores=[0.1], labels=[2], seconds=0.0, labels_aligned=True, ) self.assertEqual(row.labels, [2]) self.assertEqual(row.to_json_dict()["labels"], [2]) def test_token_prediction_schema_v2_preserves_auditable_masks(self): row = TokenPredictionRow( index=3, id="row", method="example", tokens=["a", "b", "c"], token_ids=[10, 20, 30], token_indices=[0, 1, 2], scores=[0.0, 0.5, 0.0], score_valid_mask=[True, True, False], comparison_mask=[False, True, False], labels=[1, 0, 2], seconds=0.0, labels_aligned=True, ) payload = row.to_json_dict() self.assertEqual(payload["prediction_schema_version"], 2) self.assertEqual(payload["token_ids"], [10, 20, 30]) self.assertEqual(payload["comparison_mask"], [False, True, False]) self.assertEqual(TokenPredictionRow.from_json_dict(payload), row) def test_comparison_mask_cannot_include_invalid_score(self): with self.assertRaises(ValueError): TokenPredictionRow( index=0, id=None, method="bad", tokens=["a"], scores=[0.0], labels=[0], seconds=0.0, labels_aligned=True, score_valid_mask=[False], comparison_mask=[True], ) def test_token_prediction_row_rejects_invalid_labels(self): with self.assertRaises(ValueError): TokenPredictionRow( index=0, id=None, method="bad", tokens=["Paris"], scores=[0.1], labels=[0.9], seconds=0.0, labels_aligned=True, ) def test_token_prediction_row_rejects_nonfinite_scores(self): with self.assertRaises(ValueError): TokenPredictionRow( index=0, id=None, method="bad", tokens=["Paris"], scores=[float("nan")], labels=None, seconds=0.0, labels_aligned=False, ) def test_dataclasses_are_available_from_structs_module(self): self.assertIs(BaselineDetectorName, StructBaselineDetectorName) self.assertIs(BaselineExample, StructBaselineExample) self.assertIs(HallucinatedSpan, StructHallucinatedSpan) self.assertIs(LabelCounts, StructLabelCounts) self.assertIs(LogisticConfig, StructLogisticConfig) self.assertIs(TokenEvalMetadata, StructTokenEvalMetadata) self.assertIs(TokenEvalRow, StructTokenEvalRow) self.assertIs(TokenPredictionRow, StructTokenPredictionRow) self.assertIs(TokenScoreResult, StructTokenScoreResult) def test_token_score_detector_is_abstract(self): with self.assertRaises(TypeError): TokenScoreDetector() def test_trainable_token_score_detector_is_abstract(self): with self.assertRaises(TypeError): TrainableTokenScoreDetector() def test_method_score_aggregate_keeps_labelled_score_subset(self): aggregate = MethodScoreAggregate(method="unit") aggregate.extend([0.1, 0.9], [0, 1]) aggregate.extend([0.2, 0.3], None) self.assertEqual(aggregate.scores, [0.1, 0.9, 0.2, 0.3]) self.assertEqual(aggregate.labels_or_none(), [0, 1]) self.assertEqual(aggregate.labelled_scores_or_none(), [0.1, 0.9]) with self.assertRaises(ValueError): aggregate.extend([0.4, 0.5], [1]) self.assertEqual(aggregate.scores, [0.1, 0.9, 0.2, 0.3]) self.assertEqual(aggregate.labels_or_none(), [0, 1]) self.assertEqual(aggregate.labelled_scores_or_none(), [0.1, 0.9]) aggregate.extend([0.4, 0.8, 0.6], [2, 1, 0]) self.assertEqual(aggregate.scores, [0.1, 0.9, 0.2, 0.3, 0.4, 0.8, 0.6]) self.assertEqual(aggregate.labels_or_none(), [0, 1, 1, 0]) self.assertEqual( aggregate.labelled_scores_or_none(), [0.1, 0.9, 0.8, 0.6], ) with self.assertRaises(ValueError): aggregate.extend([0.4], [0.9]) with self.assertRaises(ValueError): aggregate.extend([float("inf")], None) self.assertEqual( aggregate.scores, [0.1, 0.9, 0.2, 0.3, 0.4, 0.8, 0.6], ) def test_prompt_tokenizer_kwargs_adds_truncation_only_when_requested(self): self.assertEqual( prompt_tokenizer_kwargs(add_special_tokens=True, max_length=None), {"add_special_tokens": True, "return_tensors": "pt"}, ) self.assertEqual( prompt_tokenizer_kwargs(add_special_tokens=True, max_length=2048), { "add_special_tokens": True, "return_tensors": "pt", "truncation": True, "max_length": 2048, }, ) if __name__ == "__main__": unittest.main()