test(train): add unit tests, typing, and docstrings for train utilities

src/speculators/train/noise_transforms.py

@@ -1,25 +1,88 @@
+"""Noise transforms for augmenting tensors (e.g. hidden states) during training."""
+
 import torch
 
+__all__ = [
+    "AddGaussianNoise",
+    "AddUniformNoise",
+    "TransformTensors",
+]
+
 
 class TransformTensors:
-    def __init__(self, std=0.05, tensors=("hidden_states",)):
+    """Base class for noise transforms applied to named tensors of a batch.
+
+    Subclasses must override :meth:`transform` to define the noise operation
+    applied to each of the configured tensors.
+    """
+
+    def __init__(
+        self, std: float = 0.05, tensors: tuple[str, ...] = ("hidden_states",)
+    ):
+        """Initialize the transform.
+
+        Args:
+            std: Scale of the noise applied to the configured tensors.
+            tensors: Keys of the batch entries to transform.
+        """
         self.tensors = tensors
         self.std = std
 
-    def __call__(self, data):
+    def __call__(self, data: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        """Apply the transform to each configured tensor of the batch.
+
+        Args:
+            data: Mapping of batch entries to tensors. Entries listed in
+                ``self.tensors`` are replaced with their transformed values.
+
+        Returns:
+            The batch mapping with the configured tensors transformed.
+        """
         for tensor in self.tensors:
             data[tensor] = self.transform(data[tensor])
         return data
 
     def transform(self, tensor: torch.Tensor) -> torch.Tensor:
+        """Apply noise to a single tensor.
+
+        Args:
+            tensor: The tensor to transform.
+
+        Returns:
+            The transformed tensor.
+
+        Raises:
+            NotImplementedError: If the subclass does not implement this method.
+        """
         raise NotImplementedError("Subclasses must implement this method")
 
 
 class AddGaussianNoise(TransformTensors):
+    """Add zero-mean Gaussian noise to the configured tensors."""
+
     def transform(self, tensor: torch.Tensor) -> torch.Tensor:
+        """Add Gaussian noise with standard deviation ``self.std`` to a tensor.
+
+        Args:
+            tensor: The tensor to transform.
+
+        Returns:
+            A new tensor with element-wise Gaussian noise added.
+        """
         return tensor + torch.randn_like(tensor) * self.std
 
 
 class AddUniformNoise(TransformTensors):
+    """Add uniform noise to the configured tensors."""
+
     def transform(self, tensor: torch.Tensor) -> torch.Tensor:
+        """Add uniform noise bounded by ``self.std`` to a tensor.
+
+        Args:
+            tensor: The tensor to transform.
+
+        Returns:
+            A new tensor with element-wise uniform noise added, where each
+            noise sample lies within ``[-self.std, self.std)``.
+        """
         return tensor + 2 * (torch.rand_like(tensor) - 0.5) * self.std

src/speculators/train/vocab_mapping.py

@@ -10,22 +10,28 @@
 
 __all__ = [
     "build_vocab_mappings_from_distribution",
+    "combine_token_frequency_distributions",
+    "get_target_vocab_size",
     "save_token_frequency_distribution",
 ]
 
 
 def save_token_frequency_distribution(
     dataset: HFDataset,
     output_path: Path | str = "./token_freq.pt",
-):
+) -> None:
     """Save token frequency distribution from the dataset.
 
+    Only tokens where ``loss_mask`` is 1 (assistant tokens) are counted. If
+    ``output_path`` already exists, the dataset is skipped and the existing
+    file is left untouched.
+
     Args:
         dataset: HuggingFace dataset with input_ids and loss_mask
         output_path: Path where to save the token frequency distribution
 
     Returns:
-        Path to the saved frequency distribution file
+        None. The frequency distribution is written to ``output_path``.
     """
     path = Path(output_path)
     if path.exists():
