Adding unit test cases

Signed-off-by: Swati Allabadi <[email protected]>

Author: Swati Allabadi

QEfficient/finetune/experimental/core/model.py

@@ -9,13 +9,13 @@
 from abc import ABC, abstractmethod
 from typing import Any, Dict, Optional, Type
 
-import torch
 import torch.nn as nn
 from transformers import AutoTokenizer, BitsAndBytesConfig
 import transformers
+from transformers.utils.logging import get_logger
 
 from QEfficient.finetune.experimental.core.component_registry import registry
-from QEfficient.finetune.experimental.utils.dataset_helper import insert_pad_token
+from QEfficient.finetune.experimental.core.utils.dataset_utils import insert_pad_token
 
 logger = get_logger(__name__)
 
@@ -38,7 +38,6 @@ def create(cls, model_name: str, **model_kwargs: Any) -> "BaseModel":
         if not isinstance(module, nn.Module):
             raise TypeError(f"load_model() must return nn.Module, got {type(module)}")
         obj._model = module
-        obj.add_module("_wrapped_model", module)  # register
         return obj
 
     @abstractmethod
@@ -70,14 +69,14 @@ def forward(self, *args, **kwargs):
     def get_input_embeddings(self):
         if hasattr(self.model, "get_input_embeddings"):
             return self.model.get_input_embeddings()
-        logger.log_rank_zero(f"Model {self.model_name} does not expose input embeddings", logging.WARNING)
+        logger.info(f"Model {self.model_name} does not expose input embeddings", logging.WARNING)
         return None
 
     def resize_token_embeddings(self, new_num_tokens: int) -> None:
         if hasattr(self.model, "resize_token_embeddings"):
             self.model.resize_token_embeddings(new_num_tokens)
         else:
-            logger.log_rank_zero(f"Model {self.model_name} cannot resize token embeddings", logging.WARNING)
+            logger.info(f"Model {self.model_name} cannot resize token embeddings", logging.WARNING)
 
     # optional
     def to(self, *args, **kwargs):
@@ -134,7 +133,7 @@ def configure_model_kwargs(self) -> Dict[str, Any]:
         return extra
 
     def load_model(self) -> nn.Module:
-        logger.log_rank_zero(f"Loading HuggingFace model '{self.model_name}' via {self.auto_class.__name__}")
+        logger.info(f"Loading HuggingFace model '{self.model_name}' via {self.auto_class.__name__}")
 
         return self.auto_class.from_pretrained(
             self.model_name,
@@ -143,7 +142,7 @@ def load_model(self) -> nn.Module:
 
     def load_tokenizer(self) -> AutoTokenizer:
         """Load Hugging Face tokenizer."""
-        logger.log_rank_zero(f"Loading tokenizer '{self.tokenizer_name}'")
+        logger.info(f"Loading tokenizer '{self.tokenizer_name}'")
         tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name)
         insert_pad_token(tokenizer)
         return tokenizer

QEfficient/finetune/experimental/tests/test_model.py

@@ -0,0 +1,143 @@
+import pytest
+import torch
+import torch.nn as nn
+from unittest import mock
+
+import transformers
+from QEfficient.finetune.experimental.core import model
+from QEfficient.finetune.experimental.core.model import BaseModel, HFModel
+
+
+class TestMockModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(2, 2)
+
+    def forward(self, x):
+        return self.linear(x)
+
+
+class TestCustomModel(BaseModel):
+    def __init__(self, model_name):
+        super().__init__(model_name)
+        print("init of custom class")
+
+    def load_model(self) -> nn.Module:
+        return TestMockModel()
+
+    def load_tokenizer(self):
+        return "dummy-tokenizer"
+
+
+# BaseModel tests
+def test_model_property_errors_if_not_created():
+    m = TestCustomModel("dummy")
+    with pytest.raises(RuntimeError):
+        _ = m.model  # must call .create()
+
+
+def test_create_builds_and_registers():
+    breakpoint()
+    m = TestCustomModel.create("dummy")
+    # inner model exists and registered
+    assert "_model" in m._modules
+    assert isinstance(m.model, TestMockModel)
+    # forward works
+    out = m(torch.zeros(1, 2))
+    assert out.shape == (1, 2)
+
+
+def test_tokenizer_lazy_loading():
+    m = TestCustomModel.create("dummy")
+    assert m._tokenizer is None
+    tok = m.tokenizer
+    assert tok == "dummy-tokenizer"
+    assert m._tokenizer == tok
+
+
+def test_to_moves_inner_and_returns_self():
+    m = TestCustomModel.create("dummy")
+    with mock.patch.object(TestMockModel, "to", autospec=True) as mocked_to:
+        ret = m.to("cuda:0")
+    mocked_to.assert_called_once_with(m.model, "cuda:0")
+    assert ret is m
+
+
+def test_train_eval_sync_flags():
+    m = TestCustomModel.create("dummy")
+    m.eval()
+    assert m.training is False
+    assert m.model.training is False
+    m.train()
+    assert m.training is True
+    assert m.model.training is True
+
+
+def test_resize_token_embeddings_and_get_input_embeddings_warn(monkeypatch):
+    m = TestCustomModel.create("dummy")
+
+    # resize_token_embeddings: underlying model lacks the method, should warn and not raise
+    with mock.patch("QEfficient.finetune.experimental.core.model.logger.info") as mocked_log:
+        m.resize_token_embeddings(10)
+        mocked_log.assert_called_once()
+
+    # get_input_embeddings: underlying model lacks method, should warn and return None
+    with mock.patch("QEfficient.finetune.experimental.core.model.logger.info") as mocked_log:
+        assert m.get_input_embeddings() is None
+        mocked_log.assert_called_once()
+
+
+def test_state_dict_contains_inner_params():
+    m = TestCustomModel.create("dummy")
+    sd = m.state_dict()
+    # should contain params from TestMockModel.linear
+    assert any("linear.weight" in k for k in sd)
+    assert any("linear.bias" in k for k in sd)
+
+
+# HFModel tests
+def test_hfmodel_invalid_auto_class_raises():
+    with pytest.raises(ValueError):
+        HFModel.create("hf-name", auto_class_name="AutoDoesNotExist")
+
+
+def test_hfmodel_loads_auto_and_tokenizer(monkeypatch):
+    # fake HF Auto class
+    class FakeAuto(nn.Module):
+        @classmethod
+        def from_pretrained(cls, name, **kwargs):
+            inst = cls()
+            inst.loaded = (name, kwargs)
+            return inst
+
+        def forward(self, x):
+            return x
+
+    fake_tok = mock.Mock()
+
+    # Monkeypatch transformer classes used in HFModel
+    monkeypatch.setattr(
+        "QEfficient.finetune.experimental.core.model.transformers.AutoModelForCausalLM",
+        FakeAuto,
+        raising=False,
+    )
+    monkeypatch.setattr(
+        model,
+        "AutoTokenizer",
+        mock.Mock(from_pretrained=mock.Mock(return_value=fake_tok)),
+    )
+    monkeypatch.setattr(
+        "QEfficient.finetune.experimental.core.model.insert_pad_token",
+        mock.Mock(),
+        raising=False,
+    )
+
+    m = HFModel.create("hf-name")
+    assert isinstance(m.model, FakeAuto)
+
+    # load tokenizer
+    tok = m.load_tokenizer()
+
+    # tokenizer was loaded and pad token inserted
+    model.AutoTokenizer.from_pretrained.assert_called_once_with("hf-name")
+    model.insert_pad_token.assert_called_once_with(fake_tok)