bead 0.1.0__py3-none-any.whl

bead/active_learning/strategies.py

@@ -0,0 +1,391 @@
+"""Sampling strategies for active learning.
+
+This module implements various uncertainty quantification methods for
+active learning item selection, including entropy, margin, and least
+confidence sampling, plus a random baseline.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Literal
+
+import numpy as np
+
+
+class SamplingStrategy(ABC):
+    """Base class for active learning sampling strategies.
+
+    All sampling strategies must implement compute_scores to quantify
+    uncertainty or informativeness of predictions, and select_top_k
+    to select the most informative items.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> class MyStrategy(SamplingStrategy):
+    ...     def compute_scores(self, probabilities):
+    ...         return np.max(probabilities, axis=1)
+    >>> strategy = MyStrategy()
+    >>> probs = np.array([[0.7, 0.2, 0.1], [0.4, 0.4, 0.2]])
+    >>> scores = strategy.compute_scores(probs)
+    >>> indices = strategy.select_top_k(scores, k=1)
+    >>> len(indices)
+    1
+    """
+
+    @abstractmethod
+    def compute_scores(self, probabilities: np.ndarray) -> np.ndarray:
+        """Compute uncertainty scores from prediction probabilities.
+
+        Parameters
+        ----------
+        probabilities : np.ndarray
+            Prediction probabilities with shape (n_samples, n_classes).
+            Each row should sum to 1.0.
+
+        Returns
+        -------
+        np.ndarray
+            Uncertainty scores with shape (n_samples,).
+            Higher scores indicate more informative/uncertain items.
+
+        Examples
+        --------
+        >>> strategy = UncertaintySampling()  # doctest: +SKIP
+        >>> probs = np.array([[0.5, 0.5], [0.9, 0.1]])  # doctest: +SKIP
+        >>> scores = strategy.compute_scores(probs)  # doctest: +SKIP
+        >>> scores[0] > scores[1]  # First is more uncertain  # doctest: +SKIP
+        True
+        """
+        pass
+
+    def select_top_k(self, scores: np.ndarray, k: int) -> np.ndarray:
+        """Select top k items by score.
+
+        Parameters
+        ----------
+        scores : np.ndarray
+            Uncertainty scores with shape (n_samples,).
+        k : int
+            Number of items to select.
+
+        Returns
+        -------
+        np.ndarray
+            Indices of top k items with shape (k,).
+            If k > len(scores), returns all indices.
+            If k <= 0, returns empty array.
+
+        Examples
+        --------
+        >>> strategy = UncertaintySampling()
+        >>> scores = np.array([0.5, 0.9, 0.3, 0.7])
+        >>> indices = strategy.select_top_k(scores, k=2)
+        >>> list(indices)
+        [1, 3]
+        """
+        if k <= 0:
+            return np.array([], dtype=int)
+
+        if k >= len(scores):
+            return np.arange(len(scores))
+
+        # get indices of top k scores (descending order)
+        return np.argsort(scores)[-k:][::-1]
+
+
+class UncertaintySampling(SamplingStrategy):
+    """Entropy-based uncertainty sampling.
+
+    Selects items where the model's prediction entropy is highest,
+    indicating maximum uncertainty across all classes.
+
+    Mathematical definition:
+        H(p) = -∑(p_i * log(p_i))
+
+    where p is the probability distribution over classes.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> strategy = UncertaintySampling()
+    >>> # Uniform distribution (high entropy)
+    >>> probs = np.array([[0.33, 0.33, 0.34]])
+    >>> score = strategy.compute_scores(probs)
+    >>> score[0] > 1.0  # High uncertainty
+    True
+    >>> # Confident prediction (low entropy)
+    >>> probs = np.array([[0.9, 0.05, 0.05]])
+    >>> score = strategy.compute_scores(probs)
+    >>> score[0] < 0.5  # Low uncertainty
+    True
+    """
+
+    def compute_scores(self, probabilities: np.ndarray) -> np.ndarray:
+        """Compute entropy for each prediction.
+
+        Parameters
+        ----------
+        probabilities : np.ndarray
+            Prediction probabilities with shape (n_samples, n_classes).
+
+        Returns
+        -------
+        np.ndarray
+            Entropy scores with shape (n_samples,).
+            Higher entropy indicates more uncertainty.
+
+        Examples
+        --------
+        >>> strategy = UncertaintySampling()
+        >>> probs = np.array([[0.5, 0.5], [0.9, 0.1]])
+        >>> scores = strategy.compute_scores(probs)
+        >>> scores[0] > scores[1]  # Uniform is more uncertain
+        True
+        """
+        # add small epsilon to avoid log(0)
+        epsilon = 1e-10
+        probs_safe = np.clip(probabilities, epsilon, 1.0)
+
+        # compute entropy: -∑(p * log(p))
+        entropy = -np.sum(probs_safe * np.log(probs_safe), axis=1)
+
+        return entropy
+
+
+class MarginSampling(SamplingStrategy):
+    """Margin-based uncertainty sampling.
+
+    Selects items where the margin between the top two predicted classes
+    is smallest, indicating uncertainty between the two most likely options.
+
+    Mathematical definition:
+        margin(p) = 1 - (p₁ - p₂)
+
+    where p₁ and p₂ are the highest and second-highest probabilities.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> strategy = MarginSampling()
+    >>> # Small margin (uncertain)
+    >>> probs = np.array([[0.51, 0.49, 0.0]])
+    >>> score = strategy.compute_scores(probs)
+    >>> score[0] > 0.95  # High uncertainty
+    True
+    >>> # Large margin (confident)
+    >>> probs = np.array([[0.9, 0.05, 0.05]])
+    >>> score = strategy.compute_scores(probs)
+    >>> score[0] < 0.2  # Low uncertainty
+    True
+    """
+
+    def compute_scores(self, probabilities: np.ndarray) -> np.ndarray:
+        """Compute margin scores for each prediction.
+
+        Parameters
+        ----------
+        probabilities : np.ndarray
+            Prediction probabilities with shape (n_samples, n_classes).
+
+        Returns
+        -------
+        np.ndarray
+            Margin scores with shape (n_samples,).
+            Higher scores indicate smaller margin (more uncertainty).
+
+        Examples
+        --------
+        >>> strategy = MarginSampling()
+        >>> probs = np.array([[0.6, 0.3, 0.1], [0.8, 0.15, 0.05]])
+        >>> scores = strategy.compute_scores(probs)
+        >>> scores[0] > scores[1]  # First has smaller margin
+        True
+        """
+        # sort probabilities in descending order
+        sorted_probs = np.sort(probabilities, axis=1)
+
+        # get top 2 probabilities
+        top1 = sorted_probs[:, -1]
+        top2 = sorted_probs[:, -2]
+
+        # compute margin: 1 - (p1 - p2)
+        margin = 1.0 - (top1 - top2)
+
+        return margin
+
+
+class LeastConfidenceSampling(SamplingStrategy):
+    """Least confidence sampling.
+
+    Selects items where the model is least confident, measured as
+    1 minus the maximum predicted probability.
+
+    Mathematical definition:
+        lc(p) = 1 - max(p)
+
+    where p is the probability distribution over classes.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> strategy = LeastConfidenceSampling()
+    >>> # Low confidence
+    >>> probs = np.array([[0.4, 0.3, 0.3]])
+    >>> score = strategy.compute_scores(probs)
+    >>> score[0] == 0.6  # 1 - 0.4
+    True
+    >>> # High confidence
+    >>> probs = np.array([[0.95, 0.03, 0.02]])
+    >>> score = strategy.compute_scores(probs)
+    >>> score[0] == 0.05  # 1 - 0.95
+    True
+    """
+
+    def compute_scores(self, probabilities: np.ndarray) -> np.ndarray:
+        """Compute least confidence scores for each prediction.
+
+        Parameters
+        ----------
+        probabilities : np.ndarray
+            Prediction probabilities with shape (n_samples, n_classes).
+
+        Returns
+        -------
+        np.ndarray
+            Least confidence scores with shape (n_samples,).
+            Higher scores indicate lower confidence (more uncertainty).
+
+        Examples
+        --------
+        >>> strategy = LeastConfidenceSampling()
+        >>> probs = np.array([[0.5, 0.5], [0.9, 0.1]])
+        >>> scores = strategy.compute_scores(probs)
+        >>> scores[0] > scores[1]  # First is less confident
+        True
+        """
+        # get maximum probability for each sample
+        max_probs = np.max(probabilities, axis=1)
+
+        # compute least confidence: 1 - max(p)
+        least_confidence = 1.0 - max_probs
+
+        return least_confidence
+
+
+class RandomSampling(SamplingStrategy):
+    """Random sampling baseline.
+
+    Selects items randomly, serving as a baseline for comparison
+    with uncertainty-based methods. Uses seeded random number generation
+    for reproducibility.
+
+    Parameters
+    ----------
+    seed : int | None
+        Random seed for reproducibility. If None, uses non-deterministic seed.
+
+    Attributes
+    ----------
+    rng : np.random.Generator
+        Random number generator.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> strategy = RandomSampling(seed=42)
+    >>> probs = np.array([[0.9, 0.1], [0.5, 0.5]])
+    >>> scores = strategy.compute_scores(probs)
+    >>> len(scores) == 2
+    True
+    >>> # Scores are random, not based on probabilities
+    >>> strategy2 = RandomSampling(seed=42)
+    >>> scores2 = strategy2.compute_scores(probs)
+    >>> np.allclose(scores, scores2)  # Same seed gives same results
+    True
+    """
+
+    def __init__(self, seed: int | None = None) -> None:
+        self.rng = np.random.default_rng(seed)
+
+    def compute_scores(self, probabilities: np.ndarray) -> np.ndarray:
+        """Generate random scores for each item.
+
+        Parameters
+        ----------
+        probabilities : np.ndarray
+            Prediction probabilities with shape (n_samples, n_classes).
+            Not used in random sampling, but required by interface.
+
+        Returns
+        -------
+        np.ndarray
+            Random scores with shape (n_samples,).
+
+        Examples
+        --------
+        >>> strategy = RandomSampling(seed=123)
+        >>> probs = np.array([[0.9, 0.1], [0.1, 0.9]])
+        >>> scores = strategy.compute_scores(probs)
+        >>> len(scores) == 2
+        True
+        >>> 0.0 <= scores[0] <= 1.0
+        True
+        """
+        n_samples = probabilities.shape[0]
+        return self.rng.random(n_samples)
+
+
+# Type alias for strategy methods
+StrategyMethod = Literal["entropy", "margin", "least_confidence", "random"]
+
+
+def create_strategy(
+    method: StrategyMethod, seed: int | None = None
+) -> SamplingStrategy:
+    """Create a sampling strategy instance.
+
+    Parameters
+    ----------
+    method : StrategyMethod
+        Strategy method name ("entropy", "margin", "least_confidence", "random").
+    seed : int | None
+        Random seed for random strategy. Ignored for other strategies.
+
+    Returns
+    -------
+    SamplingStrategy
+        Instantiated sampling strategy.
+
+    Raises
+    ------
+    ValueError
+        If method is not recognized.
+
+    Examples
+    --------
+    >>> strategy = create_strategy("entropy")
+    >>> isinstance(strategy, UncertaintySampling)
+    True
+    >>> strategy = create_strategy("margin")
+    >>> isinstance(strategy, MarginSampling)
+    True
+    >>> strategy = create_strategy("least_confidence")
+    >>> isinstance(strategy, LeastConfidenceSampling)
+    True
+    >>> strategy = create_strategy("random", seed=42)
+    >>> isinstance(strategy, RandomSampling)
+    True
+    """
+    match method:
+        case "entropy":
+            return UncertaintySampling()
+        case "margin":
+            return MarginSampling()
+        case "least_confidence":
+            return LeastConfidenceSampling()
+        case "random":
+            return RandomSampling(seed=seed)
+        case _:
+            raise ValueError(f"Unknown sampling method: {method}")

bead/active_learning/trainers/__init__.py

@@ -0,0 +1,26 @@
+"""Training framework adapters.
+
+Provides trainer implementations for HuggingFace Transformers and PyTorch
+Lightning. All trainers implement the BaseTrainer interface.
+"""
+
+from __future__ import annotations
+
+from bead.active_learning.trainers.base import BaseTrainer, ModelMetadata
+from bead.active_learning.trainers.huggingface import HuggingFaceTrainer
+from bead.active_learning.trainers.lightning import PyTorchLightningTrainer
+from bead.active_learning.trainers.registry import (
+    get_trainer,
+    list_trainers,
+    register_trainer,
+)
+
+__all__ = [
+    "BaseTrainer",
+    "HuggingFaceTrainer",
+    "ModelMetadata",
+    "PyTorchLightningTrainer",
+    "get_trainer",
+    "list_trainers",
+    "register_trainer",
+]

bead/active_learning/trainers/base.py

@@ -0,0 +1,210 @@
+"""Base trainer interface for model training.
+
+This module provides the abstract base class for all trainers and the
+ModelMetadata model for tracking training results.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from bead.data.base import BeadBaseModel
+
+if TYPE_CHECKING:
+    from datasets import Dataset
+    from transformers import PreTrainedModel
+
+
+class ModelMetadata(BeadBaseModel):
+    """Training metadata.
+
+    Parameters
+    ----------
+    model_name : str
+        Model identifier.
+    framework : str
+        Training framework ("huggingface" or "pytorch_lightning").
+    training_config : dict[str, str | int | float | bool | Path | None]
+        Training configuration used.
+    training_data_path : Path
+        Path to training data.
+    eval_data_path : Path | None
+        Path to eval data if used.
+    metrics : dict[str, float]
+        Final evaluation metrics.
+    best_checkpoint : Path | None
+        Path to best checkpoint.
+    training_time : float
+        Total training time in seconds.
+    training_timestamp : str
+        ISO 8601 timestamp when training completed.
+
+    Attributes
+    ----------
+    model_name : str
+        Model identifier.
+    framework : str
+        Training framework ("huggingface" or "pytorch_lightning").
+    training_config : dict[str, str | int | float | bool | Path | None]
+        Training configuration used.
+    training_data_path : Path
+        Path to training data.
+    eval_data_path : Path | None
+        Path to eval data if used.
+    metrics : dict[str, float]
+        Final evaluation metrics.
+    best_checkpoint : Path | None
+        Path to best checkpoint.
+    training_time : float
+        Total training time in seconds.
+    training_timestamp : str
+        ISO 8601 timestamp when training completed.
+
+    Examples
+    --------
+    >>> from pathlib import Path
+    >>> metadata = ModelMetadata(
+    ...     model_name="bert-base-uncased",
+    ...     framework="huggingface",
+    ...     training_config={"epochs": 3},
+    ...     training_data_path=Path("train.json"),
+    ...     metrics={"accuracy": 0.95},
+    ...     training_time=120.5,
+    ...     training_timestamp="2025-01-17T00:00:00+00:00"
+    ... )
+    >>> metadata.framework
+    'huggingface'
+    >>> metadata.metrics["accuracy"]
+    0.95
+    """
+
+    model_name: str
+    framework: str
+    training_config: dict[str, str | int | float | bool | Path | None]
+    training_data_path: Path
+    eval_data_path: Path | None = None
+    metrics: dict[str, float]
+    best_checkpoint: Path | None = None
+    training_time: float
+    training_timestamp: str
+
+
+class BaseTrainer(ABC):
+    """Base trainer interface.
+
+    All trainers must implement the train, save_model, and load_model methods.
+    This provides a consistent interface across different training frameworks.
+
+    Parameters
+    ----------
+    config : dict[str, str | int | float | bool | Path | None] | BeadBaseModel
+        Training configuration (framework-specific).
+
+    Attributes
+    ----------
+    config : dict[str, str | int | float | bool | Path | None] | BeadBaseModel
+        Training configuration.
+
+    Examples
+    --------
+    >>> class MyTrainer(BaseTrainer):
+    ...     def train(self, train_data, eval_data=None):
+    ...         return ModelMetadata(
+    ...             model_name="test",
+    ...             framework="custom",
+    ...             training_config={},
+    ...             training_data_path=Path("train.json"),
+    ...             metrics={},
+    ...             training_time=0.0,
+    ...             training_timestamp="2025-01-17T00:00:00+00:00"
+    ...         )
+    ...     def save_model(self, output_dir, metadata):
+    ...         pass
+    ...     def load_model(self, model_dir):
+    ...         return None
+    >>> trainer = MyTrainer(config={})
+    >>> trainer.config
+    {}
+    """
+
+    def __init__(
+        self, config: dict[str, str | int | float | bool | Path | None] | BeadBaseModel
+    ) -> None:
+        self.config = config
+
+    @abstractmethod
+    def train(
+        self,
+        train_data: Dataset
+        | dict[str, str | int | float | bool | None]
+        | list[dict[str, str | int | float | bool | None]],
+        eval_data: Dataset
+        | dict[str, str | int | float | bool | None]
+        | list[dict[str, str | int | float | bool | None]]
+        | None = None,
+    ) -> ModelMetadata:
+        """Train model and return metadata.
+
+        Parameters
+        ----------
+        train_data : Dataset | dict | list
+            Training dataset (framework-specific format).
+        eval_data : Dataset | dict | list | None
+            Evaluation dataset (framework-specific format).
+
+        Returns
+        -------
+        ModelMetadata
+            Metadata about the training run.
+
+        Examples
+        --------
+        >>> trainer = MyTrainer(config={})  # doctest: +SKIP
+        >>> metadata = trainer.train(train_dataset)  # doctest: +SKIP
+        >>> metadata.framework  # doctest: +SKIP
+        'custom'
+        """
+        pass
+
+    @abstractmethod
+    def save_model(self, output_dir: Path, metadata: ModelMetadata) -> None:
+        """Save model and metadata to directory.
+
+        Parameters
+        ----------
+        output_dir : Path
+            Directory to save model and metadata.
+        metadata : ModelMetadata
+            Training metadata to save.
+
+        Examples
+        --------
+        >>> trainer = MyTrainer(config={})  # doctest: +SKIP
+        >>> trainer.save_model(Path("output"), metadata)  # doctest: +SKIP
+        """
+        pass
+
+    @abstractmethod
+    def load_model(
+        self, model_dir: Path
+    ) -> PreTrainedModel | dict[str, str | int | float | bool | None] | BeadBaseModel:
+        """Load model from directory.
+
+        Parameters
+        ----------
+        model_dir : Path
+            Directory containing saved model.
+
+        Returns
+        -------
+        PreTrainedModel | dict[str, str | int | float | bool | None] | BeadBaseModel
+            Loaded model (framework-specific type).
+
+        Examples
+        --------
+        >>> trainer = MyTrainer(config={})  # doctest: +SKIP
+        >>> model = trainer.load_model(Path("saved_model"))  # doctest: +SKIP
+        """
+        pass