Adds grid stratification, acceptability model, and full layers interop (0.8.0)

CHANGELOG.md

@@ -7,6 +7,89 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.8.0] - 2026-07-13
+
+### Added
+
+#### N-dimensional, mixed-variable grid stratification
+
+- `bead.lists.stratification` gained a grid engine that bins items along several
+  dimensions at once. `assign_grid_cells` maps each item to a tuple of
+  per-dimension bin indices, and `assign_grid_cells_by_uuid` returns the
+  row-major flattened cell id for the stand-off annotation pattern;
+  `grid_shape`, `flatten_cell`, and `unflatten_cell` convert between cell tuples
+  and ids.
+- Each dimension declares a binning strategy via the `BinningSpec` tagged union:
+  `QuantileBinning` (equal-frequency), `EqualWidthBinning` (equal-interval),
+  `ThresholdBinning` (caller-supplied cut points), and `StdDevBinning`
+  (standard-deviation offsets) for continuous values, and `CategoricalBinning`
+  (one bin per value, with optional declared categories and a catch-all) for
+  finite discrete values.
+- `GridStratificationConstraint` and its `GridDimension` axes express a grid as a
+  list constraint; `ListPartitioner` stratifies on the grid cell and
+  `QuantileBalancer.balance_by_cell` spreads each cell uniformly across lists.
+  The one-dimensional `assign_quantiles` path is now the quantile specialization
+  of this single engine.
+
+#### Acceptability-model stratification in the `eng` gallery
+
+- The `gallery/eng/argument_structure` example trains a 2AFC acceptability model
+  on MegaAcceptability (single-item Likert ratings mapped to forced-choice pairs
+  per annotator), uses it together with the language-model score to stratify
+  items across a two-dimensional grid, and seeds the active-learning loop with
+  the pretrained model. Gallery artifacts round-trip through the `layers` schema
+  via `bead.interop.layers`.
+- `AcceptabilityScorer` in `bead.items.scoring` scores a forced-choice item by a
+  trained `ForcedChoiceModel`, returning the predicted preference margin.
+
+#### More resource and response layers lenses
+
+- `FilledTemplateFillingLens` (`FILLED_TEMPLATE_FILLING`) maps a `FilledTemplate`
+  to and from a layers `filling` record (`resource.Filling`): the template
+  reference, per-slot `SlotFilling` records, rendered text, and filling strategy
+  form the view, while the bead framework identity, source template name, slot
+  requirement map, and exact lexical-item fillers travel in the lens complement.
+- `AnnotationRecordJudgmentLens` (`ANNOTATION_RECORD_JUDGMENT`) maps an
+  `AnnotationRecord` to and from a layers `judgment.Judgment`, and
+  `records_to_judgment_set` / `judgment_set_to_records` group a single
+  annotator's records into a `judgment.JudgmentSet` (with the annotator as an
+  `AgentRef`).
+- `ListConstraintLens` (`LIST_CONSTRAINT`) maps a bead `ListConstraint` to and
+  from a layers `judgment.ListConstraint`, and `ExperimentListLens`
+  (`EXPERIMENT_LIST_LAYERS`) maps an `ExperimentList` to a layers `collection`
+  with one `collectionMembership` per item plus its list constraints.
+- `ParticipantAgentLens` (`PARTICIPANT_AGENT`) maps a bead `Participant` to and
+  from a layers `AgentRef` identity, and `participant_features` renders the
+  participant's study fields (demographics, study id, sessions, consent) as the
+  `FeatureMap` that a `judgment.JudgmentSet` documents as the home for annotator
+  demographics, session metadata, and payment info. `records_to_judgment_set`
+  now takes an optional `participant`, so a judgment set carries that annotator's
+  identity and study fields natively. A layers `persona.Persona` is an annotator
+  role and interpretive framework, not a concrete enrolled participant, so it is
+  not the target of this mapping.
+
+  All of these round-trip exactly.
+
+#### Forced-choice training: dev-set early stopping
+
+- `ForcedChoiceModel` supports early stopping on a held-out dev set. A new
+  `early_stopping_patience` config field stops training after that many epochs
+  without dev cross-entropy improvement and restores the best encoder and head.
+  The `eng` gallery holds out whole sentences (never seen during training) for the
+  dev set and, with `max_pairs_per_annotator` unset, trains over every within-rater
+  same-frame pair.
+
+### Fixed
+
+- `ForcedChoiceModel` reports train and validation accuracy, precision, recall,
+  and F1 from the model's own predict path. The shared mixed-effects trainer's
+  `evaluate()` only labels cloze batches, so it could not score forced choice, and
+  train accuracy previously always read `0.0`. `compute_binary_metrics` and
+  `compute_multiclass_metrics` no longer pass `zero_division` to `evaluate`'s F1,
+  which this version rejects.
+- Forced-choice training and prediction move the per-participant random-effects
+  parameters to the model device, so the model runs on MPS and CUDA, not only CPU.
+
 ## [0.7.0] - 2026-06-29
 
 ### Added
