Added test case to improve the coverage

Author: amitsrivastava78

keras/src/quantizers/gptq_core_test.py

@@ -1,56 +1,146 @@
 import pytest
+from absl import logging
 
 from keras.src import layers
 from keras.src import models
 from keras.src.quantizers import gptq_core
 from keras.src.quantizers.gptq_config import GPTQConfig
 
-
-def _get_model_no_embedding():
-    """Returns a simple model that lacks an Embedding layer."""
-    return models.Sequential([layers.Dense(10, input_shape=(5,))])
-
-
-def _get_model_no_blocks():
-    """Returns a model with an embedding layer but no subsequent container
-    layers."""
-    return models.Sequential([layers.Embedding(100, 10, input_shape=(5,))])
+VOCAB_SIZE = 100
 
 
 class MockTokenizer:
     """A mock tokenizer that mimics the real API for testing."""
 
     def tokenize(self, text):
-        return [ord(c) for c in "".join(text)]
+        return [ord(c) % VOCAB_SIZE for c in "".join(text)]
 
     def __call__(self, text):
         return self.tokenize(text)
 
 
-architecture_test_cases = [
-    (
-        _get_model_no_embedding(),
-        "Could not automatically find an embedding layer",
-        "no_embedding_layer",
-    ),
-    (
-        _get_model_no_blocks(),
-        "Could not automatically find any transformer-like blocks",
-        "no_transformer_blocks",
-    ),
-]
+class MockEmptyBlock(layers.Layer):
+    """A mock block that contains no quantizable layers."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.ln = layers.LayerNormalization()
+
+    def call(self, inputs):
+        return self.ln(inputs)
+
+
+class MockTransformerBlock(layers.Layer):
+    """A mock transformer block with a quantizable Dense layer."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.dense = layers.Dense(128)
+
+    def call(self, inputs):
+        return self.dense(inputs)
+
+
+def _get_model_with_backbone(
+    has_transformer_layers=True, embedding_name="embedding"
+):
+    """Creates a mock KerasNLP-style model with a backbone."""
+
+    class MockBackbone(layers.Layer):
+        def __init__(self, **kwargs):
+            super().__init__(**kwargs)
+            if has_transformer_layers:
+                self.transformer_layers = [MockTransformerBlock()]
+            setattr(self, embedding_name, layers.Embedding(VOCAB_SIZE, 128))
+
+    class MockModel(models.Model):
+        def __init__(self, **kwargs):
+            super().__init__(**kwargs)
+            self.backbone = MockBackbone()
+
+        def call(self, inputs):
+            return self.backbone(inputs)
+
+    model = MockModel()
+    model.build(input_shape=(None, 10))
+    return model
 
 
 @pytest.mark.requires_trainable_backend
 class TestGPTQCore:
-    def test_get_dataloader_with_empty_dataset(self):
-        """
-        Tests that get_dataloader raises a ValueError for an empty dataset.
-        """
+    def test_get_dataloader_error_scenarios(self):
+        """Tests error cases for get_dataloader."""
         with pytest.raises(ValueError, match="Provided dataset is empty"):
             gptq_core.get_dataloader(
                 tokenizer=MockTokenizer(), seqlen=10, dataset=[], nsamples=10
             )
+        with pytest.raises(
+            TypeError,
+            match="Providing a dataset name as a "
+            "string is not supported. Please pass the "
+            "loaded dataset directly.",
+        ):
+            gptq_core.get_dataloader(
+                tokenizer=MockTokenizer(),
+                seqlen=10,
+                dataset="wikitext2",
+                nsamples=10,
+            )
+
+    def test_apply_gptq_on_multi_block_model(self):
+        """Tests quantization on a model with multiple blocks."""
+        model = models.Sequential(
+            [
+                layers.Embedding(VOCAB_SIZE, 128),
+                MockTransformerBlock(),
+                MockTransformerBlock(),
+            ]
+        )
+        model.build(input_shape=(None, 10))
+        config = GPTQConfig(
+            dataset=["test data"], tokenizer=MockTokenizer(), group_size=32
+        )
+        try:
+            model.quantize("gptq", config=config)
+        except Exception as e:
+            pytest.fail(f"Multi-block quantization failed unexpectedly: {e}")
+
+    def test_apply_gptq_with_empty_block(self, caplog):
+        """Tests that a block with no quantizable layers is skipped
+        correctly."""
+        caplog.set_level(logging.INFO)
+        model = models.Sequential(
+            [layers.Embedding(VOCAB_SIZE, 10), MockEmptyBlock()]
+        )
+        model.build(input_shape=(None, 10))
+        config = GPTQConfig(dataset=["test data"], tokenizer=MockTokenizer())
+        model.quantize("gptq", config=config)
+        assert "No Dense or EinsumDense layers found" in caplog.text
+
+    architecture_test_cases = [
+        (
+            models.Sequential([layers.Dense(10)]),
+            "Could not automatically find an embedding layer",
+            "no_embedding_layer",
+        ),
+        (
+            models.Sequential(
+                [layers.Embedding(VOCAB_SIZE, 10), layers.Dense(10)]
+            ),
+            "Could not automatically find any transformer-like blocks",
+            "no_transformer_blocks",
+        ),
+        (
+            _get_model_with_backbone(has_transformer_layers=False),
+            "backbone does not have a 'transformer_layers' attribute",
+            "backbone_no_layers",
+        ),
+        (
+            _get_model_with_backbone(embedding_name="wrong_name"),
+            "Could not automatically find an embedding layer in the model",
+            "backbone_no_embedding",
+        ),
+    ]
 
     @pytest.mark.parametrize(
         "model, match_message, test_id",
@@ -60,11 +150,11 @@ def test_get_dataloader_with_empty_dataset(self):
     def test_apply_gptq_with_unsupported_architectures(
         self, model, match_message, test_id
     ):
-        """
-        Tests that quantize fails correctly for various unsupported model
-        architectures.
-        """
-        config = GPTQConfig(dataset=["test"], tokenizer=MockTokenizer())
+        """Tests that quantize fails correctly for various unsupported
+        model architectures."""
+        if not model.built:
+            model.build(input_shape=(None, 10))
 
+        config = GPTQConfig(dataset=["test"], tokenizer=MockTokenizer())
         with pytest.raises(ValueError, match=match_message):
             model.quantize("gptq", config=config)