bead 0.1.0__py3-none-any.whl

bead/simulation/strategies/multi_select.py

@@ -0,0 +1,129 @@
+"""Multi-select simulation strategy."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from bead.simulation.strategies.base import SimulationStrategy
+
+if TYPE_CHECKING:
+    from bead.items.item import Item
+    from bead.items.item_template import ItemTemplate
+
+
+class MultiSelectStrategy(SimulationStrategy):
+    """Strategy for multi_select tasks.
+
+    Handles tasks where multiple options can be selected independently.
+    Uses model outputs to compute independent selection probabilities
+    for each option via sigmoid.
+
+    For each option i:
+        P(select option i) = sigmoid(score_i / temperature)
+
+    Parameters
+    ----------
+    threshold
+        Probability threshold for selection. Default: 0.5.
+    temperature
+        Temperature for scaling decisions. Default: 1.0.
+
+    Examples
+    --------
+    >>> strategy = MultiSelectStrategy()
+    >>> strategy.supported_task_type
+    'multi_select'
+    """
+
+    def __init__(self, threshold: float = 0.5, temperature: float = 1.0) -> None:
+        self.threshold = threshold
+        self.temperature = temperature
+
+    @property
+    def supported_task_type(self) -> str:
+        """Return 'multi_select'."""
+        return "multi_select"
+
+    def validate_item(self, item: Item, item_template: ItemTemplate) -> None:
+        """Validate item for multi-select.
+
+        Checks:
+        - task_type is 'multi_select'
+        - task_spec.options is defined
+        - At least 2 options
+
+        Parameters
+        ----------
+        item : Item
+            Item to validate.
+        item_template : ItemTemplate
+            Template defining task.
+
+        Raises
+        ------
+        ValueError
+            If validation fails.
+        """
+        if item_template.task_type != "multi_select":
+            msg = f"Expected task_type 'multi_select', got '{item_template.task_type}'"
+            raise ValueError(msg)
+
+        if not item_template.task_spec.options:
+            raise ValueError("task_spec.options must be defined for multi_select")
+
+        if len(item_template.task_spec.options) < 2:
+            raise ValueError("multi_select requires at least 2 options")
+
+    def simulate_response(
+        self,
+        item: Item,
+        item_template: ItemTemplate,
+        model_output_key: str,
+        rng: np.random.RandomState,
+    ) -> list[str]:
+        """Generate multi-select response.
+
+        Parameters
+        ----------
+        item : Item
+            Item to respond to.
+        item_template : ItemTemplate
+            Template defining task.
+        model_output_key : str
+            Key for model outputs (e.g., "lm_score").
+        rng : np.random.RandomState
+            Random number generator.
+
+        Returns
+        -------
+        list[str]
+            List of selected option names.
+        """
+        options = item_template.task_spec.options
+        assert options is not None, "options validated in validate()"
+        n_options = len(options)
+
+        # extract model outputs for each option
+        scores = self.extract_model_outputs(item, model_output_key, n_options)
+
+        if scores is None:
+            # fallback to random selection (each option has threshold probability)
+            selected = []
+            for option in options:
+                if rng.random() < self.threshold:
+                    selected.append(option)
+            return selected
+
+        # compute selection probability for each option using sigmoid
+        selected = []
+        for option, score in zip(options, scores, strict=True):
+            # sigmoid(score / temperature)
+            prob = 1.0 / (1.0 + np.exp(-score / self.temperature))
+
+            # sample selection
+            if rng.random() < prob:
+                selected.append(option)
+
+        return selected

bead/simulation/strategies/ordinal_scale.py

@@ -0,0 +1,131 @@
+"""Ordinal scale simulation strategy."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from bead.simulation.strategies.base import SimulationStrategy
+
+if TYPE_CHECKING:
+    from bead.items.item import Item
+    from bead.items.item_template import ItemTemplate
+
+
+class OrdinalScaleStrategy(SimulationStrategy):
+    """Strategy for ordinal_scale tasks (Likert scales).
+
+    Handles discrete ordinal scales (e.g., 1-7, 1-5). Maps model outputs
+    to scale positions, then samples with noise around that position.
+
+    For ordinal scales with LM score:
+        - Map score to continuous position on scale
+        - Add noise
+        - Round to nearest integer within bounds
+
+    Examples
+    --------
+    >>> strategy = OrdinalScaleStrategy()
+    >>> strategy.supported_task_type
+    'ordinal_scale'
+    """
+
+    @property
+    def supported_task_type(self) -> str:
+        """Return 'ordinal_scale'.
+
+        Returns
+        -------
+        str
+            Task type identifier.
+        """
+        return "ordinal_scale"
+
+    def validate_item(self, item: Item, item_template: ItemTemplate) -> None:
+        """Validate item for ordinal scale.
+
+        Checks:
+        - task_type is 'ordinal_scale'
+        - task_spec.scale_bounds is defined
+        - scale_bounds has valid min/max
+
+        Parameters
+        ----------
+        item : Item
+            Item to validate.
+        item_template : ItemTemplate
+            Template defining task.
+
+        Raises
+        ------
+        ValueError
+            If validation fails.
+        """
+        if item_template.task_type != "ordinal_scale":
+            msg = f"Expected task_type 'ordinal_scale', got '{item_template.task_type}'"
+            raise ValueError(msg)
+
+        if not item_template.task_spec.scale_bounds:
+            msg = "task_spec.scale_bounds must be defined for ordinal_scale"
+            raise ValueError(msg)
+
+        min_val, max_val = item_template.task_spec.scale_bounds
+        if min_val >= max_val:
+            msg = f"scale_bounds min ({min_val}) must be less than max ({max_val})"
+            raise ValueError(msg)
+
+    def simulate_response(
+        self,
+        item: Item,
+        item_template: ItemTemplate,
+        model_output_key: str,
+        rng: np.random.RandomState,
+    ) -> int:
+        """Generate ordinal scale response.
+
+        Parameters
+        ----------
+        item : Item
+            Item to respond to.
+        item_template : ItemTemplate
+            Template defining task.
+        model_output_key : str
+            Key for model outputs (e.g., "lm_score").
+        rng : np.random.RandomState
+            Random number generator.
+
+        Returns
+        -------
+        int
+            Rating on ordinal scale.
+        """
+        scale_bounds = item_template.task_spec.scale_bounds
+        if scale_bounds is None:
+            msg = "task_spec.scale_bounds must be defined"
+            raise ValueError(msg)
+
+        min_val, max_val = scale_bounds
+        scale_range = max_val - min_val
+
+        # extract model output (expecting single score)
+        scores = self.extract_model_outputs(item, model_output_key, required_count=1)
+
+        if scores is None:
+            # fallback to uniform random across scale
+            return int(rng.randint(min_val, max_val + 1))
+
+        # map LM score to scale position; use sigmoid to map unbounded score to [0, 1]
+        score = scores[0]
+        sigmoid_score = 1.0 / (1.0 + np.exp(-score))
+
+        # map [0, 1] to scale range
+        continuous_rating = min_val + sigmoid_score * scale_range
+
+        # round to nearest integer
+        rating = int(np.round(continuous_rating))
+
+        # clamp to scale bounds (in case of rounding issues)
+        rating = max(min_val, min(max_val, rating))
+
+        return rating

