From c0711b43e2c535fa9ca2e6b4683a708302c805b5 Mon Sep 17 00:00:00 2001 From: Aaron Steven White Date: Tue, 30 Jun 2026 14:00:48 -0400 Subject: [PATCH 1/9] Adds grid stratification, acceptability model, full layers interop (0.8.0) Adds N-dimensional, mixed-variable grid stratification to bead.lists: a BinningSpec tagged union (quantile, equal-width, threshold, standard-deviation, and categorical binners), assign_grid_cells / assign_grid_cells_by_uuid / grid_shape / flatten_cell / unflatten_cell, a GridStratificationConstraint, and a stratum-aware QuantileBalancer.balance_by_cell with a matching ListPartitioner path. The one-dimensional assign_quantiles is now the quantile specialization of the single engine. Adds AcceptabilityScorer (bead.items.scoring), which scores a forced-choice item by a trained ForcedChoiceModel and returns the predicted preference margin. Extends bead.interop.layers with lenses for every mappable bead model: a FilledTemplate to a resource.Filling, an AnnotationRecord to a judgment.Judgment (grouped per annotator into a judgment.JudgmentSet), a ListConstraint to a judgment.ListConstraint, and an ExperimentList to a resource.Collection of kind stimulus-pool with one membership per item and its list constraints. All round-trip exactly. A Participant is intentionally not mapped: a persona.Persona is an annotator role, not a concrete participant, whose identity is an AgentRef. Rebuilds the gallery/eng/argument_structure example: an acceptability model trained on MegaAcceptability (per-annotator Likert to 2AFC pairing with rater random effects) scores items, the 2AFC pairs are stratified across a grid of the acceptability margin, the language-model score, and the pair type, and the active-learning loop is seeded with the pretrained model. Every artifact (lexicons, templates, filled templates, cross-product items, 2AFC pairs, the training corpus, and experiment lists) is persisted through the layers schema and the bead lairs codec. Bumps the version to 0.8.0. See CHANGELOG for the full entry. --- CHANGELOG.md | 62 +- bead/interop/layers/__init__.py | 26 + bead/interop/layers/judgment_lens.py | 185 +++++ bead/interop/layers/list_lens.py | 244 +++++++ bead/interop/layers/resource_lens.py | 79 +- bead/items/scoring.py | 113 ++- bead/lists/__init__.py | 36 +- bead/lists/balancer.py | 92 ++- bead/lists/constraints.py | 219 ++++++ bead/lists/partitioner.py | 83 ++- bead/lists/stratification.py | 483 +++++++++---- docs/api/interop.md | 21 + docs/api/lists.md | 9 + gallery/eng/argument_structure/Makefile | 41 ++ gallery/eng/argument_structure/README.md | 52 +- gallery/eng/argument_structure/config.yaml | 84 ++- .../argument_structure/create_2afc_pairs.py | 275 ++++--- .../extract_generic_templates.py | 6 + .../eng/argument_structure/fill_templates.py | 6 + .../generate_cross_product.py | 18 +- .../argument_structure/generate_lexicons.py | 15 +- .../eng/argument_structure/generate_lists.py | 107 ++- .../argument_structure/generate_templates.py | 6 + gallery/eng/argument_structure/layers_io.py | 352 +++++++++ .../prepare_megaacceptability.py | 675 ++++++++++++++++++ .../eng/argument_structure/run_pipeline.py | 507 ++++--------- .../train_acceptability_model.py | 326 +++++++++ pyproject.toml | 2 +- tests/interop/test_layers_judgment.py | 87 +++ tests/interop/test_layers_list.py | 189 +++++ tests/interop/test_layers_resource.py | 63 ++ tests/items/test_acceptability_scorer.py | 89 +++ tests/lists/test_grid_stratification.py | 317 ++++++++ uv.lock | 2 +- 34 files changed, 4235 insertions(+), 636 deletions(-) create mode 100644 bead/interop/layers/judgment_lens.py create mode 100644 bead/interop/layers/list_lens.py create mode 100644 gallery/eng/argument_structure/layers_io.py create mode 100644 gallery/eng/argument_structure/prepare_megaacceptability.py create mode 100644 gallery/eng/argument_structure/train_acceptability_model.py create mode 100644 tests/interop/test_layers_judgment.py create mode 100644 tests/interop/test_layers_list.py create mode 100644 tests/items/test_acceptability_scorer.py create mode 100644 tests/lists/test_grid_stratification.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 4c663f9..2703ee2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,64 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +## [0.8.0] - 2026-06-30 + +### 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. + + All of these round-trip exactly. A bead `Participant` is intentionally not + mapped: a layers `persona.Persona` is an annotator role and interpretive + framework, not a concrete enrolled participant, whose identity is instead an + `AgentRef` and whose study fields (consent, sessions, demographics) are outside + the shared schema. + ## [0.7.0] - 2026-06-29 ### Added @@ -564,7 +622,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 diff --git a/bead/interop/layers/__init__.py b/bead/interop/layers/__init__.py index 91d7e2a..96c2156 100644 --- a/bead/interop/layers/__init__.py +++ b/bead/interop/layers/__init__.py @@ -55,6 +55,19 @@ ItemLayersLens, item_to_layers, ) +from bead.interop.layers.judgment_lens import ( + ANNOTATION_RECORD_JUDGMENT, + AnnotationRecordJudgmentLens, + judgment_set_to_records, + records_to_judgment_set, +) +from bead.interop.layers.list_lens import ( + EXPERIMENT_LIST_LAYERS, + LIST_CONSTRAINT, + ExperimentListLayers, + ExperimentListLens, + ListConstraintLens, +) from bead.interop.layers.parse_lens import ( PARSED_SENTENCE_LAYERS, ParsedSentenceLayers, @@ -62,9 +75,11 @@ parse_to_layers, ) from bead.interop.layers.resource_lens import ( + FILLED_TEMPLATE_FILLING, LEXICAL_ITEM_ENTRY, LEXICON_COLLECTION, TEMPLATE_LAYERS, + FilledTemplateFillingLens, LexicalItemEntryLens, LexiconCollectionLens, LexiconLayers, @@ -72,20 +87,29 @@ ) __all__ = [ + "ANNOTATION_RECORD_JUDGMENT", "CORPUS_GRAPH_LAYERS", + "EXPERIMENT_LIST_LAYERS", + "FILLED_TEMPLATE_FILLING", "ITEM_LAYERS", "LEXICAL_ITEM_ENTRY", "LEXICON_COLLECTION", + "LIST_CONSTRAINT", "PARSED_SENTENCE_LAYERS", "RECORD_EXPRESSION", "TEMPLATE_LAYERS", + "AnnotationRecordJudgmentLens", "BeadCodec", "CorpusGraphLayers", "CorpusGraphLayersLens", + "ExperimentListLayers", + "ExperimentListLens", + "FilledTemplateFillingLens", "ItemLayersLens", "LexicalItemEntryLens", "LexiconCollectionLens", "LexiconLayers", + "ListConstraintLens", "ParsedSentenceLayers", "ParsedSentenceLayersIso", "RecordExpressionLens", @@ -98,10 +122,12 @@ "graph_to_layers", "item_to_layers", "items_to_corpus", + "judgment_set_to_records", "load_layers_corpus", "materialize_corpus", "parse_to_layers", "publish_corpus", "record_to_expression", + "records_to_judgment_set", "save_corpus_repo", ] diff --git a/bead/interop/layers/judgment_lens.py b/bead/interop/layers/judgment_lens.py new file mode 100644 index 0000000..253d5e6 --- /dev/null +++ b/bead/interop/layers/judgment_lens.py @@ -0,0 +1,185 @@ +"""Lenses between bead annotation records and layers judgment records. + +Maps bead's reliability inputs to their canonical +:mod:`lairs.records.judgment` counterparts: + +- ``AnnotationRecord`` <-> a layers ``judgment`` +- a tuple of ``AnnotationRecord`` sharing one annotator <-> a layers + ``judgmentSet`` + +A single layers ``judgment`` has no slot for the annotator or the question name; +those belong to the parent ``judgmentSet`` and experiment. The lens therefore +keeps the bead framework identity, the annotator id, the question name, and the +raw item id / response label in the lens complement so reconstruction is exact. +The aggregate is modelled as two plain functions (a tuple is not a single +``dx.Model``), each one delegating to the per-record lens so every record +round-trips exactly. +""" + +from __future__ import annotations + +import didactic.api as dx +from lairs.records import defs, judgment + +from bead.data.base import JsonValue +from bead.evaluation.reliability import AnnotationRecord +from bead.interop.layers._convert import ( + apply_identity, + identity_of, + j_list, + j_obj, + j_str, + object_ref, +) + +# layers requires an experiment reference on every judgment set; bead annotation +# records carry none, so the view uses an empty at-uri and nothing round-trips. +_NO_EXPERIMENT_REF = "" + + +class AnnotationRecordJudgmentLens( + dx.Lens[AnnotationRecord, judgment.Judgment, JsonValue] +): + """Lossless lens ``AnnotationRecord <-> (layers judgment, complement)``. + + The layers ``judgment`` record (:class:`lairs.records.judgment.Judgment`) is + the canonical representation of a single annotator response: it carries the + item reference and the categorical value. The bead-only remainder (identity, + annotator id, question name, and the raw item id / response label) travels in + the lens complement, since a single ``judgment`` has no slot for an annotator + or a question name. + """ + + def forward(self, record: AnnotationRecord) -> tuple[judgment.Judgment, JsonValue]: + """Project an annotation record to a layers judgment and complement.""" + view = judgment.Judgment( + item=object_ref(record.item_id), + categoricalValue=record.response_label, + ) + complement: JsonValue = { + "identity": identity_of(record), + "annotator_id": record.annotator_id, + "question_name": record.question_name, + "item_id": record.item_id, + "response_label": record.response_label, + } + return view, complement + + def backward( + self, view: judgment.Judgment, complement: JsonValue + ) -> AnnotationRecord: + """Reconstruct an annotation record from its layers judgment + complement.""" + comp = j_obj(complement) + record = AnnotationRecord( + annotator_id=j_str(comp["annotator_id"]), + item_id=j_str(comp["item_id"]), + question_name=j_str(comp["question_name"]), + response_label=j_str(comp["response_label"]), + ) + return apply_identity(record, comp["identity"]) + + +ANNOTATION_RECORD_JUDGMENT = AnnotationRecordJudgmentLens() + + +def records_to_judgment_set( + records: tuple[AnnotationRecord, ...], +) -> tuple[judgment.JudgmentSet, JsonValue]: + """Project records sharing one annotator to a layers judgment set. + + Groups a homogeneous tuple of annotation records (all from one annotator) + into a layers ``judgmentSet``, projecting each record through + :data:`ANNOTATION_RECORD_JUDGMENT`. The set's ``agent`` carries the shared + annotator id and its ``createdAt`` is the earliest record creation time. The + complement carries the shared annotator id and the per-record complements, so + :func:`judgment_set_to_records` inverts it exactly. + + Parameters + ---------- + records : tuple[AnnotationRecord, ...] + Annotation records that all share one ``annotator_id``. Must be + non-empty. + + Returns + ------- + tuple[judgment.JudgmentSet, JsonValue] + The layers judgment set view and the lens complement. + + Raises + ------ + ValueError + If ``records`` is empty, or if the records do not all share the same + ``annotator_id``. + + See Also + -------- + judgment_set_to_records : Invert this projection. + """ + if not records: + raise ValueError( + "records_to_judgment_set requires at least one record; " + "an empty tuple has no annotator to attach to the judgment set." + ) + annotator_ids = {record.annotator_id for record in records} + if len(annotator_ids) != 1: + raise ValueError( + "records_to_judgment_set requires every record to share one " + f"annotator_id, got {len(annotator_ids)} distinct ids: " + f"{sorted(annotator_ids)}." + ) + annotator_id = records[0].annotator_id + + judgments: list[judgment.Judgment] = [] + record_complements: list[JsonValue] = [] + for record in records: + judgment_view, record_complement = ANNOTATION_RECORD_JUDGMENT.forward(record) + judgments.append(judgment_view) + record_complements.append(record_complement) + + view = judgment.JudgmentSet( + agent=defs.AgentRef(id=annotator_id), + createdAt=min(record.created_at for record in records), + experimentRef=_NO_EXPERIMENT_REF, + judgments=tuple(judgments), + ) + complement: JsonValue = { + "annotator_id": annotator_id, + "record_complements": tuple(record_complements), + } + return view, complement + + +def judgment_set_to_records( + view: judgment.JudgmentSet, complement: JsonValue +) -> tuple[AnnotationRecord, ...]: + """Reconstruct annotation records from a layers judgment set + complement. + + Inverts :func:`records_to_judgment_set` by zipping each layers ``judgment`` + in the set with its per-record complement and delegating to + :data:`ANNOTATION_RECORD_JUDGMENT`, so every record round-trips exactly. + + Parameters + ---------- + view : judgment.JudgmentSet + The layers judgment set produced by :func:`records_to_judgment_set`. + complement : JsonValue + The matching lens complement (per-record complements plus the shared + annotator id). + + Returns + ------- + tuple[AnnotationRecord, ...] + The reconstructed annotation records, in their original order. + + See Also + -------- + records_to_judgment_set : The forward projection. + """ + comp = j_obj(complement) + record_complements = j_list(comp["record_complements"]) + return tuple( + ANNOTATION_RECORD_JUDGMENT.backward(judgment_view, record_complement) + for judgment_view, record_complement in zip( + view.judgments, record_complements, strict=True + ) + ) diff --git a/bead/interop/layers/list_lens.py b/bead/interop/layers/list_lens.py new file mode 100644 index 0000000..28634b4 --- /dev/null +++ b/bead/interop/layers/list_lens.py @@ -0,0 +1,244 @@ +"""Lenses between bead list-composition models and layers judgment records. + +Maps bead's experimental list models to their canonical +:mod:`lairs.records` counterparts: + +- ``ListConstraint`` <-> a layers ``listConstraint`` (``judgment.ListConstraint``) +- ``ExperimentList`` <-> a layers ``collection`` with its ``collectionMembership`` + records and per-list ``listConstraint`` records, bundled as the bead-side + :class:`ExperimentListLayers` aggregate view + +The bead ``ListConstraint`` union is far richer than the four scalar fields of the +layers ``listConstraint`` record, so the lens carries the entire bead constraint +(serialized with ``model_dump_json``) in the lens complement and reconstructs the +exact subclass with ``ListConstraint.model_validate_json`` on the way back. The +``ExperimentList`` lens keeps the bead framework identity, the numeric list id, the +ordered item references, the presentation order, the list and balance metadata, the +per-constraint satisfaction records, and the per-constraint complements in its own +complement, so reconstruction is exact. +""" + +from __future__ import annotations + +from uuid import UUID + +import didactic.api as dx +from lairs.records import defs, judgment, resource + +from bead.data.base import JsonValue +from bead.interop.layers._convert import ( + apply_identity, + dumps_meta, + identity_of, + j_int, + j_list, + j_obj, + j_str, + loads_meta, +) +from bead.lists.constraints import ( + BalanceConstraint, + ConditionalUniquenessConstraint, + DiversityConstraint, + GroupedQuantileConstraint, + ListConstraint, + QuantileConstraint, + UniquenessConstraint, +) +from bead.lists.experiment_list import ConstraintSatisfaction, ExperimentList + +# An experiment list is a curated pool of stimuli; "stimulus-pool" is the +# layers collection kind documented for exactly this (the bead list number and +# identity ride in the lens complement). +_EXPERIMENT_LIST_KIND = "stimulus-pool" + + +def _constraint_kind(constraint: ListConstraint) -> str: + """Read a list constraint's discriminator tag as a plain string.""" + value = constraint.model_dump().get("constraint_type") + return value if isinstance(value, str) else "" + + +def _target_property(constraint: ListConstraint) -> str | None: + """Return a constraint's ``property_expression`` when its variant has one.""" + if isinstance( + constraint, + ( + UniquenessConstraint, + ConditionalUniquenessConstraint, + BalanceConstraint, + QuantileConstraint, + GroupedQuantileConstraint, + DiversityConstraint, + ), + ): + return constraint.property_expression + return None + + +class ListConstraintLens(dx.Lens[ListConstraint, judgment.ListConstraint, JsonValue]): + """Lossless lens ``ListConstraint <-> (layers list constraint, complement)``. + + The layers ``listConstraint`` record (``judgment.ListConstraint``) keeps only a + ``kind`` slug, an optional ``targetProperty``, an optional formal ``constraint`` + expression, and an optional ``parameters`` map. The bead union carries many more + variant-specific fields, so the projection is lossy on its own; the complement + stores the entire bead constraint as ``model_dump_json`` and reconstruction + resolves the exact subclass through the discriminated-union root. + """ + + def forward( + self, constraint: ListConstraint + ) -> tuple[judgment.ListConstraint, JsonValue]: + """Project a list constraint to a layers list constraint and complement.""" + kind = _constraint_kind(constraint) + target_property = _target_property(constraint) + view = judgment.ListConstraint( + kind=kind, + targetProperty=target_property, + constraint=defs.Constraint( + expression=target_property if target_property is not None else kind, + description=kind, + ), + ) + complement: JsonValue = {"bead": constraint.model_dump_json()} + return view, complement + + def backward( + self, view: judgment.ListConstraint, complement: JsonValue + ) -> ListConstraint: + """Reconstruct a list constraint from its complement. + + The layers view is purely a projection; the exact bead constraint comes + from the serialized model in the complement, whose discriminated-union root + resolves the correct subclass. + """ + comp = j_obj(complement) + return ListConstraint.model_validate_json(j_str(comp["bead"])) + + +LIST_CONSTRAINT = ListConstraintLens() + + +class ExperimentListLayers(dx.Model): + """A layers view of an experiment list. + + Attributes + ---------- + collection : resource.Collection + The list itself, projected to a layers ``collection`` of kind + ``stimulus-pool``. + memberships : tuple[resource.CollectionMembership, ...] + One membership per item reference, in presentation-independent item order, + carrying the item ordinal. + list_constraints : tuple[judgment.ListConstraint, ...] + The per-list constraints, projected via :data:`LIST_CONSTRAINT`. + """ + + collection: dx.Embed[resource.Collection] = dx.field() + memberships: tuple[dx.Embed[resource.CollectionMembership], ...] = dx.field( + default=() + ) + list_constraints: tuple[dx.Embed[judgment.ListConstraint], ...] = dx.field( + default=() + ) + + +class ExperimentListLens(dx.Lens[ExperimentList, ExperimentListLayers, JsonValue]): + """Lossless lens ``ExperimentList <-> (layers collection aggregate, complement)``. + + The layers view bundles a ``collection`` (the list), one + ``collectionMembership`` per item reference (preserving order via the membership + ordinal), and the per-list ``listConstraint`` records. The bead-only remainder + (framework identity, the numeric list id, the ordered item references, the + presentation order, the list and balance metadata, the constraint-satisfaction + records, and the per-constraint complements) travels in the lens complement, so + a round-trip reconstructs the original list exactly. + """ + + def forward( + self, experiment_list: ExperimentList + ) -> tuple[ExperimentListLayers, JsonValue]: + """Project an experiment list to a layers collection aggregate.""" + memberships: list[resource.CollectionMembership] = [] + for ordinal, item_id in enumerate(experiment_list.item_refs): + memberships.append( + resource.CollectionMembership( + collectionRef=str(experiment_list.id), + entryRef=str(item_id), + createdAt=experiment_list.created_at, + ordinal=ordinal, + ) + ) + constraint_views: list[judgment.ListConstraint] = [] + constraint_complements: list[JsonValue] = [] + for constraint in experiment_list.list_constraints: + constraint_view, constraint_complement = LIST_CONSTRAINT.forward(constraint) + constraint_views.append(constraint_view) + constraint_complements.append(constraint_complement) + view = ExperimentListLayers( + collection=resource.Collection( + name=experiment_list.name, + kind=_EXPERIMENT_LIST_KIND, + createdAt=experiment_list.created_at, + ), + memberships=tuple(memberships), + list_constraints=tuple(constraint_views), + ) + presentation_order: JsonValue = ( + tuple(str(item_id) for item_id in experiment_list.presentation_order) + if experiment_list.presentation_order is not None + else None + ) + complement: JsonValue = { + "identity": identity_of(experiment_list), + "list_number": experiment_list.list_number, + "item_refs": tuple(str(item_id) for item_id in experiment_list.item_refs), + "presentation_order": presentation_order, + "list_metadata": dumps_meta(experiment_list.list_metadata), + "balance_metrics": dumps_meta(experiment_list.balance_metrics), + "constraint_satisfaction": tuple( + record.model_dump_json() + for record in experiment_list.constraint_satisfaction + ), + "constraint_complements": tuple(constraint_complements), + } + return view, complement + + def backward( + self, view: ExperimentListLayers, complement: JsonValue + ) -> ExperimentList: + """Reconstruct an experiment list from its layers aggregate and complement.""" + comp = j_obj(complement) + item_refs = tuple(UUID(j_str(value)) for value in j_list(comp["item_refs"])) + order_value = comp["presentation_order"] + presentation_order = ( + None + if order_value is None + else tuple(UUID(j_str(value)) for value in j_list(order_value)) + ) + constraint_complements = j_list(comp["constraint_complements"]) + list_constraints = tuple( + LIST_CONSTRAINT.backward(constraint_view, constraint_complement) + for constraint_view, constraint_complement in zip( + view.list_constraints, constraint_complements, strict=True + ) + ) + constraint_satisfaction = tuple( + ConstraintSatisfaction.model_validate_json(j_str(value)) + for value in j_list(comp["constraint_satisfaction"]) + ) + experiment_list = ExperimentList( + name=view.collection.name, + list_number=j_int(comp["list_number"]), + item_refs=item_refs, + list_constraints=list_constraints, + constraint_satisfaction=constraint_satisfaction, + presentation_order=presentation_order, + list_metadata=loads_meta(comp["list_metadata"]), + balance_metrics=loads_meta(comp["balance_metrics"]), + ) + return apply_identity(experiment_list, comp["identity"]) + + +EXPERIMENT_LIST_LAYERS = ExperimentListLens() diff --git a/bead/interop/layers/resource_lens.py b/bead/interop/layers/resource_lens.py index 864c333..be5517e 100644 --- a/bead/interop/layers/resource_lens.py +++ b/bead/interop/layers/resource_lens.py @@ -6,15 +6,19 @@ - ``LexicalItem`` <-> a layers ``entry`` - ``Lexicon`` <-> a layers ``collection`` with its ``entry`` records - ``Template`` <-> a layers ``template`` (with its slots and constraints) +- ``FilledTemplate`` <-> a layers ``filling`` (with its per-slot fillings) Each lens produces a layers-shaped view from the generated models and keeps the fields that have no layers equivalent (the bead framework identity, tags, the -``LexicalItem`` original ``form`` / free-text ``source``, and the bead DSL -constraint context) in the lens complement, so reconstruction is exact. +``LexicalItem`` original ``form`` / free-text ``source``, the bead DSL +constraint context, and the filled-template slot requirement map) in the lens +complement, so reconstruction is exact. """ from __future__ import annotations +from typing import cast + import didactic.api as dx from lairs.records import defs, resource @@ -24,6 +28,7 @@ dumps_meta, feature_map, identity_of, + j_bool, j_list, j_obj, j_str, @@ -35,6 +40,7 @@ from bead.resources.lexical_item import LexicalItem from bead.resources.lexicon import Lexicon from bead.resources.template import Slot, Template +from bead.templates.filler import FilledTemplate _LEXICON_KIND = "lexicon" @@ -267,3 +273,72 @@ def backward(self, view: resource.Template, complement: JsonValue) -> Template: TEMPLATE_LAYERS = TemplateLayersLens() + + +class FilledTemplateFillingLens(dx.Lens[FilledTemplate, resource.Filling, JsonValue]): + """Lossless lens ``FilledTemplate <-> (layers filling, complement)``. + + The layers ``filling`` record (``resource.Filling``) is the canonical + representation of a filled template: it carries the template reference, the + per-slot fillings, the rendered text, and the filling strategy. The bead-only + remainder (identity, the source template name, the slot requirement map, and + the exact lexical-item fillers) travels in the lens complement. + """ + + def forward(self, filled: FilledTemplate) -> tuple[resource.Filling, JsonValue]: + """Project a filled template to a layers filling and complement.""" + slot_fillings: list[resource.SlotFilling] = [] + filler_complements: dict[str, JsonValue] = {} + for slot_name, item in filled.slot_fillers.items(): + slot_fillings.append( + resource.SlotFilling( + slotName=slot_name, + literalValue=item.lemma, + renderedForm=item.form if item.form is not None else item.lemma, + features=feature_map(item.features), + ) + ) + filler_complements[slot_name] = item.model_dump_json() + view = resource.Filling( + templateRef=filled.template_id, + slotFillings=tuple(slot_fillings), + renderedText=filled.rendered_text, + strategy=filled.strategy_name, + createdAt=filled.created_at, + ) + complement: JsonValue = { + "identity": identity_of(filled), + "template_id": filled.template_id, + "template_name": filled.template_name, + "template_slots": cast("dict[str, JsonValue]", dict(filled.template_slots)), + "slot_order": tuple(filled.slot_fillers), + "filler_complements": filler_complements, + } + return view, complement + + def backward(self, view: resource.Filling, complement: JsonValue) -> FilledTemplate: + """Reconstruct a filled template from its layers filling and complement.""" + comp = j_obj(complement) + filler_complements = j_obj(comp["filler_complements"]) + slot_fillers: dict[str, LexicalItem] = {} + for slot_name_value in j_list(comp["slot_order"]): + slot_name = j_str(slot_name_value) + slot_fillers[slot_name] = LexicalItem.model_validate_json( + j_str(filler_complements[slot_name]) + ) + template_slots = { + slot_name: j_bool(required) + for slot_name, required in j_obj(comp["template_slots"]).items() + } + filled = FilledTemplate( + template_id=j_str(comp["template_id"]), + template_name=j_str(comp["template_name"]), + slot_fillers=slot_fillers, + rendered_text=view.renderedText if view.renderedText is not None else "", + strategy_name=view.strategy if view.strategy is not None else "", + template_slots=template_slots, + ) + return apply_identity(filled, comp["identity"]) + + +FILLED_TEMPLATE_FILLING = FilledTemplateFillingLens() diff --git a/bead/items/scoring.py b/bead/items/scoring.py index d11bc8d..2b0f6c5 100644 --- a/bead/items/scoring.py +++ b/bead/items/scoring.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: + from bead.active_learning.models.forced_choice import ForcedChoiceModel from bead.items.adapters.huggingface import HuggingFaceLanguageModel from collections.abc import Callable from uuid import UUID, uuid4 @@ -401,7 +402,7 @@ def score(self, item: Item) -> float: return self.comparison_fn(option_scores) def _extract_precomputed_scores(self, item: Item) -> list[float] | None: - """Extract precomputed scores from item metadata if available. + """Extract precomputed option scores from item metadata if available. Looks for keys like: lm_score_0, lm_score_1, ... or lm_score_a, lm_score_b, ... @@ -446,3 +447,113 @@ def _extract_precomputed_scores(self, item: Item) -> list[float] | None: break return scores if scores else None + + +class AcceptabilityScorer(ItemScorer): + """Scorer wrapping a trained forced-choice acceptability model. + + Scores a forced-choice item by the model's predicted preference margin + ``|2 * P(prefer first option) - 1|``, a value in ``[0, 1]`` where 0 is a + near-tie (the model has no preference) and 1 is a clear winner. The margin + is the stratification signal that replaces a raw language-model score + difference. + + Parameters + ---------- + model : ForcedChoiceModel + A trained 2AFC model exposing ``predict_proba``. + + Examples + -------- + >>> from bead.items.scoring import AcceptabilityScorer + >>> scorer = AcceptabilityScorer(model) # doctest: +SKIP + >>> margin = scorer.score(item) # doctest: +SKIP + >>> 0.0 <= margin <= 1.0 # doctest: +SKIP + True + """ + + def __init__(self, model: ForcedChoiceModel) -> None: + self._model = model + + @classmethod + def from_checkpoint(cls, path: str | Path) -> AcceptabilityScorer: + """Load a trained forced-choice model from a checkpoint directory. + + Parameters + ---------- + path : str | Path + Directory written by ``ForcedChoiceModel.save``. + + Returns + ------- + AcceptabilityScorer + Scorer wrapping the loaded model. + """ + from bead.active_learning.models.forced_choice import ( # noqa: PLC0415 + ForcedChoiceModel, + ) + from bead.config.active_learning import ( # noqa: PLC0415 + ForcedChoiceModelConfig, + ) + + model = ForcedChoiceModel(ForcedChoiceModelConfig()) + model.load(str(path)) + return cls(model) + + def _participant_ids(self, n_items: int) -> list[str] | None: + """Population-level participant ids for prediction. + + Fixed-effects models take ``None``; mixed-effects models take a constant + unknown id so prediction falls back to the population mean. + """ + mode = self._model.config.mixed_effects.mode + if mode == "fixed": + return None + return ["__population__"] * n_items + + def _margins(self, items: list[Item]) -> list[float]: + """Predicted preference margins for a batch of forced-choice items.""" + proba = self._model.predict_proba(items, self._participant_ids(len(items))) + return [float(2.0 * np.max(row) - 1.0) for row in proba] + + def score(self, item: Item) -> float: + """Return the predicted preference margin for a single item.""" + return self._margins([item])[0] + + def score_batch(self, items: list[Item]) -> list[float]: + """Return predicted preference margins for multiple items.""" + if not items: + return [] + return self._margins(items) + + def score_with_metadata( + self, items: list[Item] + ) -> dict[UUID, dict[str, float | str]]: + """Score items and return the margin plus per-option probabilities. + + Parameters + ---------- + items + Forced-choice items to score. + + Returns + ------- + dict[UUID, dict[str, float | str]] + For each item: ``"score"`` and ``"acceptability_margin"`` (the + preference margin), ``"p_first"`` (probability of the first option), + and ``"predicted_option"`` (argmax option index). + """ + if not items: + return {} + + proba = self._model.predict_proba(items, self._participant_ids(len(items))) + results: dict[UUID, dict[str, float | str]] = {} + for item, row in zip(items, proba, strict=True): + margin = float(2.0 * np.max(row) - 1.0) + results[item.id] = { + "score": margin, + "acceptability_margin": margin, + "p_first": float(row[0]), + "predicted_option": int(np.argmax(row)), + } + return results diff --git a/bead/lists/__init__.py b/bead/lists/__init__.py index c4110c8..226f54f 100644 --- a/bead/lists/__init__.py +++ b/bead/lists/__init__.py @@ -2,20 +2,39 @@ Provides data models for organizing experimental items into balanced lists for presentation to participants. Includes ExperimentList, ListCollection, -and constraint types (uniqueness, balance, quantile, size, diversity, ordering). +constraint types (uniqueness, balance, quantile, grid stratification, size, +diversity, ordering), and the stratification engine that bins items along one +or more dimensions. """ from bead.lists.constraints import ( BalanceConstraint, + BinningSpec, + CategoricalBinning, DiversityConstraint, + EqualWidthBinning, + GridDimension, + GridStratificationConstraint, ListConstraint, OrderingConstraint, + QuantileBinning, QuantileConstraint, SizeConstraint, + StdDevBinning, + ThresholdBinning, UniquenessConstraint, ) from bead.lists.experiment_list import ExperimentList from bead.lists.list_collection import ListCollection +from bead.lists.stratification import ( + assign_grid_cells, + assign_grid_cells_by_uuid, + assign_quantiles, + assign_quantiles_by_uuid, + flatten_cell, + grid_shape, + unflatten_cell, +) __all__ = [ "ExperimentList", @@ -24,7 +43,22 @@ "UniquenessConstraint", "BalanceConstraint", "QuantileConstraint", + "GridStratificationConstraint", + "GridDimension", + "BinningSpec", + "QuantileBinning", + "EqualWidthBinning", + "ThresholdBinning", + "StdDevBinning", + "CategoricalBinning", "DiversityConstraint", "SizeConstraint", "OrderingConstraint", + "assign_quantiles", + "assign_quantiles_by_uuid", + "assign_grid_cells", + "assign_grid_cells_by_uuid", + "grid_shape", + "flatten_cell", + "unflatten_cell", ] diff --git a/bead/lists/balancer.py b/bead/lists/balancer.py index e83d9cd..7e54392 100644 --- a/bead/lists/balancer.py +++ b/bead/lists/balancer.py @@ -118,28 +118,88 @@ def balance( f"got {items_per_quantile_per_list}" ) - # create quantile-based strata + # create quantile-based strata and distribute across lists strata = self._create_strata(item_ids, value_func) + return self._distribute( + strata, self.n_quantiles, n_lists, items_per_quantile_per_list + ) - # initialize lists - lists: list[list[UUID]] = [[] for _ in range(n_lists)] + def balance_by_cell( + self, + item_ids: list[UUID], + cell_func: Callable[[UUID], int], + n_cells: int, + n_lists: int, + items_per_cell_per_list: int, + ) -> list[list[UUID]]: + """Balance items across lists by an arbitrary precomputed stratum id. - # distribute items from each quantile across lists - for q in range(self.n_quantiles): - q_items = strata[q] + Generalizes :meth:`balance` from one-dimensional quantiles to the cells + of an N-dimensional stratification grid. Each item's stratum is its grid + cell id; items in each cell are spread uniformly across lists. - # shuffle items in this quantile - q_items_array = np.array(q_items) - self._rng.shuffle(q_items_array) + Parameters + ---------- + item_ids : list[UUID] + UUIDs of items to balance. + cell_func : Callable[[UUID], int] + Function returning the grid cell id (0 to n_cells-1) for an item. + n_cells : int + Number of grid cells. + n_lists : int + Number of lists to create. + items_per_cell_per_list : int + Target number of items per cell per list. - # distribute to lists - for list_idx in range(n_lists): - # take items for this list - start_idx = list_idx * items_per_quantile_per_list - end_idx = start_idx + items_per_quantile_per_list - list_items = q_items_array[start_idx:end_idx].tolist() + Returns + ------- + list[list[UUID]] + Balanced lists of item UUIDs. - lists[list_idx].extend(list_items) + Raises + ------ + ValueError + If n_cells < 1, n_lists < 1, or items_per_cell_per_list < 1. + """ + if n_cells < 1: + raise ValueError(f"n_cells must be >= 1, got {n_cells}") + if n_lists < 1: + raise ValueError(f"n_lists must be >= 1, got {n_lists}") + if items_per_cell_per_list < 1: + raise ValueError( + f"items_per_cell_per_list must be >= 1, got {items_per_cell_per_list}" + ) + + strata: dict[int, list[UUID]] = {c: [] for c in range(n_cells)} + for item_id in item_ids: + strata[cell_func(item_id)].append(item_id) + + return self._distribute(strata, n_cells, n_lists, items_per_cell_per_list) + + def _distribute( + self, + strata: dict[int, list[UUID]], + n_strata: int, + n_lists: int, + items_per_stratum_per_list: int, + ) -> list[list[UUID]]: + """Distribute each stratum's items uniformly across lists. + + Within each stratum, items are shuffled and sliced contiguously per + list. If a stratum holds fewer items than ``n_lists * + items_per_stratum_per_list``, the later lists receive fewer items. + """ + lists: list[list[UUID]] = [[] for _ in range(n_lists)] + + for stratum in range(n_strata): + stratum_items = strata.get(stratum, []) + stratum_array = np.array(stratum_items, dtype=object) + self._rng.shuffle(stratum_array) + + for list_idx in range(n_lists): + start_idx = list_idx * items_per_stratum_per_list + end_idx = start_idx + items_per_stratum_per_list + lists[list_idx].extend(stratum_array[start_idx:end_idx].tolist()) return lists diff --git a/bead/lists/constraints.py b/bead/lists/constraints.py index 722551a..785e8af 100644 --- a/bead/lists/constraints.py +++ b/bead/lists/constraints.py @@ -24,11 +24,20 @@ "balance", "quantile", "grouped_quantile", + "grid_stratification", "diversity", "size", "ordering", ] +BinningType = typing.Literal[ + "quantile", + "equal_width", + "threshold", + "stddev", + "categorical", +] + BatchConstraintType = typing.Literal[ "coverage", "balance", @@ -244,6 +253,206 @@ def _check_expr(self, value: str) -> str: return _check_non_empty(type(self), value) +class BinningSpec(BeadBaseModel, dx.TaggedUnion, discriminator="binning"): + """Discriminated union root for per-dimension binning strategies. + + A binning strategy maps one dimension's value to a non-negative bin index. + Continuous strategies (``quantile``, ``equal_width``, ``threshold``, + ``stddev``) bin numeric values; the ``categorical`` strategy bins finite + discrete values. Each strategy has a fixed bin count given its parameters + and the data, which is used to flatten an N-dimensional grid cell to a + single integer id. + """ + + +class QuantileBinning(BinningSpec): + """Equal-frequency bins over the empirical quantiles of a numeric value. + + Attributes + ---------- + n_quantiles : int + Number of equal-frequency bins (>= 2). + """ + + binning: typing.Literal["quantile"] + n_quantiles: int = 5 + + __axioms__ = (dx.axiom("n_quantiles >= 2", message="n_quantiles must be >= 2"),) + + +class EqualWidthBinning(BinningSpec): + """Equal-interval bins over the range of a numeric value. + + Attributes + ---------- + n_bins : int + Number of equal-width bins (>= 2). + range_min : float | None + Lower edge of the binned range. ``None`` uses the data minimum. + range_max : float | None + Upper edge of the binned range. ``None`` uses the data maximum. + """ + + binning: typing.Literal["equal_width"] + n_bins: int = 5 + range_min: float | None = None + range_max: float | None = None + + __axioms__ = ( + dx.axiom("n_bins >= 2", message="n_bins must be >= 2"), + dx.axiom( + "range_min == None or range_max == None or range_min < range_max", + message="range_min must be < range_max", + ), + ) + + +class ThresholdBinning(BinningSpec): + """Bins delimited by caller-supplied cut points. + + ``edges`` of length ``k`` yields ``k + 1`` bins: a value below the first + edge falls in bin 0, a value at or above the last edge in bin ``k``. + + Attributes + ---------- + edges : tuple[float, ...] + Strictly increasing cut points (at least one). + """ + + binning: typing.Literal["threshold"] + edges: tuple[float, ...] + + @dx.validates("edges") + def _check_edges(self, value: tuple[float, ...]) -> tuple[float, ...]: + if len(value) < 1: + raise ValueError("edges must contain at least one cut point") + if any(b <= a for a, b in zip(value, value[1:], strict=False)): + raise ValueError("edges must be strictly increasing") + return value + + +class StdDevBinning(BinningSpec): + """Bins delimited by standard-deviation offsets from the mean. + + ``k_values`` of length ``k`` yields ``k + 1`` bins with edges at + ``mean + k_i * sd``. + + Attributes + ---------- + k_values : tuple[float, ...] + Strictly increasing standard-deviation offsets (at least one). + """ + + binning: typing.Literal["stddev"] + k_values: tuple[float, ...] = (-1.0, 0.0, 1.0) + + @dx.validates("k_values") + def _check_k_values(self, value: tuple[float, ...]) -> tuple[float, ...]: + if len(value) < 1: + raise ValueError("k_values must contain at least one offset") + if any(b <= a for a, b in zip(value, value[1:], strict=False)): + raise ValueError("k_values must be strictly increasing") + return value + + +class CategoricalBinning(BinningSpec): + """One bin per distinct discrete value. + + Attributes + ---------- + categories : tuple[str, ...] | None + Ordered category values, one bin each. ``None`` discovers the + categories from the data (sorted for determinism). + include_other : bool + Add a trailing catch-all bin for values outside ``categories``. + Ignored when ``categories`` is ``None``. + """ + + binning: typing.Literal["categorical"] + categories: tuple[str, ...] | None = None + include_other: bool = False + + @dx.validates("categories") + def _check_categories( + self, value: tuple[str, ...] | None + ) -> tuple[str, ...] | None: + if value is None: + return value + if len(value) < 1: + raise ValueError("categories must contain at least one value") + if len(set(value)) != len(value): + raise ValueError("categories must be unique") + return value + + +class GridDimension(BeadBaseModel): + """One axis of a stratification grid: a value plus a binning strategy. + + Attributes + ---------- + property_expression : str + DSL expression returning the value to bin (e.g. + ``"item['acceptability_score_diff']"``). + binning : BinningSpec + Strategy mapping the value to a bin index. + context : dict[str, ContextValue] + Extra DSL evaluation context. + """ + + property_expression: str + binning: dx.Embed[BinningSpec] + context: dict[str, ContextValue] = dx.field(default_factory=dict) + + @dx.validates("property_expression") + def _check_property_expression(self, value: str) -> str: + return _check_non_empty(type(self), value) + + +class GridStratificationConstraint(ListConstraint): + """Uniform distribution across the cells of an N-dimensional bin grid. + + Each dimension bins one value (continuous or discrete) by its own + strategy; the Cartesian product of the per-dimension bins forms a grid, + and items are distributed uniformly across the occupied cells. This is the + N-dimensional, mixed-variable generalization of :class:`QuantileConstraint` + and :class:`GroupedQuantileConstraint`. + + Attributes + ---------- + dimensions : tuple[GridDimension, ...] + Grid axes (at least one). + group_by_expression : str | None + Optional DSL expression; continuous bins are computed within each + group when set. + context : dict[str, ContextValue] + Extra DSL evaluation context. + items_per_cell : int + Target items per grid cell per list (>= 1). + priority : int + Constraint priority. + """ + + constraint_type: typing.Literal["grid_stratification"] + dimensions: tuple[dx.Embed[GridDimension], ...] + group_by_expression: str | None = None + context: dict[str, ContextValue] = dx.field(default_factory=dict) + items_per_cell: int = 2 + priority: int = 1 + + __axioms__ = ( + dx.axiom("items_per_cell >= 1", message="items_per_cell must be >= 1"), + dx.axiom("priority >= 1", message="priority must be >= 1"), + ) + + @dx.validates("dimensions") + def _check_dimensions( + self, value: tuple[GridDimension, ...] + ) -> tuple[GridDimension, ...]: + if len(value) < 1: + raise ValueError("at least one dimension is required") + return value + + class DiversityConstraint(ListConstraint): """Minimum number of unique values for a property within a list. @@ -573,11 +782,21 @@ def _check_property_expression(self, value: str) -> str: | BalanceConstraint | QuantileConstraint | GroupedQuantileConstraint + | GridStratificationConstraint | DiversityConstraint | SizeConstraint | OrderingConstraint ) +# Union of the binning-strategy variants for a stratification grid dimension. +type BinningSpecUnion = ( + QuantileBinning + | EqualWidthBinning + | ThresholdBinning + | StdDevBinning + | CategoricalBinning +) + type BatchConstraintUnion = ( BatchCoverageConstraint | BatchBalanceConstraint diff --git a/bead/lists/partitioner.py b/bead/lists/partitioner.py index 7bccea7..fc02345 100644 --- a/bead/lists/partitioner.py +++ b/bead/lists/partitioner.py @@ -8,6 +8,7 @@ from __future__ import annotations from collections import Counter, defaultdict +from collections.abc import Callable, Hashable from typing import Any from uuid import UUID @@ -22,12 +23,15 @@ BatchCoverageConstraint, BatchDiversityConstraint, BatchMinOccurrenceConstraint, + GridDimension, + GridStratificationConstraint, ListConstraint, QuantileConstraint, SizeConstraint, UniquenessConstraint, ) from bead.lists.experiment_list import ExperimentList, MetadataValue +from bead.lists.stratification import assign_grid_cells, flatten_cell, grid_shape from bead.resources.constraints import ContextValue # Type aliases for clarity @@ -262,7 +266,17 @@ def _partition_stratified( list[ExperimentList] Partitioned lists. """ - # Find quantile constraints + # Prefer an N-dimensional grid constraint when present + grid_constraints = [ + c for c in constraints if isinstance(c, GridStratificationConstraint) + ] + if grid_constraints: + balanced_lists = self._balance_grid( + items, grid_constraints[0], n_lists, metadata + ) + return self._build_lists(balanced_lists, constraints, metadata) + + # Otherwise stratify on the first one-dimensional quantile constraint quantile_constraints = [ c for c in constraints if isinstance(c, QuantileConstraint) ] @@ -292,7 +306,15 @@ def value_func(item_id: UUID) -> float: items, value_func, n_lists, qc.items_per_quantile ) - # Convert to ExperimentList objects + return self._build_lists(balanced_lists, constraints, metadata) + + def _build_lists( + self, + balanced_lists: list[list[UUID]], + constraints: list[ListConstraint], + metadata: MetadataDict, + ) -> list[ExperimentList]: + """Wrap balanced UUID lists in ``ExperimentList`` objects with metrics.""" lists: list[ExperimentList] = [] for i, item_ids in enumerate(balanced_lists): exp_list = ExperimentList( @@ -310,6 +332,63 @@ def value_func(item_id: UUID) -> float: return lists + def _balance_grid( + self, + items: list[UUID], + constraint: GridStratificationConstraint, + n_lists: int, + metadata: MetadataDict, + ) -> list[list[UUID]]: + """Distribute items across lists by their N-dimensional grid cell. + + Bins each item along every grid dimension (continuous bins computed + within ``group_by_expression`` groups when set), flattens the cell to a + single id, and spreads each cell uniformly across lists. + """ + getters = [ + self._make_dimension_getter(dim, metadata) for dim in constraint.dimensions + ] + binnings = [dim.binning for dim in constraint.dimensions] + + stratify_func: Callable[[UUID], Hashable] | None + if constraint.group_by_expression is not None: + group_expr = constraint.group_by_expression + + def group_getter(item_id: UUID) -> Hashable: + return self._extract_property_value( + item_id, group_expr, constraint.context, metadata + ) + + stratify_func = group_getter + else: + stratify_func = None + + coords = assign_grid_cells(items, getters, binnings, stratify_func) + shape = grid_shape(items, getters, binnings) + n_cells = int(np.prod(shape)) if shape else 1 + cell_ids = {item_id: flatten_cell(coords[item_id], shape) for item_id in items} + + balancer = QuantileBalancer(n_quantiles=2, random_seed=self.random_seed) + return balancer.balance_by_cell( + items, + lambda item_id: cell_ids[item_id], + n_cells, + n_lists, + constraint.items_per_cell, + ) + + def _make_dimension_getter( + self, dimension: GridDimension, metadata: MetadataDict + ) -> Callable[[UUID], float | str]: + """Build a value accessor for one grid dimension via its DSL expression.""" + + def getter(item_id: UUID) -> float | str: + return self._extract_property_value( + item_id, dimension.property_expression, dimension.context, metadata + ) + + return getter + def _find_best_list_index( self, item_id: UUID, diff --git a/bead/lists/stratification.py b/bead/lists/stratification.py index 0237eac..4a83076 100644 --- a/bead/lists/stratification.py +++ b/bead/lists/stratification.py @@ -1,170 +1,409 @@ -"""Stratification utilities for quantile-based item assignment. +"""Stratification utilities for binning items into quantile and grid strata. -This module provides utilities for assigning items to quantile bins based -on numeric properties, with optional stratification by grouping variables. +This module assigns items to bins along one or more dimensions. Each dimension +declares a binning strategy via a :class:`~bead.lists.constraints.BinningSpec` +(quantile, equal-width, threshold, standard-deviation, or categorical). The +Cartesian product of the per-dimension bins forms a grid; an item's grid cell +is the tuple of its per-dimension bin indices, optionally flattened to a single +integer id for stand-off storage and balancing. + +The one-dimensional :func:`assign_quantiles` helper is the quantile +specialization of the grid engine, kept for ergonomics. """ from __future__ import annotations from collections import defaultdict -from collections.abc import Callable, Hashable -from typing import Any +from collections.abc import Callable, Hashable, Sequence +from typing import Any, Protocol from uuid import UUID import numpy as np from bead.items.item import MetadataValue +from bead.lists.constraints import ( + BinningSpec, + CategoricalBinning, + EqualWidthBinning, + QuantileBinning, + StdDevBinning, + ThresholdBinning, +) + +# --------------------------------------------------------------------------- +# fitted binners (behavior; not didactic models, since they carry numpy state) +# --------------------------------------------------------------------------- + + +class _Binner(Protocol): + """A fitted binning function mapping a single value to a bin index.""" + + n_bins: int + + def bin(self, value: float | str) -> int: ... + + +class _QuantileBinner: + """Equal-frequency binner using empirical quantile cut points.""" + + def __init__( + self, + cut_points: np.ndarray[Any, np.dtype[np.floating[Any]]], + n_bins: int, + ) -> None: + self._cut = cut_points + self.n_bins = n_bins + + def bin(self, value: float | str) -> int: + idx = int(np.searchsorted(self._cut, float(value))) + return min(max(idx, 0), self.n_bins - 1) + + +class _EqualWidthBinner: + """Equal-interval binner over a fixed range.""" + + def __init__(self, lo: float, width: float, n_bins: int) -> None: + self._lo = lo + self._width = width + self.n_bins = n_bins + + def bin(self, value: float | str) -> int: + if self._width <= 0.0: + return 0 + idx = int((float(value) - self._lo) // self._width) + return min(max(idx, 0), self.n_bins - 1) + + +class _ThresholdBinner: + """Binner delimited by a strictly increasing array of cut points.""" + + def __init__(self, edges: Sequence[float]) -> None: + self._edges = np.asarray(edges, dtype=float) + self.n_bins = len(edges) + 1 + + def bin(self, value: float | str) -> int: + idx = int(np.searchsorted(self._edges, float(value), side="right")) + return min(max(idx, 0), self.n_bins - 1) + + +class _CategoricalBinner: + """Binner mapping each distinct discrete value to its own bin.""" + + def __init__( + self, index: dict[str, int], other_bin: int | None, n_bins: int + ) -> None: + self._index = index + self._other_bin = other_bin + self.n_bins = n_bins + + def bin(self, value: float | str) -> int: + key = str(value) + if key in self._index: + return self._index[key] + if self._other_bin is not None: + return self._other_bin + raise ValueError(f"value {key!r} is outside the declared categories") + + +def _fit_binner(spec: BinningSpec, values: Sequence[float | str]) -> _Binner: + """Fit a binner for one dimension from its spec and the dimension's values.""" + if isinstance(spec, QuantileBinning): + numeric = np.asarray([float(v) for v in values], dtype=float) + edges = np.quantile(numeric, np.linspace(0.0, 1.0, spec.n_quantiles + 1)) + return _QuantileBinner(edges[1:], spec.n_quantiles) + if isinstance(spec, EqualWidthBinning): + numeric = np.asarray([float(v) for v in values], dtype=float) + lo = spec.range_min if spec.range_min is not None else float(np.min(numeric)) + hi = spec.range_max if spec.range_max is not None else float(np.max(numeric)) + width = (hi - lo) / spec.n_bins + return _EqualWidthBinner(lo, width, spec.n_bins) + if isinstance(spec, ThresholdBinning): + return _ThresholdBinner(spec.edges) + if isinstance(spec, StdDevBinning): + numeric = np.asarray([float(v) for v in values], dtype=float) + mean = float(np.mean(numeric)) + sd = float(np.std(numeric)) + return _ThresholdBinner([mean + k * sd for k in spec.k_values]) + assert isinstance(spec, CategoricalBinning) + if spec.categories is not None: + declared = {category: i for i, category in enumerate(spec.categories)} + other = len(spec.categories) if spec.include_other else None + n_bins = len(spec.categories) + (1 if spec.include_other else 0) + return _CategoricalBinner(declared, other, n_bins) + discovered = sorted({str(v) for v in values}) + index = {category: i for i, category in enumerate(discovered)} + return _CategoricalBinner(index, None, len(discovered)) + + +def _fit_binners[T]( + items: list[T], + getters: Sequence[Callable[[T], float | str]], + binnings: Sequence[BinningSpec], +) -> list[_Binner]: + """Fit one binner per dimension over all items.""" + binners: list[_Binner] = [] + for getter, spec in zip(getters, binnings, strict=True): + binners.append(_fit_binner(spec, [getter(item) for item in items])) + return binners -def assign_quantiles[T]( +# --------------------------------------------------------------------------- +# grid stratification engine +# --------------------------------------------------------------------------- + + +def assign_grid_cells[T]( items: list[T], - property_getter: Callable[[T], float], - n_quantiles: int = 10, + getters: Sequence[Callable[[T], float | str]], + binnings: Sequence[BinningSpec], stratify_by: Callable[[T], Hashable] | None = None, -) -> dict[T, int]: - """Assign quantile bins to items based on numeric property. +) -> dict[T, tuple[int, ...]]: + """Assign each item to an N-dimensional grid cell. - Divides items into n_quantiles bins based on the distribution of - a numeric property extracted via property_getter. Optionally stratifies - by a grouping variable, computing separate quantiles for each group. + Each dimension bins one value via its strategy; the item's cell is the + tuple of per-dimension bin indices. When ``stratify_by`` is given, + continuous bins are computed within each group while categorical bin counts + stay global so the grid shape is consistent across groups. Parameters ---------- items : list[T] - List of items to assign to quantile bins. - property_getter : Callable[[T], float] - Function that extracts a numeric value from each item. - This value is used to compute quantiles. - n_quantiles : int - Number of quantile bins (default: 10 for deciles). - Must be >= 2. + Items to assign. + getters : Sequence[Callable[[T], float | str]] + One value accessor per dimension. Numeric strategies coerce to float; + the categorical strategy coerces to str. + binnings : Sequence[BinningSpec] + One binning strategy per dimension; same length as ``getters``. stratify_by : Callable[[T], Hashable] | None - Optional function that extracts a grouping variable from each item. - If provided, quantiles are computed separately for each group. - Groups must be hashable (str, int, UUID, tuple, etc.). + Optional grouping function; continuous bins are fit within each group. Returns ------- - dict[T, int] - Dictionary mapping each item to its quantile bin (0 to n_quantiles-1). + dict[T, tuple[int, ...]] + Mapping from each item to its per-dimension bin-index tuple. Raises ------ ValueError - If n_quantiles < 2 or items list is empty. - - Examples - -------- - Basic usage with simple numeric values: - >>> items = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] - >>> result = assign_quantiles( - ... items, - ... property_getter=lambda x: x, - ... n_quantiles=4 - ... ) - >>> result[1] # First item in lowest quartile - 0 - >>> result[10] # Last item in highest quartile - 3 - - With Item objects and stratification: - >>> from bead.items.item import Item - >>> from uuid import uuid4 - >>> items = [ - ... Item(item_template_id=uuid4(), item_metadata={"score": 10.5, "group": "A"}), - ... Item(item_template_id=uuid4(), item_metadata={"score": 5.2, "group": "A"}), - ... Item(item_template_id=uuid4(), item_metadata={"score": 8.1, "group": "B"}), - ... Item(item_template_id=uuid4(), item_metadata={"score": 3.3, "group": "B"}), - ... ] - >>> result = assign_quantiles( - ... items, - ... property_getter=lambda x: x.item_metadata["score"], - ... n_quantiles=2, - ... stratify_by=lambda x: x.item_metadata["group"] - ... ) # doctest: +SKIP - - With UUID keys (common pattern): - >>> from uuid import UUID - >>> item_uuids = [uuid4() for _ in range(100)] - >>> item_scores = {uid: float(i) for i, uid in enumerate(item_uuids)} - >>> result = assign_quantiles( - ... item_uuids, - ... property_getter=lambda uid: item_scores[uid], - ... n_quantiles=10 - ... ) # doctest: +SKIP + If ``items`` is empty, no dimensions are given, or ``getters`` and + ``binnings`` differ in length. """ if not items: raise ValueError("items list cannot be empty") + if len(getters) != len(binnings): + raise ValueError("getters and binnings must have the same length") + if len(binnings) < 1: + raise ValueError("at least one dimension is required") - if n_quantiles < 2: - raise ValueError(f"n_quantiles must be >= 2, got {n_quantiles}") + global_binners = _fit_binners(items, getters, binnings) - # if no stratification, compute quantiles for all items if stratify_by is None: - return _assign_quantiles_single_group(items, property_getter, n_quantiles) + return { + item: tuple( + binner.bin(getter(item)) + for binner, getter in zip(global_binners, getters, strict=True) + ) + for item in items + } - # stratified: compute quantiles separately for each group groups: dict[Hashable, list[T]] = defaultdict(list) for item in items: - group_key = stratify_by(item) - groups[group_key].append(item) + groups[stratify_by(item)].append(item) - # compute quantiles for each group - result: dict[T, int] = {} + result: dict[T, tuple[int, ...]] = {} for group_items in groups.values(): - group_result = _assign_quantiles_single_group( - group_items, property_getter, n_quantiles + group_binners: list[_Binner] = [] + for dim, (getter, spec) in enumerate(zip(getters, binnings, strict=True)): + if isinstance(spec, CategoricalBinning): + group_binners.append(global_binners[dim]) + else: + group_binners.append( + _fit_binner(spec, [getter(item) for item in group_items]) + ) + for item in group_items: + result[item] = tuple( + binner.bin(getter(item)) + for binner, getter in zip(group_binners, getters, strict=True) + ) + return result + + +def grid_shape[T]( + items: list[T], + getters: Sequence[Callable[[T], float | str]], + binnings: Sequence[BinningSpec], +) -> tuple[int, ...]: + """Return the per-dimension bin counts of the grid. + + The shape is the cardinality of each dimension fit over all items, used to + flatten grid cells to single integer ids. + """ + if not items: + raise ValueError("items list cannot be empty") + return tuple(binner.n_bins for binner in _fit_binners(items, getters, binnings)) + + +def flatten_cell(coords: Sequence[int], shape: Sequence[int]) -> int: + """Flatten a per-dimension bin-index tuple to a row-major integer id.""" + if len(coords) != len(shape): + raise ValueError("coords and shape must have the same length") + cell_id = 0 + for coord, size in zip(coords, shape, strict=True): + cell_id = cell_id * size + coord + return cell_id + + +def unflatten_cell(cell_id: int, shape: Sequence[int]) -> tuple[int, ...]: + """Invert :func:`flatten_cell`, recovering the per-dimension bin indices.""" + coords: list[int] = [] + for size in reversed(shape): + coords.append(cell_id % size) + cell_id //= size + return tuple(reversed(coords)) + + +def assign_grid_cells_by_uuid( + item_ids: list[UUID], + item_metadata: dict[UUID, dict[str, MetadataValue]], + property_keys: Sequence[str], + binnings: Sequence[BinningSpec], + stratify_by_key: str | None = None, +) -> dict[UUID, int]: + """Assign UUID-keyed items to flattened grid cells via metadata lookup. + + Convenience wrapper over :func:`assign_grid_cells` for the stand-off + annotation pattern. Each dimension reads one metadata key; the returned + cell id is the row-major flattening of the per-dimension bin indices. + + Parameters + ---------- + item_ids : list[UUID] + Item UUIDs. + item_metadata : dict[UUID, dict[str, MetadataValue]] + Metadata dict mapping UUIDs to their metadata dicts. + property_keys : Sequence[str] + One metadata key per dimension; same length as ``binnings``. + binnings : Sequence[BinningSpec] + One binning strategy per dimension. + stratify_by_key : str | None + Optional metadata key for stratification. + + Returns + ------- + dict[UUID, int] + Mapping from each UUID to its flattened grid-cell id. + + Raises + ------ + ValueError + If ``property_keys`` and ``binnings`` differ in length, a property key + is missing from some item, or the stratification key is missing. + KeyError + If any UUID is absent from ``item_metadata``. + """ + if len(property_keys) != len(binnings): + raise ValueError("property_keys and binnings must have the same length") + + for uid in item_ids: + if uid not in item_metadata: + raise KeyError(f"UUID {uid} not found in item_metadata") + for key in property_keys: + if key not in item_metadata[uid]: + raise ValueError( + f"Property '{key}' not found in metadata for UUID {uid}" + ) + if stratify_by_key is not None and any( + stratify_by_key not in item_metadata[uid] for uid in item_ids + ): + raise ValueError( + f"Stratification key '{stratify_by_key}' not found in all items" ) - result.update(group_result) - return result + def make_getter(key: str) -> Callable[[UUID], float | str]: + def getter(uid: UUID) -> float | str: + value = item_metadata[uid][key] + return value # type: ignore[return-value] + + return getter + getters = [make_getter(key) for key in property_keys] -def _assign_quantiles_single_group[T]( + stratify_func: Callable[[UUID], Hashable] | None + if stratify_by_key is not None: + strat_key = stratify_by_key + + def stratify_getter(uid: UUID) -> Hashable: + return item_metadata[uid][strat_key] # type: ignore[return-value] + + stratify_func = stratify_getter + else: + stratify_func = None + + shape = grid_shape(item_ids, getters, binnings) + coords = assign_grid_cells(item_ids, getters, binnings, stratify_func) + return {uid: flatten_cell(coords[uid], shape) for uid in item_ids} + + +# --------------------------------------------------------------------------- +# one-dimensional quantile helpers (quantile specialization of the grid engine) +# --------------------------------------------------------------------------- + + +def assign_quantiles[T]( items: list[T], property_getter: Callable[[T], float], - n_quantiles: int, + n_quantiles: int = 10, + stratify_by: Callable[[T], Hashable] | None = None, ) -> dict[T, int]: - """Assign quantiles to items within a single group. + """Assign quantile bins to items based on a numeric property. + + Divides items into ``n_quantiles`` equal-frequency bins based on the + distribution of a numeric property. Optionally stratifies by a grouping + variable, computing separate quantiles for each group. This is the + one-dimensional quantile case of :func:`assign_grid_cells`. Parameters ---------- items : list[T] - List of items in this group. + List of items to assign to quantile bins. property_getter : Callable[[T], float] - Function to extract numeric value. + Function that extracts a numeric value from each item. n_quantiles : int - Number of quantile bins. + Number of quantile bins (default: 10 for deciles). Must be >= 2. + stratify_by : Callable[[T], Hashable] | None + Optional grouping function; quantiles are computed per group. Returns ------- dict[T, int] - Mapping from item to quantile bin (0 to n_quantiles-1). + Dictionary mapping each item to its quantile bin (0 to n_quantiles-1). + + Raises + ------ + ValueError + If n_quantiles < 2 or items list is empty. + + Examples + -------- + >>> items = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + >>> result = assign_quantiles( + ... items, + ... property_getter=lambda x: x, + ... n_quantiles=4, + ... ) + >>> result[1] + 0 + >>> result[10] + 3 """ if not items: - return {} - - # extract scores - scores: np.ndarray[Any, np.dtype[np.floating[Any]]] = np.array( - [property_getter(item) for item in items] - ) - - # compute quantile edges - # linspace(0, 1, n+1) gives [0, 1/n, 2/n, ..., 1] - quantile_edges: np.ndarray[Any, np.dtype[np.floating[Any]]] = np.quantile( - scores, np.linspace(0, 1, n_quantiles + 1) - ) - - # assign each item to a quantile bin - result: dict[T, int] = {} - for item, score in zip(items, scores.tolist(), strict=True): - # searchsorted finds the index where score would be inserted; - # we use quantile_edges[1:] to exclude the 0th edge; - # this maps scores to bins [0, n_quantiles-1] - quantile_idx = int(np.searchsorted(quantile_edges[1:], float(score))) - result[item] = quantile_idx + raise ValueError("items list cannot be empty") + if n_quantiles < 2: + raise ValueError(f"n_quantiles must be >= 2, got {n_quantiles}") - return result + binning = QuantileBinning(binning="quantile", n_quantiles=n_quantiles) + coords = assign_grid_cells(items, [property_getter], [binning], stratify_by) + return {item: cell[0] for item, cell in coords.items()} def assign_quantiles_by_uuid( @@ -176,8 +415,8 @@ def assign_quantiles_by_uuid( ) -> dict[UUID, int]: """Assign quantile bins to items by UUID with metadata lookup. - Convenience function for the common pattern of working with UUIDs - and metadata dictionaries (stand-off annotation pattern). + Convenience function for the common pattern of working with UUIDs and + metadata dictionaries (stand-off annotation pattern). Parameters ---------- @@ -203,22 +442,6 @@ def assign_quantiles_by_uuid( If property_key missing from any item's metadata. KeyError If any UUID not found in item_metadata. - - Examples - -------- - >>> from uuid import uuid4 - >>> uuids = [uuid4() for _ in range(100)] - >>> metadata = { - ... uid: {"score": float(i), "group": "A" if i < 50 else "B"} - ... for i, uid in enumerate(uuids) - ... } - >>> result = assign_quantiles_by_uuid( - ... uuids, - ... metadata, - ... property_key="score", - ... n_quantiles=4, - ... stratify_by_key="group" - ... ) # doctest: +SKIP """ # validate that all items have the property for uid in item_ids: @@ -235,15 +458,17 @@ def property_getter(uid: UUID) -> float: return float(value) # type: ignore[arg-type] # create stratification getter if needed - stratify_func: Callable[[UUID], int | float | str | bool] | None + stratify_func: Callable[[UUID], Hashable] | None if stratify_by_key: if any(stratify_by_key not in item_metadata[uid] for uid in item_ids): raise ValueError( f"Stratification key '{stratify_by_key}' not found in all items" ) - def stratify_getter(uid: UUID) -> int | float | str | bool: - return item_metadata[uid][stratify_by_key] # type: ignore[return-value] + strat_key = stratify_by_key + + def stratify_getter(uid: UUID) -> Hashable: + return item_metadata[uid][strat_key] # type: ignore[return-value] stratify_func = stratify_getter else: diff --git a/docs/api/interop.md b/docs/api/interop.md index 5671fcf..830eeeb 100644 --- a/docs/api/interop.md +++ b/docs/api/interop.md @@ -33,11 +33,32 @@ usage. ## Resource lenses +Lenses between bead's resource models and their layers counterparts: lexical +items and lexicons to entries and collections, templates to layers templates, +and filled templates to layers fillings. + ::: bead.interop.layers.resource_lens options: show_root_heading: true show_source: false +## Judgment and list lenses + +Lenses between bead's response and list-composition models and their layers +counterparts: annotation records to judgments (and judgment sets), bead list +constraints to layers list constraints, and experiment lists to layers +collections with their memberships and constraints. + +::: bead.interop.layers.judgment_lens + options: + show_root_heading: true + show_source: false + +::: bead.interop.layers.list_lens + options: + show_root_heading: true + show_source: false + ## Codec and corpus I/O ::: bead.interop.layers.codec diff --git a/docs/api/lists.md b/docs/api/lists.md index aebc08d..d4dc108 100644 --- a/docs/api/lists.md +++ b/docs/api/lists.md @@ -2,6 +2,15 @@ Stage 4 of the bead pipeline: list partitioning with constraint satisfaction. +Items can be stratified along several dimensions at once. Each dimension of a +`GridStratificationConstraint` declares a binning strategy (`QuantileBinning`, +`EqualWidthBinning`, `ThresholdBinning`, or `StdDevBinning` for continuous +values; `CategoricalBinning` for finite discrete values), and items are +distributed uniformly across the cells of the resulting grid. The +`assign_grid_cells` and `assign_grid_cells_by_uuid` helpers compute the grid +cell for each item; one-dimensional quantile stratification via +`assign_quantiles` is the single-dimension quantile case of the same engine. + ## Core Classes ::: bead.lists.experiment_list diff --git a/gallery/eng/argument_structure/Makefile b/gallery/eng/argument_structure/Makefile index 9fdcad0..fcb34b1 100644 --- a/gallery/eng/argument_structure/Makefile +++ b/gallery/eng/argument_structure/Makefile @@ -16,6 +16,7 @@ # Python interpreter (use uv run with project root for dependency management) PROJECT_ROOT := $(shell cd ../../.. && pwd) PYTHON := uv run --project $(PROJECT_ROOT) python +BEAD := uv run --project $(PROJECT_ROOT) bead # Directories ITEMS_DIR := items @@ -169,6 +170,46 @@ lists: ## Partition 2AFC pairs into experiment lists $(PYTHON) generate_lists.py @echo "$(GREEN)✓ Experiment lists generated$(NC)" +# ============================================================================ +# Acceptability Model Targets (MegaAcceptability -> 2AFC classifier) +# ============================================================================ + +.PHONY: acceptability-data +acceptability-data: ## Map MegaAcceptability Likert ratings to 2AFC training pairs + @echo "$(BLUE)Preparing MegaAcceptability 2AFC pairs...$(NC)" + $(PYTHON) prepare_megaacceptability.py -y + @echo "$(GREEN)✓ Acceptability training pairs prepared$(NC)" + +.PHONY: acceptability-train +acceptability-train: ## Train the acceptability model (needs torch; slow on CPU) + @echo "$(BLUE)Training acceptability model...$(NC)" + $(PYTHON) train_acceptability_model.py -y + @echo "$(GREEN)✓ Acceptability model trained$(NC)" + +.PHONY: acceptability +acceptability: acceptability-data acceptability-train ## Prepare data and train the acceptability model + @echo "$(GREEN)✓ Acceptability model ready for stratification$(NC)" + +# ============================================================================ +# Layers / lairs Interop Targets +# ============================================================================ + +.PHONY: layers-encode +layers-encode: ## Encode the 2AFC item collection as a layers fragment + @echo "$(BLUE)Encoding 2AFC items as a layers fragment...$(NC)" + $(BEAD) layers encode items/2afc_pairs.jsonl --out items/2afc_pairs.layers.json + @echo "$(GREEN)✓ Layers fragment written$(NC)" + +.PHONY: layers-materialize +layers-materialize: ## Materialize the 2AFC items as an Arrow/Parquet corpus + @echo "$(BLUE)Materializing 2AFC items as a layers corpus...$(NC)" + $(BEAD) layers materialize items/2afc_pairs.jsonl --out items/2afc_corpus + @echo "$(GREEN)✓ Layers corpus materialized$(NC)" + +.PHONY: layers +layers: layers-encode layers-materialize ## Produce all layers artifacts for the 2AFC items + @echo "$(GREEN)✓ Layers artifacts produced$(NC)" + .PHONY: deployment deployment: ## Generate jsPsych/JATOS deployment files @echo "$(BLUE)Generating jsPsych/JATOS deployment...$(NC)" diff --git a/gallery/eng/argument_structure/README.md b/gallery/eng/argument_structure/README.md index 3bf3c83..3e97dd6 100644 --- a/gallery/eng/argument_structure/README.md +++ b/gallery/eng/argument_structure/README.md @@ -22,7 +22,7 @@ This project implements a human-in-the-loop active learning pipeline for studyin 1. **Extracts verb-specific templates** from VerbNet with detailed frame information 2. **Generates generic frame structures** by abstracting across verb-specific patterns 3. **Tests all verbs in all frame structures** using a full cross-product design -4. **Generates 2AFC (two-alternative forced choice) pairs** stratified by language model scores +4. **Generates 2AFC (two-alternative forced choice) pairs** stratified across a grid of a MegaAcceptability-trained acceptability model and language model scores 5. **Iteratively collects human judgments** through web-based experiments 6. **Trains predictive models** that converge to human inter-annotator agreement 7. **Detects convergence** using Krippendorff's alpha and other reliability metrics @@ -183,7 +183,7 @@ The pipeline consists of 10 main scripts organized into 4 stages: │ ├─ Load filled templates from previous step │ │ ├─ Score with language model (GPT-2) │ │ ├─ Create minimal pairs (same_verb, different_verb) │ -│ └─ Stratify by LM score quantiles │ +│ └─ Stratify by acceptability x LM score grid │ │ → Output: items/2afc_pairs.jsonl │ │ │ │ 7. generate_lists.py │ @@ -625,9 +625,19 @@ We create two types of pairs: same_verb pairs test alternations by using the sam This format has several advantages over Likert scales. It eliminates scale bias and increases inter-rater reliability, forces relative rather than absolute judgments (making it more sensitive to acceptability gradients), speeds up annotation by removing multi-point scale deliberation, and feels like a natural task that aligns with linguistic intuitions. -### Stratified Sampling +### Acceptability Model -Items are stratified into 10 quantiles based on language model score differences. Quantiles 1-3 contain easy pairs with large LM score differences, quantiles 4-7 have medium pairs with moderate differences, and quantiles 8-10 include hard pairs with small differences near the ceiling. This ensures the model learns across the full difficulty spectrum. +A 2AFC acceptability classifier trained on the MegaAcceptability dataset informs stratification. MegaAcceptability collects single-sentence ordinal (1-7) ratings from many annotators, so `prepare_megaacceptability.py` converts those Likert ratings into 2AFC training pairs by per-annotator within-rater pairing: for each annotator, two sentences they both rated become a forced-choice pair whose gold label is the sentence they rated higher, and the annotator id rides along as a participant id so the model uses participant random effects. `train_acceptability_model.py` then trains a `ForcedChoiceModel` on these pairs. Its predicted preference margin scores each constructed pair (`AcceptabilityScorer`), and it seeds the active-learning loop, which fine-tunes it on collected human judgments. + +### Grid Stratification + +Pairs are stratified across a grid of two signals rather than one. The acceptability-model margin and the language-model score difference are each binned into quantiles and crossed with the categorical pair type, so every experiment list spans the full difficulty space along both signals at once. The grid is expressed as a `grid_stratification` list constraint over `acceptability_score_diff` (quantile), `lm_score_diff` (quantile), and `pair_type` (categorical); each pair stores its flattened grid cell as `stratum_cell`. Grid stratification is general: a dimension can use quantile, equal-width, threshold, or standard-deviation binning for continuous values, or categorical binning for discrete values. + +### Layers Interop + +Every artifact that has a `layers` representation is persisted through the `layers` schema and the `bead` lairs codec, alongside the bead-native JSONL. Lexicons become resource collections, templates become resource templates, filled templates become resource fillings, item sets (the verb x frame cross product and the 2AFC pairs) become layers fragments plus materialized Arrow/Parquet corpora, and experiment lists become `stimulus-pool` collections with one membership per item and their list constraints; the MegaAcceptability source and the derived training pairs are emitted as corpora too. Downstream scripts reload from the canonical fragment. The codec round-trip is law-verified, so identity and construction metadata survive losslessly. + +Collected human responses map onto layers judgments (`AnnotationRecord` to `judgment.Judgment`, grouped per annotator into a `judgment.JudgmentSet`). A study `Participant` is intentionally not mapped: a layers `persona.Persona` is an annotator role and interpretive framework, not a concrete enrolled participant. ## Quick Start @@ -902,7 +912,7 @@ python create_2afc_pairs.py --limit 200 **Input:** `filled_templates/generic_frames_filled.jsonl` **Output:** `items/2afc_pairs.jsonl` -This script loads filled templates from the previous step, scores each filled sentence with a language model (GPT-2), and creates forced-choice pairs. The scoring uses `LanguageModelScorer` to compute log probabilities with caching for efficiency. Pairs are created using `create_forced_choice_items_from_groups`, generating both same_verb pairs (same verb in different frames, testing alternations) and different_verb pairs (different verbs in same frame, testing verb licensing). After pairing, score differences are computed and items are stratified into 10 quantiles using `assign_quantiles_by_uuid`, ensuring the model learns across the full difficulty spectrum. +This script loads filled templates from the previous step, scores each filled sentence with a language model (GPT-2), and creates forced-choice pairs. The scoring uses `LanguageModelScorer` to compute log probabilities with caching for efficiency. Pairs are created using `create_forced_choice_items_from_groups`, generating both same_verb pairs (same verb in different frames, testing alternations) and different_verb pairs (different verbs in same frame, testing verb licensing). Each pair is then scored by the trained acceptability model (`AcceptabilityScorer`, its predicted preference margin stored as `acceptability_score_diff`) and stratified across a grid of the acceptability margin, the language-model score difference, and the pair type using `assign_grid_cells_by_uuid`; the flattened grid cell is stored as `stratum_cell`. The pairs are saved as a bead-native JSONL and as a canonical layers fragment plus an Arrow/Parquet corpus. **Example pair:** ```json @@ -920,7 +930,9 @@ This script loads filled templates from the previous step, scores each filled se "lm_score1": -27.59, "lm_score2": -30.18, "lm_score_diff": 2.59, - "quantile": 3, + "acceptability_score_diff": 0.42, + "accept_p_prefer_a": 0.71, + "stratum_cell": 18, "template_id": "019a2bbc-4c41-7b33-befc-248335924f3f", "template_structure": "{subj} {verb}.", "verb1": "giggle", @@ -960,10 +972,15 @@ lists: target_counts: {same_verb: 50, different_verb: 50} - type: "uniqueness" property_expression: "item.metadata.verb_lemma" - - type: "grouped_quantile" - property_expression: "item.metadata.lm_score_diff" - group_by_expression: "item.metadata.pair_type" - n_quantiles: 10 + - type: "grid_stratification" + items_per_cell: 2 + dimensions: + - property_expression: "item.metadata.acceptability_score_diff" + binning: {type: "quantile", n_quantiles: 5} + - property_expression: "item.metadata.lm_score_diff" + binning: {type: "quantile", n_quantiles: 5} + - property_expression: "item.metadata.pair_type" + binning: {type: "categorical", categories: ["same_verb", "different_verb"]} batch_constraints: - type: "coverage" property_expression: "item.metadata.template_id" @@ -1202,17 +1219,22 @@ template: lists: n_lists: 8 # number of experimental lists items_per_list: 100 # items per list - quantile_bins: 10 # stratification bins + quantile_bins: 5 # per-dimension stratification bins constraints: - type: "balance" property_expression: "item.metadata.pair_type" target_counts: {same_verb: 50, different_verb: 50} - type: "uniqueness" property_expression: "item.metadata.verb_lemma" - - type: "grouped_quantile" - property_expression: "item.metadata.lm_score_diff" - group_by_expression: "item.metadata.pair_type" - n_quantiles: 10 + - type: "grid_stratification" # cross acceptability x LM score x pair type + items_per_cell: 2 + dimensions: + - property_expression: "item.metadata.acceptability_score_diff" + binning: {type: "quantile", n_quantiles: 5} + - property_expression: "item.metadata.lm_score_diff" + binning: {type: "quantile", n_quantiles: 5} + - property_expression: "item.metadata.pair_type" + binning: {type: "categorical", categories: ["same_verb", "different_verb"]} ``` ### 11. Run Production Pipeline diff --git a/gallery/eng/argument_structure/config.yaml b/gallery/eng/argument_structure/config.yaml index c7edfcf..09d4eb3 100644 --- a/gallery/eng/argument_structure/config.yaml +++ b/gallery/eng/argument_structure/config.yaml @@ -27,7 +27,12 @@ paths: filled_templates_dir: "filled_templates" cross_product_items: "items/cross_product_items.jsonl" 2afc_pairs: "items/2afc_pairs.jsonl" + # canonical layers artifacts: every item set is persisted as a layers + # fragment (bead layers encode) plus a materialized Arrow/Parquet corpus + 2afc_pairs_fragment: "items/2afc_pairs.layers.json" + 2afc_corpus_dir: "items/2afc_corpus" experiment_lists: "lists/experiment_lists.jsonl" + experiment_lists_fragment: "lists/experiment_lists.layers.json" # resources resources: @@ -151,11 +156,68 @@ items: - "template_structure" - "pair_type" - "lm_score_diff" + - "acceptability_score_diff" + - "accept_p_prefer_a" + - "stratum_cell" - "quantile" +# acceptability judgment model +# +# A 2AFC classifier trained on MegaAcceptability. Its predicted preference +# margin is one stratification dimension (paired with the language-model score +# in the grid above) and it seeds the active-learning loop, which fine-tunes it +# on collected human data. +acceptability_model: + # MegaAcceptability source: a local CSV path, or a URL fetched into cache_dir + source_path: null + source_url: null + columns: + verb: "verb" + frame: "frame" + sentence: "sentence" + participant_id: "participant" + rating: "response" + + # Likert -> 2AFC conversion: per-annotator within-rater pairing + pairing_strategy: "per_annotator" + pairing_scope: "same_frame" # all, same_verb, or same_frame + drop_ties: true + max_pairs_per_annotator: 200 + seed: 42 + + # derived training set artifacts (read by prepare/train scripts) + paths: + training_items: "items/megaacceptability_2afc.jsonl" + raw_corpus_dir: "items/megaacceptability_raw_corpus" + training_corpus_dir: "items/megaacceptability_2afc_corpus" + + # model + training + model_name: "bert-base-uncased" + encoder_mode: "single_encoder" + mixed_effects_mode: "random_intercepts" + learning_rate: 2e-5 + batch_size: 16 + epochs: 3 + device: "cpu" + checkpoint_dir: "checkpoints/acceptability_init" + +# layers / lairs interop +# +# Every item artifact is persisted as a layers fragment and a materialized +# Arrow/Parquet corpus through the bead lairs codec. +layers: + enabled: true + materialize: true + corpus_authority: "local" + # opt-in PDS publish (dry run by default) + publish: + enabled: false + to_did: null + dry_run: true + # list partitioning lists: - strategy: "quantile_balanced_minimal_pairs" + strategy: "stratified" # grid-cell stratification (see grid_stratification below) n_lists: 16 items_per_list: 50 quantile_bins: 10 @@ -170,11 +232,21 @@ lists: property_expression: "item.metadata.get('verb1', item.metadata.get('verb'))" # no verb appears multiple times in same list - - type: "grouped_quantile" - property_expression: "item.metadata.lm_score_diff" - group_by_expression: "item.metadata.pair_type" - n_quantiles: 10 - items_per_quantile: 5 + # stratify across a grid of the acceptability-model margin and the + # language-model score, split by pair type. Items are spread uniformly + # across the occupied grid cells so each list spans the full difficulty + # space along both signals. + - type: "grid_stratification" + items_per_cell: 2 + dimensions: + - property_expression: "item.metadata.acceptability_score_diff" + binning: { type: "quantile", n_quantiles: 5 } + - property_expression: "item.metadata.lm_score_diff" + binning: { type: "quantile", n_quantiles: 5 } + - property_expression: "item.metadata.pair_type" + binning: + type: "categorical" + categories: ["same_verb", "different_verb"] - type: "diversity" property_expression: "item.metadata.template_id" diff --git a/gallery/eng/argument_structure/create_2afc_pairs.py b/gallery/eng/argument_structure/create_2afc_pairs.py index fdbb71d..7273f43 100755 --- a/gallery/eng/argument_structure/create_2afc_pairs.py +++ b/gallery/eng/argument_structure/create_2afc_pairs.py @@ -18,7 +18,9 @@ from pathlib import Path from uuid import uuid4 +import layers_io import yaml +from protocol import ACCEPTABILITY_ANCHOR_NAME from bead.cli.display import ( confirm, @@ -31,11 +33,10 @@ print_warning, ) from bead.items.forced_choice import create_forced_choice_items_from_groups - -from protocol import ACCEPTABILITY_ANCHOR_NAME from bead.items.item import Item -from bead.items.scoring import LanguageModelScorer -from bead.lists.stratification import assign_quantiles_by_uuid +from bead.items.scoring import AcceptabilityScorer, LanguageModelScorer +from bead.lists.constraints import CategoricalBinning, QuantileBinning +from bead.lists.stratification import assign_grid_cells_by_uuid from bead.templates.filler import FilledTemplate @@ -132,6 +133,20 @@ def score_filled_items_with_lm( return scores +def with_metadata(item: Item, extra: dict[str, object]) -> Item: + """Return a copy of ``item`` with ``extra`` merged into its metadata. + + Items are frozen, so metadata is updated through ``with_`` rather than by + mutating the dict in place (which would not persist). ``None`` values are + skipped so absent fields do not fail metadata validation. + """ + merged: dict[str, object] = dict(item.item_metadata) + for key, value in extra.items(): + if value is not None: + merged[key] = value + return item.with_(item_metadata=merged) + + def create_forced_choice_pairs( items: list[Item], lm_scores: dict[str, float], @@ -142,20 +157,29 @@ def create_forced_choice_pairs( 1. Same-verb pairs (same verb, different frames) 2. Different-verb pairs (different verbs, same frame) """ - # Add scores to item metadata - for item in items: - item.item_metadata["lm_score"] = lm_scores.get(str(item.id), float("-inf")) - # Create lookup dict to avoid O(n) scans for each pair item_lookup = {str(item.id): item for item in items} # Helper to extract text from items def extract_text(item: Item) -> str: - return item.rendered_elements.get("text", "") + text = item.rendered_elements.get("text", "") + return text if isinstance(text, str) else "" + + def lm_diff(a: Item, b: Item) -> float: + return abs(lm_scores.get(str(a.id), 0.0) - lm_scores.get(str(b.id), 0.0)) + + def source_pair(fc_item: Item) -> tuple[Item, Item] | None: + id0 = fc_item.item_metadata.get("source_item_0_id") + id1 = fc_item.item_metadata.get("source_item_1_id") + s0 = item_lookup.get(str(id0)) if id0 is not None else None + s1 = item_lookup.get(str(id1)) if id1 is not None else None + if s0 is None or s1 is None: + return None + return s0, s1 # 1. Create same-verb pairs (group by verb_lemma) with console.status("[bold]Creating same-verb pairs...[/bold]"): - same_verb_items = create_forced_choice_items_from_groups( + same_verb_raw = create_forced_choice_items_from_groups( items=items, group_by=lambda item: item.item_metadata.get("verb_lemma", "unknown"), n_alternatives=2, @@ -163,38 +187,33 @@ def extract_text(item: Item) -> str: include_group_metadata=True, ) - # Add pair_type and additional metadata - for fc_item in same_verb_items: - item1_id = fc_item.item_metadata.get("source_item_0_id") - item2_id = fc_item.item_metadata.get("source_item_1_id") - - # Use lookup dict instead of list comprehension - source_items = [item_lookup.get(item1_id), item_lookup.get(item2_id)] - if all(source_items) and len(source_items) == 2: - fc_item.item_metadata.update( + same_verb_items: list[Item] = [] + for fc_item in same_verb_raw: + pair = source_pair(fc_item) + if pair is None: + continue + s0, s1 = pair + same_verb_items.append( + with_metadata( + fc_item, { "pair_type": "same_verb", - "verb": source_items[0].item_metadata.get("verb_lemma"), - "template1": source_items[0].item_metadata.get( - "template_structure" - ), - "template2": source_items[1].item_metadata.get( - "template_structure" - ), - "lm_score_a": lm_scores.get(str(source_items[0].id), float("-inf")), - "lm_score_b": lm_scores.get(str(source_items[1].id), float("-inf")), - "lm_score_diff": abs( - lm_scores.get(str(source_items[0].id), 0) - - lm_scores.get(str(source_items[1].id), 0) - ), - } + "verb": s0.item_metadata.get("verb_lemma"), + "template1": s0.item_metadata.get("template_structure"), + "template2": s1.item_metadata.get("template_structure"), + "lm_score_a": lm_scores.get(str(s0.id), 0.0), + "lm_score_b": lm_scores.get(str(s1.id), 0.0), + "lm_score_diff": lm_diff(s0, s1), + "anchor": ACCEPTABILITY_ANCHOR_NAME, + }, ) + ) print_success(f"Created {len(same_verb_items):,} same-verb pairs") # 2. Create different-verb pairs (group by template_id) with console.status("[bold]Creating different-verb pairs...[/bold]"): - different_verb_items = create_forced_choice_items_from_groups( + different_verb_raw = create_forced_choice_items_from_groups( items=items, group_by=lambda item: str(item.item_template_id), n_alternatives=2, @@ -202,76 +221,131 @@ def extract_text(item: Item) -> str: include_group_metadata=True, ) - # Add pair_type and additional metadata - for fc_item in different_verb_items: - item1_id = fc_item.item_metadata.get("source_item_0_id") - item2_id = fc_item.item_metadata.get("source_item_1_id") - - # Use lookup dict instead of list comprehension - source_items = [item_lookup.get(item1_id), item_lookup.get(item2_id)] - if all(source_items) and len(source_items) == 2: - fc_item.item_metadata.update( + different_verb_items: list[Item] = [] + for fc_item in different_verb_raw: + pair = source_pair(fc_item) + if pair is None: + continue + s0, s1 = pair + different_verb_items.append( + with_metadata( + fc_item, { "pair_type": "different_verb", - "template_id": str(source_items[0].item_template_id), - "template_structure": source_items[0].item_metadata.get( - "template_structure" - ), - "verb1": source_items[0].item_metadata.get("verb_lemma"), - "verb2": source_items[1].item_metadata.get("verb_lemma"), - "lm_score_a": lm_scores.get(str(source_items[0].id), float("-inf")), - "lm_score_b": lm_scores.get(str(source_items[1].id), float("-inf")), - "lm_score_diff": abs( - lm_scores.get(str(source_items[0].id), 0) - - lm_scores.get(str(source_items[1].id), 0) - ), - } + "template_id": str(s0.item_template_id), + "template_structure": s0.item_metadata.get("template_structure"), + "verb1": s0.item_metadata.get("verb_lemma"), + "verb2": s1.item_metadata.get("verb_lemma"), + "lm_score_a": lm_scores.get(str(s0.id), 0.0), + "lm_score_b": lm_scores.get(str(s1.id), 0.0), + "lm_score_diff": lm_diff(s0, s1), + "anchor": ACCEPTABILITY_ANCHOR_NAME, + }, ) + ) print_success(f"Created {len(different_verb_items):,} different-verb pairs") - # Thread the protocol anchor name onto every pair so downstream - # JATOS-result → AnnotationRecord conversion can match responses - # back to the canonical 2AFC acceptability anchor. - all_pairs = same_verb_items + different_verb_items - for fc_item in all_pairs: - fc_item.item_metadata["anchor"] = ACCEPTABILITY_ANCHOR_NAME - return all_pairs + return same_verb_items + different_verb_items -def assign_quantiles_to_pairs( +def score_pairs_with_acceptability( pair_items: list[Item], - n_quantiles: int = 10, + checkpoint_dir: Path, ) -> list[Item]: - """Assign quantile bins using bead/lists/stratification.py. + """Score 2AFC pairs with the trained acceptability model. + + Returns new pairs carrying ``acceptability_score_diff`` (the predicted + preference margin) and ``accept_p_prefer_a``. When the checkpoint is + missing, the acceptability dimension is filled with zeros so the grid still + partitions. - Stratifies by pair_type so same-verb and different-verb pairs - get separate quantile distributions. + Parameters + ---------- + pair_items : list[Item] + The 2AFC pairs to score. + checkpoint_dir : Path + Directory holding the trained ForcedChoiceModel checkpoint. + + Returns + ------- + list[Item] + The scored pairs. """ - with console.status( - "[bold]Assigning quantiles (stratified by pair_type)...[/bold]" - ): - # Build metadata dict for quantile assignment - item_metadata = {item.id: item.item_metadata for item in pair_items} + if not checkpoint_dir.exists(): + print_warning( + f"Acceptability checkpoint not found at {checkpoint_dir}; " + "filling acceptability_score_diff with 0.0. Train the model first " + "with train_acceptability_model.py for grid stratification." + ) + return [ + with_metadata( + item, + {"acceptability_score_diff": 0.0, "accept_p_prefer_a": 0.5}, + ) + for item in pair_items + ] + + with console.status("[bold]Scoring pairs with acceptability model...[/bold]"): + scorer = AcceptabilityScorer.from_checkpoint(checkpoint_dir) + scored = scorer.score_with_metadata(pair_items) + result = [ + with_metadata( + item, + { + "acceptability_score_diff": float( + scored[item.id]["acceptability_margin"] + ), + "accept_p_prefer_a": float(scored[item.id]["p_first"]), + }, + ) + for item in pair_items + ] + + print_success(f"Scored {len(pair_items):,} pairs with the acceptability model") + return result + - # Get item IDs +def assign_grid_cells_to_pairs( + pair_items: list[Item], + n_quantiles: int = 5, +) -> list[Item]: + """Stratify pairs across a grid of acceptability margin x LM score x pair type. + + Bins the acceptability-model margin and the language-model score difference + into quantiles and crosses them with the categorical pair type, storing the + flattened grid cell id as ``stratum_cell`` on each returned pair. + """ + with console.status("[bold]Assigning grid cells (acceptability x LM)...[/bold]"): + item_metadata = {item.id: dict(item.item_metadata) for item in pair_items} item_ids = [item.id for item in pair_items] - # Assign quantiles stratified by pair_type - quantile_assignments = assign_quantiles_by_uuid( + cell_ids = assign_grid_cells_by_uuid( item_ids=item_ids, item_metadata=item_metadata, - property_key="lm_score_diff", - n_quantiles=n_quantiles, - stratify_by_key="pair_type", + property_keys=[ + "acceptability_score_diff", + "lm_score_diff", + "pair_type", + ], + binnings=[ + QuantileBinning(binning="quantile", n_quantiles=n_quantiles), + QuantileBinning(binning="quantile", n_quantiles=n_quantiles), + CategoricalBinning( + binning="categorical", + categories=("same_verb", "different_verb"), + ), + ], ) - # Add quantile to each item's metadata - for item in pair_items: - item.item_metadata["quantile"] = quantile_assignments[item.id] + result = [ + with_metadata(item, {"stratum_cell": cell_ids[item.id]}) + for item in pair_items + ] - print_success(f"Assigned quantiles to {len(pair_items):,} pairs") - return pair_items + n_cells = len(set(cell_ids.values())) + print_success(f"Assigned {len(pair_items):,} pairs across {n_cells} grid cells") + return result def main( @@ -373,15 +447,19 @@ def main( console.print() - # Assign quantiles - quantile_bins = config["lists"].get("quantile_bins", 10) - pair_items = assign_quantiles_to_pairs(pair_items, n_quantiles=quantile_bins) + # Score with the acceptability model, then stratify across the grid of + # acceptability margin x language-model score x pair type. + checkpoint_dir = base_dir / config["acceptability_model"]["checkpoint_dir"] + pair_items = score_pairs_with_acceptability(pair_items, checkpoint_dir) + quantile_bins = config["lists"].get("quantile_bins", 5) + pair_items = assign_grid_cells_to_pairs(pair_items, n_quantiles=quantile_bins) console.print() except Exception as e: print_error(f"Failed to create pairs: {e}") sys.exit(1) - # Save + # Save: a bead-native JSONL plus the canonical layers fragment and an + # Arrow/Parquet corpus through the bead lairs codec. print_header("Saving Results") try: output_path.parent.mkdir(parents=True, exist_ok=True) @@ -389,8 +467,25 @@ def main( with open(output_path, "w") as f: for item in pair_items: f.write(item.model_dump_json() + "\n") - - print_success(f"Saved {len(pair_items):,} 2AFC pairs\n") + print_success(f"Saved {len(pair_items):,} 2AFC pairs to {output_path.name}") + + if config.get("layers", {}).get("enabled", True): + fragment_path = base_dir / config["paths"]["2afc_pairs_fragment"] + corpus_dir = ( + base_dir / config["paths"]["2afc_corpus_dir"] + if config["layers"].get("materialize", True) + else None + ) + with console.status("[bold]Encoding layers fragment + corpus...[/bold]"): + layers_io.write_items( + pair_items, + name="argument_structure_2afc", + fragment_path=fragment_path, + materialize_dir=corpus_dir, + ) + print_success(f"Wrote layers fragment to {fragment_path.name}") + if corpus_dir is not None: + print_success(f"Materialized layers corpus to {corpus_dir.name}\n") except Exception as e: print_error(f"Failed to save output: {e}") sys.exit(1) diff --git a/gallery/eng/argument_structure/extract_generic_templates.py b/gallery/eng/argument_structure/extract_generic_templates.py index e1a27e6..99709cd 100644 --- a/gallery/eng/argument_structure/extract_generic_templates.py +++ b/gallery/eng/argument_structure/extract_generic_templates.py @@ -14,6 +14,7 @@ from collections import defaultdict from pathlib import Path +import layers_io from utils.constraint_builder import ( build_be_participle_constraint, build_be_subject_agreement_constraint, @@ -282,6 +283,11 @@ def main( print_success(f"Saved {len(all_templates)} templates to {output_path}") + # also persist as layers resource templates + layers_path = output_path.with_suffix(".layers.json") + layers_io.write_templates_layers(list(all_templates), layers_path) + print_success(f"Wrote layers templates to {layers_path}") + # Summary statistics print_header("Summary") diff --git a/gallery/eng/argument_structure/fill_templates.py b/gallery/eng/argument_structure/fill_templates.py index 21e7f28..4181373 100755 --- a/gallery/eng/argument_structure/fill_templates.py +++ b/gallery/eng/argument_structure/fill_templates.py @@ -13,6 +13,7 @@ import sys from pathlib import Path +import layers_io import yaml from utils.renderers import OtherNounRenderer @@ -304,6 +305,11 @@ def count_noun_slots(t: Template) -> int: output_path.parent.mkdir(parents=True, exist_ok=True) write_jsonlines(filled_templates, output_path) print_success(f"Saved filled templates to {output_path}") + + # also persist as layers filling records + layers_path = output_path.with_suffix(".layers.json") + layers_io.write_fillings_layers(filled_templates, layers_path) + print_success(f"Wrote layers fillings to {layers_path}") except Exception as e: print_error(f"Failed to save output: {e}") sys.exit(1) diff --git a/gallery/eng/argument_structure/generate_cross_product.py b/gallery/eng/argument_structure/generate_cross_product.py index b273e90..505cafa 100644 --- a/gallery/eng/argument_structure/generate_cross_product.py +++ b/gallery/eng/argument_structure/generate_cross_product.py @@ -12,6 +12,8 @@ import sys from pathlib import Path +import layers_io + from bead.cli.display import ( confirm, console, @@ -105,6 +107,7 @@ def main( total_combinations = min(output_limit, total_combinations) items_generated = 0 + items_for_layers: list[Item] = [] try: with open(output_path, "w") as f: @@ -134,8 +137,9 @@ def main( }, ) - # Write to file + # Write to file and buffer for the layers corpus f.write(item.model_dump_json() + "\n") + items_for_layers.append(item) items_generated += 1 # Check limit @@ -148,6 +152,18 @@ def main( break print_success(f"Generated {items_generated:,} cross-product items\n") + + # also persist as a canonical layers fragment and Arrow/Parquet corpus + fragment_path = output_dir / "cross_product_items.layers.json" + corpus_dir = output_dir / "cross_product_corpus" + with console.status("[bold]Encoding layers fragment + corpus...[/bold]"): + layers_io.write_items( + items_for_layers, + name="argument_structure_cross_product", + fragment_path=fragment_path, + materialize_dir=corpus_dir, + ) + print_success(f"Wrote layers fragment to {fragment_path}\n") except Exception as e: print_error(f"Failed to generate items: {e}") sys.exit(1) diff --git a/gallery/eng/argument_structure/generate_lexicons.py b/gallery/eng/argument_structure/generate_lexicons.py index 8514d07..214dd41 100644 --- a/gallery/eng/argument_structure/generate_lexicons.py +++ b/gallery/eng/argument_structure/generate_lexicons.py @@ -15,6 +15,7 @@ from pathlib import Path from uuid import UUID +import layers_io import pandas as pd from utils.morphology import MorphologyExtractor from utils.verbnet_parser import VerbNetExtractor @@ -102,7 +103,7 @@ def main(verb_limit: int | None = None) -> None: ) output_path = lexicons_dir / "verbnet_verbs.jsonl" - verb_lexicon.to_jsonl(str(output_path)) + layers_io.save_lexicon(verb_lexicon, output_path) print_success(f"Saved to {output_path}") # 2. generate bleached nouns lexicon @@ -125,7 +126,7 @@ def main(verb_limit: int | None = None) -> None: ) output_path = lexicons_dir / "bleached_nouns.jsonl" - noun_lexicon.to_jsonl(str(output_path)) + layers_io.save_lexicon(noun_lexicon, output_path) print_success(f"Saved to {output_path}") # 3. generate bleached verbs lexicon @@ -147,7 +148,7 @@ def main(verb_limit: int | None = None) -> None: print_success(f"Loaded {verb_count:,} bleached verbs from {csv_path}") output_path = lexicons_dir / "bleached_verbs.jsonl" - bleached_verb_lexicon.to_jsonl(str(output_path)) + layers_io.save_lexicon(bleached_verb_lexicon, output_path) print_success(f"Saved to {output_path}") # 4. generate bleached adjectives lexicon @@ -170,7 +171,7 @@ def main(verb_limit: int | None = None) -> None: ) output_path = lexicons_dir / "bleached_adjectives.jsonl" - adj_lexicon.to_jsonl(str(output_path)) + layers_io.save_lexicon(adj_lexicon, output_path) print_success(f"Saved to {output_path}") # 5. generate prepositions lexicon @@ -193,7 +194,7 @@ def main(verb_limit: int | None = None) -> None: ) output_path = lexicons_dir / "prepositions.jsonl" - prep_lexicon.to_jsonl(str(output_path)) + layers_io.save_lexicon(prep_lexicon, output_path) print_success(f"Saved to {output_path}") # 6. generate determiners lexicon @@ -214,7 +215,7 @@ def main(verb_limit: int | None = None) -> None: print_success(f"Loaded {len(det_lexicon.items):,} determiners from {csv_path}") output_path = lexicons_dir / "determiners.jsonl" - det_lexicon.to_jsonl(str(output_path)) + layers_io.save_lexicon(det_lexicon, output_path) print_success(f"Saved to {output_path}") # 7. generate "be" verb forms lexicon @@ -253,7 +254,7 @@ def main(verb_limit: int | None = None) -> None: ) output_path = lexicons_dir / "be_forms.jsonl" - be_lexicon.to_jsonl(str(output_path)) + layers_io.save_lexicon(be_lexicon, output_path) print_success(f"Saved to {output_path}") # summary diff --git a/gallery/eng/argument_structure/generate_lists.py b/gallery/eng/argument_structure/generate_lists.py index f61b176..eaab600 100644 --- a/gallery/eng/argument_structure/generate_lists.py +++ b/gallery/eng/argument_structure/generate_lists.py @@ -13,6 +13,7 @@ from pathlib import Path from uuid import UUID, uuid4 +import layers_io import yaml from bead.cli.display import ( @@ -26,8 +27,15 @@ from bead.items.item import Item from bead.lists import ( BalanceConstraint, + CategoricalBinning, DiversityConstraint, + EqualWidthBinning, + GridDimension, + GridStratificationConstraint, ListCollection, + QuantileBinning, + StdDevBinning, + ThresholdBinning, UniquenessConstraint, ) from bead.lists.constraints import ( @@ -35,7 +43,7 @@ BatchCoverageConstraint, BatchDiversityConstraint, BatchMinOccurrenceConstraint, - GroupedQuantileConstraint, + BinningSpec, ) from bead.lists.partitioner import ListPartitioner @@ -73,6 +81,48 @@ def load_2afc_pairs(pairs_path: Path) -> tuple[list[Item], dict[UUID, dict]]: return items, metadata_dict +def build_binning(spec: dict) -> BinningSpec: + """Build a binning strategy from a config dict. + + Parameters + ---------- + spec : dict + Mapping with a ``type`` key (quantile, equal_width, threshold, stddev, + or categorical) and strategy-specific parameters. + + Returns + ------- + BinningSpec + The constructed binning strategy. + """ + binning_type = spec["type"] + if binning_type == "quantile": + return QuantileBinning( + binning="quantile", n_quantiles=spec.get("n_quantiles", 5) + ) + if binning_type == "equal_width": + return EqualWidthBinning( + binning="equal_width", + n_bins=spec.get("n_bins", 5), + range_min=spec.get("range_min"), + range_max=spec.get("range_max"), + ) + if binning_type == "threshold": + return ThresholdBinning(binning="threshold", edges=tuple(spec["edges"])) + if binning_type == "stddev": + return StdDevBinning( + binning="stddev", k_values=tuple(spec.get("k_values", (-1.0, 0.0, 1.0))) + ) + if binning_type == "categorical": + categories = spec.get("categories") + return CategoricalBinning( + binning="categorical", + categories=tuple(categories) if categories else None, + include_other=spec.get("include_other", False), + ) + raise ValueError(f"Unknown binning type: {binning_type}") + + def build_constraints(config: dict) -> tuple[list, list]: """Build list and batch constraints from config. @@ -91,28 +141,37 @@ def build_constraints(config: dict) -> tuple[list, list]: if constraint_type == "balance": list_constraints.append( BalanceConstraint( + constraint_type="balance", property_expression=constraint_spec["property_expression"], - target_counts=constraint_spec.get("target_counts", {}), + target_counts=constraint_spec.get("target_counts"), ) ) elif constraint_type == "uniqueness": list_constraints.append( UniquenessConstraint( - property_expression=constraint_spec["property_expression"] + constraint_type="uniqueness", + property_expression=constraint_spec["property_expression"], ) ) - elif constraint_type == "grouped_quantile": + elif constraint_type == "grid_stratification": list_constraints.append( - GroupedQuantileConstraint( - property_expression=constraint_spec["property_expression"], - group_by_expression=constraint_spec["group_by_expression"], - n_quantiles=constraint_spec["n_quantiles"], - items_per_quantile=constraint_spec["items_per_quantile"], + GridStratificationConstraint( + constraint_type="grid_stratification", + dimensions=tuple( + GridDimension( + property_expression=dim["property_expression"], + binning=build_binning(dim["binning"]), + ) + for dim in constraint_spec["dimensions"] + ), + group_by_expression=constraint_spec.get("group_by_expression"), + items_per_cell=constraint_spec.get("items_per_cell", 2), ) ) elif constraint_type == "diversity": list_constraints.append( DiversityConstraint( + constraint_type="diversity", property_expression=constraint_spec["property_expression"], min_unique_values=constraint_spec["min_unique_values"], ) @@ -125,14 +184,16 @@ def build_constraints(config: dict) -> tuple[list, list]: if constraint_type == "coverage": batch_constraints.append( BatchCoverageConstraint( + constraint_type="coverage", property_expression=constraint_spec["property_expression"], - target_values=constraint_spec["target_values"], + target_values=tuple(constraint_spec["target_values"]), min_coverage=constraint_spec.get("min_coverage", 1.0), ) ) elif constraint_type == "balance": batch_constraints.append( BatchBalanceConstraint( + constraint_type="balance", property_expression=constraint_spec["property_expression"], target_distribution=constraint_spec.get("target_distribution", {}), tolerance=constraint_spec.get("tolerance", 0.05), @@ -141,6 +202,7 @@ def build_constraints(config: dict) -> tuple[list, list]: elif constraint_type == "min_occurrence": batch_constraints.append( BatchMinOccurrenceConstraint( + constraint_type="min_occurrence", property_expression=constraint_spec["property_expression"], min_occurrences=constraint_spec["min_occurrences"], ) @@ -148,6 +210,7 @@ def build_constraints(config: dict) -> tuple[list, list]: elif constraint_type == "diversity": batch_constraints.append( BatchDiversityConstraint( + constraint_type="diversity", property_expression=constraint_spec["property_expression"], max_lists_per_value=constraint_spec.get("max_lists_per_value"), ) @@ -179,10 +242,18 @@ def main() -> None: console.print(f" - Items per list: [cyan]{items_per_list}[/cyan]") console.print(f" - Strategy: [cyan]{strategy}[/cyan]\n") - # Load 2AFC pairs + # Load 2AFC pairs, preferring the canonical layers fragment print_header("[2/5] Loading 2AFC Pairs") + fragment_path = base_dir / config["paths"].get("2afc_pairs_fragment", "") pairs_path = base_dir / config["paths"]["2afc_pairs"] - items, metadata_dict = load_2afc_pairs(pairs_path) + if config["paths"].get("2afc_pairs_fragment") and fragment_path.exists(): + print_success(f"Loading from layers fragment {fragment_path.name}") + items = layers_io.read_items(fragment_path) + metadata_dict = { + item.id: {"metadata": dict(item.item_metadata)} for item in items + } + else: + items, metadata_dict = load_2afc_pairs(pairs_path) # Build constraints console.print() @@ -262,7 +333,17 @@ def main() -> None: # Save as JSONL list_collection.to_jsonl(output_path) - print_success(f"Saved {len(experiment_lists)} lists to {output_path}\n") + print_success(f"Saved {len(experiment_lists)} lists to {output_path}") + + # also persist the lists as layers collection aggregates + fragment_rel = config["paths"].get("experiment_lists_fragment") + if fragment_rel: + fragment_path = base_dir / fragment_rel + layers_io.write_experiment_lists_layers( + list(experiment_lists), fragment_path + ) + print_success(f"Wrote layers list collections to {fragment_path.name}") + console.print() # Print summary table print_header("Summary") diff --git a/gallery/eng/argument_structure/generate_templates.py b/gallery/eng/argument_structure/generate_templates.py index 41a44e5..b0f4a20 100644 --- a/gallery/eng/argument_structure/generate_templates.py +++ b/gallery/eng/argument_structure/generate_templates.py @@ -8,6 +8,7 @@ import sys from pathlib import Path +import layers_io from utils.template_generator import generate_templates_for_all_verbs from utils.verbnet_parser import VerbNetExtractor @@ -84,6 +85,11 @@ def main(verb_limit: int | None = None, *, yes: bool = False) -> None: print_success(f"Saved {len(templates):,} templates to {output_path}") + # also persist as layers resource templates + layers_path = output_path.with_suffix(".layers.json") + layers_io.write_templates_layers(list(templates), layers_path) + print_success(f"Wrote layers templates to {layers_path}") + # summary statistics print_header("Template Generation Complete") diff --git a/gallery/eng/argument_structure/layers_io.py b/gallery/eng/argument_structure/layers_io.py new file mode 100644 index 0000000..a3d5a6f --- /dev/null +++ b/gallery/eng/argument_structure/layers_io.py @@ -0,0 +1,352 @@ +"""Persist gallery artifacts through the ``layers`` schema and ``lairs`` codec. + +Every item artifact in this gallery is written as a ``layers`` fragment (the +same document ``bead layers encode`` produces) and reloaded by decoding it back +through the ``bead`` lairs codec, so the canonical on-disk form is the layers +representation rather than an ad hoc JSON dump. A materialized Arrow/Parquet +corpus is emitted alongside for downstream consumers and model training. + +The codec round-trip is law-verified (``encode(decode(x)) == x``), so item +identity and all construction metadata survive the conversion. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from uuid import UUID, uuid4 + +from lairs.integrations.codecs import CorpusFragment +from lairs.records import resource + +from bead.interop.layers import ( + EXPERIMENT_LIST_LAYERS, + BeadCodec, + ExperimentListLayers, + items_to_corpus, + materialize_corpus, +) +from bead.interop.layers.resource_lens import ( + FILLED_TEMPLATE_FILLING, + LEXICON_COLLECTION, + TEMPLATE_LAYERS, + LexiconLayers, +) +from bead.items.item import Item, ItemCollection +from bead.lists.experiment_list import ExperimentList +from bead.resources.lexicon import Lexicon +from bead.resources.template import Template +from bead.templates.filler import FilledTemplate + + +def items_to_collection( + items: list[Item], + *, + name: str, + source_template_collection_id: UUID | None = None, + source_filled_collection_id: UUID | None = None, +) -> ItemCollection: + """Wrap items in an ``ItemCollection`` ready for layers conversion. + + Parameters + ---------- + items : list[Item] + Items to wrap. + name : str + Collection name. + source_template_collection_id : UUID | None + Provenance reference; a fresh id is used when omitted. + source_filled_collection_id : UUID | None + Provenance reference; a fresh id is used when omitted. + + Returns + ------- + ItemCollection + The wrapped collection. + """ + return ItemCollection( + name=name, + source_template_collection_id=source_template_collection_id or uuid4(), + source_filled_collection_id=source_filled_collection_id or uuid4(), + items=tuple(items), + ) + + +def write_fragment(collection: ItemCollection, fragment_path: Path) -> None: + """Write a collection as a ``layers`` fragment JSON document. + + Mirrors ``bead layers encode``: the bead ``ItemCollection`` is decoded into + a ``lairs`` ``CorpusFragment`` and serialized. + + Parameters + ---------- + collection : ItemCollection + Collection to encode. + fragment_path : Path + Output path for the fragment JSON. + """ + fragment = BeadCodec().decode(collection.model_dump_json()) + fragment_path.parent.mkdir(parents=True, exist_ok=True) + fragment_path.write_text(fragment.model_dump_json(), encoding="utf-8") + + +def read_items(fragment_path: Path) -> list[Item]: + """Load items from a ``layers`` fragment JSON document. + + Mirrors ``bead layers decode``: the fragment is decoded back into an + ``ItemCollection`` through the ``bead`` codec. + + Parameters + ---------- + fragment_path : Path + Path to a fragment written by :func:`write_fragment`. + + Returns + ------- + list[Item] + The reconstructed items. + """ + fragment = CorpusFragment.model_validate_json( + fragment_path.read_text(encoding="utf-8") + ) + collection = ItemCollection.model_validate_json( + BeadCodec().encode(fragment.records) + ) + return list(collection.items) + + +def materialize(collection: ItemCollection, out_dir: Path, *, name: str) -> list[Path]: + """Materialize a collection as an Arrow/Parquet ``layers`` corpus. + + Parameters + ---------- + collection : ItemCollection + Collection to materialize. + out_dir : Path + Output directory for the Parquet views. + name : str + Corpus name. + + Returns + ------- + list[Path] + The written view files. + """ + corpus = items_to_corpus(collection, corpus_name=name) + out_dir.mkdir(parents=True, exist_ok=True) + return materialize_corpus(corpus, out_dir) + + +def write_items( + items: list[Item], + *, + name: str, + fragment_path: Path, + materialize_dir: Path | None = None, +) -> ItemCollection: + """Write items as a layers fragment and an optional materialized corpus. + + Parameters + ---------- + items : list[Item] + Items to persist. + name : str + Collection / corpus name. + fragment_path : Path + Output path for the layers fragment. + materialize_dir : Path | None + When set, also materialize an Arrow/Parquet corpus into this directory. + + Returns + ------- + ItemCollection + The collection that was written. + """ + collection = items_to_collection(items, name=name) + write_fragment(collection, fragment_path) + if materialize_dir is not None: + materialize(collection, materialize_dir, name=name) + return collection + + +def write_lexicon_layers(lexicon: Lexicon, path: Path) -> None: + """Write a lexicon as a ``layers`` resource collection. + + Stores the layers view (a ``LexiconLayers`` of ``entry`` records) together + with the bead-only complement so the lexicon round-trips. + + Parameters + ---------- + lexicon : Lexicon + Lexicon to encode. + path : Path + Output path for the layers JSON. + """ + view, complement = LEXICON_COLLECTION.forward(lexicon) + payload = {"view": view.model_dump_json(), "complement": complement} + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload), encoding="utf-8") + + +def read_lexicon_layers(path: Path) -> Lexicon: + """Load a lexicon from a ``layers`` resource collection. + + Parameters + ---------- + path : Path + Path written by :func:`write_lexicon_layers`. + + Returns + ------- + Lexicon + The reconstructed lexicon. + """ + payload = json.loads(path.read_text(encoding="utf-8")) + view = LexiconLayers.model_validate_json(payload["view"]) + return LEXICON_COLLECTION.backward(view, payload["complement"]) + + +def save_lexicon(lexicon: Lexicon, jsonl_path: Path) -> None: + """Save a lexicon as bead-native JSONL and a ``layers`` resource collection. + + Writes ``jsonl_path`` plus a sibling ``.layers.json``. + + Parameters + ---------- + lexicon : Lexicon + Lexicon to save. + jsonl_path : Path + Output path for the bead-native JSONL. + """ + jsonl_path.parent.mkdir(parents=True, exist_ok=True) + lexicon.to_jsonl(str(jsonl_path)) + write_lexicon_layers(lexicon, jsonl_path.with_suffix(".layers.json")) + + +def write_templates_layers(templates: list[Template], path: Path) -> None: + """Write templates as ``layers`` resource templates. + + Each template is mapped through the ``TEMPLATE_LAYERS`` lens to a layers + ``template`` record; the bead-only complement travels alongside so the + template round-trips. + + Parameters + ---------- + templates : list[Template] + Templates to encode. + path : Path + Output path for the layers JSON. + """ + records: list[dict[str, object]] = [] + for template in templates: + view, complement = TEMPLATE_LAYERS.forward(template) + records.append({"view": view.model_dump_json(), "complement": complement}) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(records), encoding="utf-8") + + +def read_templates_layers(path: Path) -> list[Template]: + """Load templates from a ``layers`` resource-template document. + + Parameters + ---------- + path : Path + Path written by :func:`write_templates_layers`. + + Returns + ------- + list[Template] + The reconstructed templates. + """ + records = json.loads(path.read_text(encoding="utf-8")) + templates: list[Template] = [] + for record in records: + view = resource.Template.model_validate_json(record["view"]) + templates.append(TEMPLATE_LAYERS.backward(view, record["complement"])) + return templates + + +def write_fillings_layers(filled: list[FilledTemplate], path: Path) -> None: + """Write filled templates as ``layers`` filling records. + + Each filled template is mapped through ``FILLED_TEMPLATE_FILLING`` to a + layers ``filling`` record; the bead-only complement travels alongside so the + filled template round-trips. + + Parameters + ---------- + filled : list[FilledTemplate] + Filled templates to encode. + path : Path + Output path for the layers JSON. + """ + records: list[dict[str, object]] = [] + for item in filled: + view, complement = FILLED_TEMPLATE_FILLING.forward(item) + records.append({"view": view.model_dump_json(), "complement": complement}) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(records), encoding="utf-8") + + +def read_fillings_layers(path: Path) -> list[FilledTemplate]: + """Load filled templates from a ``layers`` filling document. + + Parameters + ---------- + path : Path + Path written by :func:`write_fillings_layers`. + + Returns + ------- + list[FilledTemplate] + The reconstructed filled templates. + """ + records = json.loads(path.read_text(encoding="utf-8")) + filled: list[FilledTemplate] = [] + for record in records: + view = resource.Filling.model_validate_json(record["view"]) + filled.append(FILLED_TEMPLATE_FILLING.backward(view, record["complement"])) + return filled + + +def write_experiment_lists_layers(lists: list[ExperimentList], path: Path) -> None: + """Write experiment lists as ``layers`` collection aggregates. + + Each list is mapped through ``EXPERIMENT_LIST_LAYERS`` to a layers + ``collection`` with one membership per item and its list constraints; the + bead-only complement travels alongside so the list round-trips. + + Parameters + ---------- + lists : list[ExperimentList] + Experiment lists to encode. + path : Path + Output path for the layers JSON. + """ + records: list[dict[str, object]] = [] + for experiment_list in lists: + view, complement = EXPERIMENT_LIST_LAYERS.forward(experiment_list) + records.append({"view": view.model_dump_json(), "complement": complement}) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(records), encoding="utf-8") + + +def read_experiment_lists_layers(path: Path) -> list[ExperimentList]: + """Load experiment lists from a ``layers`` collection-aggregate document. + + Parameters + ---------- + path : Path + Path written by :func:`write_experiment_lists_layers`. + + Returns + ------- + list[ExperimentList] + The reconstructed experiment lists. + """ + records = json.loads(path.read_text(encoding="utf-8")) + lists: list[ExperimentList] = [] + for record in records: + view = ExperimentListLayers.model_validate_json(record["view"]) + lists.append(EXPERIMENT_LIST_LAYERS.backward(view, record["complement"])) + return lists diff --git a/gallery/eng/argument_structure/prepare_megaacceptability.py b/gallery/eng/argument_structure/prepare_megaacceptability.py new file mode 100644 index 0000000..6242398 --- /dev/null +++ b/gallery/eng/argument_structure/prepare_megaacceptability.py @@ -0,0 +1,675 @@ +#!/usr/bin/env python3 +"""Prepare MegaAcceptability data and derive a per-annotator 2AFC training set. + +MegaAcceptability collects single-sentence ordinal (1-7 Likert) ratings from many +annotators. The 2AFC acceptability experiment in this gallery instead asks which of +two sentences sounds more natural. This script bridges the two via per-annotator +within-rater pairing: for each annotator, every pair of sentences they both rated +becomes one forced-choice training item whose gold label is the option holding the +sentence they rated higher. The annotator id rides along as ``participant_id`` so a +downstream ForcedChoiceModel can fit participant random effects. + +The script: + +1. Loads MegaAcceptability rows from a local CSV (or downloads one named in config). +2. Emits the raw ratings as a layers corpus (guarded; never fatal). +3. Builds per-annotator 2AFC training items via :func:`build_per_annotator_pairs`. +4. Writes the derived items as JSONL and as a layers corpus (guarded). + +Run ``--self-test`` to exercise the pairing logic on a small in-memory dataset +without any network access, file IO, or heavy model dependencies. +""" + +from __future__ import annotations + +import argparse +import csv +import shutil +import sys +import urllib.error +import urllib.request +from collections.abc import Iterator +from dataclasses import dataclass +from itertools import combinations +from pathlib import Path +from random import Random +from typing import Literal +from uuid import uuid4 + +import yaml +from protocol import ACCEPTABILITY_ANCHOR_NAME + +from bead.cli.display import ( + confirm, + console, + create_summary_table, + print_error, + print_header, + print_info, + print_success, + print_warning, +) +from bead.items.item import Item, ItemCollection, MetadataValue + +type PairingScope = Literal["all", "same_verb", "same_frame"] + +OPTION_A = "option_a" +OPTION_B = "option_b" +DEFAULT_SEED = 42 +DEFAULT_SCOPE: PairingScope = "all" +VALID_SCOPES: tuple[PairingScope, ...] = ("all", "same_verb", "same_frame") + + +@dataclass(frozen=True) +class RatingRow: + """A single MegaAcceptability rating by one annotator. + + Attributes + ---------- + participant_id + Identifier of the annotator who produced the rating. + verb + Verb associated with the sentence. + frame + Syntactic frame associated with the sentence. + sentence + Surface text that was rated. + rating + Ordinal acceptability rating, coerced to float (1-7 Likert). + """ + + participant_id: str + verb: str + frame: str + sentence: str + rating: float + + +# --- pure pairing logic ----------------------------------------------------- + + +def _pair_type(row_a: RatingRow, row_b: RatingRow) -> str: + """Classify a pair as ``same_verb``, ``same_frame``, or ``cross``.""" + if row_a.verb == row_b.verb: + return "same_verb" + if row_a.frame == row_b.frame: + return "same_frame" + return "cross" + + +def _group_rows( + rows: list[RatingRow], key: Literal["verb", "frame"] +) -> list[list[RatingRow]]: + """Group rows by verb or frame, returning groups in sorted key order.""" + groups: dict[str, list[RatingRow]] = {} + for row in rows: + group_key = row.verb if key == "verb" else row.frame + groups.setdefault(group_key, []).append(row) + return [groups[group_key] for group_key in sorted(groups)] + + +def _candidate_pairs( + rows: list[RatingRow], scope: PairingScope +) -> Iterator[tuple[RatingRow, RatingRow]]: + """Yield within-scope unordered pairs of an annotator's rated sentences.""" + if scope == "same_verb": + groups = _group_rows(rows, "verb") + elif scope == "same_frame": + groups = _group_rows(rows, "frame") + else: + groups = [rows] + for group in groups: + yield from combinations(group, 2) + + +def _build_pair_item( + row_i: RatingRow, + row_j: RatingRow, + *, + participant_id: str, + rng: Random, +) -> Item: + """Build one forced-choice item from a within-rater sentence pair. + + The orientation (which sentence becomes ``option_a``) is drawn from ``rng`` so + that, across many pairs, the gold label is balanced between the two options. + + Parameters + ---------- + row_i, row_j + The two rows rated by the same annotator. + participant_id + Identifier of the annotator, attached as item metadata. + rng + Seeded random generator driving deterministic orientation. + + Returns + ------- + Item + A 2AFC item with ``option_a`` / ``option_b`` rendered elements and a + ``label`` metadata key naming the higher-rated option. + """ + if rng.random() < 0.5: + option_a, option_b = row_i, row_j + else: + option_a, option_b = row_j, row_i + + if option_a.rating > option_b.rating: + label = OPTION_A + elif option_b.rating > option_a.rating: + label = OPTION_B + else: + # Exact tie: only reached when drop_ties is False. Pick a deterministic + # option and flag the pair so callers can filter ties downstream. + label = OPTION_A + + is_tie = option_a.rating == option_b.rating + shared_verb = option_a.verb if option_a.verb == option_b.verb else None + shared_frame = option_a.frame if option_a.frame == option_b.frame else None + + metadata: dict[str, MetadataValue] = { + "participant_id": participant_id, + "verb": shared_verb, + "frame": shared_frame, + "verb_option_a": option_a.verb, + "verb_option_b": option_b.verb, + "frame_option_a": option_a.frame, + "frame_option_b": option_b.frame, + "rating_option_a": option_a.rating, + "rating_option_b": option_b.rating, + "pair_type": _pair_type(option_a, option_b), + "label": label, + "is_tie": is_tie, + "anchor": ACCEPTABILITY_ANCHOR_NAME, + } + + return Item( + item_template_id=uuid4(), + rendered_elements={OPTION_A: option_a.sentence, OPTION_B: option_b.sentence}, + options=(option_a.sentence, option_b.sentence), + item_metadata=metadata, + ) + + +def build_per_annotator_pairs( + rows: list[RatingRow], + *, + scope: PairingScope = DEFAULT_SCOPE, + drop_ties: bool = True, + max_per_annotator: int | None = None, + seed: int = DEFAULT_SEED, +) -> list[Item]: + """Derive 2AFC training items by within-rater pairing of MegaAcceptability rows. + + For each annotator, every within-scope pair of sentences they both rated yields + one forced-choice item. The gold label names the option holding the sentence the + annotator rated higher; pairs with equal ratings are dropped unless ``drop_ties`` + is False. Pairs per annotator are capped at ``max_per_annotator`` by shuffling and + keeping the first N. All randomness derives from ``seed`` so the output is + reproducible. + + Parameters + ---------- + rows + Flat list of per-annotator ratings. + scope + Which pairs to form: ``all`` pairs, only ``same_verb`` pairs, or only + ``same_frame`` pairs. + drop_ties + Whether to drop pairs with equal ratings. When False, ties are kept with a + deterministic label and an ``is_tie`` metadata flag. + max_per_annotator + Optional cap on the number of pairs emitted per annotator. + seed + Seed for the deterministic random generator. + + Returns + ------- + list[Item] + Forced-choice training items, one per retained sentence pair. + + Examples + -------- + >>> rows = [ + ... RatingRow("p1", "give", "NP_NP", "She gave him a book.", 7.0), + ... RatingRow("p1", "give", "NP_PP", "She gave a book to him.", 5.0), + ... ] + >>> items = build_per_annotator_pairs(rows, scope="same_verb") + >>> items[0].item_metadata["participant_id"] + 'p1' + """ + by_participant: dict[str, list[RatingRow]] = {} + for row in rows: + by_participant.setdefault(row.participant_id, []).append(row) + + rng = Random(seed) + items: list[Item] = [] + for participant_id in sorted(by_participant): + participant_rows = sorted( + by_participant[participant_id], key=lambda row: row.sentence + ) + participant_items: list[Item] = [] + for row_i, row_j in _candidate_pairs(participant_rows, scope): + if drop_ties and row_i.rating == row_j.rating: + continue + participant_items.append( + _build_pair_item(row_i, row_j, participant_id=participant_id, rng=rng) + ) + if max_per_annotator is not None and len(participant_items) > max_per_annotator: + rng.shuffle(participant_items) + participant_items = participant_items[:max_per_annotator] + items.extend(participant_items) + return items + + +# --- raw item construction -------------------------------------------------- + + +def build_raw_items(rows: list[RatingRow]) -> list[Item]: + """Build one raw item per rating, carrying the rating in item metadata.""" + items: list[Item] = [] + for row in rows: + metadata: dict[str, MetadataValue] = { + "participant_id": row.participant_id, + "verb": row.verb, + "frame": row.frame, + "rating": row.rating, + "anchor": ACCEPTABILITY_ANCHOR_NAME, + } + items.append( + Item( + item_template_id=uuid4(), + rendered_elements={"text": row.sentence}, + item_metadata=metadata, + ) + ) + return items + + +# --- IO helpers ------------------------------------------------------------- + + +def load_config(config_path: Path) -> dict[str, object]: + """Load the YAML configuration file.""" + with open(config_path, encoding="utf-8") as handle: + return yaml.safe_load(handle) + + +def get_acceptability_config(config: dict[str, object]) -> dict[str, object]: + """Return the ``acceptability_model`` section, or an empty dict if absent.""" + section = config.get("acceptability_model") + return section if isinstance(section, dict) else {} + + +def resolve_columns(section: dict[str, object]) -> dict[str, str]: + """Resolve CSV column names from config, falling back to sensible defaults.""" + columns = section.get("columns") + columns = columns if isinstance(columns, dict) else {} + return { + "verb": str(columns.get("verb", "verb")), + "frame": str(columns.get("frame", "frame")), + "sentence": str(columns.get("sentence", "sentence")), + "participant_id": str(columns.get("participant_id", "participant")), + "rating": str(columns.get("rating", "rating")), + } + + +def resolve_scope(section: dict[str, object]) -> PairingScope: + """Resolve the pairing scope, defaulting to ``all`` for unknown values.""" + scope = section.get("pairing_scope", DEFAULT_SCOPE) + if scope in VALID_SCOPES: + return scope # type: ignore[return-value] + print_warning(f"Unknown pairing_scope '{scope}'; using '{DEFAULT_SCOPE}'") + return DEFAULT_SCOPE + + +def download_csv(url: str, target: Path) -> Path | None: + """Download a CSV from ``url`` into ``target``, returning None on failure.""" + target.parent.mkdir(parents=True, exist_ok=True) + try: + with console.status(f"[bold]Downloading {url}...[/bold]"): + with urllib.request.urlopen(url) as response, open(target, "wb") as out: + shutil.copyfileobj(response, out) + except (urllib.error.URLError, OSError, ValueError) as exc: + print_error(f"Could not download {url}: {exc}") + return None + return target + + +def ensure_source_csv( + *, + source_path: Path | None, + source_url: str | None, + cache_dir: Path, +) -> Path | None: + """Resolve the MegaAcceptability CSV, downloading into the cache if needed.""" + if source_path is not None and source_path.exists(): + return source_path + if source_url is None: + print_error( + "No MegaAcceptability source found. Set acceptability_model.source_path " + "to a local CSV, or acceptability_model.source_url to download one." + ) + return None + cached = cache_dir / "megaacceptability.csv" + if cached.exists(): + return cached + return download_csv(source_url, cached) + + +def load_rating_rows( + csv_path: Path, + *, + columns: dict[str, str], + limit: int | None = None, +) -> list[RatingRow]: + """Read MegaAcceptability rows from CSV, coercing ratings to float.""" + rows: list[RatingRow] = [] + with open(csv_path, newline="", encoding="utf-8") as handle: + reader = csv.DictReader(handle) + for index, record in enumerate(reader): + if limit is not None and index >= limit: + break + raw_rating = record.get(columns["rating"]) + try: + rating = float(raw_rating) + except (TypeError, ValueError): + continue + rows.append( + RatingRow( + participant_id=str(record.get(columns["participant_id"], "")), + verb=str(record.get(columns["verb"], "")), + frame=str(record.get(columns["frame"], "")), + sentence=str(record.get(columns["sentence"], "")), + rating=rating, + ) + ) + return rows + + +def write_items_jsonl(items: list[Item], path: Path) -> None: + """Write items to JSONL, one ``model_dump_json`` line each.""" + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, "w", encoding="utf-8") as handle: + for item in items: + handle.write(item.model_dump_json() + "\n") + + +def emit_layers_corpus(items: list[Item], *, corpus_name: str, out_dir: Path) -> None: + """Emit items as a materialized layers corpus, warning (not failing) on error.""" + try: + from bead.interop.layers import ( # noqa: PLC0415 + items_to_corpus, + materialize_corpus, + ) + except Exception as exc: # noqa: BLE001 + print_warning(f"Skipping layers corpus '{corpus_name}': {exc}") + return + try: + collection = ItemCollection( + name=corpus_name, + source_template_collection_id=uuid4(), + source_filled_collection_id=uuid4(), + items=tuple(items), + ) + out_dir.mkdir(parents=True, exist_ok=True) + corpus = items_to_corpus(collection, corpus_name=corpus_name) + written = materialize_corpus(corpus, out_dir) + print_success( + f"Wrote layers corpus '{corpus_name}' ({len(written)} files) to {out_dir}" + ) + except Exception as exc: # noqa: BLE001 + print_warning(f"Could not emit layers corpus '{corpus_name}': {exc}") + + +# --- self-test -------------------------------------------------------------- + + +def _synthetic_rows() -> list[RatingRow]: + """Build a tiny synthetic dataset: 3 annotators each rating 4 sentences.""" + verbs = ("give", "send", "throw") + frames = ("NP_NP", "NP_PP") + rows: list[RatingRow] = [] + for participant_index in range(3): + participant_id = f"annotator_{participant_index}" + for sentence_index in range(4): + verb = verbs[sentence_index % len(verbs)] + frame = frames[sentence_index % len(frames)] + sentence = f"The {verb} sentence number {sentence_index}." + rating = float(1 + (sentence_index + participant_index) % 7) + rows.append(RatingRow(participant_id, verb, frame, sentence, rating)) + return rows + + +def run_self_test() -> int: + """Exercise the pairing logic on synthetic data without IO or heavy imports.""" + print_header("prepare_megaacceptability self-test") + rows = _synthetic_rows() + expected_participants = {row.participant_id for row in rows} + + pairs = build_per_annotator_pairs( + rows, + scope="all", + drop_ties=True, + max_per_annotator=None, + seed=DEFAULT_SEED, + ) + + assert pairs, "expected at least one derived 2AFC pair" + + seen_participants: set[str] = set() + for item in pairs: + metadata = item.item_metadata + seen_participants.add(str(metadata["participant_id"])) + + rating_a = metadata["rating_option_a"] + rating_b = metadata["rating_option_b"] + higher_option = OPTION_A if rating_a > rating_b else OPTION_B + assert metadata["label"] == higher_option, ( + "gold label must point to the higher-rated sentence " + f"(option_a={rating_a}, option_b={rating_b}, label={metadata['label']})" + ) + assert item.rendered_elements[OPTION_A], "option_a text must be present" + assert item.rendered_elements[OPTION_B], "option_b text must be present" + + assert seen_participants == expected_participants, ( + "participant_ids did not propagate to the derived pairs: " + f"saw {sorted(seen_participants)}, expected {sorted(expected_participants)}" + ) + + print_success( + f"Self-test passed: {len(pairs)} pairs from " + f"{len(expected_participants)} annotators, all labels and " + "participant_ids verified" + ) + return 0 + + +# --- main ------------------------------------------------------------------- + + +def main( + config_path: Path = Path("config.yaml"), + item_limit: int | None = None, + *, + self_test: bool = False, + yes: bool = False, +) -> None: + """Prepare MegaAcceptability data and derive a per-annotator 2AFC training set. + + Parameters + ---------- + config_path + Path to the gallery configuration file. + item_limit + Optional cap on the number of CSV rows to read (for quick testing). + self_test + When True, run the in-memory pairing self-test and exit. + yes + Skip overwrite confirmation prompts for non-interactive use. + """ + if self_test: + sys.exit(run_self_test()) + + base_dir = Path(__file__).parent + if not config_path.exists(): + config_path = base_dir / config_path.name + if not config_path.exists(): + print_error(f"Config file not found: {config_path}") + sys.exit(1) + + config = load_config(config_path) + section = get_acceptability_config(config) + columns = resolve_columns(section) + scope = resolve_scope(section) + drop_ties = bool(section.get("drop_ties", True)) + seed = int(section.get("seed", DEFAULT_SEED)) + max_per_annotator = section.get("max_pairs_per_annotator") + if max_per_annotator is not None: + max_per_annotator = int(max_per_annotator) + + paths_section = section.get("paths") + paths_section = paths_section if isinstance(paths_section, dict) else {} + base_paths = config.get("paths") + base_paths = base_paths if isinstance(base_paths, dict) else {} + + cache_dir = base_dir / str(base_paths.get("cache_dir", ".cache")) + training_items_path = base_dir / str( + paths_section.get("training_items", "items/megaacceptability_2afc.jsonl") + ) + raw_corpus_dir = base_dir / str( + paths_section.get("raw_corpus_dir", "items/megaacceptability_raw_corpus") + ) + training_corpus_dir = base_dir / str( + paths_section.get("training_corpus_dir", "items/megaacceptability_2afc_corpus") + ) + + source_path_value = section.get("source_path") + source_path = base_dir / str(source_path_value) if source_path_value else None + source_url = section.get("source_url") + source_url = str(source_url) if source_url else None + + print_header("MegaAcceptability Preparation") + console.print(f"Base directory: [cyan]{base_dir}[/cyan]") + console.print(f"Pairing scope: [cyan]{scope}[/cyan]") + console.print(f"Output items: [cyan]{training_items_path}[/cyan]\n") + + if training_items_path.exists() and not yes: + if not confirm(f"Overwrite {training_items_path}?", default=False): + print_info("Operation cancelled.") + return + + if item_limit: + print_warning(f"Test mode: limiting to {item_limit:,} rows\n") + + # 1. Resolve and load the raw dataset + print_header("1/4 Loading MegaAcceptability") + csv_path = ensure_source_csv( + source_path=source_path, source_url=source_url, cache_dir=cache_dir + ) + if csv_path is None: + sys.exit(1) + try: + rows = load_rating_rows(csv_path, columns=columns, limit=item_limit) + except Exception as exc: # noqa: BLE001 + print_error(f"Failed to read {csv_path}: {exc}") + sys.exit(1) + if not rows: + print_error(f"No usable rows read from {csv_path}.") + sys.exit(1) + n_annotators = len({row.participant_id for row in rows}) + print_success(f"Loaded {len(rows):,} ratings from {n_annotators:,} annotators\n") + + # 2. Emit the raw dataset as a layers corpus (guarded) + print_header("2/4 Emitting Raw Layers Corpus") + emit_layers_corpus( + build_raw_items(rows), + corpus_name="megaacceptability_raw", + out_dir=raw_corpus_dir, + ) + console.print() + + # 3. Derive per-annotator 2AFC pairs + print_header("3/4 Building Per-Annotator 2AFC Pairs") + with console.status("[bold]Pairing within-rater sentences...[/bold]"): + pair_items = build_per_annotator_pairs( + rows, + scope=scope, + drop_ties=drop_ties, + max_per_annotator=max_per_annotator, + seed=seed, + ) + if not pair_items: + print_error("No 2AFC pairs were derived. Exiting.") + sys.exit(1) + print_success(f"Derived {len(pair_items):,} forced-choice training items\n") + + # 4. Persist the derived training set (JSONL + layers corpus) + print_header("4/4 Persisting Derived Training Set") + try: + write_items_jsonl(pair_items, training_items_path) + print_success(f"Wrote {len(pair_items):,} items to {training_items_path}") + except Exception as exc: # noqa: BLE001 + print_error(f"Failed to write {training_items_path}: {exc}") + sys.exit(1) + emit_layers_corpus( + pair_items, + corpus_name="megaacceptability_2afc", + out_dir=training_corpus_dir, + ) + console.print() + + # Summary + print_header("Summary") + n_option_a = sum( + 1 for item in pair_items if item.item_metadata.get("label") == OPTION_A + ) + n_option_b = len(pair_items) - n_option_a + table = create_summary_table( + { + "Ratings loaded": f"{len(rows):,}", + "Annotators": f"{n_annotators:,}", + "Derived pairs": f"{len(pair_items):,}", + "Gold option_a / option_b": f"{n_option_a:,} / {n_option_b:,}", + "Training items (JSONL)": str(training_items_path), + } + ) + console.print(table) + print_info("Next: run train_acceptability_model.py to fit the initial model") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Prepare MegaAcceptability and derive per-annotator 2AFC pairs" + ) + parser.add_argument( + "--config", + type=Path, + default=Path("config.yaml"), + help="Path to configuration file", + ) + parser.add_argument( + "--limit", + type=int, + default=None, + help="Limit number of CSV rows to read (default: all)", + ) + parser.add_argument( + "--self-test", + action="store_true", + help="Run the in-memory pairing self-test and exit (no network or IO)", + ) + parser.add_argument( + "-y", + "--yes", + action="store_true", + help="Skip confirmation prompts (for non-interactive use)", + ) + args = parser.parse_args() + + main( + config_path=args.config, + item_limit=args.limit, + self_test=args.self_test, + yes=args.yes, + ) diff --git a/gallery/eng/argument_structure/run_pipeline.py b/gallery/eng/argument_structure/run_pipeline.py index 8d474a4..308befa 100755 --- a/gallery/eng/argument_structure/run_pipeline.py +++ b/gallery/eng/argument_structure/run_pipeline.py @@ -1,16 +1,15 @@ #!/usr/bin/env python3 -"""Run the complete argument structure active learning pipeline. - -This script orchestrates the full pipeline: -1. Load configuration -2. Load 2AFC pairs -3. Set up convergence detection -4. Set up active learning components -5. Run active learning loop with convergence detection -6. Report results and save outputs - -The pipeline continues until the model converges to human-level inter-annotator -agreement, as measured by Krippendorff's alpha or other configured metrics. +"""Run the argument-structure active-learning pipeline. + +Seeds the active-learning loop with the acceptability model pretrained on +MegaAcceptability (see train_acceptability_model.py) and fine-tunes it on +collected human 2AFC judgments until the model converges to human-level +inter-annotator agreement. Pairs are loaded from the canonical layers fragment +produced by create_2afc_pairs.py. + +The active-learning loop trains transformer models, so a real run needs torch, +transformers, and a trained checkpoint. Use --dry-run to validate the wiring +(item template, loop config, selector, checkpoint presence) without training. """ from __future__ import annotations @@ -22,12 +21,12 @@ from pathlib import Path from typing import Any -import numpy as np +import layers_io import yaml +from protocol import acceptability_family, build_protocol from bead.active_learning.loop import ActiveLearningLoop from bead.active_learning.selection import UncertaintySampler -from bead.active_learning.trainers.base import ModelMetadata from bead.cli.display import ( console, create_summary_table, @@ -37,23 +36,18 @@ print_success, print_warning, ) +from bead.config.active_learning import ( + ActiveLearningLoopConfig, + UncertaintySamplerConfig, +) from bead.evaluation.convergence import ConvergenceDetector from bead.items.item import Item +from bead.items.item_template import ItemTemplate +from bead.protocol.items import family_to_item_template def load_config(config_path: Path) -> dict[str, Any]: - """Load configuration from YAML file. - - Parameters - ---------- - config_path : Path - Path to configuration YAML file. - - Returns - ------- - dict[str, Any] - Configuration dictionary. - """ + """Load configuration from a YAML file.""" print_info(f"Loading configuration from {config_path}") with open(config_path) as f: config = yaml.safe_load(f) @@ -61,253 +55,124 @@ def load_config(config_path: Path) -> dict[str, Any]: return config -def load_2afc_pairs(path: Path, limit: int | None = None, skip: int = 0) -> list[Item]: - """Load 2AFC pairs from JSONL file. +def load_pairs(config: dict[str, Any], base_dir: Path) -> list[Item]: + """Load 2AFC pairs from the canonical layers fragment, else the JSONL. Parameters ---------- - path : Path - Path to 2AFC pairs JSONL file. - limit : int | None - Maximum number of items to load. If None, load all. - skip : int - Number of items to skip from the beginning. + config : dict[str, Any] + Loaded configuration. + base_dir : Path + Directory the config paths are relative to. Returns ------- list[Item] - List of 2AFC item pairs. + The 2AFC pairs. """ - print_info(f"Loading 2AFC pairs from {path}") - items = [] - with open(path) as f: - for i, line in enumerate(f): - if i < skip: - continue - if limit and len(items) >= limit: - break - data = json.loads(line) - items.append(Item(**data)) - - print_success(f"Loaded {len(items)} 2AFC pairs") - if skip > 0: - console.print(f" (skipped first {skip} items)") + fragment_rel = config["paths"].get("2afc_pairs_fragment") + fragment_path = base_dir / fragment_rel if fragment_rel else None + if fragment_path is not None and fragment_path.exists(): + print_info(f"Loading pairs from layers fragment {fragment_path.name}") + return layers_io.read_items(fragment_path) + + pairs_path = base_dir / config["paths"]["2afc_pairs"] + print_info(f"Loading pairs from {pairs_path.name}") + items: list[Item] = [] + with open(pairs_path) as f: + for line in f: + items.append(Item(**json.loads(line))) return items -def setup_convergence_detector(config: dict[str, Any]) -> ConvergenceDetector: - """Set up convergence detector from configuration. +def build_item_template(config_path: Path) -> ItemTemplate: + """Build the 2AFC ItemTemplate from the protocol declared in config.yaml.""" + family = acceptability_family(build_protocol(config_path)) + return family_to_item_template(family, judgment_type="acceptability") - Parameters - ---------- - config : dict[str, Any] - Configuration dictionary. - Returns - ------- - ConvergenceDetector - Configured convergence detector. - """ - print_info("Setting up convergence detection...") - conv_config = config["training"]["convergence"] - - detector = ConvergenceDetector( - human_agreement_metric=conv_config["metric"], - convergence_threshold=conv_config["threshold"], - min_iterations=conv_config["min_iterations"], - alpha=conv_config.get("alpha", 0.05), +def build_loop_config(config: dict[str, Any]) -> ActiveLearningLoopConfig: + """Build the active-learning loop config from the configuration.""" + al = config["active_learning"] + conv = config["training"]["convergence"] + return ActiveLearningLoopConfig( + max_iterations=al.get("max_iterations", 10), + budget_per_iteration=al.get("budget_per_iteration", 100), + stopping_criterion=al.get("stopping_criterion", "convergence"), + metric_name=conv.get("metric", "accuracy"), ) - print_success("Convergence detector initialized") - console.print(f" - Metric: {conv_config['metric']}") - console.print(f" - Threshold: {conv_config['threshold']}") - console.print(f" - Min iterations: {conv_config['min_iterations']}") - - return detector +def load_acceptability_model(checkpoint_dir: Path, device: str) -> Any: + """Load the pretrained acceptability ForcedChoiceModel from a checkpoint. -def setup_active_learning(config: dict[str, Any]) -> tuple[UncertaintySampler, Any]: - """Set up active learning components from configuration. + Imports torch-backed modules lazily so the rest of the pipeline (and + --dry-run) stays light. Parameters ---------- - config : dict[str, Any] - Configuration dictionary. + checkpoint_dir : Path + Directory written by train_acceptability_model.py. + device : str + Device to load the model onto. Returns ------- - tuple[UncertaintySampler, Any] - Tuple of (item selector, trainer). - Trainer is None for now (placeholder for future implementation). + ForcedChoiceModel + The pretrained model, ready to seed the active-learning loop. """ - print_info("Setting up active learning components...") - al_config = config["active_learning"] - - # Set up item selector - if al_config["strategy"] == "uncertainty_sampling": - selector = UncertaintySampler(method=al_config["method"]) - print_success(f"Item selector: {al_config['strategy']}") - console.print(f" - Method: {al_config['method']}") - else: - raise ValueError(f"Unknown AL strategy: {al_config['strategy']}") - - # Trainer placeholder (Phase 22, not yet implemented) - trainer = None - console.print(" - Trainer: Not implemented (using placeholder)") - - return selector, trainer - - -def predict_2afc(model: Any, item: Item) -> np.ndarray: - """Prediction function for 2AFC items. - - This is a placeholder implementation that returns random probabilities. - In a real implementation, this would use the trained model to predict - which option is more acceptable. - - Parameters - ---------- - model : Any - Trained model (currently unused). - item : Item - 2AFC item to predict on. + from bead.active_learning.models.forced_choice import ( # noqa: PLC0415 + ForcedChoiceModel, + ) + from bead.config.active_learning import ForcedChoiceModelConfig # noqa: PLC0415 - Returns - ------- - np.ndarray - Array of shape (2,) with probabilities for each option. - """ - # Placeholder: return random probabilities - # In real implementation, would use model to score both options - return np.array([0.5, 0.5]) + model = ForcedChoiceModel(ForcedChoiceModelConfig(device=device)) # type: ignore[arg-type] + model.load(str(checkpoint_dir)) + return model -def load_human_ratings( - human_ratings_path: Path | None, -) -> dict[str, list[Any]] | None: - """Load human ratings for computing baseline agreement. +def load_human_ratings(path: Path) -> dict[str, str] | None: + """Load human 2AFC ratings mapping item id to the chosen option name. Parameters ---------- - human_ratings_path : Path | None - Path to human ratings file (JSONL format). - If None, returns None. + path : Path + JSONL file of ``{"item_id": ..., "option": "option_a"|"option_b"}``. Returns ------- - dict[str, list[Any]] | None - Dictionary mapping rater IDs to their ratings. - Returns None if no human ratings available. + dict[str, str] | None + Mapping from item id to chosen option, or None when absent. """ - if human_ratings_path is None or not human_ratings_path.exists(): - console.print( - " - Human ratings: Not available (will be collected during deployment)" - ) + if not path.exists(): + console.print(" - Human ratings: not available (collected during deployment)") return None + ratings: dict[str, str] = {} + with open(path) as f: + for line in f: + record = json.loads(line) + ratings[str(record["item_id"])] = str(record["option"]) + print_success(f"Loaded {len(ratings)} human ratings") + return ratings - print_info(f"Loading human ratings from {human_ratings_path}") - # Placeholder implementation - # Real implementation would parse JSONL with format: - # {"rater_id": "r1", "item_id": "...", "response": 0 or 1} - return None - - -def print_results(results: list[ModelMetadata]) -> None: - """Print final results from active learning. - Parameters - ---------- - results : list[ModelMetadata] - List of model metadata from each iteration. - """ +def print_results(results: list[Any]) -> None: + """Print per-iteration active-learning results.""" print_header("Active Learning Results") - if not results: print_info("No training iterations completed.") return - console.print(f"\nCompleted {len(results)} iterations") - - # Print per-iteration results - console.print("\nIteration Summary:") - console.print( - f"{'Iter':<6} {'Accuracy':<12} {'Precision':<12} {'Recall':<12} {'F1':<12}" - ) - console.print("-" * 60) - for i, metadata in enumerate(results, 1): metrics = metadata.metrics acc = metrics.get("accuracy", 0.0) - prec = metrics.get("precision", 0.0) - rec = metrics.get("recall", 0.0) - f1 = metrics.get("f1", 0.0) - console.print(f"{i:<6} {acc:<12.4f} {prec:<12.4f} {rec:<12.4f} {f1:<12.4f}") - - # Print final metrics - print_header("Final Model Performance") - final_metrics = results[-1].metrics - for metric, value in sorted(final_metrics.items()): - if isinstance(value, float): - console.print(f" {metric}: {value:.4f}") - else: - console.print(f" {metric}: {value}") - - -def save_results( - results: list[ModelMetadata], output_path: Path, config: dict[str, Any] -) -> None: - """Save results to JSON file. - - Parameters - ---------- - results : list[ModelMetadata] - List of model metadata from each iteration. - output_path : Path - Path to save results JSON file. - config : dict[str, Any] - Configuration dictionary. - """ - print_header("Saving Results") - - # Prepare results for serialization - results_data = { - "config": { - "project": config["project"], - "active_learning": config["active_learning"], - "convergence": config["training"]["convergence"], - }, - "iterations": [ - { - "iteration": i + 1, - "metrics": metadata.metrics, - "model_info": { - "n_train_samples": metadata.n_train_samples, - "n_val_samples": metadata.n_val_samples, - "training_time": metadata.training_time, - }, - } - for i, metadata in enumerate(results) - ], - "final_metrics": results[-1].metrics if results else {}, - } - - with open(output_path, "w") as f: - json.dump(results_data, f, indent=2) - - print_success(f"Results saved to {output_path}") + console.print(f" Iteration {i}: accuracy={acc:.4f}") def main(args: argparse.Namespace) -> None: - """Run the complete active learning pipeline. - - Parameters - ---------- - args : argparse.Namespace - Command-line arguments. - """ + """Run the active-learning pipeline.""" print_header("Argument Structure Active Learning Pipeline") - # Determine base directory if args.config: config_path = Path(args.config) base_dir = config_path.parent @@ -315,164 +180,106 @@ def main(args: argparse.Namespace) -> None: base_dir = Path(__file__).parent config_path = base_dir / "config.yaml" - console.print(f"\nBase directory: {base_dir}") - console.print(f"Configuration: {config_path}") - - # Load configuration - print_header("[1/7] Loading Configuration") - try: - config = load_config(config_path) - except Exception as e: - print_error(f"Error loading configuration: {e}") - sys.exit(1) + print_header("[1/6] Loading Configuration") + config = load_config(config_path) - # Set up convergence detection - print_header("[2/7] Setting Up Convergence Detection") - try: - convergence_detector = setup_convergence_detector(config) - except Exception as e: - print_error(f"Error setting up convergence detector: {e}") - sys.exit(1) + print_header("[2/6] Building Item Template + Loop Config") + item_template = build_item_template(config_path) + loop_config = build_loop_config(config) + sampler_config = UncertaintySamplerConfig( + method=config["active_learning"].get("method", "entropy") + ) + selector = UncertaintySampler(config=sampler_config) + loop = ActiveLearningLoop(item_selector=selector, config=loop_config) + print_success( + f"Loop ready: {loop_config.max_iterations} iterations, " + f"stop on {loop_config.stopping_criterion}" + ) - # Set up active learning components - print_header("[3/7] Setting Up Active Learning") - try: - selector, trainer = setup_active_learning(config) - except Exception as e: - print_error(f"Error setting up active learning: {e}") - sys.exit(1) + print_header("[3/6] Loading 2AFC Pairs") + pairs = load_pairs(config, base_dir) + initial_size = args.initial_size or config["active_learning"].get( + "initial_training_size", 100 + ) + initial_items = pairs[:initial_size] + unlabeled_pool = pairs[initial_size:] + print_success( + f"{len(initial_items)} initial items, {len(unlabeled_pool)} in the pool" + ) - # Create active learning loop - al_config = config["active_learning"] - loop = ActiveLearningLoop( - item_selector=selector, - trainer=trainer, - predict_fn=predict_2afc, - max_iterations=al_config["max_iterations"], - budget_per_iteration=al_config["budget_per_iteration"], + print_header("[4/6] Seeding Pretrained Acceptability Model") + checkpoint_dir = base_dir / config["acceptability_model"]["checkpoint_dir"] + if not checkpoint_dir.exists(): + print_warning( + f"No acceptability checkpoint at {checkpoint_dir}. " + "Run prepare_megaacceptability.py then train_acceptability_model.py." + ) + else: + print_success(f"Found pretrained checkpoint at {checkpoint_dir}") + + print_header("[5/6] Convergence Detection") + conv = config["training"]["convergence"] + convergence_detector = ConvergenceDetector( + human_agreement_metric=conv["metric"], + convergence_threshold=conv["threshold"], + min_iterations=conv["min_iterations"], + alpha=conv.get("alpha", 0.05), ) - print_success("Active learning loop initialized") - console.print(f" - Max iterations: {al_config['max_iterations']}") - console.print(f" - Budget per iteration: {al_config['budget_per_iteration']}") + human_ratings = load_human_ratings(base_dir / "data" / "human_ratings.jsonl") - # Load 2AFC pairs - print_header("[4/7] Loading 2AFC Pairs") - pairs_path = base_dir / config["paths"]["2afc_pairs"] + print_header("[6/6] Running Active Learning Loop") + if args.dry_run: + print_warning("DRY RUN: validated wiring; not training.") + print_success("Wiring OK (template, loop, selector, convergence detector).") + return - # Use command-line overrides if provided - initial_size = args.initial_size or al_config.get("initial_training_size", 100) - unlabeled_size = args.unlabeled_size or 500 + if not checkpoint_dir.exists(): + print_error("Cannot run without a pretrained checkpoint. See step [4/6].") + sys.exit(1) + + device = config["acceptability_model"].get("device", "cpu") + initial_model = load_acceptability_model(checkpoint_dir, device) try: - initial_items = load_2afc_pairs(pairs_path, limit=initial_size, skip=0) - unlabeled_pool = load_2afc_pairs( - pairs_path, limit=unlabeled_size, skip=initial_size + results = loop.run( + initial_items=initial_items, + initial_model=initial_model, + item_template=item_template, + unlabeled_pool=unlabeled_pool, + human_ratings=human_ratings, + convergence_detector=convergence_detector, ) except Exception as e: - print_error(f"Error loading 2AFC pairs: {e}") + print_error(f"Error during active learning: {e}") + traceback.print_exc() sys.exit(1) - console.print("\nData split:") - console.print(f" - Initial training set: {len(initial_items)} items") - console.print(f" - Unlabeled pool: {len(unlabeled_pool)} items") - - # Load human ratings (if available) - print_header("[5/7] Loading Human Ratings") - human_ratings_path = base_dir / "data" / "human_ratings.jsonl" - human_ratings = load_human_ratings( - human_ratings_path if human_ratings_path.exists() else None - ) - - # Run active learning loop - print_header("[6/7] Running Active Learning Loop") - console.print("\nStarting active learning with convergence detection...") - console.print("=" * 70) - - if args.dry_run: - print_warning("DRY RUN MODE: No actual training will occur") - console.print("=" * 70) - results = [] - else: - try: - results = loop.run( - initial_items=initial_items, - initial_model=None, - unlabeled_pool=unlabeled_pool, - human_ratings=human_ratings, - convergence_detector=convergence_detector, - stopping_criterion=al_config["stopping_criterion"], - metric_name=config["training"]["convergence"]["metric"], - ) - except Exception as e: - print_error(f"Error during active learning: {e}") - traceback.print_exc() - sys.exit(1) - - console.print("\n" + "=" * 70) - print_success("Active learning complete!") - - # Print and save results - print_header("[7/7] Results") - if not args.dry_run: - print_results(results) - - # Save results - output_path = base_dir / "results" / "pipeline_results.json" - output_path.parent.mkdir(parents=True, exist_ok=True) - save_results(results, output_path, config) - - # Print summary - print_header("Pipeline Complete") - summary = loop.get_summary() - + print_results(results) table = create_summary_table( { - "Total iterations": str(summary["total_iterations"]), - "Total items selected": str(summary["total_items_selected"]), + "Iterations": str(len(results)), + "Initial items": str(len(initial_items)), + "Pool size": str(len(unlabeled_pool)), } ) console.print(table) - if not args.dry_run and results: - final_acc = results[-1].metrics.get("accuracy", 0.0) - console.print(f"\nFinal accuracy: {final_acc:.4f}") - if __name__ == "__main__": parser = argparse.ArgumentParser( description="Run the argument structure active learning pipeline" ) - parser.add_argument( "--config", type=str, - help="Path to configuration YAML file (default: config.yaml in script dir)", - ) - - parser.add_argument( - "--initial-size", - type=int, - help="Size of initial training set (overrides config)", + help="Path to config.yaml (default: alongside this script)", ) - - parser.add_argument( - "--unlabeled-size", - type=int, - help="Size of unlabeled pool (overrides config)", - ) - parser.add_argument( - "--dry-run", - action="store_true", - help="Run in dry-run mode (load data but don't train)", + "--initial-size", type=int, help="Size of the initial training set" ) - parser.add_argument( - "--verbose", - action="store_true", - help="Enable verbose output", + "--dry-run", action="store_true", help="Validate wiring without training" ) - args = parser.parse_args() try: diff --git a/gallery/eng/argument_structure/train_acceptability_model.py b/gallery/eng/argument_structure/train_acceptability_model.py new file mode 100644 index 0000000..f3c9857 --- /dev/null +++ b/gallery/eng/argument_structure/train_acceptability_model.py @@ -0,0 +1,326 @@ +#!/usr/bin/env python3 +"""Train the initial acceptability model on MegaAcceptability-derived 2AFC pairs. + +This script consumes the forced-choice training items written by +``prepare_megaacceptability.py`` and fits a :class:`ForcedChoiceModel` with +participant random effects. Each item carries its gold ``label`` (``option_a`` or +``option_b``) and the annotator ``participant_id`` in metadata, so the model can fit +a per-annotator random intercept on top of the shared acceptability classifier. The +trained model is saved to a checkpoint directory for use as the active-learning +initializer. + +Real training pulls torch, transformers, and BERT weights and is slow on CPU, so the +heavy imports are deferred into the training path. Run ``--self-test`` to validate +the config wiring and confirm a model config can be built without loading any model +weights or running training. +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import TYPE_CHECKING + +import yaml + +from bead.cli.display import ( + confirm, + console, + create_summary_table, + print_error, + print_header, + print_info, + print_success, + print_warning, +) +from bead.items.item import Item + +if TYPE_CHECKING: + from bead.config.active_learning import ForcedChoiceModelConfig + +DEFAULT_MIXED_EFFECTS_MODE = "random_intercepts" +DEFAULT_MODEL_NAME = "bert-base-uncased" +DEFAULT_DEVICE = "cpu" +DEFAULT_LEARNING_RATE = 2e-5 +DEFAULT_BATCH_SIZE = 16 +DEFAULT_EPOCHS = 3 +DEFAULT_CHECKPOINT_DIR = "checkpoints/acceptability_init" +DEFAULT_TRAINING_ITEMS = "items/megaacceptability_2afc.jsonl" + + +def load_config(config_path: Path) -> dict[str, object]: + """Load the YAML configuration file.""" + with open(config_path, encoding="utf-8") as handle: + return yaml.safe_load(handle) + + +def get_acceptability_config(config: dict[str, object]) -> dict[str, object]: + """Return the ``acceptability_model`` section, or an empty dict if absent.""" + section = config.get("acceptability_model") + return section if isinstance(section, dict) else {} + + +def load_training_items(path: Path, *, limit: int | None = None) -> list[Item]: + """Load derived 2AFC training items from JSONL written by the prepare script.""" + items: list[Item] = [] + with open(path, encoding="utf-8") as handle: + for index, line in enumerate(handle): + if limit is not None and index >= limit: + break + line = line.strip() + if not line: + continue + items.append(Item(**json.loads(line))) + return items + + +def extract_labels_and_participants( + items: list[Item], +) -> tuple[list[str], list[str]]: + """Pull gold labels and participant ids from each item's metadata. + + Parameters + ---------- + items + Forced-choice training items, each with ``label`` and ``participant_id`` + keys in ``item_metadata``. + + Returns + ------- + tuple[list[str], list[str]] + Parallel lists of option-name labels and participant identifiers. + + Raises + ------ + ValueError + If any item is missing a ``label`` or ``participant_id``. + """ + labels: list[str] = [] + participant_ids: list[str] = [] + for item in items: + label = item.item_metadata.get("label") + participant_id = item.item_metadata.get("participant_id") + if label is None or participant_id is None: + raise ValueError( + "Each training item must carry 'label' and 'participant_id' in " + "item_metadata. Regenerate items with prepare_megaacceptability.py." + ) + labels.append(str(label)) + participant_ids.append(str(participant_id)) + return labels, participant_ids + + +def build_model_config(section: dict[str, object]) -> ForcedChoiceModelConfig: + """Build a ForcedChoiceModelConfig from the acceptability config section. + + Parameters + ---------- + section + The ``acceptability_model`` config section (possibly empty), read with + defensive defaults. + + Returns + ------- + ForcedChoiceModelConfig + Config wired for single-encoder 2AFC training with participant random + effects. + """ + from bead.active_learning.config import MixedEffectsConfig # noqa: PLC0415 + from bead.config.active_learning import ForcedChoiceModelConfig # noqa: PLC0415 + + mode = str(section.get("mixed_effects_mode", DEFAULT_MIXED_EFFECTS_MODE)) + return ForcedChoiceModelConfig( + model_name=str(section.get("model_name", DEFAULT_MODEL_NAME)), + encoder_mode="single_encoder", + learning_rate=float(section.get("learning_rate", DEFAULT_LEARNING_RATE)), + batch_size=int(section.get("batch_size", DEFAULT_BATCH_SIZE)), + num_epochs=int(section.get("epochs", DEFAULT_EPOCHS)), + device=str(section.get("device", DEFAULT_DEVICE)), + mixed_effects=MixedEffectsConfig(mode=mode), + ) + + +def run_self_test(config_path: Path) -> int: + """Validate config wiring by building a model config without loading weights.""" + print_header("train_acceptability_model self-test") + + config: dict[str, object] = {} + if config_path.exists(): + config = load_config(config_path) + print_info(f"Loaded config from {config_path}") + else: + print_warning(f"Config not found at {config_path}; using built-in defaults") + + section = get_acceptability_config(config) + model_config = build_model_config(section) + + table = create_summary_table( + { + "model_name": model_config.model_name, + "encoder_mode": model_config.encoder_mode, + "learning_rate": str(model_config.learning_rate), + "batch_size": str(model_config.batch_size), + "num_epochs": str(model_config.num_epochs), + "device": model_config.device, + "mixed_effects.mode": model_config.mixed_effects.mode, + } + ) + console.print(table) + print_success("Self-test passed: model config built without loading weights") + return 0 + + +def main( + config_path: Path = Path("config.yaml"), + item_limit: int | None = None, + *, + self_test: bool = False, + yes: bool = False, +) -> None: + """Train the initial acceptability model on derived 2AFC pairs. + + Parameters + ---------- + config_path + Path to the gallery configuration file. + item_limit + Optional cap on the number of training items to load (for quick testing). + self_test + When True, validate config wiring and exit without loading model weights. + yes + Skip overwrite confirmation prompts for non-interactive use. + """ + base_dir = Path(__file__).parent + if not config_path.exists(): + config_path = base_dir / config_path.name + + if self_test: + sys.exit(run_self_test(config_path)) + + if not config_path.exists(): + print_error(f"Config file not found: {config_path}") + sys.exit(1) + + config = load_config(config_path) + section = get_acceptability_config(config) + + paths_section = section.get("paths") + paths_section = paths_section if isinstance(paths_section, dict) else {} + training_items_path = base_dir / str( + paths_section.get("training_items", DEFAULT_TRAINING_ITEMS) + ) + checkpoint_dir = base_dir / str( + section.get("checkpoint_dir", DEFAULT_CHECKPOINT_DIR) + ) + + print_header("Acceptability Model Training") + console.print(f"Training items: [cyan]{training_items_path}[/cyan]") + console.print(f"Checkpoint dir: [cyan]{checkpoint_dir}[/cyan]\n") + + if not training_items_path.exists(): + print_error( + f"Training items not found at {training_items_path}. " + "Run prepare_megaacceptability.py first." + ) + sys.exit(1) + + if checkpoint_dir.exists() and not yes: + if not confirm(f"Overwrite checkpoint at {checkpoint_dir}?", default=False): + print_info("Operation cancelled.") + return + + # 1. Load derived training items + print_header("1/3 Loading Training Items") + try: + items = load_training_items(training_items_path, limit=item_limit) + labels, participant_ids = extract_labels_and_participants(items) + except Exception as exc: # noqa: BLE001 + print_error(f"Failed to load training items: {exc}") + sys.exit(1) + if not items: + print_error("No training items loaded. Exiting.") + sys.exit(1) + n_participants = len(set(participant_ids)) + print_success(f"Loaded {len(items):,} items from {n_participants:,} annotators\n") + + # 2. Build config and train (heavy imports deferred to here) + print_header("2/3 Training Forced-Choice Model") + model_config = build_model_config(section) + console.print( + f"Model: [cyan]{model_config.model_name}[/cyan] | " + f"mode: [cyan]{model_config.mixed_effects.mode}[/cyan] | " + f"epochs: [cyan]{model_config.num_epochs}[/cyan]\n" + ) + try: + from bead.active_learning.models.forced_choice import ( # noqa: PLC0415 + ForcedChoiceModel, + ) + + model = ForcedChoiceModel(model_config) + with console.status("[bold]Training (this can be slow on CPU)...[/bold]"): + metrics = model.train(items, labels, participant_ids=participant_ids) + except Exception as exc: # noqa: BLE001 + print_error(f"Training failed: {exc}") + sys.exit(1) + print_success("Training complete\n") + + # 3. Save the checkpoint + print_header("3/3 Saving Checkpoint") + try: + checkpoint_dir.parent.mkdir(parents=True, exist_ok=True) + model.save(str(checkpoint_dir)) + print_success(f"Saved model to {checkpoint_dir}\n") + except Exception as exc: # noqa: BLE001 + print_error(f"Failed to save checkpoint: {exc}") + sys.exit(1) + + # Summary + print_header("Summary") + summary: dict[str, str] = { + "Training items": f"{len(items):,}", + "Annotators": f"{n_participants:,}", + "Checkpoint": str(checkpoint_dir), + } + for name, value in metrics.items(): + summary[name] = f"{value:.4f}" if isinstance(value, float) else str(value) + console.print(create_summary_table(summary)) + print_info("Next: use this checkpoint to seed the active-learning loop") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Train the initial acceptability model on derived 2AFC pairs" + ) + parser.add_argument( + "--config", + type=Path, + default=Path("config.yaml"), + help="Path to configuration file", + ) + parser.add_argument( + "--limit", + type=int, + default=None, + help="Limit number of training items to load (default: all)", + ) + parser.add_argument( + "--self-test", + action="store_true", + help="Validate config wiring without loading model weights, then exit", + ) + parser.add_argument( + "-y", + "--yes", + action="store_true", + help="Skip confirmation prompts (for non-interactive use)", + ) + args = parser.parse_args() + + main( + config_path=args.config, + item_limit=args.limit, + self_test=args.self_test, + yes=args.yes, + ) diff --git a/pyproject.toml b/pyproject.toml index df967ae..566a460 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "bead" -version = "0.7.0" +version = "0.8.0" description = "Lexicon and Template Collection Construction Pipeline for Acceptability and Inference Judgment Data" authors = [{name = "Aaron Steven White", email = "aaron.white@rochester.edu"}] readme = "README.md" diff --git a/tests/interop/test_layers_judgment.py b/tests/interop/test_layers_judgment.py new file mode 100644 index 0000000..c75d72d --- /dev/null +++ b/tests/interop/test_layers_judgment.py @@ -0,0 +1,87 @@ +"""Round-trip law tests for the judgment overlap lenses.""" + +from __future__ import annotations + +import pytest +from lairs.records import judgment + +from bead.evaluation.reliability import AnnotationRecord +from bead.interop.layers._convert import from_object_ref +from bead.interop.layers.judgment_lens import ( + ANNOTATION_RECORD_JUDGMENT, + judgment_set_to_records, + records_to_judgment_set, +) + + +def _record( + *, + annotator_id: str = "ann_1", + item_id: str = "item_1", + question_name: str = "completion", + response_label: str = "yes", +) -> AnnotationRecord: + return AnnotationRecord( + annotator_id=annotator_id, + item_id=item_id, + question_name=question_name, + response_label=response_label, + ) + + +class TestAnnotationRecordJudgment: + """AnnotationRecord <-> layers judgment.""" + + def test_view_shape(self) -> None: + record = _record() + view, _ = ANNOTATION_RECORD_JUDGMENT.forward(record) + assert from_object_ref(view.item) == record.item_id + assert view.categoricalValue == record.response_label + + def test_roundtrip_exact(self) -> None: + record = _record() + view, complement = ANNOTATION_RECORD_JUDGMENT.forward(record) + assert ANNOTATION_RECORD_JUDGMENT.backward(view, complement) == record + + def test_roundtrip_through_serialization(self) -> None: + record = _record() + view, complement = ANNOTATION_RECORD_JUDGMENT.forward(record) + view2 = judgment.Judgment.model_validate_json(view.model_dump_json()) + assert ANNOTATION_RECORD_JUDGMENT.backward(view2, complement) == record + + +class TestAnnotationSetJudgment: + """tuple[AnnotationRecord, ...] <-> layers judgment set.""" + + def test_roundtrip_exact(self) -> None: + records = ( + _record(item_id="item_1", response_label="yes"), + _record(item_id="item_2", response_label="no"), + _record(item_id="item_3", response_label="yes"), + ) + view, complement = records_to_judgment_set(records) + assert view.agent is not None + assert view.agent.id == "ann_1" + assert len(view.judgments) == 3 + assert judgment_set_to_records(view, complement) == records + + def test_roundtrip_through_serialization(self) -> None: + records = ( + _record(item_id="item_1", response_label="yes"), + _record(item_id="item_2", response_label="no"), + ) + view, complement = records_to_judgment_set(records) + view2 = judgment.JudgmentSet.model_validate_json(view.model_dump_json()) + assert judgment_set_to_records(view2, complement) == records + + def test_mixed_annotator_raises(self) -> None: + records = ( + _record(annotator_id="ann_1"), + _record(annotator_id="ann_2"), + ) + with pytest.raises(ValueError, match="share one"): + records_to_judgment_set(records) + + def test_empty_raises(self) -> None: + with pytest.raises(ValueError, match="at least one record"): + records_to_judgment_set(()) diff --git a/tests/interop/test_layers_list.py b/tests/interop/test_layers_list.py new file mode 100644 index 0000000..d44886c --- /dev/null +++ b/tests/interop/test_layers_list.py @@ -0,0 +1,189 @@ +"""Round-trip law tests for the list-composition overlap lenses.""" + +from __future__ import annotations + +from uuid import uuid4 + +from lairs.records import judgment + +from bead.interop.layers.list_lens import ( + EXPERIMENT_LIST_LAYERS, + LIST_CONSTRAINT, +) +from bead.lists.constraints import ( + BalanceConstraint, + GridDimension, + GridStratificationConstraint, + ListConstraint, + QuantileBinning, + UniquenessConstraint, +) +from bead.lists.experiment_list import ConstraintSatisfaction, ExperimentList + + +def _uniqueness() -> UniquenessConstraint: + return UniquenessConstraint( + constraint_type="uniqueness", + property_expression="item['lemma']", + allow_null=True, + priority=3, + ) + + +def _balance() -> BalanceConstraint: + return BalanceConstraint( + constraint_type="balance", + property_expression="item['condition']", + target_counts={"grammatical": 4, "ungrammatical": 4}, + tolerance=0.2, + priority=2, + ) + + +def _grid() -> GridStratificationConstraint: + return GridStratificationConstraint( + constraint_type="grid_stratification", + dimensions=( + GridDimension( + property_expression="item['acceptability_score_diff']", + binning=QuantileBinning(binning="quantile", n_quantiles=4), + ), + GridDimension( + property_expression="item['length']", + binning=QuantileBinning(binning="quantile", n_quantiles=3), + ), + ), + items_per_cell=3, + priority=5, + ) + + +class TestListConstraint: + """ListConstraint <-> layers list constraint.""" + + def test_uniqueness_roundtrip_exact(self) -> None: + constraint = _uniqueness() + view, complement = LIST_CONSTRAINT.forward(constraint) + assert view.kind == "uniqueness" + assert view.targetProperty == "item['lemma']" + restored = LIST_CONSTRAINT.backward(view, complement) + assert restored == constraint + assert isinstance(restored, UniquenessConstraint) + + def test_balance_roundtrip_exact(self) -> None: + constraint = _balance() + view, complement = LIST_CONSTRAINT.forward(constraint) + assert view.kind == "balance" + assert view.targetProperty == "item['condition']" + restored = LIST_CONSTRAINT.backward(view, complement) + assert restored == constraint + assert isinstance(restored, BalanceConstraint) + + def test_grid_stratification_roundtrip_exact(self) -> None: + constraint = _grid() + view, complement = LIST_CONSTRAINT.forward(constraint) + assert view.kind == "grid_stratification" + # grid stratification has no single property_expression + assert view.targetProperty is None + restored = LIST_CONSTRAINT.backward(view, complement) + assert restored == constraint + assert isinstance(restored, GridStratificationConstraint) + + def test_uniqueness_roundtrip_through_serialization(self) -> None: + constraint = _uniqueness() + view, complement = LIST_CONSTRAINT.forward(constraint) + view2 = judgment.ListConstraint.model_validate_json(view.model_dump_json()) + assert LIST_CONSTRAINT.backward(view2, complement) == constraint + + def test_balance_roundtrip_through_serialization(self) -> None: + constraint = _balance() + view, complement = LIST_CONSTRAINT.forward(constraint) + view2 = judgment.ListConstraint.model_validate_json(view.model_dump_json()) + assert LIST_CONSTRAINT.backward(view2, complement) == constraint + + def test_grid_roundtrip_through_serialization(self) -> None: + constraint = _grid() + view, complement = LIST_CONSTRAINT.forward(constraint) + view2 = judgment.ListConstraint.model_validate_json(view.model_dump_json()) + restored = LIST_CONSTRAINT.backward(view2, complement) + assert restored == constraint + + def test_backward_returns_union_root_type(self) -> None: + constraint: ListConstraint = _balance() + view, complement = LIST_CONSTRAINT.forward(constraint) + restored = LIST_CONSTRAINT.backward(view, complement) + assert type(restored) is BalanceConstraint + + +class TestExperimentListLayers: + """ExperimentList <-> layers collection aggregate.""" + + def _experiment_list(self) -> tuple[ExperimentList, tuple[str, ...]]: + item_a = uuid4() + item_b = uuid4() + item_c = uuid4() + uniqueness = _uniqueness() + balance = _balance() + experiment_list = ExperimentList( + name="list_0", + list_number=2, + item_refs=(item_a, item_b, item_c), + list_constraints=(uniqueness, balance), + constraint_satisfaction=( + ConstraintSatisfaction(constraint_id=uniqueness.id, satisfied=True), + ConstraintSatisfaction(constraint_id=balance.id, satisfied=False), + ), + presentation_order=(item_c, item_a, item_b), + list_metadata={"block": 1, "notes": ("a", "b")}, + balance_metrics={"entropy": 0.87, "balanced": True}, + ) + entry_refs = (str(item_a), str(item_b), str(item_c)) + return experiment_list, entry_refs + + def test_roundtrip_exact(self) -> None: + experiment_list, _ = self._experiment_list() + view, complement = EXPERIMENT_LIST_LAYERS.forward(experiment_list) + assert view.collection.kind == "stimulus-pool" + assert view.collection.name == "list_0" + assert len(view.list_constraints) == 2 + restored = EXPERIMENT_LIST_LAYERS.backward(view, complement) + assert restored == experiment_list + + def test_membership_count_and_order(self) -> None: + experiment_list, entry_refs = self._experiment_list() + view, _ = EXPERIMENT_LIST_LAYERS.forward(experiment_list) + assert len(view.memberships) == len(experiment_list.item_refs) + assert tuple(m.entryRef for m in view.memberships) == entry_refs + assert tuple(m.ordinal for m in view.memberships) == (0, 1, 2) + for membership in view.memberships: + assert membership.collectionRef == str(experiment_list.id) + + def test_roundtrip_through_serialization(self) -> None: + experiment_list, _ = self._experiment_list() + view, complement = EXPERIMENT_LIST_LAYERS.forward(experiment_list) + collection2 = view.collection.__class__.model_validate_json( + view.collection.model_dump_json() + ) + view2 = view.with_(collection=collection2) + restored = EXPERIMENT_LIST_LAYERS.backward(view2, complement) + assert restored == experiment_list + + def test_empty_roundtrip(self) -> None: + experiment_list = ExperimentList(name="empty", list_number=0) + view, complement = EXPERIMENT_LIST_LAYERS.forward(experiment_list) + assert view.memberships == () + assert view.list_constraints == () + restored = EXPERIMENT_LIST_LAYERS.backward(view, complement) + assert restored == experiment_list + + def test_presentation_order_none_recovered(self) -> None: + item_a = uuid4() + experiment_list = ExperimentList( + name="list_1", + list_number=1, + item_refs=(item_a,), + ) + view, complement = EXPERIMENT_LIST_LAYERS.forward(experiment_list) + restored = EXPERIMENT_LIST_LAYERS.backward(view, complement) + assert restored.presentation_order is None + assert restored == experiment_list diff --git a/tests/interop/test_layers_resource.py b/tests/interop/test_layers_resource.py index 28a3916..a1169af 100644 --- a/tests/interop/test_layers_resource.py +++ b/tests/interop/test_layers_resource.py @@ -2,7 +2,10 @@ from __future__ import annotations +from lairs.records import resource + from bead.interop.layers.resource_lens import ( + FILLED_TEMPLATE_FILLING, LEXICAL_ITEM_ENTRY, LEXICON_COLLECTION, TEMPLATE_LAYERS, @@ -11,6 +14,7 @@ from bead.resources.lexical_item import LexicalItem from bead.resources.lexicon import Lexicon from bead.resources.template import Slot, Template +from bead.templates.filler import FilledTemplate class TestLexicalItemEntry: @@ -105,3 +109,62 @@ def test_minimal(self) -> None: template = Template(name="t", template_string="{x}") view, complement = TEMPLATE_LAYERS.forward(template) assert TEMPLATE_LAYERS.backward(view, complement) == template + + +class TestFilledTemplateFilling: + """FilledTemplate <-> layers filling.""" + + def _filled(self) -> FilledTemplate: + return FilledTemplate( + template_id="tmpl-019a", + template_name="transitive", + slot_fillers={ + "subj": LexicalItem( + lemma="dog", + form="dogs", + language_code="eng", + features={"pos": "NOUN", "number": "plural"}, + source="bleached", + ), + "verb": LexicalItem( + lemma="chase", form="chased", language_code="eng", + features={"pos": "VERB"}, + ), + }, + rendered_text="Dogs chased.", + strategy_name="mlm", + template_slots={"subj": True, "verb": False}, + ) + + def test_view_shape(self) -> None: + view, _ = FILLED_TEMPLATE_FILLING.forward(self._filled()) + assert view.templateRef == "tmpl-019a" + assert view.strategy == "mlm" + assert view.renderedText == "Dogs chased." + names = {s.slotName for s in view.slotFillings} + assert names == {"subj", "verb"} + + def test_roundtrip_exact(self) -> None: + filled = self._filled() + view, complement = FILLED_TEMPLATE_FILLING.forward(filled) + assert FILLED_TEMPLATE_FILLING.backward(view, complement) == filled + + def test_roundtrip_through_serialization(self) -> None: + filled = self._filled() + view, complement = FILLED_TEMPLATE_FILLING.forward(filled) + view2 = resource.Filling.model_validate_json(view.model_dump_json()) + assert FILLED_TEMPLATE_FILLING.backward(view2, complement) == filled + + def test_none_form_recovered(self) -> None: + filled = FilledTemplate( + template_id="t1", + template_name="intransitive", + slot_fillers={"subj": LexicalItem(lemma="rain", language_code="eng")}, + rendered_text="Rain.", + strategy_name="manual", + template_slots={"subj": True}, + ) + view, complement = FILLED_TEMPLATE_FILLING.forward(filled) + restored = FILLED_TEMPLATE_FILLING.backward(view, complement) + assert restored.slot_fillers["subj"].form is None + assert restored == filled diff --git a/tests/items/test_acceptability_scorer.py b/tests/items/test_acceptability_scorer.py new file mode 100644 index 0000000..34e11d0 --- /dev/null +++ b/tests/items/test_acceptability_scorer.py @@ -0,0 +1,89 @@ +"""Tests for AcceptabilityScorer using a stub forced-choice model.""" + +from __future__ import annotations + +from uuid import uuid4 + +import numpy as np + +from bead.items.item import Item +from bead.items.scoring import AcceptabilityScorer + + +class _StubMixedEffects: + mode = "fixed" + + +class _StubConfig: + mixed_effects = _StubMixedEffects() + + +class _StubForcedChoiceModel: + """Minimal stand-in for ForcedChoiceModel exposing predict_proba.""" + + def __init__(self, proba: np.ndarray) -> None: + self._proba = proba + self.config = _StubConfig() + + def predict_proba( + self, items: list[Item], participant_ids: list[str] | None = None + ) -> np.ndarray: + # fixed-effects model is asked with participant_ids=None + assert participant_ids is None + return self._proba[: len(items)] + + +def _item(text: str = "a vs b") -> Item: + return Item( + item_template_id=uuid4(), + rendered_elements={"option_a": "A", "option_b": "B"}, + options=("first", "second"), + item_metadata={"text": text}, + ) + + +def test_score_returns_preference_margin() -> None: + # p_first=0.7 -> margin = 2*0.7-1 = 0.4 + model = _StubForcedChoiceModel(np.array([[0.7, 0.3]])) + scorer = AcceptabilityScorer(model) # type: ignore[arg-type] + assert abs(scorer.score(_item()) - 0.4) < 1e-9 + + +def test_margin_is_column_agnostic() -> None: + # whichever option dominates, margin reflects |p_a - p_b| + model = _StubForcedChoiceModel(np.array([[0.2, 0.8]])) + scorer = AcceptabilityScorer(model) # type: ignore[arg-type] + assert abs(scorer.score(_item()) - 0.6) < 1e-9 + + +def test_tie_gives_zero_margin() -> None: + model = _StubForcedChoiceModel(np.array([[0.5, 0.5]])) + scorer = AcceptabilityScorer(model) # type: ignore[arg-type] + assert abs(scorer.score(_item())) < 1e-9 + + +def test_score_batch() -> None: + proba = np.array([[0.9, 0.1], [0.5, 0.5], [0.3, 0.7]]) + scorer = AcceptabilityScorer(_StubForcedChoiceModel(proba)) # type: ignore[arg-type] + margins = scorer.score_batch([_item(), _item(), _item()]) + assert len(margins) == 3 + assert abs(margins[0] - 0.8) < 1e-9 + assert abs(margins[1]) < 1e-9 + assert abs(margins[2] - 0.4) < 1e-9 + + +def test_score_batch_empty() -> None: + scorer = AcceptabilityScorer(_StubForcedChoiceModel(np.zeros((0, 2)))) # type: ignore[arg-type] + assert scorer.score_batch([]) == [] + + +def test_score_with_metadata() -> None: + proba = np.array([[0.7, 0.3]]) + item = _item() + scorer = AcceptabilityScorer(_StubForcedChoiceModel(proba)) # type: ignore[arg-type] + out = scorer.score_with_metadata([item]) + entry = out[item.id] + assert abs(entry["score"] - 0.4) < 1e-9 # type: ignore[operator] + assert abs(entry["acceptability_margin"] - 0.4) < 1e-9 # type: ignore[operator] + assert abs(entry["p_first"] - 0.7) < 1e-9 # type: ignore[operator] + assert entry["predicted_option"] == 0 diff --git a/tests/lists/test_grid_stratification.py b/tests/lists/test_grid_stratification.py new file mode 100644 index 0000000..698d52c --- /dev/null +++ b/tests/lists/test_grid_stratification.py @@ -0,0 +1,317 @@ +"""Tests for N-dimensional mixed-variable grid stratification.""" + +from __future__ import annotations + +from uuid import uuid4 + +import didactic.api as dx +import pytest + +from bead.lists.balancer import QuantileBalancer +from bead.lists.constraints import ( + CategoricalBinning, + EqualWidthBinning, + GridDimension, + GridStratificationConstraint, + QuantileBinning, + StdDevBinning, + ThresholdBinning, +) +from bead.lists.partitioner import ListPartitioner +from bead.lists.stratification import ( + assign_grid_cells, + assign_grid_cells_by_uuid, + assign_quantiles, + flatten_cell, + grid_shape, + unflatten_cell, +) + + +def _q(n: int) -> QuantileBinning: + return QuantileBinning(binning="quantile", n_quantiles=n) + + +def _c(**kwargs: object) -> CategoricalBinning: + return CategoricalBinning(binning="categorical", **kwargs) # type: ignore[arg-type] + + +class TestBinningStrategies: + """Each binning strategy maps values to the expected bin indices.""" + + def test_quantile_equal_frequency(self) -> None: + items = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + result = assign_grid_cells(items, [lambda x: float(x)], [_q(4)]) + assert result[1] == (0,) + assert result[10] == (3,) + assert {c[0] for c in result.values()} == {0, 1, 2, 3} + + def test_equal_width(self) -> None: + items = [0.0, 0.4, 0.6, 1.0] + binning = EqualWidthBinning( + binning="equal_width", n_bins=2, range_min=0.0, range_max=1.0 + ) + result = assign_grid_cells(items, [lambda x: x], [binning]) + assert result[0.0] == (0,) + assert result[0.4] == (0,) + assert result[0.6] == (1,) + assert result[1.0] == (1,) + + def test_threshold(self) -> None: + items = [0.2, 0.5, 0.8] + binning = ThresholdBinning(binning="threshold", edges=(0.5,)) + result = assign_grid_cells(items, [lambda x: x], [binning]) + assert result[0.2] == (0,) + assert result[0.5] == (1,) # at the edge -> upper bin + assert result[0.8] == (1,) + + def test_threshold_multiple_edges(self) -> None: + items = [-1.0, 0.5, 1.5, 3.0] + binning = ThresholdBinning(binning="threshold", edges=(0.0, 1.0, 2.0)) + result = assign_grid_cells(items, [lambda x: x], [binning]) + assert result[-1.0] == (0,) + assert result[0.5] == (1,) + assert result[1.5] == (2,) + assert result[3.0] == (3,) + + def test_stddev(self) -> None: + items = [0.0, 1.0, 2.0, 3.0, 4.0] + binning = StdDevBinning(binning="stddev", k_values=(-1.0, 1.0)) + result = assign_grid_cells(items, [lambda x: x], [binning]) + # mean=2, sd~1.41; below mean-sd -> 0, around mean -> 1, above mean+sd -> 2 + assert result[0.0] == (0,) + assert result[2.0] == (1,) + assert result[4.0] == (2,) + + def test_categorical_declared(self) -> None: + items = ["a", "b", "a", "c"] + binning = _c(categories=("a", "b"), include_other=True) + result = assign_grid_cells(items, [lambda x: x], [binning]) + assert result["a"] == (0,) + assert result["b"] == (1,) + assert result["c"] == (2,) # catch-all bin + + def test_categorical_no_other_raises(self) -> None: + items = ["a", "b", "c"] + binning = _c(categories=("a", "b"), include_other=False) + with pytest.raises(ValueError, match="outside the declared categories"): + assign_grid_cells(items, [lambda x: x], [binning]) + + def test_categorical_auto_discovery(self) -> None: + items = ["z", "a", "m", "a"] + result = assign_grid_cells(items, [lambda x: x], [_c()]) + # categories discovered and sorted: a=0, m=1, z=2 + assert result["a"] == (0,) + assert result["m"] == (1,) + assert result["z"] == (2,) + + +class TestAssignGridCells: + """Grid assignment composes dimensions and validates inputs.""" + + def test_reduces_to_quantiles_at_n1(self) -> None: + items = list(range(100)) + grid = assign_grid_cells(items, [lambda x: float(x)], [_q(4)]) + quant = assign_quantiles(items, lambda x: float(x), n_quantiles=4) + assert all(grid[i][0] == quant[i] for i in items) + + def test_mixed_continuous_and_discrete(self) -> None: + data = [(0.1, "a"), (0.9, "a"), (0.2, "b"), (0.8, "b")] + getters = [lambda x: x[0], lambda x: x[1]] + binnings = [_q(2), _c()] + result = assign_grid_cells(list(data), getters, binnings) + assert result[(0.1, "a")] == (0, 0) + assert result[(0.9, "a")] == (1, 0) + assert result[(0.2, "b")] == (0, 1) + assert result[(0.8, "b")] == (1, 1) + + def test_stratify_by_groups_continuous(self) -> None: + data = [(10, "A"), (20, "A"), (30, "A"), (40, "A"), (5, "B"), (35, "B")] + result = assign_grid_cells( + list(data), + [lambda x: float(x[0])], + [_q(2)], + stratify_by=lambda x: x[1], + ) + # within each group both quantiles appear + a = {result[item][0] for item in data if item[1] == "A"} + b = {result[item][0] for item in data if item[1] == "B"} + assert a == {0, 1} + assert b == {0, 1} + + def test_empty_items_raises(self) -> None: + with pytest.raises(ValueError, match="items list cannot be empty"): + assign_grid_cells([], [lambda x: x], [_q(2)]) + + def test_length_mismatch_raises(self) -> None: + with pytest.raises(ValueError, match="same length"): + assign_grid_cells([1, 2], [lambda x: float(x)], [_q(2), _q(2)]) + + +class TestFlattenCell: + """Flatten/unflatten round-trips with heterogeneous cardinalities.""" + + def test_round_trip(self) -> None: + shape = (3, 4, 2) + seen = set() + for i in range(3): + for j in range(4): + for k in range(2): + cell = flatten_cell((i, j, k), shape) + assert unflatten_cell(cell, shape) == (i, j, k) + seen.add(cell) + # bijection onto 0..23 + assert seen == set(range(24)) + + def test_length_mismatch_raises(self) -> None: + with pytest.raises(ValueError, match="same length"): + flatten_cell((1, 2), (3,)) + + +class TestAssignGridCellsByUUID: + """The UUID convenience returns flattened cell ids and validates inputs.""" + + def test_basic(self) -> None: + uuids = [uuid4() for _ in range(8)] + meta = { + uid: { + "acc": float(i) / 8, + "lm": float(7 - i) / 8, + "pair_type": "x" if i % 2 else "y", + } + for i, uid in enumerate(uuids) + } + cells = assign_grid_cells_by_uuid( + uuids, + meta, + ["acc", "lm", "pair_type"], + [_q(2), _q(2), _c(categories=("x", "y"))], + ) + assert set(cells.keys()) == set(uuids) + assert all(0 <= c < 2 * 2 * 2 for c in cells.values()) + + def test_shape_matches_grid_shape(self) -> None: + uuids = [uuid4() for _ in range(6)] + meta = {uid: {"a": float(i)} for i, uid in enumerate(uuids)} + getters = [lambda uid: meta[uid]["a"]] + shape = grid_shape(uuids, getters, [_q(3)]) + assert shape == (3,) + + def test_missing_property_raises(self) -> None: + uuids = [uuid4() for _ in range(4)] + meta = {uid: {"a": float(i)} for i, uid in enumerate(uuids)} + with pytest.raises(ValueError, match="Property 'b' not found"): + assign_grid_cells_by_uuid(uuids, meta, ["b"], [_q(2)]) + + def test_missing_uuid_raises(self) -> None: + uuids = [uuid4() for _ in range(4)] + meta = {uid: {"a": float(i)} for i, uid in enumerate(uuids[:-1])} + with pytest.raises(KeyError): + assign_grid_cells_by_uuid(uuids, meta, ["a"], [_q(2)]) + + def test_missing_stratify_key_raises(self) -> None: + uuids = [uuid4() for _ in range(4)] + meta = {uid: {"a": float(i)} for i, uid in enumerate(uuids)} + with pytest.raises(ValueError, match="Stratification key"): + assign_grid_cells_by_uuid( + uuids, meta, ["a"], [_q(2)], stratify_by_key="grp" + ) + + +class TestGridStratificationConstraint: + """The constraint validates its dimensions and parameters.""" + + def test_valid(self) -> None: + dim = GridDimension( + property_expression="item['a']", binning=_q(3) + ) + constraint = GridStratificationConstraint( + constraint_type="grid_stratification", + dimensions=(dim,), + items_per_cell=2, + ) + assert constraint.constraint_type == "grid_stratification" + assert len(constraint.dimensions) == 1 + + def test_empty_dimensions_raises(self) -> None: + with pytest.raises((ValueError, dx.ValidationError)): + GridStratificationConstraint( + constraint_type="grid_stratification", dimensions=() + ) + + def test_items_per_cell_axiom(self) -> None: + dim = GridDimension(property_expression="item['a']", binning=_q(2)) + with pytest.raises((ValueError, dx.ValidationError)): + GridStratificationConstraint( + constraint_type="grid_stratification", + dimensions=(dim,), + items_per_cell=0, + ) + + def test_threshold_edges_must_increase(self) -> None: + with pytest.raises((ValueError, dx.ValidationError)): + ThresholdBinning(binning="threshold", edges=(1.0, 0.5)) + + +class TestBalancerByCell: + """The balancer spreads each grid cell uniformly across lists.""" + + def test_distribute_across_cells(self) -> None: + uuids = [uuid4() for _ in range(40)] + # 4 cells, 10 items each + cell_of = {uid: i % 4 for i, uid in enumerate(uuids)} + balancer = QuantileBalancer(n_quantiles=2, random_seed=0) + lists = balancer.balance_by_cell( + uuids, lambda uid: cell_of[uid], n_cells=4, n_lists=2, + items_per_cell_per_list=5, + ) + assert len(lists) == 2 + # each list gets 5 items per cell * 4 cells = 20 + assert all(len(lst) == 20 for lst in lists) + + def test_invalid_n_cells(self) -> None: + balancer = QuantileBalancer(n_quantiles=2) + with pytest.raises(ValueError, match="n_cells must be >= 1"): + balancer.balance_by_cell([], lambda uid: 0, 0, 2, 1) + + +class TestPartitionerGrid: + """The partitioner uses the grid constraint when present.""" + + def test_grid_partition(self) -> None: + items = [uuid4() for _ in range(40)] + meta = { + uid: { + "metadata": { + "acc": float(i % 10) / 10.0, + "lm": float((i * 3) % 10) / 10.0, + "pair_type": "same" if i % 2 == 0 else "diff", + } + } + for i, uid in enumerate(items) + } + constraint = GridStratificationConstraint( + constraint_type="grid_stratification", + dimensions=( + GridDimension( + property_expression="item['metadata']['acc']", binning=_q(2) + ), + GridDimension( + property_expression="item['metadata']['lm']", binning=_q(2) + ), + GridDimension( + property_expression="item['metadata']['pair_type']", + binning=_c(categories=("same", "diff")), + ), + ), + items_per_cell=1, + ) + partitioner = ListPartitioner(random_seed=42) + lists = partitioner.partition( + items, n_lists=4, constraints=[constraint], strategy="stratified", + metadata=meta, + ) + assert len(lists) == 4 + # no item appears in two lists + all_refs = [ref for lst in lists for ref in lst.item_refs] + assert len(all_refs) == len(set(all_refs)) diff --git a/uv.lock b/uv.lock index 03abc52..2e6242f 100644 --- a/uv.lock +++ b/uv.lock @@ -170,7 +170,7 @@ wheels = [ [[package]] name = "bead" -version = "0.7.0" +version = "0.8.0" source = { editable = "." } dependencies = [ { name = "accelerate" }, From 7b612bfb341aec3e88a767f034d4a1ea74b5b976 Mon Sep 17 00:00:00 2001 From: Aaron Steven White Date: Tue, 30 Jun 2026 14:15:19 -0400 Subject: [PATCH 2/9] Fixes the gallery active-learning loop so it actually trains Running the loop surfaced two blockers in run_pipeline.py. First, the model validates each item against the template's task_spec.options and reads those option names from rendered_elements, but the protocol template leaves options unset (FORCED_CHOICE uses per-item options). build_item_template now sets ("option_a", "option_b") on the model-facing template so the 2AFC pairs validate. Second, the loop trains without participant ids, which a random-intercepts model rejects; load_acceptability_model now switches the seeded model to fixed effects (the pretrained encoder carries over) so the loop fine-tunes at the population level. Verified end to end with bert-base-uncased on CPU: initial training (ForcedChoiceModel.train with random intercepts), AcceptabilityScorer scoring via the population fallback, and a full loop.run iteration with the loaded model. --- .../eng/argument_structure/run_pipeline.py | 20 +++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/gallery/eng/argument_structure/run_pipeline.py b/gallery/eng/argument_structure/run_pipeline.py index 308befa..7cb2f11 100755 --- a/gallery/eng/argument_structure/run_pipeline.py +++ b/gallery/eng/argument_structure/run_pipeline.py @@ -86,9 +86,19 @@ def load_pairs(config: dict[str, Any], base_dir: Path) -> list[Item]: def build_item_template(config_path: Path) -> ItemTemplate: - """Build the 2AFC ItemTemplate from the protocol declared in config.yaml.""" + """Build the 2AFC ItemTemplate that the forced-choice model validates against. + + The protocol declares per-item (FORCED_CHOICE) options, so its task spec + leaves ``options`` unset. The active-learning model, however, validates each + item against the template's ``task_spec.options`` and reads those option + names out of ``rendered_elements``. The 2AFC pairs use ``option_a`` / + ``option_b``, so we set those on the model-facing template. + """ family = acceptability_family(build_protocol(config_path)) - return family_to_item_template(family, judgment_type="acceptability") + template = family_to_item_template(family, judgment_type="acceptability") + return template.with_( + task_spec=template.task_spec.with_(options=("option_a", "option_b")) + ) def build_loop_config(config: dict[str, Any]) -> ActiveLearningLoopConfig: @@ -121,6 +131,7 @@ def load_acceptability_model(checkpoint_dir: Path, device: str) -> Any: ForcedChoiceModel The pretrained model, ready to seed the active-learning loop. """ + from bead.active_learning.config import MixedEffectsConfig # noqa: PLC0415 from bead.active_learning.models.forced_choice import ( # noqa: PLC0415 ForcedChoiceModel, ) @@ -128,6 +139,11 @@ def load_acceptability_model(checkpoint_dir: Path, device: str) -> Any: model = ForcedChoiceModel(ForcedChoiceModelConfig(device=device)) # type: ignore[arg-type] model.load(str(checkpoint_dir)) + # The model was pretrained with per-annotator random intercepts. The loop + # fine-tunes at the population level (it trains without participant ids), so + # switch the seeded model to fixed effects; the pretrained encoder carries + # over. + model.config = model.config.with_(mixed_effects=MixedEffectsConfig(mode="fixed")) return model From daf6f66c9a7ef6b06d19cec33815d1c2f7379b80 Mon Sep 17 00:00:00 2001 From: Aaron Steven White Date: Tue, 30 Jun 2026 15:02:27 -0400 Subject: [PATCH 3/9] Fixes training metrics and adds dev-set early stopping to the acceptability model Running real training surfaced several issues, fixed here. ForcedChoiceModel reported train_accuracy as 0.0 because it read accuracy from the shared mixed-effects trainer's evaluate(), whose prediction step only labels cloze batches, so for forced choice it returns no usable accuracy or loss. Train and dev metrics are now computed from the model's own predict path (which applies the fitted random effects), so they reflect the deployed model. compute_binary_metrics and compute_multiclass_metrics no longer pass zero_division to evaluate's f1, which this version rejects. The model now supports early stopping on a held-out dev set via a callback that scores the dev set through the predict path at each epoch end, stops after early_stopping_patience epochs without dev cross-entropy improvement, and restores the best encoder and head. A new early_stopping_patience config field controls it. The eng gallery builds the dev set by holding out whole sentences: prepare_megaacceptability.py splits unique sentences (and all their ratings) into train and dev so the dev set contains only sentences never seen during training, and writes both; train_acceptability_model.py loads the dev set and passes it for early stopping. The gallery JSONL loaders now use Item.model_validate_json, since Item(**json.loads(line)) no longer coerces string UUIDs under didactic. Verified end to end on bert-base-uncased: training learns (loss drops, train and dev accuracy above chance), early stopping engages, and the existing forced-choice, categorical, and binary model tests pass. --- bead/active_learning/models/forced_choice.py | 213 ++++++++++++++---- bead/active_learning/trainers/metrics.py | 9 +- bead/config/active_learning.py | 1 + gallery/eng/argument_structure/config.yaml | 9 +- .../eng/argument_structure/generate_lists.py | 11 +- .../prepare_megaacceptability.py | 68 +++++- .../eng/argument_structure/run_pipeline.py | 4 +- .../train_acceptability_model.py | 31 ++- 8 files changed, 281 insertions(+), 65 deletions(-) diff --git a/bead/active_learning/models/forced_choice.py b/bead/active_learning/models/forced_choice.py index 1ca09da..4b5a979 100644 --- a/bead/active_learning/models/forced_choice.py +++ b/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,86 @@ __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, + ) -> TrainerControl: + 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 + return control + + def on_train_end( + self, + args: TrainingArguments, + state: TrainerState, + control: TrainerControl, + **kwargs: object, + ) -> TrainerControl: + 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) + return control + + class ForcedChoiceModel(ActiveLearningModel): """Model for forced_choice tasks with n alternatives. @@ -479,9 +567,12 @@ 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 = ( @@ -489,13 +580,6 @@ def _train_with_huggingface_trainer( 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 +589,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 +601,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 +618,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 +679,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], diff --git a/bead/active_learning/trainers/metrics.py b/bead/active_learning/trainers/metrics.py index 9f17a06..8030f2d 100644 --- a/bead/active_learning/trainers/metrics.py +++ b/bead/active_learning/trainers/metrics.py @@ -72,9 +72,9 @@ def compute_binary_metrics(eval_pred: EvalPrediction) -> dict[str, float]: recall = recall_metric.compute( predictions=preds, references=labels, average="binary", zero_division=0 )["recall"] - f1 = f1_metric.compute( - predictions=preds, references=labels, average="binary", zero_division=0 - )["f1"] + # f1's compute() does not accept zero_division in this evaluate version; + # the underlying sklearn f1 defaults undefined cases to 0.0. + f1 = f1_metric.compute(predictions=preds, references=labels, average="binary")["f1"] return { "accuracy": accuracy, @@ -214,11 +214,12 @@ def compute_multiclass_metrics( average="macro", zero_division=0, )["recall"] + # f1's compute() does not accept zero_division in this evaluate version; + # the underlying sklearn f1 defaults undefined cases to 0.0. f1 = f1_metric.compute( predictions=preds, references=labels, average="macro", - zero_division=0, )["f1"] return { diff --git a/bead/config/active_learning.py b/bead/config/active_learning.py index 46ae226..874cab8 100644 --- a/bead/config/active_learning.py +++ b/bead/config/active_learning.py @@ -48,6 +48,7 @@ class BaseEncoderModelConfig(dx.Model): batch_size: int = 16 num_epochs: int = 3 device: Literal["cpu", "cuda", "mps"] = "cpu" + early_stopping_patience: int = 2 mixed_effects: dx.Embed[MixedEffectsConfig] = dx.field( default_factory=_default_mixed_effects ) diff --git a/gallery/eng/argument_structure/config.yaml b/gallery/eng/argument_structure/config.yaml index 09d4eb3..b71809d 100644 --- a/gallery/eng/argument_structure/config.yaml +++ b/gallery/eng/argument_structure/config.yaml @@ -185,9 +185,15 @@ acceptability_model: max_pairs_per_annotator: 200 seed: 42 + # dev split: hold out this fraction of unique sentences (and all their + # ratings) so the dev set contains only sentences unseen during training; + # used for early stopping + dev_fraction: 0.15 + # derived training set artifacts (read by prepare/train scripts) paths: training_items: "items/megaacceptability_2afc.jsonl" + dev_items: "items/megaacceptability_2afc_dev.jsonl" raw_corpus_dir: "items/megaacceptability_raw_corpus" training_corpus_dir: "items/megaacceptability_2afc_corpus" @@ -197,7 +203,8 @@ acceptability_model: mixed_effects_mode: "random_intercepts" learning_rate: 2e-5 batch_size: 16 - epochs: 3 + epochs: 10 + early_stopping_patience: 2 # stop after N dev epochs without improvement device: "cpu" checkpoint_dir: "checkpoints/acceptability_init" diff --git a/gallery/eng/argument_structure/generate_lists.py b/gallery/eng/argument_structure/generate_lists.py index eaab600..ec98c7e 100644 --- a/gallery/eng/argument_structure/generate_lists.py +++ b/gallery/eng/argument_structure/generate_lists.py @@ -7,7 +7,6 @@ from __future__ import annotations -import json import sys import traceback from pathlib import Path @@ -68,8 +67,10 @@ def load_2afc_pairs(pairs_path: Path) -> tuple[list[Item], dict[UUID, dict]]: with open(pairs_path) as f: for line in f: - data = json.loads(line) - item = Item(**data) + line = line.strip() + if not line: + continue + item = Item.model_validate_json(line) items.append(item) # Build metadata dict for constraint checking @@ -339,9 +340,7 @@ def main() -> None: fragment_rel = config["paths"].get("experiment_lists_fragment") if fragment_rel: fragment_path = base_dir / fragment_rel - layers_io.write_experiment_lists_layers( - list(experiment_lists), fragment_path - ) + layers_io.write_experiment_lists_layers(list(experiment_lists), fragment_path) print_success(f"Wrote layers list collections to {fragment_path.name}") console.print() diff --git a/gallery/eng/argument_structure/prepare_megaacceptability.py b/gallery/eng/argument_structure/prepare_megaacceptability.py index 6242398..09524fd 100644 --- a/gallery/eng/argument_structure/prepare_megaacceptability.py +++ b/gallery/eng/argument_structure/prepare_megaacceptability.py @@ -262,6 +262,42 @@ def build_per_annotator_pairs( return items +def split_rows_by_sentence( + rows: list[RatingRow], dev_fraction: float, seed: int +) -> tuple[list[RatingRow], list[RatingRow]]: + """Split rows into train and dev so no sentence appears in both. + + Unique sentences are partitioned: a ``dev_fraction`` of them (and every + rating of those sentences) is held out as the dev split, the rest form the + train split. Holding out whole sentences guarantees the dev set contains + only sentences never seen during training, so dev performance measures + generalization to new stimuli rather than new annotator pairings. + + Parameters + ---------- + rows : list[RatingRow] + All rating rows. + dev_fraction : float + Fraction of unique sentences to hold out for the dev split. + seed : int + Seed for the deterministic sentence shuffle. + + Returns + ------- + tuple[list[RatingRow], list[RatingRow]] + The train rows and dev rows. + """ + if dev_fraction <= 0.0: + return list(rows), [] + sentences = sorted({row.sentence for row in rows}) + Random(seed).shuffle(sentences) + n_dev = max(1, round(len(sentences) * dev_fraction)) + dev_sentences = set(sentences[:n_dev]) + train_rows = [row for row in rows if row.sentence not in dev_sentences] + dev_rows = [row for row in rows if row.sentence in dev_sentences] + return train_rows, dev_rows + + # --- raw item construction -------------------------------------------------- @@ -533,10 +569,15 @@ def main( base_paths = config.get("paths") base_paths = base_paths if isinstance(base_paths, dict) else {} + dev_fraction = float(section.get("dev_fraction", 0.15)) + cache_dir = base_dir / str(base_paths.get("cache_dir", ".cache")) training_items_path = base_dir / str( paths_section.get("training_items", "items/megaacceptability_2afc.jsonl") ) + dev_items_path = base_dir / str( + paths_section.get("dev_items", "items/megaacceptability_2afc_dev.jsonl") + ) raw_corpus_dir = base_dir / str( paths_section.get("raw_corpus_dir", "items/megaacceptability_raw_corpus") ) @@ -589,11 +630,20 @@ def main( ) console.print() - # 3. Derive per-annotator 2AFC pairs + # 3. Split by sentence (dev sentences are unseen during training), then + # derive per-annotator 2AFC pairs within each split print_header("3/4 Building Per-Annotator 2AFC Pairs") + train_rows, dev_rows = split_rows_by_sentence(rows, dev_fraction, seed) with console.status("[bold]Pairing within-rater sentences...[/bold]"): pair_items = build_per_annotator_pairs( - rows, + train_rows, + scope=scope, + drop_ties=drop_ties, + max_per_annotator=max_per_annotator, + seed=seed, + ) + dev_items = build_per_annotator_pairs( + dev_rows, scope=scope, drop_ties=drop_ties, max_per_annotator=max_per_annotator, @@ -602,15 +652,21 @@ def main( if not pair_items: print_error("No 2AFC pairs were derived. Exiting.") sys.exit(1) - print_success(f"Derived {len(pair_items):,} forced-choice training items\n") + print_success( + f"Derived {len(pair_items):,} train and {len(dev_items):,} dev items " + f"(dev sentences held out)\n" + ) - # 4. Persist the derived training set (JSONL + layers corpus) + # 4. Persist the derived train and dev sets (JSONL + layers corpus) print_header("4/4 Persisting Derived Training Set") try: write_items_jsonl(pair_items, training_items_path) - print_success(f"Wrote {len(pair_items):,} items to {training_items_path}") + print_success(f"Wrote {len(pair_items):,} train items to {training_items_path}") + if dev_items: + write_items_jsonl(dev_items, dev_items_path) + print_success(f"Wrote {len(dev_items):,} dev items to {dev_items_path}") except Exception as exc: # noqa: BLE001 - print_error(f"Failed to write {training_items_path}: {exc}") + print_error(f"Failed to write items: {exc}") sys.exit(1) emit_layers_corpus( pair_items, diff --git a/gallery/eng/argument_structure/run_pipeline.py b/gallery/eng/argument_structure/run_pipeline.py index 7cb2f11..ceff651 100755 --- a/gallery/eng/argument_structure/run_pipeline.py +++ b/gallery/eng/argument_structure/run_pipeline.py @@ -81,7 +81,9 @@ def load_pairs(config: dict[str, Any], base_dir: Path) -> list[Item]: items: list[Item] = [] with open(pairs_path) as f: for line in f: - items.append(Item(**json.loads(line))) + line = line.strip() + if line: + items.append(Item.model_validate_json(line)) return items diff --git a/gallery/eng/argument_structure/train_acceptability_model.py b/gallery/eng/argument_structure/train_acceptability_model.py index f3c9857..709d586 100644 --- a/gallery/eng/argument_structure/train_acceptability_model.py +++ b/gallery/eng/argument_structure/train_acceptability_model.py @@ -18,7 +18,6 @@ from __future__ import annotations import argparse -import json import sys from pathlib import Path from typing import TYPE_CHECKING @@ -72,7 +71,7 @@ def load_training_items(path: Path, *, limit: int | None = None) -> list[Item]: line = line.strip() if not line: continue - items.append(Item(**json.loads(line))) + items.append(Item.model_validate_json(line)) return items @@ -138,6 +137,7 @@ def build_model_config(section: dict[str, object]) -> ForcedChoiceModelConfig: batch_size=int(section.get("batch_size", DEFAULT_BATCH_SIZE)), num_epochs=int(section.get("epochs", DEFAULT_EPOCHS)), device=str(section.get("device", DEFAULT_DEVICE)), + early_stopping_patience=int(section.get("early_stopping_patience", 2)), mixed_effects=MixedEffectsConfig(mode=mode), ) @@ -211,6 +211,9 @@ def main( training_items_path = base_dir / str( paths_section.get("training_items", DEFAULT_TRAINING_ITEMS) ) + dev_items_path = base_dir / str( + paths_section.get("dev_items", "items/megaacceptability_2afc_dev.jsonl") + ) checkpoint_dir = base_dir / str( section.get("checkpoint_dir", DEFAULT_CHECKPOINT_DIR) ) @@ -243,7 +246,21 @@ def main( print_error("No training items loaded. Exiting.") sys.exit(1) n_participants = len(set(participant_ids)) - print_success(f"Loaded {len(items):,} items from {n_participants:,} annotators\n") + print_success(f"Loaded {len(items):,} items from {n_participants:,} annotators") + + # Load the held-out dev set (unseen sentences) for early stopping, if present + dev_items: list[Item] = [] + dev_labels: list[str] = [] + if dev_items_path.exists(): + dev_items = load_training_items(dev_items_path, limit=item_limit) + dev_labels, _ = extract_labels_and_participants(dev_items) + print_success(f"Loaded {len(dev_items):,} dev items from {dev_items_path.name}") + else: + print_warning( + "No dev set found; training without early stopping. " + "Run prepare_megaacceptability.py to create one." + ) + console.print() # 2. Build config and train (heavy imports deferred to here) print_header("2/3 Training Forced-Choice Model") @@ -260,7 +277,13 @@ def main( model = ForcedChoiceModel(model_config) with console.status("[bold]Training (this can be slow on CPU)...[/bold]"): - metrics = model.train(items, labels, participant_ids=participant_ids) + metrics = model.train( + items, + labels, + participant_ids=participant_ids, + validation_items=dev_items or None, + validation_labels=dev_labels or None, + ) except Exception as exc: # noqa: BLE001 print_error(f"Training failed: {exc}") sys.exit(1) From db9e27ab5391a2cbd0c2503e997863fe88115c87 Mon Sep 17 00:00:00 2001 From: Aaron Steven White Date: Mon, 13 Jul 2026 15:13:27 -0400 Subject: [PATCH 4/9] Removes the acceptability pairing cap to use all within-rater same-frame pairs MegaAcceptability lists are structured so no verb or frame repeats within a list, so the only within-rater controlled contrast is same-frame (different verbs, one frame), which a rater sees once per list across their lists. Capping at 200 pairs per annotator discarded most of these, so max_pairs_per_annotator now defaults to null (no cap). On the full dataset this raises the derived same-frame training set from roughly 145k pairs to roughly 760k (minus ties and the held-out dev split). --- gallery/eng/argument_structure/config.yaml | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/gallery/eng/argument_structure/config.yaml b/gallery/eng/argument_structure/config.yaml index b71809d..3daf92c 100644 --- a/gallery/eng/argument_structure/config.yaml +++ b/gallery/eng/argument_structure/config.yaml @@ -182,7 +182,10 @@ acceptability_model: pairing_strategy: "per_annotator" pairing_scope: "same_frame" # all, same_verb, or same_frame drop_ties: true - max_pairs_per_annotator: 200 + # no cap: use every within-rater same-frame pair (different verbs, one frame). + # MegaAcceptability lists have no repeated verb or frame, so same-frame pairs + # come from a rater seeing each frame once per list across their lists. + max_pairs_per_annotator: null seed: 42 # dev split: hold out this fraction of unique sentences (and all their From 08cfbb1424e34356b69eee57fe244f9b1c268da6 Mon Sep 17 00:00:00 2001 From: Aaron Steven White Date: Mon, 13 Jul 2026 15:57:06 -0400 Subject: [PATCH 5/9] Fixes random-effects device placement so forced choice trains on MPS and CUDA Switching the acceptability model to device: mps surfaced a device mismatch: _do_predict_proba adds the per-participant random-effects intercept, a CPU tensor, to the encoder logits on the accelerator, raising "Expected all tensors to be on the same device". The random-effects bias (and the random-slopes head) are now moved to the logits / embedding device, which is a no-op on CPU. The early-stopping callback methods now return None and mutate control in place, matching the TrainerCallback signature. The eng gallery config sets acceptability_model.device to mps, since the uncapped same-frame training set (~760k pairs) is impractical on CPU. Verified: BERT trains and scores on MPS, and the forced-choice model tests still pass on CPU. --- bead/active_learning/models/forced_choice.py | 12 ++++++------ gallery/eng/argument_structure/config.yaml | 2 +- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/bead/active_learning/models/forced_choice.py b/bead/active_learning/models/forced_choice.py index 4b5a979..8b147d1 100644 --- a/bead/active_learning/models/forced_choice.py +++ b/bead/active_learning/models/forced_choice.py @@ -77,7 +77,7 @@ def on_epoch_end( state: TrainerState, control: TrainerControl, **kwargs: object, - ) -> TrainerControl: + ) -> None: loss = self._dev_cross_entropy() if loss < self._best_loss - 1e-4: self._best_loss = loss @@ -96,7 +96,6 @@ def on_epoch_end( self._wait += 1 if self._wait >= self._patience: control.should_training_stop = True - return control def on_train_end( self, @@ -104,12 +103,11 @@ def on_train_end( state: TrainerState, control: TrainerControl, **kwargs: object, - ) -> TrainerControl: + ) -> 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) - return control class ForcedChoiceModel(ActiveLearningModel): @@ -970,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) diff --git a/gallery/eng/argument_structure/config.yaml b/gallery/eng/argument_structure/config.yaml index 3daf92c..822b5a9 100644 --- a/gallery/eng/argument_structure/config.yaml +++ b/gallery/eng/argument_structure/config.yaml @@ -208,7 +208,7 @@ acceptability_model: batch_size: 16 epochs: 10 early_stopping_patience: 2 # stop after N dev epochs without improvement - device: "cpu" + device: "mps" # Apple Silicon GPU; the ~760k-pair set is impractical on CPU checkpoint_dir: "checkpoints/acceptability_init" # layers / lairs interop From 3143629080ef91eb8c273ae02189d12caceb9559 Mon Sep 17 00:00:00 2001 From: Aaron Steven White Date: Mon, 13 Jul 2026 16:01:22 -0400 Subject: [PATCH 6/9] Updates the 0.8.0 changelog with training fixes and dev-set early stopping --- CHANGELOG.md | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2703ee2..1151e2c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,7 +7,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] -## [0.8.0] - 2026-06-30 +## [0.8.0] - 2026-07-02 ### Added @@ -65,6 +65,26 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `AgentRef` and whose study fields (consent, sessions, demographics) are outside the shared schema. +#### 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 From ef6ed6c7e11bdf7fe96ec2cfe818ecf3d1e2f390 Mon Sep 17 00:00:00 2001 From: Aaron Steven White Date: Mon, 13 Jul 2026 16:08:43 -0400 Subject: [PATCH 7/9] Formats the new grid-stratification and layers-resource tests --- tests/interop/test_layers_resource.py | 4 +++- tests/lists/test_grid_stratification.py | 14 +++++++++----- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/tests/interop/test_layers_resource.py b/tests/interop/test_layers_resource.py index a1169af..fe08787 100644 --- a/tests/interop/test_layers_resource.py +++ b/tests/interop/test_layers_resource.py @@ -127,7 +127,9 @@ def _filled(self) -> FilledTemplate: source="bleached", ), "verb": LexicalItem( - lemma="chase", form="chased", language_code="eng", + lemma="chase", + form="chased", + language_code="eng", features={"pos": "VERB"}, ), }, diff --git a/tests/lists/test_grid_stratification.py b/tests/lists/test_grid_stratification.py index 698d52c..86cc249 100644 --- a/tests/lists/test_grid_stratification.py +++ b/tests/lists/test_grid_stratification.py @@ -222,9 +222,7 @@ class TestGridStratificationConstraint: """The constraint validates its dimensions and parameters.""" def test_valid(self) -> None: - dim = GridDimension( - property_expression="item['a']", binning=_q(3) - ) + dim = GridDimension(property_expression="item['a']", binning=_q(3)) constraint = GridStratificationConstraint( constraint_type="grid_stratification", dimensions=(dim,), @@ -262,7 +260,10 @@ def test_distribute_across_cells(self) -> None: cell_of = {uid: i % 4 for i, uid in enumerate(uuids)} balancer = QuantileBalancer(n_quantiles=2, random_seed=0) lists = balancer.balance_by_cell( - uuids, lambda uid: cell_of[uid], n_cells=4, n_lists=2, + uuids, + lambda uid: cell_of[uid], + n_cells=4, + n_lists=2, items_per_cell_per_list=5, ) assert len(lists) == 2 @@ -308,7 +309,10 @@ def test_grid_partition(self) -> None: ) partitioner = ListPartitioner(random_seed=42) lists = partitioner.partition( - items, n_lists=4, constraints=[constraint], strategy="stratified", + items, + n_lists=4, + constraints=[constraint], + strategy="stratified", metadata=meta, ) assert len(lists) == 4 From 89a399f011c486a27914cc4bc93c5b98272b37fd Mon Sep 17 00:00:00 2001 From: Aaron Steven White Date: Mon, 13 Jul 2026 16:30:35 -0400 Subject: [PATCH 8/9] Maps bead Participant to a layers AgentRef and judgment-set features layers represents a data producer as an AgentRef (identity), separate from the persona (role/framework) and tool (software), and documents JudgmentSet.features as the home for annotator demographics, session metadata, and payment info. Map a bead Participant onto that: PARTICIPANT_AGENT lenses Participant <-> AgentRef losslessly, participant_features renders the study fields as a FeatureMap, and records_to_judgment_set takes an optional participant so a judgment set carries the annotator's identity and study fields natively. --- CHANGELOG.md | 17 ++- bead/interop/layers/__init__.py | 12 ++ bead/interop/layers/judgment_lens.py | 24 +++- bead/interop/layers/participant_lens.py | 133 +++++++++++++++++++++++ tests/interop/test_layers_participant.py | 114 +++++++++++++++++++ 5 files changed, 293 insertions(+), 7 deletions(-) create mode 100644 bead/interop/layers/participant_lens.py create mode 100644 tests/interop/test_layers_participant.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 1151e2c..3c11188 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -58,12 +58,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 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. - - All of these round-trip exactly. A bead `Participant` is intentionally not - mapped: a layers `persona.Persona` is an annotator role and interpretive - framework, not a concrete enrolled participant, whose identity is instead an - `AgentRef` and whose study fields (consent, sessions, demographics) are outside - the shared schema. +- `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 diff --git a/bead/interop/layers/__init__.py b/bead/interop/layers/__init__.py index 96c2156..121412a 100644 --- a/bead/interop/layers/__init__.py +++ b/bead/interop/layers/__init__.py @@ -22,6 +22,8 @@ - bead's resources to their layers counterparts: ``LexicalItem`` to an ``entry``, ``Lexicon`` to a ``collection``, and ``Template`` to a ``template`` (see ``resource_lens``). +- a ``Participant`` to a layers ``agentRef`` identity, with its study fields + rendered as the ``featureMap`` of a ``judgmentSet`` (see ``participant_lens``). """ from __future__ import annotations @@ -74,6 +76,12 @@ ParsedSentenceLayersIso, parse_to_layers, ) +from bead.interop.layers.participant_lens import ( + PARTICIPANT_AGENT, + ParticipantAgentLens, + agent_ref_of, + participant_features, +) from bead.interop.layers.resource_lens import ( FILLED_TEMPLATE_FILLING, LEXICAL_ITEM_ENTRY, @@ -96,6 +104,7 @@ "LEXICON_COLLECTION", "LIST_CONSTRAINT", "PARSED_SENTENCE_LAYERS", + "PARTICIPANT_AGENT", "RECORD_EXPRESSION", "TEMPLATE_LAYERS", "AnnotationRecordJudgmentLens", @@ -112,8 +121,10 @@ "ListConstraintLens", "ParsedSentenceLayers", "ParsedSentenceLayersIso", + "ParticipantAgentLens", "RecordExpressionLens", "TemplateLayersLens", + "agent_ref_of", "corpus_to_graph", "corpus_to_items", "corpus_to_records", @@ -126,6 +137,7 @@ "load_layers_corpus", "materialize_corpus", "parse_to_layers", + "participant_features", "publish_corpus", "record_to_expression", "records_to_judgment_set", diff --git a/bead/interop/layers/judgment_lens.py b/bead/interop/layers/judgment_lens.py index 253d5e6..9831be4 100644 --- a/bead/interop/layers/judgment_lens.py +++ b/bead/interop/layers/judgment_lens.py @@ -18,6 +18,8 @@ from __future__ import annotations +from typing import TYPE_CHECKING + import didactic.api as dx from lairs.records import defs, judgment @@ -31,6 +33,10 @@ j_str, object_ref, ) +from bead.interop.layers.participant_lens import agent_ref_of, participant_features + +if TYPE_CHECKING: + from bead.participants.models import Participant # layers requires an experiment reference on every judgment set; bead annotation # records carry none, so the view uses an empty at-uri and nothing round-trips. @@ -84,6 +90,7 @@ def backward( def records_to_judgment_set( records: tuple[AnnotationRecord, ...], + participant: Participant | None = None, ) -> tuple[judgment.JudgmentSet, JsonValue]: """Project records sharing one annotator to a layers judgment set. @@ -94,11 +101,23 @@ def records_to_judgment_set( complement carries the shared annotator id and the per-record complements, so :func:`judgment_set_to_records` inverts it exactly. + When the annotator's :class:`~bead.participants.models.Participant` is + supplied, the set's ``agent`` is that participant's ``agentRef`` and its + ``features`` carries the participant's study fields (demographics, sessions, + consent), which the ``judgmentSet`` schema documents as the home for + "annotator demographics, session metadata, completion time, payment info". + The participant is supplementary provenance on the view: the records + reconstruct from the complement alone, so ``judgment_set_to_records`` is + unaffected. + Parameters ---------- records : tuple[AnnotationRecord, ...] Annotation records that all share one ``annotator_id``. Must be non-empty. + participant : Participant | None, optional + The annotator behind ``records``. When given, its identity and study + fields enrich the judgment set's ``agent`` and ``features``. Returns ------- @@ -136,10 +155,13 @@ def records_to_judgment_set( judgments.append(judgment_view) record_complements.append(record_complement) + agent = agent_ref_of(participant) if participant else defs.AgentRef(id=annotator_id) + features = participant_features(participant) if participant else None view = judgment.JudgmentSet( - agent=defs.AgentRef(id=annotator_id), + agent=agent, createdAt=min(record.created_at for record in records), experimentRef=_NO_EXPERIMENT_REF, + features=features, judgments=tuple(judgments), ) complement: JsonValue = { diff --git a/bead/interop/layers/participant_lens.py b/bead/interop/layers/participant_lens.py new file mode 100644 index 0000000..50fbe91 --- /dev/null +++ b/bead/interop/layers/participant_lens.py @@ -0,0 +1,133 @@ +"""Lens between a bead ``Participant`` and its layers representation. + +layers has no dedicated participant record. It represents whoever produced data +as an ``agentRef`` (:class:`lairs.records.defs.AgentRef`) whose docstring is +explicit: it "separates the identity of the producer from the interpretive +framework (persona) and the software used (tool)". A participant's study-level +attributes (demographics, session metadata, consent, payment) have their own +native home: the ``featureMap`` on the ``judgmentSet`` that binds one annotator +to one experiment, documented to hold "annotator demographics, session metadata, +completion time, payment info". A ``persona`` is the annotator *role* and +interpretive framework, not a concrete enrolled participant, so it is not the +target here. + +This module maps a bead :class:`~bead.participants.models.Participant` onto that +representation: + +- :data:`PARTICIPANT_AGENT` is a lossless lens + ``Participant <-> (agentRef, complement)``: the participant's UUID becomes the + ``agentRef`` identity and the bead framework identity plus the study fields ride + in the lens complement so reconstruction is exact. +- :func:`participant_features` renders a participant's study fields as a layers + ``featureMap`` for attaching to a ``judgmentSet``, the schema's native slot for + participant-level attributes. +- :func:`agent_ref_of` is the identity half on its own. +""" + +from __future__ import annotations + +import json +from datetime import datetime + +import didactic.api as dx +from lairs.records import defs + +from bead.data.base import JsonValue +from bead.interop.layers._convert import ( + apply_identity, + identity_of, + j_list, + j_obj, + j_str, + j_str_or_none, +) +from bead.participants.models import Participant + + +def agent_ref_of(participant: Participant) -> defs.AgentRef: + """Build the layers ``agentRef`` identity for a participant. + + The internal UUID becomes the ``agentRef`` ``id``; the layers docstring + blesses ``id`` for "anonymized crowdworker ID, platform username" and similar + opaque identifiers, which is exactly a bead participant UUID. + """ + return defs.AgentRef(id=str(participant.id)) + + +def participant_features(participant: Participant) -> defs.FeatureMap | None: + """Render a participant's study fields as a layers ``featureMap``. + + Collects ``participant_metadata`` (demographics) together with ``study_id``, + ``session_ids``, ``consent_timestamp``, and ``notes`` into the ``featureMap`` + that a ``judgmentSet`` documents as the home for "annotator demographics, + session metadata, completion time, payment info". Each value is serialized + with ``json.dumps`` (the shared feature-map convention), so + :func:`bead.interop.layers._convert.read_feature_map` decodes it exactly. + Returns ``None`` when the participant carries no study fields (a faithful + layers view omits empty optionals). + """ + entries: dict[str, JsonValue] = dict(participant.participant_metadata) + if participant.study_id is not None: + entries["study_id"] = participant.study_id + if participant.session_ids: + entries["session_ids"] = participant.session_ids + if participant.consent_timestamp is not None: + entries["consent_timestamp"] = participant.consent_timestamp.isoformat() + if participant.notes is not None: + entries["notes"] = participant.notes + if not entries: + return None + return defs.FeatureMap( + entries=tuple( + defs.Feature(key=key, value=json.dumps(entries[key])) for key in entries + ) + ) + + +class ParticipantAgentLens(dx.Lens[Participant, defs.AgentRef, JsonValue]): + """Lossless lens ``Participant <-> (layers agentRef, complement)``. + + The layers ``agentRef`` (:class:`lairs.records.defs.AgentRef`) is the + canonical identity of a data producer, so a participant projects to the + ``agentRef`` carrying its UUID. Everything else, the bead framework identity + and the study fields (demographics, study id, sessions, consent, notes), + rides in the lens complement, since a standalone ``agentRef`` has no slot for + study-level attributes (those attach to a ``judgmentSet`` via + :func:`participant_features`). The GetPut/PutGet laws guarantee an exact + round-trip. + """ + + def forward(self, participant: Participant) -> tuple[defs.AgentRef, JsonValue]: + """Project a participant to a layers agent reference and complement.""" + view = agent_ref_of(participant) + complement: JsonValue = { + "identity": identity_of(participant), + "participant_metadata": dict(participant.participant_metadata), + "study_id": participant.study_id, + "session_ids": participant.session_ids, + "consent_timestamp": ( + participant.consent_timestamp.isoformat() + if participant.consent_timestamp is not None + else None + ), + "notes": participant.notes, + } + return view, complement + + def backward(self, view: defs.AgentRef, complement: JsonValue) -> Participant: + """Reconstruct a participant from its layers agent reference + complement.""" + comp = j_obj(complement) + consent = comp["consent_timestamp"] + participant = Participant( + participant_metadata=j_obj(comp["participant_metadata"]), + study_id=j_str_or_none(comp["study_id"]), + session_ids=tuple(j_str(value) for value in j_list(comp["session_ids"])), + consent_timestamp=( + datetime.fromisoformat(j_str(consent)) if consent is not None else None + ), + notes=j_str_or_none(comp["notes"]), + ) + return apply_identity(participant, comp["identity"]) + + +PARTICIPANT_AGENT = ParticipantAgentLens() diff --git a/tests/interop/test_layers_participant.py b/tests/interop/test_layers_participant.py new file mode 100644 index 0000000..df1bcdd --- /dev/null +++ b/tests/interop/test_layers_participant.py @@ -0,0 +1,114 @@ +"""Round-trip law tests for the participant / agent-reference lens.""" + +from __future__ import annotations + +from datetime import datetime + +from lairs.records import defs + +from bead.evaluation.reliability import AnnotationRecord +from bead.interop.layers._convert import read_feature_map +from bead.interop.layers.judgment_lens import ( + judgment_set_to_records, + records_to_judgment_set, +) +from bead.interop.layers.participant_lens import ( + PARTICIPANT_AGENT, + agent_ref_of, + participant_features, +) +from bead.participants.models import Participant + + +def _record( + *, annotator_id: str = "ann_1", item_id: str = "item_1" +) -> AnnotationRecord: + return AnnotationRecord( + annotator_id=annotator_id, + item_id=item_id, + question_name="acceptability", + response_label="option_a", + ) + + +def _participant() -> Participant: + return Participant( + participant_metadata={"age": 34, "native_speaker": True, "l1": "eng"}, + study_id="megaacceptability", + session_ids=("session_a", "session_b"), + consent_timestamp=datetime(2026, 6, 1, 9, 30, 0), + notes="recruited via prolific", + ) + + +class TestParticipantAgent: + """Participant <-> layers agentRef.""" + + def test_view_is_agent_ref(self) -> None: + participant = _participant() + view, _ = PARTICIPANT_AGENT.forward(participant) + assert isinstance(view, defs.AgentRef) + assert view.id == str(participant.id) + + def test_roundtrip_exact(self) -> None: + participant = _participant() + view, complement = PARTICIPANT_AGENT.forward(participant) + assert PARTICIPANT_AGENT.backward(view, complement) == participant + + def test_roundtrip_through_serialization(self) -> None: + participant = _participant() + view, complement = PARTICIPANT_AGENT.forward(participant) + view2 = defs.AgentRef.model_validate_json(view.model_dump_json()) + assert PARTICIPANT_AGENT.backward(view2, complement) == participant + + def test_roundtrip_minimal(self) -> None: + participant = Participant() + view, complement = PARTICIPANT_AGENT.forward(participant) + assert PARTICIPANT_AGENT.backward(view, complement) == participant + + +class TestParticipantFeatures: + """A participant's study fields render to a layers featureMap.""" + + def test_features_decode_to_study_fields(self) -> None: + participant = _participant() + features = participant_features(participant) + assert features is not None + decoded = read_feature_map(features) + assert decoded["age"] == 34 + assert decoded["native_speaker"] is True + assert decoded["study_id"] == "megaacceptability" + assert decoded["session_ids"] == ("session_a", "session_b") + assert decoded["consent_timestamp"] == "2026-06-01T09:30:00" + assert decoded["notes"] == "recruited via prolific" + + def test_features_none_when_empty(self) -> None: + assert participant_features(Participant()) is None + + def test_agent_ref_of_carries_uuid(self) -> None: + participant = _participant() + assert agent_ref_of(participant).id == str(participant.id) + + +class TestJudgmentSetWithParticipant: + """A participant enriches the judgment set's agent and features.""" + + def test_participant_populates_agent_and_features(self) -> None: + participant = _participant() + records = ( + _record(item_id="item_1"), + _record(item_id="item_2"), + ) + view, _ = records_to_judgment_set(records, participant) + assert view.agent is not None + assert view.agent.id == str(participant.id) + assert view.features is not None + assert read_feature_map(view.features)["study_id"] == "megaacceptability" + + def test_records_still_roundtrip_with_participant(self) -> None: + records = ( + _record(item_id="item_1"), + _record(item_id="item_2"), + ) + view, complement = records_to_judgment_set(records, _participant()) + assert judgment_set_to_records(view, complement) == records From db03f5352e69a84ccfc48045777e55f85d0b51bd Mon Sep 17 00:00:00 2001 From: Aaron Steven White Date: Mon, 13 Jul 2026 18:13:53 -0400 Subject: [PATCH 9/9] Sets the 0.8.0 changelog date to 2026-07-13 --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3c11188..b2df7b8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,7 +7,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] -## [0.8.0] - 2026-07-02 +## [0.8.0] - 2026-07-13 ### Added