Merge pull request #4 from trustyai-explainability/hf_detector_tests

Further tests for the HF detector class

Author: m-misiura

tests/detectors/huggingface/test_method_get_probabilities.py

@@ -0,0 +1,57 @@
+# third-party imports
+import pytest
+import torch
+from unittest.mock import Mock
+from transformers import PreTrainedTokenizer
+
+# relative imports
+from detectors.huggingface.detector import Detector
+
+
+class TestGetProbabilities:
+    @pytest.fixture
+    def detector(self):
+        detector = Detector.__new__(Detector)
+        detector.tokenizer = Mock(spec=PreTrainedTokenizer)
+        return detector
+
+    def test_normal_case(self, detector):
+        # Setup
+        logprobs = [
+            Mock(values=torch.tensor([[0.0, -1.0]]), indices=torch.tensor([[1, 2]]))
+        ]
+        detector.tokenizer.convert_ids_to_tokens.side_effect = lambda x: (
+            "safe" if x == 1 else "unsafe"
+        )
+        result = detector.get_probabilities(logprobs, "safe", "unsafe")
+        assert isinstance(result, torch.Tensor)
+        assert len(result) == 2
+        assert torch.allclose(result.sum(), torch.tensor(1.0))
+        assert result[0] > result[1]  # Safe token has higher probability
+
+    def test_empty_logprobs(self, detector):
+        result = detector.get_probabilities([], "safe", "unsafe")
+        assert torch.allclose(result, torch.tensor([0.5, 0.5]))
+
+    def test_very_small_probabilities(self, detector):
+        logprobs = [
+            Mock(values=torch.tensor([[-50.0, -50.0]]), indices=torch.tensor([[1, 2]]))
+        ]
+        detector.tokenizer.convert_ids_to_tokens.side_effect = lambda x: (
+            "safe" if x == 1 else "unsafe"
+        )
+        result = detector.get_probabilities(logprobs, "safe", "unsafe")
+        assert torch.allclose(result.sum(), torch.tensor(1.0))
+        assert torch.allclose(result[0], result[1])  # Should be equal probabilities
+
+    def test_case_sensitivity(self, detector):
+        logprobs = [Mock(values=torch.tensor([[0.0]]), indices=torch.tensor([[1]]))]
+        detector.tokenizer.convert_ids_to_tokens.return_value = "SAFE"
+        result = detector.get_probabilities(logprobs, "safe", "unsafe")
+        assert result[0] > result[1]
+
+    def test_invalid_tokens(self, detector):
+        logprobs = [Mock(values=torch.tensor([[0.0]]), indices=torch.tensor([[1]]))]
+        detector.tokenizer.convert_ids_to_tokens.return_value = "invalid"
+        result = detector.get_probabilities(logprobs, "safe", "unsafe")
+        assert torch.allclose(result, torch.tensor([0.5, 0.5]))

tests/test_detector_device.py renamed to tests/detectors/huggingface/test_method_initalize_device.py

@@ -1,4 +1,4 @@
-# global imports
+# third-party imports
 import os
 import pytest
 
@@ -12,7 +12,9 @@ def setup_environment():
     """
     Setup the required environment variable for the model directory.
     """
-    os.environ["MODEL_DIR"] = os.path.join(os.path.dirname(__file__), "dummy_models")
+    current_dir = os.path.dirname(__file__)
+    parent_dir = os.path.dirname(os.path.dirname(current_dir))
+    os.environ["MODEL_DIR"] = os.path.join(parent_dir, "dummy_models")
 
 
 # tests to check the model initialization

tests/test_detector_initalisation.py renamed to tests/detectors/huggingface/test_method_initialize_model.py

@@ -0,0 +1,89 @@
+# third-party imports
+import numpy as np
+import pytest
+import torch
+from unittest.mock import Mock, patch
+from transformers import PreTrainedTokenizer
+
+# relative imports
+from detectors.huggingface.detector import Detector, ContentAnalysisResponse
+
+
+class TestDetectorParseOutput:
+    @pytest.fixture
+    def detector(self):
+        detector = Detector.__new__(Detector)
+        detector.tokenizer = Mock(spec=PreTrainedTokenizer)
+        detector.get_probabilities = Mock()
+        return detector
+
+    @pytest.fixture
+    def mock_output(self):
+        output = Mock()
+        output.sequences = torch.tensor([[1, 2, 3, 4]])
+        output.scores = [torch.randn(1, 5)]
+        return output
+
+    @pytest.fixture
+    def default_params(self):
+        return {
+            "input_len": 2,
+            "nlogprobs": 5,
+            "safe_token": "Safe",
+            "unsafe_token": "Unsafe",
+        }
+
+    def test_parse_output_safe_classification(
+        self, detector, mock_output, default_params
+    ):
+        """Test safe token classification with probabilities"""
+        detector.tokenizer.decode.return_value = "safe"
+        detector.get_probabilities.return_value = torch.tensor([0.7, 0.3])
+
+        label, prob = detector.parse_output(output=mock_output, **default_params)
+
+        assert label == "Safe"
+        assert isinstance(prob, float)
+        np.testing.assert_almost_equal(prob, 0.3, decimal=5)
+
+    def test_parse_output_unsafe_classification(
+        self, detector, mock_output, default_params
+    ):
+        """Test unsafe token classification with probabilities"""
+        detector.tokenizer.decode.return_value = "unsafe"
+        detector.get_probabilities.return_value = torch.tensor([0.3, 0.7])
+
+        label, prob = detector.parse_output(output=mock_output, **default_params)
+
+        assert label == "Unsafe"
+        assert isinstance(prob, float)
+        np.testing.assert_almost_equal(prob, 0.7, decimal=5)
+
+    def test_parse_output_failed_classification(
+        self, detector, mock_output, default_params
+    ):
+        """Test when decoded token doesn't match safe/unsafe"""
+        detector.tokenizer.decode.return_value = "invalid"
+        detector.get_probabilities.return_value = torch.tensor([0.5, 0.5])
+
+        label, prob = detector.parse_output(output=mock_output, **default_params)
+
+        assert label == "failed"
+        assert prob == 0.5
+
+    def test_parse_output_empty_sequence(self, detector, default_params):
+        """Test with empty sequence"""
+        mock_output = Mock()
+        mock_output.sequences = torch.tensor([[]])
+        detector.tokenizer.decode.return_value = ""
+
+        label, prob = detector.parse_output(
+            output=mock_output,
+            input_len=0,
+            nlogprobs=0,
+            safe_token="Safe",
+            unsafe_token="Unsafe",
+        )
+
+        assert label == "failed"
+        assert prob is None