bead/templates/__init__.py

@@ -0,0 +1,27 @@
+"""Template filling functionality.
+
+Provides template filling strategies (exhaustive, random, stratified) and
+constraint resolution for generating experimental stimuli.
+"""
+
+from __future__ import annotations
+
+from bead.templates.filler import CSPFiller, FilledTemplate, TemplateFiller
+from bead.templates.resolver import ConstraintResolver
+from bead.templates.strategies import (
+    ExhaustiveStrategy,
+    RandomStrategy,
+    StrategyFiller,
+    StratifiedStrategy,
+)
+
+__all__ = [
+    "TemplateFiller",  # ABC
+    "CSPFiller",
+    "StrategyFiller",
+    "FilledTemplate",
+    "ConstraintResolver",
+    "ExhaustiveStrategy",
+    "RandomStrategy",
+    "StratifiedStrategy",
+]

bead/templates/adapters/__init__.py

@@ -0,0 +1,17 @@
+"""Template filling model adapters.
+
+Provides masked language model adapters for template filling (Stage 2).
+Separate from judgment prediction models (Stage 3).
+"""
+
+from __future__ import annotations
+
+from bead.templates.adapters.base import TemplateFillingModelAdapter
+from bead.templates.adapters.cache import ModelOutputCache
+from bead.templates.adapters.huggingface import HuggingFaceMLMAdapter
+
+__all__ = [
+    "TemplateFillingModelAdapter",
+    "ModelOutputCache",
+    "HuggingFaceMLMAdapter",
+]

