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