@@ -49,12 +55,22 @@ def save_token_frequency_distribution(
 def combine_token_frequency_distributions(
     token_freq_paths: list[str | Path],
     output_path: str | Path,
-):
-    """Combine multiple token frequency distributions into a single file."""
-    token_freq_dicts = [
+) -> None:
+    """Combine multiple token frequency distributions into a single file.
+
+    Args:
+        token_freq_paths: Paths of token frequency files, as written by
+            :func:`save_token_frequency_distribution`.
+        output_path: Path where to save the combined frequency distribution.
+
+    Returns:
+        None. The combined frequency distribution is written to
+        ``output_path``.
+    """
+    token_freq_dicts: list[dict[int, int]] = [
         torch.load(path, weights_only=True) for path in token_freq_paths
     ]
-    combined_token_freq: Counter[str] = Counter()
+    combined_token_freq: Counter[int] = Counter()
     for token_freq_dict in token_freq_dicts:
         combined_token_freq.update(token_freq_dict)
     combined_token_freq_dict = dict(combined_token_freq)
@@ -97,22 +113,39 @@ def build_vocab_mappings_from_distribution(
 
 
 def get_target_vocab_size(
-    target_vocab_size,
-    target_model_path,
+    target_vocab_size: int | None,
+    target_model_path: str | Path | None,
     trust_remote_code: bool = False,
-):
-    has_vocab = target_vocab_size is not None
-    has_model = target_model_path is not None
+) -> int:
+    """Resolve the vocabulary size of the target (verifier) model.
 
-    if has_vocab and has_model:
-        raise ValueError("Cannot specify both target-vocab-size and target-model-path")
+    Exactly one of ``target_vocab_size`` and ``target_model_path`` must be
+    provided. When a model path is given, the vocabulary size is read from
+    the model config, unwrapping ``text_config`` for multimodal models.
 
-    if not has_vocab and not has_model:
-        raise ValueError("Must specify either target-vocab-size or target-model-path")
+    Args:
+        target_vocab_size: Explicit vocabulary size of the target model.
+        target_model_path: Path or model name of the target model to load
+            the config from.
+        trust_remote_code: Whether to trust remote code when loading the
+            model config.
+
+    Returns:
+        The target model's vocabulary size.
+
+    Raises:
+        ValueError: If both or neither of ``target_vocab_size`` and
+            ``target_model_path`` are provided.
+    """
+    if target_vocab_size is not None and target_model_path is not None:
+        raise ValueError("Cannot specify both target-vocab-size and target-model-path")
 
-    if has_vocab:
+    if target_vocab_size is not None:
         return target_vocab_size
 
+    if target_model_path is None:
+        raise ValueError("Must specify either target-vocab-size or target-model-path")
+
     config = AutoConfig.from_pretrained(
         target_model_path,
         trust_remote_code=trust_remote_code,

tests/unit/train/test_noise_transforms.py

@@ -0,0 +1,120 @@
+"""Unit tests for the tensor noise transforms used during training."""
+
+import pytest
+import torch
+
+from speculators.train.noise_transforms import (
+    AddGaussianNoise,
+    AddUniformNoise,
+    TransformTensors,
+)
+
+
+def test_base_transform_raises_not_implemented_error():
+    transform = TransformTensors()
+
+    with pytest.raises(NotImplementedError, match="Subclasses must implement"):
+        transform.transform(torch.zeros(2, 2))
+
+
+def test_base_call_raises_not_implemented_error():
+    transform = TransformTensors()
+
+    with pytest.raises(NotImplementedError, match="Subclasses must implement"):
+        transform({"hidden_states": torch.zeros(2, 2)})
+
+
+def test_add_gaussian_noise_transforms_only_configured_keys(seed):
+    data = {
+        "hidden_states": torch.zeros(4, 8),
+        "labels": torch.arange(4),
+    }
+    original_hidden = data["hidden_states"].clone()
+    original_labels = data["labels"].clone()
+
+    result = AddGaussianNoise(std=0.1)(data)
+
+    assert result is data
+    assert not torch.equal(data["hidden_states"], original_hidden)
+    assert torch.equal(data["labels"], original_labels)
+
+
+def test_add_gaussian_noise_supports_custom_tensor_keys(seed):
+    data = {
+        "hidden_states": torch.zeros(2, 2),
+        "positions": torch.zeros(2, 2),
+    }
+
+    AddGaussianNoise(std=0.1, tensors=("positions",))(data)
+
+    assert torch.equal(data["hidden_states"], torch.zeros(2, 2))
+    assert not torch.equal(data["positions"], torch.zeros(2, 2))
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
+def test_add_gaussian_noise_preserves_shape_dtype_and_device(seed, dtype):
+    tensor = torch.randn(3, 5, 7, dtype=dtype)
+
+    noisy = AddGaussianNoise(std=0.05).transform(tensor)
+
+    assert noisy.shape == tensor.shape
+    assert noisy.dtype == tensor.dtype
+    assert noisy.device == tensor.device
+
+
+def test_add_gaussian_noise_with_zero_std_is_identity(seed):
+    tensor = torch.randn(6, 6)
+
+    assert torch.equal(AddGaussianNoise(std=0.0).transform(tensor), tensor)
+
+
+def test_add_gaussian_noise_matches_requested_std(seed):
+    std = 0.25
+
+    noise = AddGaussianNoise(std=std).transform(torch.zeros(100, 100))
+
+    assert noise.std().item() == pytest.approx(std, rel=0.05)
+
+
+def test_add_uniform_noise_is_strictly_bounded(seed):
+    std = 0.3
+
+    noise = AddUniformNoise(std=std).transform(torch.zeros(10, 10))
+
+    # torch.rand samples lie in [0, 1), so the noise lies in [-std, std).
+    assert (noise >= -std).all()
+    assert (noise < std).all()
+    # The noise should span most of the bounded interval.
+    assert noise.abs().max() > 0.9 * std
+
+
+def test_add_uniform_noise_transforms_only_configured_keys(seed):
+    data = {
+        "hidden_states": torch.zeros(4, 8),
+        "labels": torch.arange(4),
+    }
+    original_hidden = data["hidden_states"].clone()
+    original_labels = data["labels"].clone()
+
+    result = AddUniformNoise(std=0.1)(data)
+
+    assert result is data
+    assert not torch.equal(data["hidden_states"], original_hidden)
+    assert torch.equal(data["labels"], original_labels)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
+def test_add_uniform_noise_preserves_shape_dtype_and_device(seed, dtype):
+    tensor = torch.randn(2, 3, 4, dtype=dtype)
+
+    noisy = AddUniformNoise(std=0.1).transform(tensor)
+
+    assert noisy.shape == tensor.shape
+    assert noisy.dtype == tensor.dtype
+    assert noisy.device == tensor.device
+
+
+def test_add_uniform_noise_with_zero_std_is_identity(seed):
+    tensor = torch.randn(6, 6)
+
+    assert torch.equal(AddUniformNoise(std=0.0).transform(tensor), tensor)