@@ -564,7 +647,9 @@ guards as type-checkers.
 - CI/CD: GitHub Actions for testing, docs, PyPI publishing
 - Read the Docs integration
 
-[Unreleased]: https://github.com/FACTSlab/bead/compare/v0.6.0...HEAD
+[Unreleased]: https://github.com/FACTSlab/bead/compare/v0.8.0...HEAD
+[0.8.0]: https://github.com/FACTSlab/bead/compare/v0.7.0...v0.8.0
+[0.7.0]: https://github.com/FACTSlab/bead/compare/v0.6.0...v0.7.0
 [0.6.0]: https://github.com/FACTSlab/bead/compare/v0.5.0...v0.6.0
 [0.5.0]: https://github.com/FACTSlab/bead/compare/v0.4.0...v0.5.0
 [0.4.0]: https://github.com/FACTSlab/bead/compare/v0.3.0...v0.4.0

bead/active_learning/models/forced_choice.py

@@ -3,12 +3,20 @@
 from __future__ import annotations
 
 import tempfile
+from collections.abc import Callable
 from pathlib import Path
 
 import numpy as np
 import torch
 import torch.nn as nn
-from transformers import AutoModel, AutoTokenizer, TrainingArguments
+from transformers import (
+    AutoModel,
+    AutoTokenizer,
+    TrainerCallback,
+    TrainerControl,
+    TrainerState,
+    TrainingArguments,
+)
 
 from bead.active_learning.config import MixedEffectsConfig, VarianceComponents
 from bead.active_learning.models.base import ActiveLearningModel, ModelPrediction
@@ -24,6 +32,84 @@
 __all__ = ["ForcedChoiceModel"]
 
 