bead/templates/adapters/base.py

@@ -0,0 +1,128 @@
+"""Base adapter for template filling models.
+
+This module defines the abstract interface for models used in template filling.
+These adapters are SEPARATE from judgment prediction model adapters (Stage 6).
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from pathlib import Path
+
+
+class TemplateFillingModelAdapter(ABC):
+    """Base adapter for models used in template filling.
+
+    This is SEPARATE from judgment prediction model adapters,
+    which are used later in the pipeline for predicting human judgments.
+
+    Parameters
+    ----------
+    model_name : str
+        Model identifier (e.g., "bert-base-uncased")
+    device : str
+        Computation device ("cpu", "cuda", "mps")
+    cache_dir : Path | None
+        Directory for caching model files
+
+    Examples
+    --------
+    >>> from bead.templates.adapters import TemplateFillingModelAdapter
+    >>> # Implemented by HuggingFaceMLMAdapter
+    >>> adapter = HuggingFaceMLMAdapter("bert-base-uncased", device="cpu")
+    >>> adapter.load_model()
+    >>> predictions = adapter.predict_masked_token(
+    ...     text="The cat [MASK] on the mat",
+    ...     mask_position=2,
+    ...     top_k=5
+    ... )
+    >>> adapter.unload_model()
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        device: str = "cpu",
+        cache_dir: Path | None = None,
+    ) -> None:
+        self.model_name = model_name
+        self.device = device
+        self.cache_dir = cache_dir
+        self._model_loaded = False
+
+    @abstractmethod
+    def load_model(self) -> None:
+        """Load model into memory.
+
+        Raises
+        ------
+        RuntimeError
+            If model loading fails
+        """
+        pass
+
+    @abstractmethod
+    def unload_model(self) -> None:
+        """Unload model from memory to free resources."""
+        pass
+
+    @abstractmethod
+    def predict_masked_token(
+        self,
+        text: str,
+        mask_position: int,
+        top_k: int = 10,
+    ) -> list[tuple[str, float]]:
+        """Predict masked token at specified position.
+
+        Parameters
+        ----------
+        text : str
+            Text with mask token (e.g., "The cat [MASK] quickly")
+        mask_position : int
+            Token position of mask (0-indexed)
+        top_k : int
+            Number of top predictions to return
+
+        Returns
+        -------
+        list[tuple[str, float]]
+            List of (token, log_probability) tuples, sorted by probability
+
+        Raises
+        ------
+        RuntimeError
+            If model is not loaded
+        ValueError
+            If mask_position is invalid
+
+        Examples
+        --------
+        >>> predictions = adapter.predict_masked_token(
+        ...     text="The cat [MASK] on the mat",
+        ...     mask_position=2,
+        ...     top_k=3
+        ... )
+        >>> predictions
+        [("sat", -0.5), ("slept", -1.2), ("jumped", -1.5)]
+        """
+        pass
+
+    def is_loaded(self) -> bool:
+        """Check if model is loaded.
+
+        Returns
+        -------
+        bool
+            True if model is loaded in memory
+        """
+        return self._model_loaded
+
+    def __enter__(self) -> TemplateFillingModelAdapter:
+        """Context manager entry."""
+        self.load_model()
+        return self
+
+    def __exit__(self, *args: object) -> None:
+        """Context manager exit."""
+        self.unload_model()

bead/templates/adapters/cache.py

