Add IndependentSpeculator w/ unit tests

Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>

Author: fynnsu

src/speculators/models/__init__.py

@@ -1,10 +1,11 @@
 from .eagle import EagleSpeculator, EagleSpeculatorConfig
-from .independent import IndependentSpeculatorConfig
+from .independent import IndependentSpeculator, IndependentSpeculatorConfig
 from .mlp import MLPSpeculatorConfig
 
 __all__ = [
     "EagleSpeculator",
     "EagleSpeculatorConfig",
+    "IndependentSpeculator",
     "IndependentSpeculatorConfig",
     "MLPSpeculatorConfig",
 ]

src/speculators/models/independent.py

@@ -1,6 +1,13 @@
-from transformers import PretrainedConfig
+import os
+from typing import ClassVar, Literal, Optional, Union
+
+import torch
+from pydantic.fields import Field
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
 
 from speculators import SpeculatorModelConfig, SpeculatorsConfig
+from speculators.model import SpeculatorModel
 
 __all__ = ["IndependentSpeculatorConfig"]
 
@@ -16,16 +23,76 @@ def from_pretrained_config(
 
         return cls(**pretrained_dict, speculators_config=speculators_config)
 
-    speculators_model_type: str = "independent"
+    speculators_model_type: Literal["independent"] = "independent"
+    architectures: list[str] = Field(
+        default_factory=lambda: ["LlamaForCausalLM"],
+        description=("List of model architectures that can be used with the model "),
+    )
+    draft_model: str = Field(
+        default="",
+        description=(
+            "The name or path to the draft model to use for the speculator. "
+            "Must be a model that is compatible with the speculator."
+        ),
+    )
+
+
+@SpeculatorModel.register("independent")
+class IndependentSpeculator(SpeculatorModel):
+    config_class: ClassVar[type[IndependentSpeculatorConfig]] = (  # type: ignore[misc]
+        IndependentSpeculatorConfig
+    )
+    _keys_to_ignore_on_load_missing: ClassVar[list[str]] = [  # type: ignore[misc]
+        "verifier*",
+    ]
+    _keys_to_ignore_on_save: ClassVar[list[str]] = [  # type: ignore[assignment,misc]
+        "verifier*",
+    ]
+
+    def __init__(
+        self,
+        config: IndependentSpeculatorConfig,
+        verifier: Optional[Union[str, os.PathLike, PreTrainedModel]] = None,
+        verifier_attachment_mode: Optional[
+            Literal["detached", "full", "train_only"]
+        ] = None,
+    ):
+        if not isinstance(config, IndependentSpeculatorConfig):
+            raise ValueError(
+                "config must be an instance of IndependentSpeculatorConfig, "
+                f"got {type(config)} instead."
+            )
 
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
+        super().__init__(
+            config=config,
+            verifier=verifier,
+            verifier_attachment_mode=verifier_attachment_mode,
+        )
 
-        # ensure we set the model_type to the one from the original config
-        self._model_type = kwargs.get("model_type")
+        self.draft_model: PreTrainedModel = self.resolve_verifier(config.draft_model)
+        self.post_init()
 
-    def to_dict(self):
-        config_dict = super().to_dict()
-        config_dict["model_type"] = self._model_type
-        del config_dict["_model_type"]
-        return config_dict
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[tuple[tuple[torch.FloatTensor]]] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,  # noqa: ARG002
+        return_dict: Optional[bool] = None,
+    ) -> Union[torch.FloatTensor, CausalLMOutputWithPast]:
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+        return self.draft_model(
+            input_ids,
+            attention_mask,
+            position_ids,
+            past_key_values,
+            use_cache,
+            output_attentions,
+            output_hidden_states,
+            return_dict,
+        )

tests/unit/models/test_independent_config.py