tests/detectors/huggingface/test_method_parse_output.py

@@ -0,0 +1,106 @@
+# third-party imports
+import os
+import pytest
+import torch
+from unittest.mock import Mock, patch
+
+# relative imports
+from detectors.huggingface.detector import Detector, ContentAnalysisResponse
+
+
+class MockGraniteOutput:
+    def __init__(self):
+        self.sequences = torch.tensor([[1, 2, 3, 4]])
+        self.scores = [torch.randn(1, 5)]
+
+
+@pytest.fixture
+def setup_environment():
+    """Setup the required environment variable for the model directory."""
+    current_dir = os.path.dirname(__file__)
+    parent_dir = os.path.dirname(os.path.dirname(current_dir))
+    os.environ["MODEL_DIR"] = os.path.join(parent_dir, "dummy_models")
+
+
+class TestDetector:
+    @pytest.fixture(autouse=True)
+    def setup(self, setup_environment):
+        pass
+
+    @pytest.fixture
+    def detector_instance(self):
+        with patch.dict("os.environ", {"MODEL_DIR": "/dummy/path"}):
+            detector = Detector.__new__(Detector)
+
+            detector.tokenizer = Mock()
+            detector.tokenizer.apply_chat_template = Mock(
+                return_value=torch.tensor([[1, 2, 3]])
+            )
+            detector.tokenizer.decode = Mock(return_value="Yes")
+
+            detector.model = Mock()
+            detector.model.device = torch.device("cpu")
+            detector.model.generate = Mock(return_value=MockGraniteOutput())
+
+            detector.model_name = "causal_lm"
+            detector.is_causal_lm = True
+            detector.cuda_device = None
+            detector.risk_names = ["harm", "bias"]
+
+            return detector
+
+    def validate_results(self, results, input_text, detector):
+        """Helper method to validate the classification results"""
+        assert len(results) == len(detector.risk_names)
+
+        for result in results:
+            expected_fields = [
+                "start",
+                "end",
+                "detection",
+                "detection_type",
+                "score",
+                "sequence_classification",
+                "sequence_probability",
+                "token_classifications",
+                "token_probabilities",
+                "text",
+                "evidences",
+            ]
+
+            for field in expected_fields:
+                assert hasattr(
+                    result, field
+                ), f"Missing '{field}' in ContentAnalysisResponse"
+
+            assert isinstance(result, ContentAnalysisResponse)
+            assert isinstance(result.start, int)
+            assert isinstance(result.end, int)
+            assert isinstance(result.detection, str)
+            assert isinstance(result.detection_type, str)
+            assert isinstance(result.score, float)
+            assert isinstance(result.sequence_classification, str)
+            assert isinstance(result.sequence_probability, float)
+            assert isinstance(result.text, str)
+            assert isinstance(result.evidences, list)
+
+            assert 0 <= result.start <= len(input_text)
+            assert 0 <= result.end <= len(input_text)
+            assert 0.0 <= result.score <= 1.0
+            assert 0.0 <= result.sequence_probability <= 1.0
+            assert result.sequence_classification in detector.risk_names
+
+    def test_process_causal_lm_single_short_input(self, detector_instance):
+        text = "This is a test."
+        results = detector_instance.process_causal_lm(text)
+        self.validate_results(results, text, detector_instance)
+
+    def test_process_causal_lm_single_long_input(self, detector_instance):
+        text = "This is a test." * 1_000
+        results = detector_instance.process_causal_lm(text)
+        self.validate_results(results, text, detector_instance)
+
+    def test_process_causal_lm_single_empty_input(self, detector_instance):
+        text = ""
+        results = detector_instance.process_causal_lm(text)
+        self.validate_results(results, text, detector_instance)

tests/detectors/huggingface/test_method_process_causal_lm.py

@@ -11,7 +11,9 @@ def setup_environment():
     """
     Setup the required environment variable for the model directory.
     """
-    os.environ["MODEL_DIR"] = os.path.join(os.path.dirname(__file__), "dummy_models")
+    current_dir = os.path.dirname(__file__)
+    parent_dir = os.path.dirname(os.path.dirname(current_dir))
+    os.environ["MODEL_DIR"] = os.path.join(parent_dir, "dummy_models")
 
 
 # tests to check the detector output