bead 0.1.0__py3-none-any.whl

bead/templates/adapters/huggingface.py

@@ -0,0 +1,312 @@
+"""HuggingFace masked language model adapter for template filling."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+from transformers import (
+    AutoModelForMaskedLM,
+    AutoTokenizer,
+    PreTrainedModel,
+    PreTrainedTokenizer,
+)
+
+from bead.adapters.huggingface import DeviceType, HuggingFaceAdapterMixin
+from bead.templates.adapters.base import TemplateFillingModelAdapter
+
+
+class HuggingFaceMLMAdapter(HuggingFaceAdapterMixin, TemplateFillingModelAdapter):
+    """Adapter for HuggingFace masked language models.
+
+    Supports BERT, RoBERTa, ALBERT, and other MLM architectures.
+
+    Parameters
+    ----------
+    model_name : str
+        HuggingFace model identifier (e.g., "bert-base-uncased")
+    device : DeviceType
+        Computation device ("cpu", "cuda", "mps")
+    cache_dir : Path | None
+        Directory for caching model files
+
+    Examples
+    --------
+    >>> adapter = HuggingFaceMLMAdapter("bert-base-uncased", device="cpu")
+    >>> adapter.load_model()
+    >>> predictions = adapter.predict_masked_token(
+    ...     text="The cat sat on the mat",
+    ...     mask_position=2,
+    ...     top_k=5
+    ... )
+    >>> for token, log_prob in predictions:
+    ...     print(f"{token}: {log_prob:.2f}")
+    >>> adapter.unload_model()
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        device: DeviceType = "cpu",
+        cache_dir: Path | None = None,
+    ) -> None:
+        # validate device before passing to parent
+        validated_device = self._validate_device(device)
+        super().__init__(model_name, validated_device, cache_dir)
+        self.model: PreTrainedModel | None = None
+        self.tokenizer: PreTrainedTokenizer | None = None
+
+    def load_model(self) -> None:
+        """Load model and tokenizer from HuggingFace.
+
+        Raises
+        ------
+        RuntimeError
+            If model loading fails
+        """
+        if self._model_loaded:
+            return
+
+        try:
+            # load tokenizer
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                self.model_name,
+                cache_dir=self.cache_dir,
+            )
+
+            # load model
+            self.model = AutoModelForMaskedLM.from_pretrained(
+                self.model_name,
+                cache_dir=self.cache_dir,
+            )
+
+            # move to device
+            self.model.to(self.device)
+
+            # set to evaluation mode
+            self.model.eval()
+
+            self._model_loaded = True
+
+        except Exception as e:
+            raise RuntimeError(f"Failed to load model {self.model_name}: {e}") from e
+
+    def unload_model(self) -> None:
+        """Unload model from memory."""
+        if not self._model_loaded:
+            return
+
+        # move model to CPU and delete
+        if self.model is not None:
+            self.model.to("cpu")
+            del self.model
+            self.model = None
+
+        del self.tokenizer
+        self.tokenizer = None
+
+        self._model_loaded = False
+
+        # clear CUDA cache if using GPU
+        if self.device == "cuda":
+            torch.cuda.empty_cache()
+
+    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 or text has no mask token
+        """
+        if not self._model_loaded:
+            raise RuntimeError("Model not loaded. Call load_model() first.")
+
+        # tokenize input
+        inputs = self.tokenizer(text, return_tensors="pt")
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+
+        # find mask token ID
+        mask_token_id = self.tokenizer.mask_token_id
+        if mask_token_id is None:
+            raise ValueError(f"Model {self.model_name} does not have a mask token")
+
+        # find mask position in tokenized input
+        input_ids = inputs["input_ids"][0]
+        mask_positions = (input_ids == mask_token_id).nonzero(as_tuple=True)[0]
+
+        if len(mask_positions) == 0:
+            raise ValueError(f"No mask token found in text: {text}")
+
+        if mask_position >= len(mask_positions):
+            raise ValueError(
+                f"mask_position {mask_position} out of range. "
+                f"Found {len(mask_positions)} mask tokens in text."
+            )
+
+        # get actual token index
+        mask_idx = mask_positions[mask_position].item()
+
+        # forward pass
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+            logits = outputs.logits
+
+        # get predictions for mask position
+        mask_logits = logits[0, mask_idx]
+
+        # convert to log probabilities
+        log_probs = torch.log_softmax(mask_logits, dim=0)
+
+        # get top-k predictions
+        top_log_probs, top_indices = torch.topk(log_probs, k=min(top_k, len(log_probs)))
+
+        # convert to tokens
+        predictions: list[tuple[str, float]] = []
+        for log_prob, idx in zip(top_log_probs.cpu(), top_indices.cpu(), strict=True):
+            token = self.tokenizer.decode([idx], skip_special_tokens=True).strip()
+            predictions.append((token, float(log_prob)))
+
+        return predictions
+
+    def predict_masked_token_batch(
+        self,
+        texts: list[str],
+        mask_position: int = 0,
+        top_k: int = 10,
+    ) -> list[list[tuple[str, float]]]:
+        """Predict masked tokens for multiple texts in a single batch.
+
+        Parameters
+        ----------
+        texts : list[str]
+            List of texts with mask tokens
+        mask_position : int
+            Token position of mask (0-indexed, relative to mask tokens found)
+        top_k : int
+            Number of top predictions to return per text
+
+        Returns
+        -------
+        list[list[tuple[str, float]]]
+            List of predictions for each text. Each element is a list of
+            (token, log_probability) tuples.
+
+        Raises
+        ------
+        RuntimeError
+            If model is not loaded
+        ValueError
+            If any text has no mask token
+        """
+        if not self._model_loaded:
+            raise RuntimeError("Model not loaded. Call load_model() first.")
+
+        if not texts:
+            return []
+
+        # tokenize all texts with padding
+        inputs = self.tokenizer(
+            texts,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+        )
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+
+        # find mask token ID
+        mask_token_id = self.tokenizer.mask_token_id
+        if mask_token_id is None:
+            raise ValueError(f"Model {self.model_name} does not have a mask token")
+
+        # forward pass for entire batch
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+            logits = outputs.logits  # shape: (batch_size, seq_len, vocab_size)
+
+        # process each text in batch
+        results: list[list[tuple[str, float]]] = []
+        for i, text in enumerate(texts):
+            # find mask position in this text
+            input_ids = inputs["input_ids"][i]
+            mask_positions = (input_ids == mask_token_id).nonzero(as_tuple=True)[0]
+
+            if len(mask_positions) == 0:
+                raise ValueError(f"No mask token found in text: {text}")
+
+            if mask_position >= len(mask_positions):
+                raise ValueError(
+                    f"mask_position {mask_position} out of range. "
+                    f"Found {len(mask_positions)} mask tokens in text."
+                )
+
+            # get actual token index
+            mask_idx = mask_positions[mask_position].item()
+
+            # get predictions for this mask position
+            mask_logits = logits[i, mask_idx]
+
+            # convert to log probabilities
+            log_probs = torch.log_softmax(mask_logits, dim=0)
+
+            # get top-k predictions
+            top_log_probs, top_indices = torch.topk(
+                log_probs, k=min(top_k, len(log_probs))
+            )
+
+            # convert to tokens
+            predictions: list[tuple[str, float]] = []
+            for log_prob, idx in zip(
+                top_log_probs.cpu(), top_indices.cpu(), strict=True
+            ):
+                token = self.tokenizer.decode([idx], skip_special_tokens=True).strip()
+                predictions.append((token, float(log_prob)))
+
+            results.append(predictions)
+
+        return results
+
+    def get_mask_token(self) -> str:
+        """Get the mask token for this model.
+
+        Returns
+        -------
+        str
+            Mask token string (e.g., "[MASK]" for BERT)
+
+        Raises
+        ------
+        RuntimeError
+            If model is not loaded
+        ValueError
+            If model has no mask token
+        """
+        if not self._model_loaded:
+            raise RuntimeError("Model not loaded. Call load_model() first.")
+
+        mask_token = self.tokenizer.mask_token
+        if mask_token is None:
+            raise ValueError(f"Model {self.model_name} does not have a mask token")
+
+        return mask_token

bead/templates/combinatorics.py

@@ -0,0 +1,103 @@
+"""Combinatorial utilities for template filling."""
+
+from __future__ import annotations
+
+import itertools
+import random
+from collections.abc import Iterator
+from typing import TypeVar
+
+T = TypeVar("T")
+
+
+def cartesian_product[T](*iterables: list[T]) -> Iterator[tuple[T, ...]]:
+    """Generate Cartesian product of iterables.
+
+    Equivalent to itertools.product but with explicit type hints
+    and documentation for template filling use case.
+
+    Parameters
+    ----------
+    *iterables : list[T]
+        Variable number of iterables to combine.
+
+    Yields
+    ------
+    tuple[T, ...]
+        Each combination from the Cartesian product.
+
+    Examples
+    --------
+    >>> list(cartesian_product([1, 2], ['a', 'b']))
+    [(1, 'a'), (1, 'b'), (2, 'a'), (2, 'b')]
+    """
+    return itertools.product(*iterables)
+
+
+def count_combinations[T](*iterables: list[T]) -> int:
+    """Count total combinations without generating them.
+
+    Calculate the size of the Cartesian product space efficiently
+    without actually generating combinations.
+
+    Parameters
+    ----------
+    *iterables : list[Any]
+        Variable number of iterables.
+
+    Returns
+    -------
+    int
+        Total number of combinations.
+
+    Examples
+    --------
+    >>> count_combinations([1, 2], ['a', 'b'], [True, False])
+    8
+    """
+    count = 1
+    for iterable in iterables:
+        count *= len(iterable)
+    return count
+
+
+def stratified_sample[T](
+    groups: dict[str, list[T]],
+    n_per_group: int,
+    seed: int | None = None,
+) -> list[T]:
+    """Sample items from groups with balanced representation.
+
+    Ensure each group is represented approximately equally in the sample.
+
+    Parameters
+    ----------
+    groups : dict[str, list[T]]
+        Dictionary mapping group names to lists of items.
+    n_per_group : int
+        Number of items to sample from each group.
+    seed : int | None
+        Random seed for reproducibility.
+
+    Returns
+    -------
+    list[T]
+        Sampled items, balanced across groups.
+
+    Examples
+    --------
+    >>> groups = {"verbs": [v1, v2, v3], "nouns": [n1, n2, n3]}
+    >>> sample = stratified_sample(groups, n_per_group=2, seed=42)
+    >>> len(sample)
+    4
+    """
+    if seed is not None:
+        random.seed(seed)
+
+    sampled: list[T] = []
+    for group_items in groups.values():
+        # Sample with replacement if n_per_group > len(group_items)
+        k = min(n_per_group, len(group_items))
+        sampled.extend(random.sample(group_items, k))
+
+    return sampled