@@ -0,0 +1,226 @@
+"""
+Unit tests for the eagle model module in the Speculators library.
+"""
+
+import pytest
+from pydantic import BaseModel, ValidationError
+
+from speculators import (
+    SpeculatorModelConfig,
+    SpeculatorsConfig,
+    VerifierConfig,
+)
+from speculators.models import IndependentSpeculatorConfig
+from speculators.proposals import GreedyTokenProposalConfig
+
+# ===== Fixtures =====
+
+
+@pytest.fixture
+def sample_verifier_config():
+    return VerifierConfig(
+        name_or_path="test/verifier",
+        architectures=["LlamaForCausalLM"],
+    )
+
+
+@pytest.fixture
+def sample_token_proposal_config():
+    return GreedyTokenProposalConfig(
+        speculative_tokens=5,
+        verifier_accept_k=1,
+        accept_tolerance=0.0,
+    )
+
+
+@pytest.fixture
+def sample_speculators_config(sample_token_proposal_config, sample_verifier_config):
+    return SpeculatorsConfig(
+        algorithm="independent",
+        proposal_methods=[sample_token_proposal_config],
+        default_proposal_method="greedy",
+        verifier=sample_verifier_config,
+    )
+
+
+@pytest.fixture
+def independent_config_dict():
+    return {
+        "speculators_model_type": "independent",
+        "architectures": ["LlamaForCausalLM"],
+        "draft_model": "test/draft",
+        "speculators_config": {
+            "algorithm": "independent",
+            "proposal_methods": [
+                {
+                    "proposal_type": "greedy",
+                    "speculative_tokens": 5,
+                    "verifier_accept_k": 1,
+                    "accept_tolerance": 0.0,
+                }
+            ],
+            "default_proposal_method": "greedy",
+            "verifier": {
+                "name_or_path": "test/verifier",
+                "architectures": ["LlamaForCausalLM"],
+                "hidden_size": 768,
+                "intermediate_size": 3072,
+                "vocab_size": 32000,
+                "max_position_embeddings": 2048,
+                "bos_token_id": 1,
+                "eos_token_id": 2,
+            },
+        },
+    }
+
+
+# ===== EagleSpeculatorConfig Tests =====
+
+
+@pytest.mark.smoke
+def test_independent_speculator_config_initialization():
+    """Test default initialization of IndependentSpeculatorConfig."""
+    config = IndependentSpeculatorConfig()
+
+    # Verify Independent-specific defaults
+    assert config.speculators_model_type == "independent"
+    assert config.architectures == ["LlamaForCausalLM"]
+    assert config.draft_model == ""
+
+    # Verify base class defaults
+    assert config.model_type == "speculator_model"
+    assert config.speculators_config is None
+
+
+@pytest.mark.smoke
+def test_independent_speculator_config_custom_initialization(sample_speculators_config):
+    """Test custom initialization of IndependentSpeculatorConfig."""
+    config = IndependentSpeculatorConfig(
+        architectures=["CustomIndependentSpeculator"],
+        draft_model="test/draft",
+        speculators_config=sample_speculators_config,
+    )
+
+    # Verify custom values
+    assert config.speculators_model_type == "independent"
+    assert "CustomIndependentSpeculator" in config.architectures
+    assert config.draft_model == "test/draft"
+    assert config.speculators_config == sample_speculators_config
+
+
+@pytest.mark.smoke
+def test_independent_speculator_config_base_initialization(sample_speculators_config):
+    # Create IndependentSpeculatorConfig with custom values
+    original_config = IndependentSpeculatorConfig(
+        architectures=["CustomIndependentSpeculator"],
+        draft_model="test/draft",
+        speculators_config=sample_speculators_config,
+    )
+
+    # Convert to dict and validate through base class
+    config_dict = original_config.model_dump()
+    recreated_config = SpeculatorModelConfig.model_validate(config_dict)
+
+    # Verify type and values preservation
+    assert isinstance(recreated_config, IndependentSpeculatorConfig)
+    assert recreated_config.speculators_model_type == "independent"
+    assert "CustomIndependentSpeculator" in recreated_config.architectures
+    assert recreated_config.draft_model == "test/draft"
+    assert recreated_config.speculators_config == sample_speculators_config
+
+
+@pytest.mark.regression
+def test_eagle_speculator_config_nested_initialization():
+    class ParentModel(BaseModel):
+        single_config: IndependentSpeculatorConfig
+        config_list: list[IndependentSpeculatorConfig]
+        config_dict: dict[str, IndependentSpeculatorConfig]
+
+    parent = ParentModel(
+        single_config=IndependentSpeculatorConfig(draft_model="test/draft"),
+        config_list=[
+            IndependentSpeculatorConfig(draft_model="test/draft1"),
+            IndependentSpeculatorConfig(draft_model="test/draft2"),
+        ],
+        config_dict={
+            "draft1": IndependentSpeculatorConfig(draft_model="test/draft1"),
+            "draft2": IndependentSpeculatorConfig(draft_model="test/draft2"),
+        },
+    )
+
+    # Verify single config
+    assert isinstance(parent.single_config, IndependentSpeculatorConfig)
+    assert parent.single_config.draft_model == "test/draft"
+
+    # Verify config list
+    assert len(parent.config_list) == 2
+    assert all(isinstance(c, IndependentSpeculatorConfig) for c in parent.config_list)
+    assert parent.config_list[0].draft_model == "test/draft1"
+    assert parent.config_list[1].draft_model == "test/draft2"
+
+    # Verify config dict
+    assert len(parent.config_dict) == 2
+    assert all(
+        isinstance(c, IndependentSpeculatorConfig) for c in parent.config_dict.values()
+    )
+    assert parent.config_dict["draft1"].draft_model == "test/draft1"
+    assert parent.config_dict["draft2"].draft_model == "test/draft2"
+
+
+@pytest.mark.smoke
+def test_independent_speculator_config_invalid_initialization():
+    # Test invalid speculators_model_type
+    with pytest.raises(ValidationError) as exc_info:
+        IndependentSpeculatorConfig(speculators_model_type="invalid")  # type: ignore[arg-type]
+    assert "speculators_model_type" in str(exc_info.value)
+
+    # Test invalid architectures type
+    with pytest.raises(ValidationError) as exc_info:
+        IndependentSpeculatorConfig(architectures="not_a_list")  # type: ignore[arg-type]
+    assert "architectures" in str(exc_info.value)
+
+    # Test invalid draft_model type
+    with pytest.raises(ValidationError) as exc_info:
+        IndependentSpeculatorConfig(draft_model=123)  # type: ignore[arg-type]
+    assert "draft_model" in str(exc_info.value)
+
+
+@pytest.mark.smoke
+def test_independent_speculator_config_auto_registry():
+    registered_classes = SpeculatorModelConfig.registered_classes()
+    class_names = [cls.__name__ for cls in registered_classes]
+
+    # Verify IndependentSpeculatorConfig is registered
+    assert "IndependentSpeculatorConfig" in class_names
+
+    # Verify registry key mapping
+    assert SpeculatorModelConfig.registry is not None
+    assert "independent" in SpeculatorModelConfig.registry
+    assert SpeculatorModelConfig.registry["independent"] == IndependentSpeculatorConfig
+
+
+@pytest.mark.smoke
+def test_independent_speculator_config_marshalling(sample_speculators_config):
+    original_config = IndependentSpeculatorConfig(
+        draft_model="test/draft",
+        speculators_config=sample_speculators_config,
+    )
+
+    # Test model_dump()
+    config_dict = original_config.model_dump()
+    assert isinstance(config_dict, dict)
+    assert config_dict["speculators_model_type"] == "independent"
+    assert config_dict["draft_model"] == "test/draft"
+    assert config_dict["speculators_config"] == sample_speculators_config.model_dump()
+
+    # Test model_validate() on base class
+    recreated_base = SpeculatorModelConfig.model_validate(config_dict)
+    assert isinstance(recreated_base, IndependentSpeculatorConfig)
+    assert recreated_base.draft_model == "test/draft"
+    assert recreated_base.speculators_config == sample_speculators_config
+
+    # Test model_validate() on derived class
+    recreated_derived = IndependentSpeculatorConfig.model_validate(config_dict)
+    assert isinstance(recreated_derived, IndependentSpeculatorConfig)
+    assert recreated_derived.draft_model == "test/draft"
+    assert recreated_derived.speculators_config == sample_speculators_config