distil-trainer 0.1.10__py3-none-any.whl

distil_trainer/evaluation/evaluators.py

@@ -0,0 +1,343 @@
+"""Evaluators for distillation quality assessment."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Callable
+
+import numpy as np
+import torch
+from scipy.stats import pearsonr, spearmanr
+from tqdm import tqdm
+
+from sentence_transformers import SentenceTransformer
+
+logger = logging.getLogger(__name__)
+
+
+class EmbeddingSimilarityEvaluator:
+    """
+    Evaluate sentence embedding quality on STS tasks.
+
+    Computes correlation between model similarity and human scores.
+
+    Example:
+        >>> evaluator = EmbeddingSimilarityEvaluator(
+        ...     sentences1=["Hello", "Hi"],
+        ...     sentences2=["World", "Earth"],
+        ...     scores=[0.5, 0.8],
+        ... )
+        >>> results = evaluator(model)
+    """
+
+    def __init__(
+        self,
+        sentences1: list[str],
+        sentences2: list[str],
+        scores: list[float],
+        batch_size: int = 32,
+        name: str = "sts-eval",
+        main_similarity: str = "cosine",
+        show_progress_bar: bool = True,
+    ):
+        """
+        Initialize the evaluator.
+
+        Args:
+            sentences1: First set of sentences.
+            sentences2: Second set of sentences.
+            scores: Ground truth similarity scores.
+            batch_size: Batch size for encoding.
+            name: Name for logging.
+            main_similarity: Similarity function to use.
+            show_progress_bar: Whether to show progress bar.
+        """
+        self.sentences1 = sentences1
+        self.sentences2 = sentences2
+        self.scores = scores
+        self.batch_size = batch_size
+        self.name = name
+        self.main_similarity = main_similarity
+        self.show_progress_bar = show_progress_bar
+
+    def __call__(self, model: SentenceTransformer) -> dict[str, float]:
+        """
+        Evaluate the model.
+
+        Args:
+            model: Model to evaluate.
+
+        Returns:
+            Dictionary with correlation scores.
+        """
+        logger.info(f"Running {self.name} evaluation")
+
+        # Encode sentences
+        embeddings1 = model.encode(
+            self.sentences1,
+            batch_size=self.batch_size,
+            show_progress_bar=self.show_progress_bar,
+            convert_to_tensor=True,
+        )
+        embeddings2 = model.encode(
+            self.sentences2,
+            batch_size=self.batch_size,
+            show_progress_bar=self.show_progress_bar,
+            convert_to_tensor=True,
+        )
+
+        # Compute similarities
+        if self.main_similarity == "cosine":
+            similarities = torch.nn.functional.cosine_similarity(
+                embeddings1, embeddings2
+            ).cpu().numpy()
+        else:
+            # Euclidean distance (convert to similarity)
+            distances = torch.pairwise_distance(embeddings1, embeddings2).cpu().numpy()
+            similarities = 1 / (1 + distances)
+
+        # Compute correlations
+        pearson = pearsonr(similarities, self.scores)[0]
+        spearman = spearmanr(similarities, self.scores)[0]
+
+        results = {
+            f"{self.name}_pearson_{self.main_similarity}": pearson,
+            f"{self.name}_spearman_{self.main_similarity}": spearman,
+        }
+
+        logger.info(f"{self.name}: Pearson={pearson:.4f}, Spearman={spearman:.4f}")
+
+        return results
+
+
+class MSEEvaluator:
+    """
+    Evaluate MSE between student and teacher embeddings.
+
+    Example:
+        >>> evaluator = MSEEvaluator(
+        ...     source_sentences=["Hello", "World"],
+        ...     target_sentences=["Hello", "World"],
+        ...     teacher_model=teacher,
+        ... )
+        >>> results = evaluator(student_model)
+    """
+
+    def __init__(
+        self,
+        source_sentences: list[str],
+        target_sentences: list[str],
+        teacher_model: SentenceTransformer,
+        batch_size: int = 32,
+        name: str = "mse-eval",
+        show_progress_bar: bool = True,
+    ):
+        """
+        Initialize the evaluator.
+
+        Args:
+            source_sentences: Source sentences.
+            target_sentences: Target sentences (same as source for distillation).
+            teacher_model: Teacher model for comparison.
+            batch_size: Batch size for encoding.
+            name: Name for logging.
+            show_progress_bar: Whether to show progress bar.
+        """
+        self.source_sentences = source_sentences
+        self.target_sentences = target_sentences
+        self.teacher_model = teacher_model
+        self.batch_size = batch_size
+        self.name = name
+        self.show_progress_bar = show_progress_bar
+
+        # Precompute teacher embeddings
+        with torch.no_grad():
+            self.teacher_embeddings = teacher_model.encode(
+                source_sentences,
+                batch_size=batch_size,
+                show_progress_bar=show_progress_bar,
+                convert_to_tensor=True,
+            )
+
+    def __call__(self, model: SentenceTransformer) -> dict[str, float]:
+        """
+        Evaluate the model.
+
+        Args:
+            model: Student model to evaluate.
+
+        Returns:
+            Dictionary with MSE score.
+        """
+        logger.info(f"Running {self.name} evaluation")
+
+        # Encode with student
+        student_embeddings = model.encode(
+            self.target_sentences,
+            batch_size=self.batch_size,
+            show_progress_bar=self.show_progress_bar,
+            convert_to_tensor=True,
+        )
+
+        # Compute MSE
+        mse = torch.nn.functional.mse_loss(
+            student_embeddings, self.teacher_embeddings
+        ).item()
+
+        # Compute cosine similarity
+        cosine_sim = torch.nn.functional.cosine_similarity(
+            student_embeddings, self.teacher_embeddings
+        ).mean().item()
+
+        results = {
+            f"{self.name}_mse": mse,
+            f"{self.name}_cosine_similarity": cosine_sim,
+        }
+
+        logger.info(f"{self.name}: MSE={mse:.6f}, Cosine={cosine_sim:.4f}")
+
+        return results
+
+
+class TranslationEvaluator:
+    """
+    Evaluate multilingual alignment via translation retrieval.
+
+    Checks if source[i] embedding is closest to target[i].
+
+    Example:
+        >>> evaluator = TranslationEvaluator(
+        ...     source_sentences=["Hello", "World"],
+        ...     target_sentences=["Hallo", "Welt"],
+        ... )
+        >>> results = evaluator(model)
+    """
+
+    def __init__(
+        self,
+        source_sentences: list[str],
+        target_sentences: list[str],
+        batch_size: int = 32,
+        name: str = "translation-eval",
+        show_progress_bar: bool = True,
+    ):
+        """
+        Initialize the evaluator.
+
+        Args:
+            source_sentences: Source language sentences.
+            target_sentences: Target language sentences (parallel).
+            batch_size: Batch size for encoding.
+            name: Name for logging.
+            show_progress_bar: Whether to show progress bar.
+        """
+        self.source_sentences = source_sentences
+        self.target_sentences = target_sentences
+        self.batch_size = batch_size
+        self.name = name
+        self.show_progress_bar = show_progress_bar
+
+    def __call__(self, model: SentenceTransformer) -> dict[str, float]:
+        """
+        Evaluate the model.
+
+        Args:
+            model: Model to evaluate.
+
+        Returns:
+            Dictionary with retrieval accuracy scores.
+        """
+        logger.info(f"Running {self.name} evaluation")
+
+        # Encode both sets
+        source_embeddings = model.encode(
+            self.source_sentences,
+            batch_size=self.batch_size,
+            show_progress_bar=self.show_progress_bar,
+            convert_to_tensor=True,
+        )
+        target_embeddings = model.encode(
+            self.target_sentences,
+            batch_size=self.batch_size,
+            show_progress_bar=self.show_progress_bar,
+            convert_to_tensor=True,
+        )
+
+        # Normalize for cosine similarity
+        source_embeddings = torch.nn.functional.normalize(source_embeddings, p=2, dim=1)
+        target_embeddings = torch.nn.functional.normalize(target_embeddings, p=2, dim=1)
+
+        # Compute similarity matrix
+        similarity_matrix = torch.mm(source_embeddings, target_embeddings.t())
+
+        # Get rankings
+        correct_indices = torch.arange(len(self.source_sentences), device=similarity_matrix.device)
+        rankings = (similarity_matrix.argsort(dim=1, descending=True) == correct_indices.unsqueeze(1)).nonzero()[:, 1]
+
+        # Compute metrics
+        acc_at_1 = (rankings == 0).float().mean().item()
+        acc_at_5 = (rankings < 5).float().mean().item()
+        acc_at_10 = (rankings < 10).float().mean().item()
+        mrr = (1.0 / (rankings.float() + 1)).mean().item()
+
+        results = {
+            f"{self.name}_accuracy@1": acc_at_1,
+            f"{self.name}_accuracy@5": acc_at_5,
+            f"{self.name}_accuracy@10": acc_at_10,
+            f"{self.name}_mrr": mrr,
+        }
+
+        logger.info(f"{self.name}: Acc@1={acc_at_1:.4f}, MRR={mrr:.4f}")
+
+        return results
+
+
+class SequentialEvaluator:
+    """
+    Run multiple evaluators sequentially.
+
+    Example:
+        >>> evaluator = SequentialEvaluator([evaluator1, evaluator2])
+        >>> results = evaluator(model)
+    """
+
+    def __init__(
+        self,
+        evaluators: list[Any],
+        main_score_function: Callable[[dict], float] | None = None,
+    ):
+        """
+        Initialize the sequential evaluator.
+
+        Args:
+            evaluators: List of evaluators to run.
+            main_score_function: Function to compute main score from results.
+        """
+        self.evaluators = evaluators
+        self.main_score_function = main_score_function or (
+            lambda x: np.mean([v for v in x.values() if isinstance(v, (int, float))])
+        )
+
+    def __call__(self, model: SentenceTransformer) -> dict[str, float]:
+        """
+        Run all evaluators.
+
+        Args:
+            model: Model to evaluate.
+
+        Returns:
+            Combined dictionary of all evaluation results.
+        """
+        all_results = {}
+
+        for evaluator in self.evaluators:
+            try:
+                results = evaluator(model)
+                all_results.update(results)
+            except Exception as e:
+                logger.warning(f"Evaluator {evaluator} failed: {e}")
+
+        # Compute main score
+        all_results["main_score"] = self.main_score_function(all_results)
+
+        return all_results

distil_trainer/evaluation/metrics.py

@@ -0,0 +1,75 @@
+"""Metrics for evaluating distillation quality."""
+
+from __future__ import annotations
+
+import torch
+from transformers import PreTrainedModel
+
+from sentence_transformers import SentenceTransformer
+
+
+class DistillationMetrics:
+    """Metrics for evaluating distillation quality."""
+
+    @staticmethod
+    def embedding_mse(
+        student_embeddings: torch.Tensor,
+        teacher_embeddings: torch.Tensor,
+    ) -> float:
+        """Mean squared error between embeddings."""
+        return torch.nn.functional.mse_loss(student_embeddings, teacher_embeddings).item()
+
+    @staticmethod
+    def embedding_cosine_similarity(
+        student_embeddings: torch.Tensor,
+        teacher_embeddings: torch.Tensor,
+    ) -> float:
+        """Average cosine similarity between embeddings."""
+        return torch.nn.functional.cosine_similarity(
+            student_embeddings, teacher_embeddings, dim=-1
+        ).mean().item()
+
+    @staticmethod
+    def compression_ratio(
+        student_model: SentenceTransformer | PreTrainedModel,
+        teacher_model: SentenceTransformer | PreTrainedModel,
+    ) -> float:
+        """Parameter count ratio (teacher / student)."""
+        student_params = sum(p.numel() for p in student_model.parameters())
+        teacher_params = sum(p.numel() for p in teacher_model.parameters())
+        return teacher_params / student_params
+
+    @staticmethod
+    def speedup_factor(
+        student_model: SentenceTransformer | PreTrainedModel,
+        teacher_model: SentenceTransformer | PreTrainedModel,
+        input_batch: dict | list[str],
+        num_runs: int = 10,
+    ) -> float:
+        """Inference speedup of student over teacher."""
+        import time
+
+        # Warm up
+        if isinstance(student_model, SentenceTransformer):
+            student_model.encode(input_batch[:1])
+            teacher_model.encode(input_batch[:1])
+
+        # Time student
+        start = time.time()
+        for _ in range(num_runs):
+            if isinstance(student_model, SentenceTransformer):
+                student_model.encode(input_batch)
+            else:
+                student_model(**input_batch)
+        student_time = time.time() - start
+
+        # Time teacher
+        start = time.time()
+        for _ in range(num_runs):
+            if isinstance(teacher_model, SentenceTransformer):
+                teacher_model.encode(input_batch)
+            else:
+                teacher_model(**input_batch)
+        teacher_time = time.time() - start
+
+        return teacher_time / student_time

distil_trainer/models/__init__.py

@@ -0,0 +1,5 @@
+"""Models module for distil_trainer."""
+
+from distil_trainer.models.layers import DenseProjection
+
+__all__ = ["DenseProjection"]

distil_trainer/models/layers.py

@@ -0,0 +1,115 @@
+"""Custom layers for distillation models."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+
+class DenseProjection(nn.Module):
+    """
+    Dense projection layer for dimension reduction.
+
+    Used to project teacher embeddings to student dimension via PCA weights.
+
+    Example:
+        >>> projection = DenseProjection(in_features=768, out_features=256)
+        >>> reduced = projection(embeddings)
+    """
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = False,
+        weights: torch.Tensor | None = None,
+    ):
+        """
+        Initialize the projection layer.
+
+        Args:
+            in_features: Input dimension (teacher).
+            out_features: Output dimension (student).
+            bias: Whether to include bias term.
+            weights: Optional initial weights (e.g., from PCA).
+        """
+        super().__init__()
+
+        self.in_features = in_features
+        self.out_features = out_features
+
+        self.linear = nn.Linear(in_features, out_features, bias=bias)
+
+        if weights is not None:
+            with torch.no_grad():
+                self.linear.weight.copy_(weights)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply projection."""
+        return self.linear(x)
+
+
+class PoolingLayer(nn.Module):
+    """
+    Pooling layer for converting token embeddings to sentence embeddings.
+
+    Supports multiple pooling strategies.
+
+    Example:
+        >>> pooler = PoolingLayer(pooling_mode="mean")
+        >>> sentence_embedding = pooler(token_embeddings, attention_mask)
+    """
+
+    def __init__(
+        self,
+        pooling_mode: str = "mean",
+    ):
+        """
+        Initialize the pooling layer.
+
+        Args:
+            pooling_mode: One of "mean", "cls", "max", "weighted_mean".
+        """
+        super().__init__()
+        self.pooling_mode = pooling_mode
+
+    def forward(
+        self,
+        token_embeddings: torch.Tensor,
+        attention_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Apply pooling to token embeddings.
+
+        Args:
+            token_embeddings: Token embeddings [batch, seq_len, dim].
+            attention_mask: Attention mask [batch, seq_len].
+
+        Returns:
+            Pooled embeddings [batch, dim].
+        """
+        if self.pooling_mode == "cls":
+            return token_embeddings[:, 0]
+
+        elif self.pooling_mode == "max":
+            # Mask out padding tokens
+            token_embeddings = token_embeddings.masked_fill(
+                ~attention_mask.unsqueeze(-1).bool(), float("-inf")
+            )
+            return token_embeddings.max(dim=1).values
+
+        elif self.pooling_mode == "weighted_mean":
+            # Position-weighted mean
+            weights = torch.arange(
+                1, token_embeddings.size(1) + 1,
+                device=token_embeddings.device,
+                dtype=token_embeddings.dtype,
+            )
+            weights = weights.unsqueeze(0).unsqueeze(-1) * attention_mask.unsqueeze(-1)
+            return (token_embeddings * weights).sum(dim=1) / weights.sum(dim=1)
+
+        else:  # mean pooling
+            input_mask_expanded = attention_mask.unsqueeze(-1).float()
+            sum_embeddings = (token_embeddings * input_mask_expanded).sum(dim=1)
+            sum_mask = input_mask_expanded.sum(dim=1).clamp(min=1e-9)
+            return sum_embeddings / sum_mask