@@ -0,0 +1,178 @@
+"""Content-addressable cache for model predictions.
+
+This module implements caching for template filling model predictions
+using SHA256-based content addressing.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from pathlib import Path
+
+
+class ModelOutputCache:
+    """Content-addressable cache for model predictions.
+
+    Uses SHA256 hashing to create deterministic cache keys based on:
+    - Model name
+    - Input text
+    - Mask position
+    - Top-K parameter
+
+    Parameters
+    ----------
+    cache_dir : Path
+        Directory for cache storage
+    enabled : bool
+        Enable/disable caching
+
+    Examples
+    --------
+    >>> cache = ModelOutputCache(cache_dir=Path("/tmp/cache"), enabled=True)
+    >>> key_args = ("bert-base-uncased", "The cat [MASK]", 2, 10)
+    >>> predictions = cache.get(*key_args)
+    >>> if predictions is None:
+    ...     predictions = model.predict(...)
+    ...     cache.set(*key_args, predictions)
+    """
+
+    def __init__(self, cache_dir: Path, enabled: bool = True) -> None:
+        self.cache_dir = cache_dir
+        self.enabled = enabled
+
+        if self.enabled:
+            self.cache_dir.mkdir(parents=True, exist_ok=True)
+
+    def _compute_key(
+        self,
+        model_name: str,
+        input_text: str,
+        mask_position: int,
+        top_k: int,
+    ) -> str:
+        """Compute cache key from inputs.
+
+        Parameters
+        ----------
+        model_name : str
+            Model identifier
+        input_text : str
+            Input text with mask
+        mask_position : int
+            Position of mask token
+        top_k : int
+            Number of predictions
+
+        Returns
+        -------
+        str
+            SHA256 hex digest
+        """
+        # create deterministic key
+        key_data = {
+            "model_name": model_name,
+            "input_text": input_text,
+            "mask_position": mask_position,
+            "top_k": top_k,
+        }
+
+        # serialize to JSON with sorted keys for determinism
+        key_json = json.dumps(key_data, sort_keys=True)
+
+        # hash with SHA256
+        return hashlib.sha256(key_json.encode("utf-8")).hexdigest()
+
+    def get(
+        self,
+        model_name: str,
+        input_text: str,
+        mask_position: int,
+        top_k: int,
+    ) -> list[tuple[str, float]] | None:
+        """Get cached predictions.
+
+        Parameters
+        ----------
+        model_name : str
+            Model identifier
+        input_text : str
+            Input text
+        mask_position : int
+            Mask position
+        top_k : int
+            Number of predictions
+
+        Returns
+        -------
+        list[tuple[str, float]] | None
+            Cached predictions or None if not found
+        """
+        if not self.enabled:
+            return None
+
+        cache_key = self._compute_key(model_name, input_text, mask_position, top_k)
+        cache_file = self.cache_dir / f"{cache_key}.json"
+
+        if not cache_file.exists():
+            return None
+
+        try:
+            with open(cache_file) as f:
+                data = json.load(f)
+                return [(item["token"], item["log_prob"]) for item in data]
+        except (json.JSONDecodeError, KeyError, OSError):
+            # cache corruption; return None
+            return None
+
+    def set(
+        self,
+        model_name: str,
+        input_text: str,
+        mask_position: int,
+        top_k: int,
+        predictions: list[tuple[str, float]],
+    ) -> None:
+        """Store predictions in cache.
+
+        Parameters
+        ----------
+        model_name : str
+            Model identifier
+        input_text : str
+            Input text
+        mask_position : int
+            Mask position
+        top_k : int
+            Number of predictions
+        predictions : list[tuple[str, float]]
+            Predictions to cache
+        """
+        if not self.enabled:
+            return
+
+        cache_key = self._compute_key(model_name, input_text, mask_position, top_k)
+        cache_file = self.cache_dir / f"{cache_key}.json"
+
+        # convert to serializable format
+        data = [
+            {"token": token, "log_prob": log_prob} for token, log_prob in predictions
+        ]
+
+        try:
+            with open(cache_file, "w") as f:
+                json.dump(data, f, indent=2)
+        except OSError:
+            # silently fail on cache write errors
+            pass
+
+    def clear(self) -> None:
+        """Clear all cached predictions."""
+        if not self.enabled or not self.cache_dir.exists():
+            return
+
+        for cache_file in self.cache_dir.glob("*.json"):
+            try:
+                cache_file.unlink()
+            except OSError:
+                pass