+class _DevEarlyStopping(TrainerCallback):
+    """Early-stop forced-choice training on a held-out dev set.
+
+    The shared mixed-effects trainer's evaluation emits no usable loss or
+    accuracy for forced choice, so the dev set is scored through the model's own
+    predict path at the end of each epoch. Training stops after ``patience``
+    consecutive epochs without dev cross-entropy improvement, and the best
+    encoder and classifier head seen so far are restored.
+    """
+
+    def __init__(
+        self,
+        *,
+        predict_proba: Callable[[list[Item], list[str]], np.ndarray],
+        encoder: nn.Module,
+        classifier_head: nn.Module,
+        dev_items: list[Item],
+        dev_participant_ids: list[str],
+        dev_y: list[int],
+        patience: int,
+    ) -> None:
+        self._predict_proba = predict_proba
+        self._encoder = encoder
+        self._classifier_head = classifier_head
+        self._dev_items = dev_items
+        self._dev_participant_ids = dev_participant_ids
+        self._dev_y = np.asarray(dev_y)
+        self._patience = patience
+        self._best_loss = float("inf")
+        self._best_state: (
+            tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]] | None
+        ) = None
+        self._wait = 0
+
+    def _dev_cross_entropy(self) -> float:
+        proba = self._predict_proba(self._dev_items, self._dev_participant_ids)
+        chosen = np.clip(proba[np.arange(len(self._dev_y)), self._dev_y], 1e-9, 1.0)
+        return float(-np.mean(np.log(chosen)))
+
+    def on_epoch_end(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs: object,
+    ) -> None:
+        loss = self._dev_cross_entropy()
+        if loss < self._best_loss - 1e-4:
+            self._best_loss = loss
+            self._best_state = (
+                {
+                    k: v.detach().cpu().clone()
+                    for k, v in self._encoder.state_dict().items()
+                },
+                {
+                    k: v.detach().cpu().clone()
+                    for k, v in self._classifier_head.state_dict().items()
+                },
+            )
+            self._wait = 0
+        else:
+            self._wait += 1
+            if self._wait >= self._patience:
+                control.should_training_stop = True
+
+    def on_train_end(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs: object,
+    ) -> None:
+        if self._best_state is not None:
+            encoder_state, head_state = self._best_state
+            self._encoder.load_state_dict(encoder_state)
+            self._classifier_head.load_state_dict(head_state)
+
+
 class ForcedChoiceModel(ActiveLearningModel):
     """Model for forced_choice tasks with n alternatives.
 
@@ -479,23 +565,19 @@ def _train_with_huggingface_trainer(
             max_length=self.config.max_length,
         )
 
-        # Create validation dataset if provided
-        eval_dataset = None
-        if validation_items is not None and validation_labels is not None:
+        # Prepare a held-out dev set for early stopping, if provided
+        has_dev = validation_items is not None and validation_labels is not None
+        val_y_numeric: list[int] = []
+        val_participant_ids: list[str] = []
+        if has_dev:
+            assert validation_items is not None and validation_labels is not None
             label_to_idx = {label: idx for idx, label in enumerate(self.option_names)}
             val_y_numeric = [label_to_idx[label] for label in validation_labels]
             val_participant_ids = (
                 ["_validation_"] * len(validation_items)
                 if self.config.mixed_effects.mode != "fixed"
                 else ["_fixed_"] * len(validation_items)
             )
-            eval_dataset = items_to_dataset(
-                items=validation_items,
-                labels=val_y_numeric,
-                participant_ids=val_participant_ids,
-                tokenizer=self.tokenizer,
-                max_length=self.config.max_length,
-            )
 
         # Create wrapper model for Trainer
         wrapped_model = EncoderClassifierWrapper(
@@ -505,10 +587,6 @@ def _train_with_huggingface_trainer(
         # Create data collator
         data_collator = MixedEffectsDataCollator(tokenizer=self.tokenizer)
 
-        # Create metrics computation function
-        def compute_metrics_fn(eval_pred: object) -> dict[str, float]:
-            return compute_multiclass_metrics(eval_pred, num_labels=self.num_classes)
-
         # Create training arguments with checkpointing
         with tempfile.TemporaryDirectory() as tmpdir:
             checkpoint_dir = Path(tmpdir) / "checkpoints"
@@ -521,10 +599,8 @@ def compute_metrics_fn(eval_pred: object) -> dict[str, float]:
                 per_device_eval_batch_size=self.config.batch_size,
                 learning_rate=self.config.learning_rate,
                 logging_steps=10,
-                eval_strategy="epoch" if eval_dataset is not None else "no",
-                save_strategy="epoch",
-                save_total_limit=1,
-                load_best_model_at_end=False,
+                eval_strategy="no",
+                save_strategy="no",
                 report_to="none",
                 remove_unused_columns=False,
                 use_cpu=self.config.device == "cpu",
@@ -540,37 +616,54 @@ def compute_metrics_fn(eval_pred: object) -> dict[str, float]:
                 model=wrapped_model,
                 args=training_args,
                 train_dataset=train_dataset,
-                eval_dataset=eval_dataset,
+                eval_dataset=None,
                 data_collator=data_collator,
                 tokenizer=self.tokenizer,
                 random_effects_manager=self.random_effects,
-                compute_metrics=compute_metrics_fn,
             )
+            # Early-stop on the dev set, scored through the model's own predict
+            # path (the shared trainer's evaluate emits no usable loss or
+            # accuracy for forced choice); the best encoder and head are restored.
+            if has_dev:
+                assert validation_items is not None
+                trainer.add_callback(
+                    _DevEarlyStopping(
+                        predict_proba=self._do_predict_proba,
+                        encoder=self.encoder,
+                        classifier_head=self.classifier_head,
+                        dev_items=validation_items,
+                        dev_participant_ids=val_participant_ids,
+                        dev_y=val_y_numeric,
+                        patience=self.config.early_stopping_patience,
+                    )
+                )
 
-            # Train
             train_result = trainer.train()
 
-            # Get training metrics
-            train_metrics = trainer.evaluate(eval_dataset=train_dataset)
-            metrics: dict[str, float] = {
-                "train_loss": float(train_result.training_loss),
-                "train_accuracy": train_metrics.get("eval_accuracy", 0.0),
-                "train_precision": train_metrics.get("eval_precision", 0.0),
-                "train_recall": train_metrics.get("eval_recall", 0.0),
-                "train_f1": train_metrics.get("eval_f1", 0.0),
-            }
-
-            # Get validation metrics if eval_dataset was provided
-            if eval_dataset is not None:
-                val_metrics = trainer.evaluate(eval_dataset=eval_dataset)
-                metrics.update(
-                    {
-                        "val_accuracy": val_metrics.get("eval_accuracy", 0.0),
-                        "val_precision": val_metrics.get("eval_precision", 0.0),
-                        "val_recall": val_metrics.get("eval_recall", 0.0),
-                        "val_f1": val_metrics.get("eval_f1", 0.0),
-                    }
-                )
+        # Score train (and dev) metrics through the model's own predict path,
+        # which applies the fitted random effects. trainer.evaluate cannot score
+        # accuracy here (its shared prediction step only labels cloze batches).
+        train_scores = self._classification_metrics(items, participant_ids, y_numeric)
+        metrics: dict[str, float] = {
+            "train_loss": float(train_result.training_loss),
+            "train_accuracy": train_scores["accuracy"],
+            "train_precision": train_scores["precision"],
+            "train_recall": train_scores["recall"],
+            "train_f1": train_scores["f1"],
+        }
+        if has_dev:
+            assert validation_items is not None
+            val_scores = self._classification_metrics(
+                validation_items, val_participant_ids, val_y_numeric
+            )
+            metrics.update(
+                {
+                    "val_accuracy": val_scores["accuracy"],
+                    "val_precision": val_scores["precision"],
+                    "val_recall": val_scores["recall"],
+                    "val_f1": val_scores["f1"],
+                }
+            )
 
         # Estimate variance components
         if self.config.mixed_effects.estimate_variance_components:
@@ -584,6 +677,38 @@ def compute_metrics_fn(eval_pred: object) -> dict[str, float]:
 
         return metrics
 
+    def _classification_metrics(
+        self,
+        items: list[Item],
+        participant_ids: list[str],
+        y_numeric: list[int],
+    ) -> dict[str, float]:
+        """Accuracy, precision, recall, and F1 from the model's own predictions.
+
+        Predictions come from the model's predict path (which applies the fitted
+        random effects), so the reported metrics reflect the deployed model
+        rather than the shared trainer's cloze-only evaluation.
+
+        Parameters
+        ----------
+        items : list[Item]
+            Items to score.
+        participant_ids : list[str]
+            Participant id per item.
+        y_numeric : list[int]
+            True class indices.
+
+        Returns
+        -------
+        dict[str, float]
+            Keys ``accuracy``, ``precision``, ``recall``, ``f1``.
+        """
+        from transformers import EvalPrediction  # noqa: PLC0415
+
+        proba = self._do_predict_proba(items, participant_ids)
+        eval_pred = EvalPrediction(predictions=proba, label_ids=np.asarray(y_numeric))
+        return compute_multiclass_metrics(eval_pred, num_labels=self.num_classes)
+
     def _train_with_custom_loop(
         self,
         items: list[Item],
@@ -843,14 +968,16 @@ def _do_predict_proba(
                         param_name="mu",
                         create_if_missing=False,
                     )
-                    logits[i] = logits[i] + bias
+                    # Random-effects params live on CPU; move to the logits
+                    # device (matters on mps/cuda).
+                    logits[i] = logits[i] + bias.to(logits.device)
 
             elif self.config.mixed_effects.mode == "random_slopes":
                 logits_list = []
                 for i, pid in enumerate(participant_ids):
                     participant_head = self.random_effects.get_slopes(
                         pid, fixed_head=self.classifier_head, create_if_missing=False
-                    )
+                    ).to(embeddings.device)
                     logits_i = participant_head(embeddings[i : i + 1])
                     logits_list.append(logits_i)
                 logits = torch.cat(logits_list, dim=0)