distil_trainer/pruning/__init__.py

@@ -0,0 +1,13 @@
+"""Pruning module for model compression."""
+
+from distil_trainer.pruning.depth_pruning import DepthPruner
+from distil_trainer.pruning.width_pruning import WidthPruner
+from distil_trainer.pruning.combined_pruning import CombinedPruner
+from distil_trainer.pruning.importance import ImportanceEstimator
+
+__all__ = [
+    "DepthPruner",
+    "WidthPruner",
+    "CombinedPruner",
+    "ImportanceEstimator",
+]

distil_trainer/pruning/combined_pruning.py

@@ -0,0 +1,122 @@
+"""Combined depth and width pruning."""
+
+from __future__ import annotations
+
+import logging
+from typing import Literal
+
+from transformers import PreTrainedModel
+
+from sentence_transformers import SentenceTransformer
+
+from distil_trainer.core.config import CombinedPruningConfig, LayerReductionConfig, WidthPruningConfig
+from distil_trainer.pruning.depth_pruning import DepthPruner
+from distil_trainer.pruning.width_pruning import WidthPruner
+
+logger = logging.getLogger(__name__)
+
+
+class CombinedPruner:
+    """
+    Combines depth and width pruning for maximum compression.
+
+    Example:
+        >>> config = CombinedPruningConfig(
+        ...     depth_config=LayerReductionConfig(num_layers_to_keep=8),
+        ...     width_config=WidthPruningConfig(target_hidden_size=3072),
+        ...     pruning_order="depth_first"
+        ... )
+        >>> pruner = CombinedPruner(model)
+        >>> pruned_model = pruner.prune(config)
+    """
+
+    def __init__(
+        self,
+        model: SentenceTransformer | PreTrainedModel,
+        calibration_data: list[str] | None = None,
+    ):
+        """
+        Initialize the CombinedPruner.
+
+        Args:
+            model: The model to prune.
+            calibration_data: Optional calibration data for importance estimation.
+        """
+        self.model = model
+        self.calibration_data = calibration_data
+
+    def prune(
+        self,
+        config: CombinedPruningConfig,
+        calibration_data: list[str] | None = None,
+    ) -> SentenceTransformer | PreTrainedModel:
+        """
+        Apply combined pruning based on configuration.
+
+        Args:
+            config: Combined pruning configuration.
+            calibration_data: Optional calibration data.
+
+        Returns:
+            New model with both depth and width reduced.
+        """
+        data = calibration_data or self.calibration_data
+        model = self.model
+
+        if config.pruning_order == "depth_first":
+            model = self._apply_depth_pruning(model, config.depth_config, data)
+            model = self._apply_width_pruning(model, config.width_config, data)
+        elif config.pruning_order == "width_first":
+            model = self._apply_width_pruning(model, config.width_config, data)
+            model = self._apply_depth_pruning(model, config.depth_config, data)
+        elif config.pruning_order == "interleaved":
+            # Apply in iterations, alternating
+            for i in range(config.num_iterations):
+                logger.info(f"Pruning iteration {i + 1}/{config.num_iterations}")
+                if i % 2 == 0:
+                    model = self._apply_depth_pruning(model, config.depth_config, data)
+                else:
+                    model = self._apply_width_pruning(model, config.width_config, data)
+        else:
+            raise ValueError(f"Unknown pruning order: {config.pruning_order}")
+
+        # Log final statistics
+        original_params = sum(p.numel() for p in self.model.parameters())
+        final_params = sum(p.numel() for p in model.parameters())
+        reduction = 1 - (final_params / original_params)
+        logger.info(f"Total parameter reduction: {original_params:,} -> {final_params:,} ({reduction:.1%})")
+
+        return model
+
+    def _apply_depth_pruning(
+        self,
+        model: SentenceTransformer | PreTrainedModel,
+        config: LayerReductionConfig | None,
+        calibration_data: list[str] | None,
+    ) -> SentenceTransformer | PreTrainedModel:
+        """Apply depth pruning if configured."""
+        if config is None:
+            return model
+
+        logger.info("Applying depth pruning...")
+        pruner = DepthPruner(model, calibration_data)
+        return pruner.prune(
+            layers_to_keep=config.layers_to_keep,
+            num_layers_to_keep=config.num_layers_to_keep,
+            layers_to_drop=config.layers_to_drop,
+            layer_selection=config.layer_selection,
+        )
+
+    def _apply_width_pruning(
+        self,
+        model: SentenceTransformer | PreTrainedModel,
+        config: WidthPruningConfig | None,
+        calibration_data: list[str] | None,
+    ) -> SentenceTransformer | PreTrainedModel:
+        """Apply width pruning if configured."""
+        if config is None:
+            return model
+
+        logger.info("Applying width pruning...")
+        pruner = WidthPruner(model, calibration_data)
+        return pruner.prune(config, calibration_data)