From 933b8f790ead2f7e3592ddb497faf480d3614cb2 Mon Sep 17 00:00:00 2001 From: Kartik Pradeepan Date: Wed, 4 Nov 2026 06:24:11 -0500 Subject: [PATCH 1/8] feat(sessions): record neural regions during a behavioral task in one presentation A BrainScoreModel session may now request behavior together with neural:. Each stimulus is shown once: with answer, extractors record the prompt pass through an optional record(layers); with a readout, region layers are extracted in the same call as the readout layer. A neural-only session after a configured task now returns activity instead of the behavioral assembly. --- brainscore_core/_model/behavioral.py | 151 +++++++++++++++++++++- brainscore_core/_sessions/drivers.py | 38 ++++++ brainscore_core/_sessions/helpers.py | 1 + brainscore_core/capabilities/sessions.py | 19 ++- tests/test_joint_behavior_neural.py | 158 +++++++++++++++++++++++ 5 files changed, 358 insertions(+), 9 deletions(-) create mode 100644 tests/test_joint_behavior_neural.py diff --git a/brainscore_core/_model/behavioral.py b/brainscore_core/_model/behavioral.py index 7d5e2ae8..0e0b4467 100644 --- a/brainscore_core/_model/behavioral.py +++ b/brainscore_core/_model/behavioral.py @@ -1,6 +1,7 @@ """Behavioral readout and generation support for BrainScoreModel.""" -from typing import Any, Optional +from contextlib import ExitStack, contextmanager +from typing import Any, Callable, Dict, Optional, Sequence, Tuple from ..contract import TaskContext @@ -247,8 +248,124 @@ def predict_probabilities(self, stimuli): features = self.owner._extract_behavioral_features(stimuli) return self.readout_classifier.predict_proba(features) - def generate_predictions(self, stimuli): - """Call the registered answer function per stimulus, return one-hot BehavioralAssembly.""" + @contextmanager + def suspended(self): + """Keep the configured task, but let it answer nothing for a while. + + A session that asks only for ``neural:`` must get activity, even + when an earlier session configured a task on this model (rule 4.1). + """ + saved = self.task_context + self.task_context = None + try: + yield + finally: + self.task_context = saved + + def respond_and_record(self, stimuli, regions: Sequence[str]): + """Answer every stimulus once and record ``regions`` from that same pass. + + Returns ``(BehavioralAssembly, {region: NeuroidAssembly})``. With + ``answer``, the recording is the first forward pass of each answer (the + prompt, before any generated token), taken through each extractor's + ``record(layers)``. With a readout, the region layers are extracted in + the same call as the readout layer. Unsupported requests fail before + any compute. + """ + layers = {region: self._region_layer(region) for region in regions} + if self.use_generation_for_task: + return self._generate_and_record(stimuli, layers) + return self._read_out_and_record(stimuli, layers) + + def _region_layer(self, region: str) -> str: + from ..selection import CompositeSelector + selector = self.owner._region_layer_selectors.get(region) + if selector is None: + raise ValueError(f'Region {region!r} is not in region_layer_map.') + if isinstance(selector, CompositeSelector): + # ponytail: composite regions need per-unit gathering; add when a study asks. + raise NotImplementedError( + f'Region {region!r} combines several layers; record it with a ' + 'neural-only session.') + return selector.layer_path + + def _region_modality(self, region: str) -> str: + from ..io_catalog import canonical_modality + modality = self.owner._region_modality_map.get(region) + if modality is not None: + return canonical_modality(modality) + modalities = {canonical_modality(key) for key in self.owner._extractors} + if len(modalities) != 1: + raise ValueError( + f'Region {region!r} does not say which input it reads. Map it with ' + "LayerSelector(name=..., modality=...) to record it during a task.") + return next(iter(modalities)) + + def _generate_and_record(self, stimuli, layers: Dict[str, str]): + from ..io_catalog import canonical_modality + owner = self.owner + by_modality: Dict[str, list] = {} + for region, layer in layers.items(): + by_modality.setdefault(self._region_modality(region), []).append(layer) + recordings = {} + for modality, modality_layers in by_modality.items(): + extractor = next((e for key, e in owner._extractors.items() + if canonical_modality(key) == modality), None) + if extractor is None: + raise NotImplementedError(f'No {modality!r} extractor to record from.') + try: + record = extractor.record + except AttributeError: + raise NotImplementedError( + f'The {modality!r} extractor cannot record while the model answers: ' + 'it needs a record(layers) method. Use a readout task, or a ' + 'neural-only session.') from None + recordings[modality] = record(list(dict.fromkeys(modality_layers))) + trials = [] + with ExitStack() as stack: + live = {m: stack.enter_context(r) for m, r in recordings.items()} + behavior = self.generate_predictions( + stimuli, after_each=lambda: trials.append( + {m: recording.take() for m, recording in live.items()})) + neural = {} + for region, layer in layers.items(): + modality = self._region_modality(region) + neural[region] = _trials_to_assembly( + [trial[modality] for trial in trials], stimuli, layer, region) + return behavior, neural + + def _read_out_and_record(self, stimuli, layers: Dict[str, str]): + owner = self.owner + readout_layer = owner._behavioral_readout_layer + modality = owner._pick_modality(owner._detect_modalities(stimuli)) + for region in layers: + if owner._region_modality_map.get(region) and self._region_modality(region) != modality: + raise ValueError( + f'Region {region!r} reads {self._region_modality(region)!r} input, but ' + f'the readout presents {modality!r} input; one presentation cannot ' + 'record it.') + all_layers = list(dict.fromkeys([readout_layer, *layers.values()])) + extracted = owner._recorder.ensure_layer_provenance( + owner._extract_for_modality(stimuli, modality, all_layers), all_layers, + description='Readout extraction') + layer_of = extracted['layer'].values + features = extracted.isel(neuroid=layer_of == readout_layer) + if 'presentation' in features.dims and 'neuroid' in features.dims: + features = features.transpose('presentation', 'neuroid') + behavior = self.readout_classifier.predict_proba(features) + neural = {} + for region, layer in layers.items(): + selected = extracted.isel(neuroid=layer_of == layer) + neural[region] = selected.assign_coords( + region=('neuroid', [region] * selected.sizes['neuroid'])) + return behavior, neural + + def generate_predictions(self, stimuli, after_each: Optional[Callable[[], None]] = None): + """Call the registered answer function per stimulus, return one-hot BehavioralAssembly. + + ``after_each`` runs after every answer, before the next stimulus; the + joint behavior-and-neural session uses it to close each recording. + """ import numpy as np from brainscore_core.supported_data_standards.brainio.assemblies import ( BehavioralAssembly, @@ -275,6 +392,8 @@ def generate_predictions(self, stimuli): instruction=instruction, label_set=label_set, ) + if after_each is not None: + after_each() raw_answers.append(str(predicted)) valid.append(predicted in label_to_idx) if valid[-1]: @@ -311,6 +430,32 @@ def generate_predictions(self, stimuli): ) +def _trials_to_assembly(trials, stimuli, layer: str, region: str): + """Stack one-presentation recordings into a NeuroidAssembly for ``region``.""" + import numpy as np + import pandas as pd + from brainscore_core.supported_data_standards.brainio.assemblies import NeuroidAssembly + + rows = [] + for index, trial in enumerate(trials): + selected = trial.isel(neuroid=trial['layer'].values == layer) + if selected.sizes['neuroid'] == 0: + raise ValueError( + f'No activity reached layer {layer!r} while answering stimulus ' + f'{stimuli["stimulus_id"].iloc[index]!r}; that input may not use it.') + rows.append(np.asarray(selected.values).reshape(-1)) + first = trials[0].isel(neuroid=trials[0]['layer'].values == layer) + index = first.indexes.get('neuroid') + names = list(index.names) if isinstance(index, pd.MultiIndex) else [] + names += [name for name, coord in first.coords.items() + if coord.dims == ('neuroid',) and name != 'neuroid' and name not in names] + coords = {name: ('neuroid', np.asarray(first[name].values)) for name in names} + coords['region'] = ('neuroid', [region] * len(rows[0])) + for column in stimuli.columns: + coords[column] = ('presentation', list(stimuli[column].values)) + return NeuroidAssembly(np.stack(rows), coords=coords, dims=['presentation', 'neuroid']) + + # Keep existing pickle references and public type names stable. for _definition in ( BehavioralReadout, diff --git a/brainscore_core/_sessions/drivers.py b/brainscore_core/_sessions/drivers.py index f630340b..f9e25e1f 100644 --- a/brainscore_core/_sessions/drivers.py +++ b/brainscore_core/_sessions/drivers.py @@ -381,6 +381,44 @@ def drive_behavior_session( )) +def drive_behavior_and_neural_session(model, session, driver: str = "interact") -> None: + """Answer ``behavior`` and record ``neural:`` from one presentation. + + Each stimulus is shown once: the regions are recorded during the same + forward pass that produced the answer (see + ``BehavioralReadout.respond_and_record``). Only ``BrainScoreModel`` runs + this; it is reached through its session capability. + """ + task_context = getattr(session, "task_context", None) + if task_context is None: + raise ValueError("A behavior session needs session.task_context.") + if getattr(session, "streaming", False): + raise NotImplementedError( + "Recording during a task presents the stimulus set at once; " + "use streaming=False.") + time_bins = getattr(session, "time_bins", None) + if time_bins is not None and len(time_bins) > 1: + raise NotImplementedError( + "Recording during a task returns one time bin; request at most one.") + regions = [] + for channel in _requested_output_channels(session): + family, region = parse_channel(channel) + if family == "neural": + regions.append(region) + _refuse_variance(model, session) + model.start_task(task_context) + stimuli = _behavior_stimuli_from_session(session, task_context) + if stimuli is None: + raise ValueError("Recording during a task needs stimuli to present.") + behavior, neural = model._behavioral.respond_and_record(stimuli, regions) + session.emit(StreamEvent(channel="behavior", payload=behavior, t_ms=0.0, + meta={"driver": driver, "task_type": task_context.task_type})) + for region in regions: + session.emit(StreamEvent(channel=f"neural:{region}", payload=neural[region], t_ms=0.0, + meta={"driver": driver, "record": region, + "during": "behavior"})) + + _drive_neural_session_via_process = drive_neural_session diff --git a/brainscore_core/_sessions/helpers.py b/brainscore_core/_sessions/helpers.py index f333025c..c9f48e96 100644 --- a/brainscore_core/_sessions/helpers.py +++ b/brainscore_core/_sessions/helpers.py @@ -70,6 +70,7 @@ _drive_via_process, _refuse_variance, _start_recording, + drive_behavior_and_neural_session, drive_behavior_session, drive_neural_session, drive_typed_session, diff --git a/brainscore_core/capabilities/sessions.py b/brainscore_core/capabilities/sessions.py index cf4cbb18..8f3b7bf0 100644 --- a/brainscore_core/capabilities/sessions.py +++ b/brainscore_core/capabilities/sessions.py @@ -41,10 +41,12 @@ def supports_session(self, model, channels) -> bool: return bool(channels) and all(c.startswith('neural:') for c in channels) def participate(self, model, session) -> None: - if getattr(session, 'streaming', False): - _drivers._drive_neural_session_streaming(model, session) - else: - _drivers.drive_neural_session(model, session, allow_multi_modality=True) + # A neural-only session records activity even after a task was configured. + with model._behavioral.suspended(): + if getattr(session, 'streaming', False): + _drivers._drive_neural_session_streaming(model, session) + else: + _drivers.drive_neural_session(model, session, allow_multi_modality=True) class BehaviorSessions(_SessionCapability): @@ -61,10 +63,15 @@ def declared_output_channels(self, model): return {'behavior'} def supports_session(self, model, channels) -> bool: - return set(channels) == {'behavior'} + # Behavior alone, or behavior plus regions recorded during the same answers. + recorded = set(channels) - {'behavior'} + return 'behavior' in channels and recorded <= { + f'neural:{region}' for region in model.region_layer_map} def participate(self, model, session) -> None: - if getattr(session, 'streaming', False): + if set(_drivers._requested_output_channels(session)) != {'behavior'}: + _drivers.drive_behavior_and_neural_session(model, session) + elif getattr(session, 'streaming', False): _drivers._drive_behavior_session_streaming(model, session) else: _drivers.drive_behavior_session(model, session) diff --git a/tests/test_joint_behavior_neural.py b/tests/test_joint_behavior_neural.py new file mode 100644 index 00000000..4b546f54 --- /dev/null +++ b/tests/test_joint_behavior_neural.py @@ -0,0 +1,158 @@ +"""One presentation, two measurements: behavior and neural: together. + +A toy network counts its forward passes, so these tests check that every +stimulus is presented once, and that the recorded activity is the activity of +the pass that produced the answer. +""" +from contextlib import contextmanager + +import numpy as np +import pytest + +from brainscore_core.brainscore_model import BrainScoreModel +from brainscore_core.model_interface import TaskContext +from brainscore_core.selection import LayerSelector +from brainscore_core.streaming import BufferedSession +from brainscore_core.supported_data_standards.brainio.assemblies import NeuroidAssembly +from brainscore_core.supported_data_standards.brainio.stimuli import StimulusSet + +pytestmark = pytest.mark.unit + +LAYERS = {'early': 3, 'late': 2} # layer -> number of units + + +class ToyNet: + """Activity depends on the stimulus and on what else is in the prompt.""" + + def __init__(self): + self.passes, self.hooks = 0, [] + + def forward(self, stimulus_id, prompt=''): + self.passes += 1 + base = float(sum(map(ord, stimulus_id)) + len(prompt)) + out = {layer: base + np.arange(n, dtype=float) * (i + 1) + for i, (layer, n) in enumerate(LAYERS.items())} + for hook in self.hooks: + hook(out) + return out + + +def _assembly(rows, stimulus_ids, layers): + values = np.array([np.concatenate([row[layer] for layer in layers]) for row in rows]) + return NeuroidAssembly(values, dims=['presentation', 'neuroid'], coords={ + 'stimulus_id': ('presentation', list(stimulus_ids)), + 'neuroid_id': ('neuroid', [f'{l}.{u}' for l in layers for u in range(LAYERS[l])]), + 'layer': ('neuroid', [l for l in layers for _ in range(LAYERS[l])])}) + + +class ToyExtractor: + identifier = backbone_id = 'toy' + + def __init__(self, net): + self.net = net + + def __call__(self, stimuli, layers, **kwargs): + ids = list(stimuli['stimulus_id']) + return _assembly([self.net.forward(i) for i in ids], ids, layers) + + @contextmanager + def record(self, layers): + captured = [] + recording = type('Recording', (), {})() + + def take(): + first = captured[0] + captured.clear() + return _assembly([first], ['?'], layers) + recording.take = take + self.net.hooks.append(captured.append) + try: + yield recording + finally: + self.net.hooks.remove(captured.append) + + +class NoRecordExtractor(ToyExtractor): + record = property(lambda self: (_ for _ in ()).throw(AttributeError('record'))) + + +STIMULI = StimulusSet({'stimulus_id': ['a', 'b', 'c'], 'sentence': ['x', 'y', 'z'], + 'label': ['cat', 'dog', 'cat']}) +STIMULI.identifier = 'joint-toy' + + +def _model(net, extractor_type=ToyExtractor, **kwargs): + def answer(stimulus_row, instruction, label_set): + out = net.forward(stimulus_row['stimulus_id'], prompt=instruction) + return label_set[int(out['late'][0]) % 2] + return BrainScoreModel(identifier='toy', extractor={'text': extractor_type(net)}, + region_layer_map={'A': LayerSelector('early', modality='text'), + 'B': LayerSelector('late', modality='text')}, + answer=answer, **kwargs) + + +def _run(model, task, channels): + session = BufferedSession(stimulus_set=STIMULI, task_context=task, + requested_output_channels=channels) + model.participate(session) + return session + + +GENERATE = TaskContext(task_type='label', label_set=['cat', 'dog'], + instruction='Name the animal.', prefer_path='generation') + + +def test_answers_and_recording_come_from_one_presentation(): + net = ToyNet() + model = _model(net) + session = _run(model, GENERATE, ['behavior', 'neural:A', 'neural:B']) + assert net.passes == len(STIMULI) # each stimulus shown once + behavior = session.collect('behavior') + model.reset() + alone = _run(_model(ToyNet()), GENERATE, ['behavior']).collect('behavior') + np.testing.assert_array_equal(behavior.values, alone.values) + recorded = session.collect('neural:A') + assert list(recorded['stimulus_id'].values) == ['a', 'b', 'c'] + assert set(recorded['region'].values) == {'A'} + # The prompt is part of the pass: activity differs from passive viewing. + expected = [ToyNet().forward(s, prompt='Name the animal.')['early'] for s in 'abc'] + np.testing.assert_array_equal(recorded.values, expected) + + +def test_readout_records_in_the_same_extraction(): + net = ToyNet() + model = _model(net, behavioral_readout_layer='late') + task = TaskContext(task_type='probabilities', label_set=['cat', 'dog'], + fitting_stimuli=STIMULI, prefer_path='readout') + model.start_task(task) + fitted = net.passes + session = _run(model, task, ['behavior', 'neural:A']) + assert net.passes - fitted == len(STIMULI) # one extraction for both + model.reset() + reference = _model(ToyNet(), behavioral_readout_layer='late') + alone = _run(reference, task, ['behavior']).collect('behavior') + np.testing.assert_allclose(session.collect('behavior').values, alone.values) + recorded = session.collect('neural:A') + np.testing.assert_array_equal(recorded.values, + [ToyNet().forward(s)['early'] for s in 'abc']) + + +def test_extractor_that_cannot_record_fails_before_any_answer(): + net = ToyNet() + model = _model(net, extractor_type=NoRecordExtractor) + with pytest.raises(NotImplementedError, match='record'): + _run(model, GENERATE, ['behavior', 'neural:A']) + assert net.passes == 0 + + +def test_undeclared_region_is_not_supported(): + with pytest.raises(NotImplementedError): + _run(_model(ToyNet()), GENERATE, ['behavior', 'neural:IT']) + + +def test_neural_session_after_a_task_returns_activity(): + model = _model(ToyNet()) + model.start_task(GENERATE) + session = BufferedSession(stimulus_set=STIMULI, requested_output_channels=['neural:A']) + model.participate(session) + assert isinstance(session.collect('neural:A'), NeuroidAssembly) From 22053d6c0c89d6a4126ff74c8fdbe8521c70b6c5 Mon Sep 17 00:00:00 2001 From: Kartik Pradeepan Date: Wed, 4 Nov 2026 06:42:19 -0500 Subject: [PATCH 2/8] feat(sessions): let the task decide what a subject does and fit readouts on the intact network A session with a non-passive task_context that requests only neural channels now records during the task. TaskContext gains readout_fit: 'intact' fits the readout with every active change undone and tests the changed network with it. Recordings made during a task carry attrs['recording']; plain neural sessions are unchanged. --- brainscore_core/_model/behavioral.py | 19 +++++-- brainscore_core/_model/perturbation.py | 21 +++++++ brainscore_core/_sessions/drivers.py | 9 ++- brainscore_core/capabilities/sessions.py | 8 ++- brainscore_core/contract.py | 7 +++ tests/test_joint_behavior_neural.py | 72 +++++++++++++++++++++++- 6 files changed, 125 insertions(+), 11 deletions(-) diff --git a/brainscore_core/_model/behavioral.py b/brainscore_core/_model/behavioral.py index 0e0b4467..fd54122e 100644 --- a/brainscore_core/_model/behavioral.py +++ b/brainscore_core/_model/behavioral.py @@ -1,6 +1,6 @@ """Behavioral readout and generation support for BrainScoreModel.""" -from contextlib import ExitStack, contextmanager +from contextlib import ExitStack, contextmanager, nullcontext from typing import Any, Callable, Dict, Optional, Sequence, Tuple from ..contract import TaskContext @@ -92,6 +92,8 @@ def start_task(self, task_context_or_task, fitting_stimuli=None, self.task_context = task_context task_type = task_context.task_type + if task_context.readout_fit not in ('current', 'intact'): + raise ValueError(f"readout_fit must be 'current' or 'intact', not {task_context.readout_fit!r}.") if task_type != 'passive' and not owner._requires_behavioral_readout(task_type): raise ValueError( f"Unknown task_type {task_type!r}. BrainScoreModel accepts " @@ -196,7 +198,9 @@ def fit_behavioral_readout(self, fitting_stimuli) -> None: ) labels = self.extract_labels(fitting_stimuli) - features = owner._extract_behavioral_features(fitting_stimuli) + intact = (self.task_context is not None and self.task_context.readout_fit == 'intact') + with (owner._perturbations.intact() if intact else nullcontext()): + features = owner._extract_behavioral_features(fitting_stimuli) self.readout_classifier = ProbabilitiesClassifier() self.readout_classifier.fit(features, labels) @@ -331,9 +335,15 @@ def _generate_and_record(self, stimuli, layers: Dict[str, str]): for region, layer in layers.items(): modality = self._region_modality(region) neural[region] = _trials_to_assembly( - [trial[modality] for trial in trials], stimuli, layer, region) + [trial[modality] for trial in trials], stimuli, layer, region + ).assign_attrs(recording=self._stamp('answer', 'prompt')) return behavior, neural + def _stamp(self, path, forward_pass): + """How a recording was made, so task and passive recordings stay distinct.""" + return {'task_type': self.task_context.task_type, 'path': path, + 'forward_pass': forward_pass, 'readout_fit': self.task_context.readout_fit} + def _read_out_and_record(self, stimuli, layers: Dict[str, str]): owner = self.owner readout_layer = owner._behavioral_readout_layer @@ -357,7 +367,8 @@ def _read_out_and_record(self, stimuli, layers: Dict[str, str]): for region, layer in layers.items(): selected = extracted.isel(neuroid=layer_of == layer) neural[region] = selected.assign_coords( - region=('neuroid', [region] * selected.sizes['neuroid'])) + region=('neuroid', [region] * selected.sizes['neuroid']) + ).assign_attrs(recording=self._stamp('readout', 'readout extraction')) return behavior, neural def generate_predictions(self, stimuli, after_each: Optional[Callable[[], None]] = None): diff --git a/brainscore_core/_model/perturbation.py b/brainscore_core/_model/perturbation.py index 2faffcd0..a63840cf 100644 --- a/brainscore_core/_model/perturbation.py +++ b/brainscore_core/_model/perturbation.py @@ -1,5 +1,6 @@ """State-change and perturbation lifecycle support for BrainScoreModel.""" +from contextlib import contextmanager from dataclasses import replace from typing import Callable, Dict @@ -13,6 +14,23 @@ class PerturbationManager: def __init__(self, owner) -> None: self.owner = owner self.active_perturbations: Dict[str, Callable[[], None]] = {} + self._changes: Dict[str, object] = {} # handle -> resolved StateChange + + @contextmanager + def intact(self): + """Undo every active change for the block, then apply each again. + + Handles keep their identity, so a later reset by handle still works. + """ + handles = list(self.active_perturbations) + for handle in reversed(handles): + self.active_perturbations[handle]() + try: + yield + finally: + for handle in handles: + _, cleanup = self.owner._state_change_fn(self._changes[handle]) + self.active_perturbations[handle] = cleanup def reset(self) -> None: """Undo every active perturbation. @@ -32,6 +50,7 @@ def reset(self) -> None: failures.append((handle_id, error)) else: self.active_perturbations.pop(handle_id, None) + self._changes.pop(handle_id, None) if failures: import warnings detail = '; '.join(f'{handle!r}: {type(err).__name__}: {err}' @@ -79,6 +98,7 @@ def dispatch_state_change(self, state_change): # model perturbed with no way to retry. cleanup() self.active_perturbations.pop(handle_id, None) + self._changes.pop(handle_id, None) return None if owner._state_change_fn is None: @@ -107,6 +127,7 @@ def dispatch_state_change(self, state_change): f"state_change_fn must produce unique handle_ids per call." ) self.active_perturbations[applied.handle_id] = cleanup + self._changes[applied.handle_id] = state_change return applied diff --git a/brainscore_core/_sessions/drivers.py b/brainscore_core/_sessions/drivers.py index f9e25e1f..3a139ccc 100644 --- a/brainscore_core/_sessions/drivers.py +++ b/brainscore_core/_sessions/drivers.py @@ -382,7 +382,9 @@ def drive_behavior_session( def drive_behavior_and_neural_session(model, session, driver: str = "interact") -> None: - """Answer ``behavior`` and record ``neural:`` from one presentation. + """Record ``neural:`` while the subject does the session's task. + + ``behavior`` is sent too when requested. Each stimulus is shown once: the regions are recorded during the same forward pass that produced the answer (see @@ -411,8 +413,9 @@ def drive_behavior_and_neural_session(model, session, driver: str = "interact") if stimuli is None: raise ValueError("Recording during a task needs stimuli to present.") behavior, neural = model._behavioral.respond_and_record(stimuli, regions) - session.emit(StreamEvent(channel="behavior", payload=behavior, t_ms=0.0, - meta={"driver": driver, "task_type": task_context.task_type})) + if "behavior" in _requested_output_channels(session): + session.emit(StreamEvent(channel="behavior", payload=behavior, t_ms=0.0, + meta={"driver": driver, "task_type": task_context.task_type})) for region in regions: session.emit(StreamEvent(channel=f"neural:{region}", payload=neural[region], t_ms=0.0, meta={"driver": driver, "record": region, diff --git a/brainscore_core/capabilities/sessions.py b/brainscore_core/capabilities/sessions.py index 8f3b7bf0..39598a9c 100644 --- a/brainscore_core/capabilities/sessions.py +++ b/brainscore_core/capabilities/sessions.py @@ -41,7 +41,13 @@ def supports_session(self, model, channels) -> bool: return bool(channels) and all(c.startswith('neural:') for c in channels) def participate(self, model, session) -> None: - # A neural-only session records activity even after a task was configured. + task = getattr(session, 'task_context', None) + if task is not None and task.task_type != 'passive': + # The session's task decides what the subject does (rule 4.1): + # record during the task; the answers are not sent. + _drivers.drive_behavior_and_neural_session(model, session) + return + # A passive session records activity even after a task was configured. with model._behavioral.suspended(): if getattr(session, 'streaming', False): _drivers._drive_neural_session_streaming(model, session) diff --git a/brainscore_core/contract.py b/brainscore_core/contract.py index a8bd48a7..c9ce144b 100644 --- a/brainscore_core/contract.py +++ b/brainscore_core/contract.py @@ -38,6 +38,12 @@ class TaskContext: ``'probabilities'``, ``'classification'``, ``'label'``. Put the stimuli to answer (for example a held-out test set) in ``metadata['stimulus_set']`` when using ``behavioral_response``. + + ``readout_fit`` says which network a readout is fitted on while a change + (a lesion, a stimulation) is active: ``'current'`` (default) fits on the + changed network; ``'intact'`` fits on the network with every change undone, + then tests the changed network with that readout, as an animal trained + before a lesion is tested after it. """ task_type: str label_set: Optional[List[str]] = None @@ -45,6 +51,7 @@ class TaskContext: instruction: Optional[str] = None prefer_path: Optional[str] = None # 'auto' | 'generation' | 'readout' metadata: Dict[str, Any] = field(default_factory=dict) + readout_fit: str = 'current' # 'current' | 'intact' diff --git a/tests/test_joint_behavior_neural.py b/tests/test_joint_behavior_neural.py index 4b546f54..75e45ec4 100644 --- a/tests/test_joint_behavior_neural.py +++ b/tests/test_joint_behavior_neural.py @@ -10,7 +10,11 @@ import pytest from brainscore_core.brainscore_model import BrainScoreModel +from dataclasses import replace + +from brainscore_core.events import Perturbation, PerturbationApplied, Selection, StateChange from brainscore_core.model_interface import TaskContext +from brainscore_core.streaming_helpers import apply_state_change from brainscore_core.selection import LayerSelector from brainscore_core.streaming import BufferedSession from brainscore_core.supported_data_standards.brainio.assemblies import NeuroidAssembly @@ -25,11 +29,12 @@ class ToyNet: """Activity depends on the stimulus and on what else is in the prompt.""" def __init__(self): - self.passes, self.hooks = 0, [] + self.passes, self.hooks, self.silenced, self.log = 0, [], False, [] def forward(self, stimulus_id, prompt=''): self.passes += 1 - base = float(sum(map(ord, stimulus_id)) + len(prompt)) + self.log.append(self.silenced) + base = 0.0 if self.silenced else float(sum(map(ord, stimulus_id)) + len(prompt)) out = {layer: base + np.arange(n, dtype=float) * (i + 1) for i, (layer, n) in enumerate(LAYERS.items())} for hook in self.hooks: @@ -85,10 +90,18 @@ def _model(net, extractor_type=ToyExtractor, **kwargs): def answer(stimulus_row, instruction, label_set): out = net.forward(stimulus_row['stimulus_id'], prompt=instruction) return label_set[int(out['late'][0]) % 2] + + def intervention(change): + net.silenced = True + + def cleanup(): + net.silenced = False + return PerturbationApplied(handle_id=f'h{len(net.log)}', target=change.target, + perturbation=change.perturbation, applied_at=0), cleanup return BrainScoreModel(identifier='toy', extractor={'text': extractor_type(net)}, region_layer_map={'A': LayerSelector('early', modality='text'), 'B': LayerSelector('late', modality='text')}, - answer=answer, **kwargs) + answer=answer, intervention=intervention, **kwargs) def _run(model, task, channels): @@ -156,3 +169,56 @@ def test_neural_session_after_a_task_returns_activity(): session = BufferedSession(stimulus_set=STIMULI, requested_output_channels=['neural:A']) model.participate(session) assert isinstance(session.collect('neural:A'), NeuroidAssembly) + + +SILENCE = StateChange(kind='ablation', target=Selection(layer='early'), + perturbation=Perturbation(kind='zero')) + + +def test_neural_only_session_under_a_task_records_during_the_task(): + net = ToyNet() + session = _run(_model(net), GENERATE, ['neural:A']) + assert net.passes == len(STIMULI) + assert [e.channel for e in session.emitted] == ['neural:A'] # answers not sent + recorded = session.collect('neural:A') + expected = [ToyNet().forward(s, prompt='Name the animal.')['early'] for s in 'abc'] + np.testing.assert_array_equal(recorded.values, expected) + assert recorded.attrs['recording']['task_type'] == 'label' + assert recorded.attrs['recording']['forward_pass'] == 'prompt' + + +def test_passive_task_records_plain_viewing_unchanged(): + net = ToyNet() + passive = TaskContext(task_type='passive') + session = BufferedSession(stimulus_set=STIMULI, task_context=passive, + requested_output_channels=['neural:A']) + _model(net).participate(session) + recorded = session.collect('neural:A') + np.testing.assert_array_equal(recorded.values, [ToyNet().forward(s)['early'] for s in 'abc']) + assert 'recording' not in recorded.attrs # existing neural outputs are unchanged + + +def _readout_task(fit): + return TaskContext(task_type='probabilities', label_set=['cat', 'dog'], + fitting_stimuli=STIMULI, prefer_path='readout', readout_fit=fit) + + +@pytest.mark.parametrize('fit, fitted_silenced', [('intact', False), ('current', True)]) +def test_readout_fit_chooses_the_network_it_learns_on(fit, fitted_silenced): + net = ToyNet() + model = _model(net, behavioral_readout_layer='late') + applied = apply_state_change(model, SILENCE) + session = _run(model, _readout_task(fit), ['behavior', 'neural:A']) + fitting, scoring = net.log[:len(STIMULI)], net.log[len(STIMULI):] + assert fitting == [fitted_silenced] * len(STIMULI) + assert scoring == [True] * len(STIMULI) # tested on the changed network + np.testing.assert_array_equal(session.collect('neural:A').values, [np.arange(3.)] * 3) # silenced + handle = applied.handle_id + apply_state_change(model, StateChange(kind='reset', handle_id=handle)) # handle survives + assert not net.silenced + + +def test_readout_fit_is_validated(): + with pytest.raises(ValueError, match='readout_fit'): + _run(_model(ToyNet(), behavioral_readout_layer='late'), _readout_task('sometimes'), + ['behavior']) From 3799a622cbf65c74fc45c08945cb3ccf09d980a3 Mon Sep 17 00:00:00 2001 From: Kartik Pradeepan Date: Wed, 4 Nov 2026 06:42:41 -0500 Subject: [PATCH 3/8] ci: run unified tests from the single-presentation recording branch --- .github/workflows/umi-v2.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/umi-v2.yml b/.github/workflows/umi-v2.yml index 78eb9631..7248b236 100644 --- a/.github/workflows/umi-v2.yml +++ b/.github/workflows/umi-v2.yml @@ -11,9 +11,9 @@ jobs: coordinated: name: UMI v2 # Update the workflow and runner revision together. - uses: KartikP/brainscore-unified/.github/workflows/offline-tests.yml@4c76a4757fe6bbc1751638ce6df7e23b7ffdcf2b + uses: KartikP/brainscore-unified/.github/workflows/offline-tests.yml@af5309500c48b9cc27a6a468ebd2db93f7ec883f with: candidate: core candidate_repository: ${{ github.event.pull_request.head.repo.full_name || github.repository }} candidate_revision: ${{ github.event.pull_request.head.sha || github.sha }} - unified_revision: 4c76a4757fe6bbc1751638ce6df7e23b7ffdcf2b + unified_revision: af5309500c48b9cc27a6a468ebd2db93f7ec883f From 272283e97ddbf4532e44cc11eef153c93f015387 Mon Sep 17 00:00:00 2001 From: Kartik Pradeepan Date: Wed, 4 Nov 2026 07:07:51 -0500 Subject: [PATCH 4/8] feat(tasks): describe the answer's form with TaskContext.response TaskContext gains response: Choice (optionally with implied probabilities or reported confidence), Scale for ratings, and Text for written answers. Each response phrases its directive and parses replies; answers arrive on addressed behavior: channels that each part declares, so compatibility is decided before compute. BrainScoreModel gains reply and likelihood parts, and a ridge readout gives ratings for feature models. task_type values keep their meaning on plain behavior. --- brainscore_core/_model/behavior.py | 22 ++ brainscore_core/_model/behavioral.py | 265 ++++++++++++++++++++++- brainscore_core/_sessions/drivers.py | 21 +- brainscore_core/brainscore_model.py | 13 +- brainscore_core/capabilities/sessions.py | 28 ++- brainscore_core/contract.py | 25 ++- brainscore_core/io_catalog.py | 2 + brainscore_core/responses.py | 120 ++++++++++ tests/test_compatibility.py | 4 +- tests/test_subject_channels.py | 3 +- tests/test_task_responses.py | 143 ++++++++++++ 11 files changed, 623 insertions(+), 23 deletions(-) create mode 100644 brainscore_core/responses.py create mode 100644 tests/test_task_responses.py diff --git a/brainscore_core/_model/behavior.py b/brainscore_core/_model/behavior.py index 2522fbf7..aa16442a 100644 --- a/brainscore_core/_model/behavior.py +++ b/brainscore_core/_model/behavior.py @@ -24,6 +24,28 @@ ) +class RatingRegressor: + """Ridge regression readout for numeric answers (ratings). + + Scales features, fits ridge regression on numeric labels, predicts one + value per presentation. The rating counterpart of ProbabilitiesClassifier. + """ + + def __init__(self, alpha: float = 1.0): + import sklearn.linear_model + import sklearn.preprocessing + self._regressor = sklearn.linear_model.Ridge(alpha=alpha) + self._scaler = sklearn.preprocessing.StandardScaler() + + def fit(self, features, labels: Iterable) -> 'RatingRegressor': + X = self._scaler.fit_transform(np.asarray(features)) + self._regressor.fit(X, np.asarray(list(labels), dtype=float)) + return self + + def predict(self, features) -> np.ndarray: + return self._regressor.predict(self._scaler.transform(np.asarray(features))) + + class ProbabilitiesClassifier: """Logistic regression readout that produces label probabilities. diff --git a/brainscore_core/_model/behavioral.py b/brainscore_core/_model/behavioral.py index fd54122e..7ecff63a 100644 --- a/brainscore_core/_model/behavioral.py +++ b/brainscore_core/_model/behavioral.py @@ -5,6 +5,8 @@ from ..contract import TaskContext +import numpy as np + class BehavioralReadout: """Owns behavioral task state for a BrainScoreModel instance.""" @@ -23,6 +25,9 @@ def __init__(self, owner) -> None: self.configured_from: Optional[TaskContext] = None self._configured_as: Optional[set] = None # fingerprints of the fitted configuration self._stale = False + # Response tasks (TaskContext.response): how the answer and probabilities are produced. + self.response_path: Optional[str] = None + self.probability_source: Optional[str] = None @classmethod def requires_behavioral_readout(cls, task_type: str) -> bool: @@ -94,6 +99,9 @@ def start_task(self, task_context_or_task, fitting_stimuli=None, task_type = task_context.task_type if task_context.readout_fit not in ('current', 'intact'): raise ValueError(f"readout_fit must be 'current' or 'intact', not {task_context.readout_fit!r}.") + if task_context.response is not None: + self._start_response_task(task_context) + return if task_type != 'passive' and not owner._requires_behavioral_readout(task_type): raise ValueError( f"Unknown task_type {task_type!r}. BrainScoreModel accepts " @@ -173,6 +181,7 @@ def start_task(self, task_context_or_task, fitting_stimuli=None, ) def reset(self) -> None: + self.response_path = self.probability_source = None self.task_context = None self.readout_classifier = None self.use_generation_for_task = False @@ -202,8 +211,14 @@ def fit_behavioral_readout(self, fitting_stimuli) -> None: with (owner._perturbations.intact() if intact else nullcontext()): features = owner._extract_behavioral_features(fitting_stimuli) - self.readout_classifier = ProbabilitiesClassifier() - self.readout_classifier.fit(features, labels) + from ..responses import Scale + response = self.task_context.response if self.task_context is not None else None + if isinstance(response, Scale): + from brainscore_core._model.behavior import RatingRegressor + self.readout_classifier = RatingRegressor().fit(features, labels) + else: + self.readout_classifier = ProbabilitiesClassifier() + self.readout_classifier.fit(features, labels) @staticmethod def extract_labels(stimuli): @@ -252,6 +267,172 @@ def predict_probabilities(self, stimuli): features = self.owner._extract_behavioral_features(stimuli) return self.readout_classifier.predict_proba(features) + def _start_response_task(self, task_context) -> None: + """Choose how a Choice, Scale or Text answer is produced, before any compute.""" + from ..responses import Choice, Scale, Text + owner, response = self.owner, task_context.response + prefer = task_context.prefer_path or 'auto' + if prefer not in ('auto', 'generation', 'likelihood', 'readout'): + raise ValueError("prefer_path must be 'auto', 'generation', 'likelihood' or 'readout'; " + f'got {prefer!r}.') + readout_wired = (owner._behavioral_readout_layer is not None + and bool(owner._extractors or owner._layer_extractors) + and task_context.fitting_stimuli is not None) + reported_confidence = isinstance(response, Choice) and response.confidence is not None + can = { + 'generation': owner._reply_fn is not None or ( + owner._generation_fn is not None and isinstance(response, Choice) + and not reported_confidence), + 'likelihood': (owner._likelihood_fn is not None and isinstance(response, Choice) + and not reported_confidence), + 'readout': readout_wired and isinstance(response, (Choice, Scale)) and not reported_confidence, + } + order = ['generation', 'likelihood', 'readout'] if prefer == 'auto' else [prefer] + path = next((p for p in order if can[p]), None) + if path is None: + needs = {'generation': 'reply= (or answer= for a plain Choice)', + 'likelihood': 'likelihood= and a Choice without reported confidence', + 'readout': 'behavioral_readout_layer=, an extractor and fitting_stimuli'} + raise ValueError( + f"Model '{owner.identifier}' cannot give a {type(response).__name__} answer" + + (' with reported confidence (needs reply=)' if reported_confidence else '') + + '. Paths tried: ' + '; '.join(f'{p} needs {needs[p]}' for p in order)) + source = None + if isinstance(response, Choice) and response.probabilities: + if path != 'readout' and owner._likelihood_fn is not None: + source = 'likelihood' + elif readout_wired: + source = 'readout' + else: + raise ValueError(f"Model '{owner.identifier}' cannot give implied probabilities: " + 'it needs likelihood= or a readout (behavioral_readout_layer=, ' + 'fitting_stimuli).') + self.response_path, self.probability_source = path, source + self.use_generation_for_task = False # response tasks run through sessions only + self.readout_classifier = None + if 'readout' in (path, source): + owner._fit_behavioral_readout(task_context.fitting_stimuli) + + def respond(self, stimuli, channels, regions=()): + """Answer a response task; return ``({channel: assembly}, {region: NeuroidAssembly})``. + + Each stimulus is presented once on the answer's path. Regions are + recorded during that presentation: the prompt pass of a reply, or the + readout's own extraction. + """ + from ..responses import Choice, Scale, Text, labels_of, prompt_for + owner, task = self.owner, self.task_context + response, path = task.response, self.response_path + layers = {region: self._region_layer(region) for region in regions} + if layers and path == 'likelihood': + raise NotImplementedError('Recording during likelihood scoring is not supported; ' + "use prefer_path='generation' or 'readout'.") + if layers and path == 'generation' and self.probability_source is not None: + raise NotImplementedError('Implied probabilities need extra forward passes, so they ' + 'cannot be recorded in one presentation with a generated ' + "answer; use prefer_path='readout' or drop probabilities.") + labels = labels_of(response) + prompt = prompt_for(task.instruction, response) + neural, outputs = {}, {} + + if path == 'readout': + values, neural = self._read_out_and_record(stimuli, layers, raw=True) + if isinstance(response, Scale): + outputs['behavior:rating'] = _numbers(values, stimuli, 'readout') + else: + outputs['behavior:choice'] = _one_hot( + [str(values['choice'].values[i]) for i in np.asarray(values.values).argmax(1)], + [True] * len(stimuli), stimuli, labels, 'readout', raw=None) + if self.probability_source == 'readout': + outputs['behavior:probabilities'] = _implied(values, 'readout') + else: + if path == 'likelihood': + def call(row): + return owner._likelihood_fn(stimulus_row=row, prompt=prompt, + candidates=list(labels)) + elif owner._reply_fn is not None: + def call(row): + return owner._reply_fn(stimulus_row=row, prompt=prompt) + else: # a registration's answer: it phrases and parses a Choice itself + def call(row): + return owner._generation_fn(stimulus_row=row, instruction=task.instruction, + label_set=list(labels)) + replies, neural = self._per_stimulus(stimuli, layers, call) + if path == 'likelihood': + probabilities = _softmax_rows(replies, labels) + chosen = [labels[i] for i in probabilities.argmax(1)] + outputs['behavior:choice'] = _one_hot(chosen, [True] * len(stimuli), stimuli, + labels, 'likelihood', raw=None) + if self.probability_source == 'likelihood': + outputs['behavior:probabilities'] = _implied_array( + probabilities, stimuli, labels, 'likelihood') + elif isinstance(response, Text): + parsed = [response.parse(reply) for reply in replies] + outputs['behavior:text'] = _texts(parsed, replies, stimuli) + elif isinstance(response, Scale): + parsed = [response.parse(reply) for reply in replies] + outputs['behavior:rating'] = _numbers(parsed, stimuli, 'reply', raw=replies) + else: + parsed = [response.parse(str(reply)) for reply in replies] + source = 'reply' if owner._reply_fn is not None else 'answer' + outputs['behavior:choice'] = _one_hot( + [label for label, _ in parsed], [label is not None for label, _ in parsed], + stimuli, labels, source, raw=[str(r) for r in replies], + errors=[error for _, error in parsed]) + if response.confidence is not None: + outputs['behavior:confidence'] = _numbers( + [response.parse_confidence(reply) for reply in replies], stimuli, + 'reply', raw=replies) + if self.probability_source is not None: + extra = self._probabilities_after(stimuli, prompt, labels) + outputs['behavior:probabilities'] = extra + outputs['behavior'] = outputs[response.channels[0]] + stamp = self._stamp(path, 'prompt' if path == 'generation' else 'readout extraction') + return ({c: outputs[c] for c in channels}, + {r: a.assign_attrs(recording=stamp) for r, a in neural.items()}) + + def _probabilities_after(self, stimuli, prompt, labels): + """Implied probabilities alongside a generated answer (extra passes, not recorded).""" + owner = self.owner + if self.probability_source == 'likelihood': + scores = [owner._likelihood_fn(stimulus_row=row, prompt=prompt, candidates=list(labels)) + for _, row in stimuli.iterrows()] + return _implied_array(_softmax_rows(scores, labels), stimuli, labels, 'likelihood') + return _implied(self.readout_classifier.predict_proba( + self.owner._extract_behavioral_features(stimuli)), 'readout') + + def _per_stimulus(self, stimuli, layers, call): + """Call ``call(row)`` once per stimulus, recording ``layers`` on each first pass.""" + from ..io_catalog import canonical_modality + owner = self.owner + by_modality: Dict[str, list] = {} + for region, layer in layers.items(): + by_modality.setdefault(self._region_modality(region), []).append(layer) + recordings = {} + for modality, modality_layers in by_modality.items(): + extractor = next((e for key, e in owner._extractors.items() + if canonical_modality(key) == modality), None) + if extractor is None: + raise NotImplementedError(f'No {modality!r} extractor to record from.') + try: + record = extractor.record + except AttributeError: + raise NotImplementedError( + f'The {modality!r} extractor cannot record while the model answers: ' + 'it needs a record(layers) method. Use a readout task, or a ' + 'neural-only session.') from None + recordings[modality] = record(list(dict.fromkeys(modality_layers))) + replies, trials = [], [] + with ExitStack() as stack: + live = {m: stack.enter_context(r) for m, r in recordings.items()} + for _, row in stimuli.iterrows(): + replies.append(call(row)) + trials.append({m: recording.take() for m, recording in live.items()}) + neural = {region: _trials_to_assembly([t[self._region_modality(region)] for t in trials], + stimuli, layer, region) + for region, layer in layers.items()} + return replies, neural + @contextmanager def suspended(self): """Keep the configured task, but let it answer nothing for a while. @@ -342,9 +523,12 @@ def _generate_and_record(self, stimuli, layers: Dict[str, str]): def _stamp(self, path, forward_pass): """How a recording was made, so task and passive recordings stay distinct.""" return {'task_type': self.task_context.task_type, 'path': path, + 'response': type(self.task_context.response).__name__ + if self.task_context.response is not None else None, 'forward_pass': forward_pass, 'readout_fit': self.task_context.readout_fit} - def _read_out_and_record(self, stimuli, layers: Dict[str, str]): + def _read_out_and_record(self, stimuli, layers: Dict[str, str], raw: bool = False): + from ..responses import Scale owner = self.owner readout_layer = owner._behavioral_readout_layer modality = owner._pick_modality(owner._detect_modalities(stimuli)) @@ -362,13 +546,19 @@ def _read_out_and_record(self, stimuli, layers: Dict[str, str]): features = extracted.isel(neuroid=layer_of == readout_layer) if 'presentation' in features.dims and 'neuroid' in features.dims: features = features.transpose('presentation', 'neuroid') - behavior = self.readout_classifier.predict_proba(features) + if raw and isinstance(self.task_context.response, Scale): + behavior = [(float(v), None) for v in self.readout_classifier.predict(features)] + else: + behavior = self.readout_classifier.predict_proba(features) neural = {} for region, layer in layers.items(): selected = extracted.isel(neuroid=layer_of == layer) neural[region] = selected.assign_coords( region=('neuroid', [region] * selected.sizes['neuroid']) - ).assign_attrs(recording=self._stamp('readout', 'readout extraction')) + ) + if not raw: + neural[region] = neural[region].assign_attrs( + recording=self._stamp('readout', 'readout extraction')) return behavior, neural def generate_predictions(self, stimuli, after_each: Optional[Callable[[], None]] = None): @@ -441,6 +631,71 @@ def generate_predictions(self, stimuli, after_each: Optional[Callable[[], None]] ) +def _presentation(stimuli): + return {column: ('presentation', list(stimuli[column].values)) for column in stimuli.columns} + + +def _one_hot(chosen, valid, stimuli, labels, source, raw=None, errors=None): + """One row per stimulus, 1 for the chosen label; invalid rows are NaN (rule 6.2).""" + from brainscore_core.supported_data_standards.brainio.assemblies import BehavioralAssembly + values = np.zeros((len(stimuli), len(labels)), dtype=np.float32) + for i, (label, ok) in enumerate(zip(chosen, valid)): + if ok: + values[i, list(labels).index(label)] = 1.0 + else: + values[i] = np.nan + coords = {**_presentation(stimuli), 'choice': list(labels), + 'answer_valid': ('presentation', list(valid)), + 'raw_answer': ('presentation', [str(r) for r in (raw or chosen)])} + if errors is not None: + coords['parse_error'] = ('presentation', [e or '' for e in errors]) + return BehavioralAssembly(values, coords=coords, dims=['presentation', 'choice'], + attrs={'measure': {'kind': 'reported' if source in ('reply', 'answer') + else 'implied', 'source': source}}) + + +def _numbers(parsed, stimuli, source, raw=None): + """One number per stimulus (a rating or a confidence); invalid ones are NaN.""" + from brainscore_core.supported_data_standards.brainio.assemblies import BehavioralAssembly + values = np.array([np.nan if v is None else v for v, _ in parsed], dtype=float) + coords = {**_presentation(stimuli), + 'answer_valid': ('presentation', [v is not None for v, _ in parsed]), + 'parse_error': ('presentation', [e or '' for _, e in parsed])} + if raw is not None: + coords['raw_answer'] = ('presentation', [str(r) for r in raw]) + return BehavioralAssembly(values, coords=coords, dims=['presentation'], + attrs={'measure': {'kind': 'reported' if source == 'reply' + else 'implied', 'source': source}}) + + +def _texts(parsed, replies, stimuli): + from brainscore_core.supported_data_standards.brainio.assemblies import BehavioralAssembly + values = np.array([v if v is not None else '' for v, _ in parsed], dtype=object) + coords = {**_presentation(stimuli), + 'answer_valid': ('presentation', [v is not None for v, _ in parsed])} + return BehavioralAssembly(values, coords=coords, dims=['presentation'], + attrs={'measure': {'kind': 'reported', 'source': 'reply'}}) + + +def _softmax_rows(scores, labels): + logp = np.array([[float(row[label]) for label in labels] for row in scores], dtype=float) + logp -= logp.max(axis=1, keepdims=True) + p = np.exp(logp) + return p / p.sum(axis=1, keepdims=True) + + +def _implied_array(probabilities, stimuli, labels, source): + from brainscore_core.supported_data_standards.brainio.assemblies import BehavioralAssembly + return BehavioralAssembly(probabilities.astype(np.float32), + coords={**_presentation(stimuli), 'choice': list(labels)}, + dims=['presentation', 'choice'], + attrs={'measure': {'kind': 'implied', 'source': source}}) + + +def _implied(assembly, source): + return assembly.assign_attrs(measure={'kind': 'implied', 'source': source}) + + def _trials_to_assembly(trials, stimuli, layer: str, region: str): """Stack one-presentation recordings into a NeuroidAssembly for ``region``.""" import numpy as np diff --git a/brainscore_core/_sessions/drivers.py b/brainscore_core/_sessions/drivers.py index 3a139ccc..c2b1c141 100644 --- a/brainscore_core/_sessions/drivers.py +++ b/brainscore_core/_sessions/drivers.py @@ -402,19 +402,24 @@ def drive_behavior_and_neural_session(model, session, driver: str = "interact") if time_bins is not None and len(time_bins) > 1: raise NotImplementedError( "Recording during a task returns one time bin; request at most one.") - regions = [] - for channel in _requested_output_channels(session): - family, region = parse_channel(channel) - if family == "neural": - regions.append(region) + requested = _requested_output_channels(session) + regions = [parse_channel(c)[1] for c in requested if parse_channel(c)[0] == "neural"] + answers = [c for c in requested if parse_channel(c)[0] == "behavior"] _refuse_variance(model, session) model.start_task(task_context) stimuli = _behavior_stimuli_from_session(session, task_context) if stimuli is None: raise ValueError("Recording during a task needs stimuli to present.") - behavior, neural = model._behavioral.respond_and_record(stimuli, regions) - if "behavior" in _requested_output_channels(session): - session.emit(StreamEvent(channel="behavior", payload=behavior, t_ms=0.0, + if task_context.response is not None: + behavior, neural = model._behavioral.respond(stimuli, answers, regions) + else: + if set(answers) - {"behavior"}: + raise ValueError(f"{answers} need TaskContext(response=...); task_type " + f"{task_context.task_type!r} answers on 'behavior'.") + plain, neural = model._behavioral.respond_and_record(stimuli, regions) + behavior = {"behavior": plain} if answers else {} + for channel in answers: + session.emit(StreamEvent(channel=channel, payload=behavior[channel], t_ms=0.0, meta={"driver": driver, "task_type": task_context.task_type})) for region in regions: session.emit(StreamEvent(channel=f"neural:{region}", payload=neural[region], t_ms=0.0, diff --git a/brainscore_core/brainscore_model.py b/brainscore_core/brainscore_model.py index c378a0aa..5f715284 100644 --- a/brainscore_core/brainscore_model.py +++ b/brainscore_core/brainscore_model.py @@ -94,6 +94,8 @@ def __init__( answer: Optional[Callable] = None, policy: Optional[Callable] = None, intervention: Optional[Callable] = None, + reply: Optional[Callable] = None, + likelihood: Optional[Callable] = None, behavioral_readout_layer: Optional[str] = None, generation_fn: Optional[Callable] = None, action_fn: Optional[Callable] = None, @@ -130,6 +132,13 @@ def __init__( :param answer: ``(stimulus_row, instruction, label_set) -> label``; generates and parses a behavioral answer for instruction-following tasks. + :param reply: ``(stimulus_row, prompt) -> str``; the model's raw reply to + a prompt the task builds (instruction plus the response's directive). + Parsing is the task's. Gives choices, ratings, reported confidence + and written answers. + :param likelihood: ``(stimulus_row, prompt, candidates) -> {candidate: + log probability}``; scores candidate replies without generating. + Gives choices and their implied probabilities. :param policy: ``EnvironmentStep -> EnvironmentResponse``; chooses an action for closed-loop / embodied evaluation. :param intervention: ``StateChange -> (PerturbationApplied, cleanup)``; @@ -223,6 +232,7 @@ def __init__( self._behavioral_readout_layer = capability_values.get( 'behavioral_readout_layer') self._generation_fn = capability_values.get('generation_fn') + self._reply_fn, self._likelihood_fn = reply, likelihood self._action_fn = capability_values.get('action_fn') self._state_change_fn = capability_values.get('state_change_fn') @@ -613,7 +623,8 @@ def reset(self) -> None: failures.append(error) # Callable capabilities and extractors may own episode/history state. # Reset each provider once, including bound-method registrations. - providers = [self._action_fn, self._generation_fn, self._activations_model, + providers = [self._action_fn, self._generation_fn, self._reply_fn, self._likelihood_fn, + self._activations_model, *self._preprocessors.values(), *self._extractors.values()] seen = set() from .capabilities import capability_registry diff --git a/brainscore_core/capabilities/sessions.py b/brainscore_core/capabilities/sessions.py index 39598a9c..a0ae6265 100644 --- a/brainscore_core/capabilities/sessions.py +++ b/brainscore_core/capabilities/sessions.py @@ -60,22 +60,38 @@ class BehaviorSessions(_SessionCapability): def enabled_for(self, model) -> bool: return (model._behavioral_readout_layer is not None - or model._generation_fn is not None) + or model._generation_fn is not None + or model._reply_fn is not None or model._likelihood_fn is not None) def declared_input_channels(self, model): - return {'instruction'} if model._generation_fn is not None else set() + replies = (model._generation_fn, model._reply_fn, model._likelihood_fn) + return {'instruction'} if any(part is not None for part in replies) else set() def declared_output_channels(self, model): - return {'behavior'} + # Each part declares the kinds of answer it can give (rule 2.4). + channels = {'behavior'} + if model._behavioral_readout_layer is not None: + channels |= {'behavior:choice', 'behavior:probabilities', 'behavior:rating'} + if model._generation_fn is not None: + channels |= {'behavior:choice'} + if model._reply_fn is not None: + channels |= {'behavior:choice', 'behavior:rating', 'behavior:confidence', + 'behavior:text'} + if model._likelihood_fn is not None: + channels |= {'behavior:choice', 'behavior:probabilities'} + return channels def supports_session(self, model, channels) -> bool: # Behavior alone, or behavior plus regions recorded during the same answers. - recorded = set(channels) - {'behavior'} - return 'behavior' in channels and recorded <= { + behavior = {c for c in channels if c == 'behavior' or c.startswith('behavior:')} + recorded = set(channels) - behavior + return bool(behavior) and behavior <= self.declared_output_channels(model) and recorded <= { f'neural:{region}' for region in model.region_layer_map} def participate(self, model, session) -> None: - if set(_drivers._requested_output_channels(session)) != {'behavior'}: + requested = set(_drivers._requested_output_channels(session)) + task = getattr(session, 'task_context', None) + if requested != {'behavior'} or (task is not None and task.response is not None): _drivers.drive_behavior_and_neural_session(model, session) elif getattr(session, 'streaming', False): _drivers._drive_behavior_session_streaming(model, session) diff --git a/brainscore_core/contract.py b/brainscore_core/contract.py index c9ce144b..aaf8889c 100644 --- a/brainscore_core/contract.py +++ b/brainscore_core/contract.py @@ -39,19 +39,42 @@ class TaskContext: answer (for example a held-out test set) in ``metadata['stimulus_set']`` when using ``behavioral_response``. + ``response`` is the form of the answer (``brainscore_core.responses``): + ``Choice``, ``Scale`` or ``Text``; ``task_type='passive'`` means no answer. + With a response, outputs are requested on ``behavior:`` channels. + The older ``task_type`` values ``'label'``, ``'classification'`` and + ``'probabilities'`` keep their meaning on the plain ``behavior`` channel. + ``readout_fit`` says which network a readout is fitted on while a change (a lesion, a stimulation) is active: ``'current'`` (default) fits on the changed network; ``'intact'`` fits on the network with every change undone, then tests the changed network with that readout, as an animal trained before a lesion is tested after it. """ - task_type: str + task_type: Optional[str] = None label_set: Optional[List[str]] = None fitting_stimuli: Optional[Any] = None instruction: Optional[str] = None prefer_path: Optional[str] = None # 'auto' | 'generation' | 'readout' metadata: Dict[str, Any] = field(default_factory=dict) readout_fit: str = 'current' # 'current' | 'intact' + response: Any = None # a brainscore_core.responses Choice, Scale or Text + + def __post_init__(self): + # The response, when given, decides the task type; label lists follow it. + kinds = {'Choice': 'choice', 'Scale': 'rating', 'Text': 'text'} + if self.response is not None: + kind = kinds.get(type(self.response).__name__) + if kind is None: + raise TypeError('response must be a Choice, Scale or Text from brainscore_core.responses') + if self.task_type is None: + self.task_type = kind + elif self.task_type != kind: + raise ValueError(f'task_type {self.task_type!r} contradicts a {type(self.response).__name__} response') + if kind == 'choice' and self.label_set is None: + self.label_set = list(self.response.labels) + elif self.task_type is None: + raise ValueError("Give a response (Choice, Scale, Text) or task_type='passive'.") diff --git a/brainscore_core/io_catalog.py b/brainscore_core/io_catalog.py index d2158aed..ea7859e8 100644 --- a/brainscore_core/io_catalog.py +++ b/brainscore_core/io_catalog.py @@ -359,6 +359,8 @@ def validate( "a label, a probability vector, or generated text", "a fitted readout or the model's generation path", meta_keys=("label_set",), + addressing=("address is the kind of answer a TaskContext response asks for: " + "choice, probabilities, confidence, rating or text"), materializers=("start_task", "BehavioralAssembly")), # --- outputs: motor (carried by EnvironmentResponse, harness-defined) --- CatalogEntry("motor", OUTPUT, "EnvironmentResponse.action", diff --git a/brainscore_core/responses.py b/brainscore_core/responses.py new file mode 100644 index 00000000..f2e62c42 --- /dev/null +++ b/brainscore_core/responses.py @@ -0,0 +1,120 @@ +"""The form a behavioral answer takes. + +A ``TaskContext`` says what the subject is asked to do; its ``response`` says +what form the answer takes. A response knows how to phrase itself after the +instruction (``directive``), how to read a reply (``parse``), and which output +channels it produces. A benchmark may subclass a response to change either; +prompt wording and parsing are the benchmark's science, not the model's. + +Reported answers come from the model's reply. Implied measures come from the +model's internals by a stated link: a readout's or a language model's +probabilities over the labels (``behavior:probabilities``). Every result says +which it is. +""" +import re +from dataclasses import dataclass +from typing import Optional, Sequence, Tuple + + +@dataclass(frozen=True) +class Scale: + """A number on a scale, such as a 1 to 7 rating. ``step=None`` allows any value.""" + low: float + high: float + step: Optional[float] = None + + channels = ('behavior:rating',) + + def directive(self) -> str: + whole = self.step is not None and float(self.step).is_integer() + bounds = (f'{int(self.low)} to {int(self.high)}' if whole + else f'{self.low} to {self.high}') + return f'Answer with a number from {bounds}. Reply with only that number.' + + def parse(self, reply: str) -> Tuple[Optional[float], Optional[str]]: + found = re.search(r'-?\d+(?:\.\d+)?', reply or '') + if found is None: + return None, 'no number in the reply' + value = float(found.group()) + if not self.low <= value <= self.high: + return None, f'{value} is outside {self.low} to {self.high}' + if self.step is not None and abs((value - self.low) / self.step + - round((value - self.low) / self.step)) > 1e-9: + return None, f'{value} is not on the scale step {self.step}' + return value, None + + +@dataclass(frozen=True) +class Choice: + """One label from ``labels``. + + ``probabilities=True`` also returns the implied probability of each label + (from a readout or a language model's likelihoods). ``confidence=Scale(...)`` + also asks the subject to report how sure it is, in the same reply. + """ + labels: Tuple[str, ...] + probabilities: bool = False + confidence: Optional[Scale] = None + + def __post_init__(self): + if isinstance(self.labels, str) or len(self.labels) < 2: + raise ValueError('Choice needs at least two labels') + object.__setattr__(self, 'labels', tuple(self.labels)) + + @property + def channels(self) -> Tuple[str, ...]: + return (('behavior:choice',) + + (('behavior:probabilities',) if self.probabilities else ()) + + (('behavior:confidence',) if self.confidence is not None else ())) + + def directive(self) -> str: + text = f'Answer with exactly one of: {", ".join(self.labels)}.' + if self.confidence is not None: + text += (f' Then, on a new line, rate your confidence from ' + f'{self.confidence.low:g} to {self.confidence.high:g}.') + return text + + def parse(self, reply: str) -> Tuple[Optional[str], Optional[str]]: + first = (reply or '').strip().splitlines()[0].strip().lower() if (reply or '').strip() else '' + for label in self.labels: + if first == label.lower(): + return label, None + found = [label for label in self.labels + if re.search(rf'\b{re.escape(label.lower())}\b', first)] + if len(found) == 1: + return found[0], None + return None, ('no label in the reply' if not found + else f'several labels in the reply: {found}') + + def parse_confidence(self, reply: str) -> Tuple[Optional[float], Optional[str]]: + lines = (reply or '').strip().splitlines() + if len(lines) < 2: + return None, 'no confidence line in the reply' + return self.confidence.parse(' '.join(lines[1:])) + + +@dataclass(frozen=True) +class Text: + """A written answer, kept as written. Brain-Score never scores text itself: + the benchmark brings its own parser, judge or comparison.""" + channels = ('behavior:text',) + + def directive(self) -> str: + return '' + + def parse(self, reply: str) -> Tuple[Optional[str], Optional[str]]: + return (reply, None) if isinstance(reply, str) else (None, 'the reply is not text') + + +def response_channels(response) -> Tuple[str, ...]: + return () if response is None else tuple(response.channels) + + +def prompt_for(instruction: Optional[str], response) -> str: + """The text a replying model receives: the instruction, then the directive.""" + parts = [part for part in ((instruction or '').strip(), response.directive()) if part] + return '\n\n'.join(parts) + + +def labels_of(response) -> Optional[Sequence[str]]: + return response.labels if isinstance(response, Choice) else None diff --git a/tests/test_compatibility.py b/tests/test_compatibility.py index e308e5b4..e586bd77 100644 --- a/tests/test_compatibility.py +++ b/tests/test_compatibility.py @@ -394,7 +394,9 @@ def _model(self): def test_every_part_declares_registered_channels(self): model = self._model() - assert model.out_channels == {'neural:IT', 'behavior', 'motor', 'perturbation'} + assert model.out_channels == {'neural:IT', 'behavior', 'behavior:choice', + 'behavior:probabilities', 'behavior:rating', + 'motor', 'perturbation'} assert {'vision', 'observation', 'lesion', 'stimulation'} <= model.in_channels for output in sorted(model.out_channels): bench = FakeBenchmark(identifier=output) diff --git a/tests/test_subject_channels.py b/tests/test_subject_channels.py index 03869497..37053837 100644 --- a/tests/test_subject_channels.py +++ b/tests/test_subject_channels.py @@ -99,7 +99,8 @@ def test_brainscore_model_gets_channel_identity_from_preprocessors(): ) assert model.in_channels == {"vision"} - assert model.out_channels == {"neural:IT", "behavior"} + assert model.out_channels == {"neural:IT", "behavior", "behavior:choice", + "behavior:probabilities", "behavior:rating"} assert model.required_channels == set() diff --git a/tests/test_task_responses.py b/tests/test_task_responses.py new file mode 100644 index 00000000..99b880fe --- /dev/null +++ b/tests/test_task_responses.py @@ -0,0 +1,143 @@ +"""TaskContext(response=...): choices, confidence, ratings and written answers. + +Reported answers come from a reply and are parsed by the response; implied +measures come from a likelihood part or a readout. Each result says which. +""" +import numpy as np +import pytest + +from brainscore_core.brainscore_model import BrainScoreModel +from brainscore_core.model_interface import TaskContext +from brainscore_core.responses import Choice, Scale, Text +from brainscore_core.streaming import BufferedSession +from brainscore_core.supported_data_standards.brainio.assemblies import NeuroidAssembly +from brainscore_core.supported_data_standards.brainio.stimuli import StimulusSet + +pytestmark = pytest.mark.unit + +STIMULI = StimulusSet({'stimulus_id': ['a', 'b', 'c', 'd'], 'sentence': ['w', 'x', 'y', 'z'], + 'label': ['cat', 'dog', 'cat', 'dog']}) +STIMULI.identifier = 'responses-toy' + + +def _features(stimuli, layers, **kwargs): + ids = list(stimuli['stimulus_id']) + values = np.array([[1.0, 0.0] if i in 'ac' else [0.0, 1.0] for i in ids]) + return NeuroidAssembly(values, dims=['presentation', 'neuroid'], coords={ + 'stimulus_id': ('presentation', ids), 'neuroid_id': ('neuroid', ['f0', 'f1']), + 'layer': ('neuroid', [layers[0]] * 2)}) + + +def _model(replies=None, scores=None, readout=False): + prompts = [] + + def reply(stimulus_row, prompt): + prompts.append(prompt) + return replies[stimulus_row['stimulus_id']] + + def likelihood(stimulus_row, prompt, candidates): + return {c: scores[stimulus_row['stimulus_id']][c] for c in candidates} + + model = BrainScoreModel(identifier='toy', extractor={'text': _features}, + reply=reply if replies else None, + likelihood=likelihood if scores else None, + behavioral_readout_layer='late' if readout else None) + return model, prompts + + +def _ask(model, task, channels): + session = BufferedSession(stimulus_set=STIMULI, task_context=task, + requested_output_channels=channels) + model.participate(session) + return session + + +def test_choice_reply_is_parsed_and_invalid_replies_are_kept_as_invalid(): + model, prompts = _model(replies={'a': 'Cat', 'b': 'a dog.', 'c': 'cat or dog', 'd': 'bird'}) + task = TaskContext(instruction='Which animal?', response=Choice(['cat', 'dog'])) + choice = _ask(model, task, ['behavior:choice']).collect('behavior:choice') + np.testing.assert_array_equal(choice.values[:2], [[1, 0], [0, 1]]) + assert np.isnan(choice.values[2:]).all() + assert list(choice['answer_valid'].values) == [True, True, False, False] + assert list(choice['raw_answer'].values) == ['Cat', 'a dog.', 'cat or dog', 'bird'] + assert choice.attrs['measure'] == {'kind': 'reported', 'source': 'reply'} + assert prompts[0] == 'Which animal?\n\nAnswer with exactly one of: cat, dog.' + + +def test_reported_confidence_comes_from_the_same_reply(): + model, _ = _model(replies={'a': 'cat\n4', 'b': 'dog\n9', 'c': 'cat', 'd': 'dog\n2'}) + task = TaskContext(instruction='Which?', response=Choice(['cat', 'dog'], confidence=Scale(1, 5, 1))) + session = _ask(model, task, ['behavior:choice', 'behavior:confidence']) + confidence = session.collect('behavior:confidence') + np.testing.assert_array_equal(np.isnan(confidence.values), [False, True, True, False]) + assert confidence.values[0] == 4 and confidence.values[3] == 2 + assert 'outside' in confidence['parse_error'].values[1] + + +def test_rating_from_a_reply_and_from_a_readout(): + model, _ = _model(replies={'a': '6', 'b': 'about 2', 'c': '7.5', 'd': 'none'}) + rating = _ask(model, TaskContext(instruction='How cute?', response=Scale(1, 7, 1)), + ['behavior:rating']).collect('behavior:rating') + assert list(rating.values[:2]) == [6, 2] and np.isnan(rating.values[2:]).all() + + fitting = STIMULI.assign(label=[6.0, 2.0, 6.0, 2.0]) + fitting.identifier = 'responses-ratings' + readout, _ = _model(readout=True) + task = TaskContext(response=Scale(1, 7), fitting_stimuli=fitting) + predicted = _ask(readout, task, ['behavior:rating']).collect('behavior:rating') + np.testing.assert_allclose(predicted.values, [6, 2, 6, 2], atol=0.5) + assert predicted.attrs['measure'] == {'kind': 'implied', 'source': 'readout'} + + +def test_written_answers_are_kept_as_written(): + model, prompts = _model(replies={s: f'A story about {s}.' for s in 'abcd'}) + text = _ask(model, TaskContext(instruction='Describe it.', response=Text()), + ['behavior:text']).collect('behavior:text') + assert list(text.values) == [f'A story about {s}.' for s in 'abcd'] + assert prompts[0] == 'Describe it.' + + +def test_likelihood_gives_a_choice_and_its_implied_probabilities(): + scores = {s: {'cat': -1.0, 'dog': -3.0} if s in 'ac' else {'cat': -2.0, 'dog': -0.5} + for s in 'abcd'} + model, _ = _model(scores=scores) + task = TaskContext(instruction='Which?', response=Choice(['cat', 'dog'], probabilities=True)) + session = _ask(model, task, ['behavior:choice', 'behavior:probabilities']) + probabilities = session.collect('behavior:probabilities') + np.testing.assert_allclose(probabilities.values[0], [1 / (1 + np.exp(-2)), 1 - 1 / (1 + np.exp(-2))], + rtol=1e-6) + assert probabilities.attrs['measure'] == {'kind': 'implied', 'source': 'likelihood'} + np.testing.assert_array_equal(session.collect('behavior:choice').values.argmax(1), [0, 1, 0, 1]) + + +def test_readout_choice_with_probabilities(): + model, _ = _model(readout=True) + task = TaskContext(response=Choice(['cat', 'dog'], probabilities=True), fitting_stimuli=STIMULI) + session = _ask(model, task, ['behavior:choice', 'behavior:probabilities']) + np.testing.assert_array_equal(session.collect('behavior:choice').values.argmax(1), [0, 1, 0, 1]) + assert session.collect('behavior:probabilities').values.shape == (4, 2) + + +@pytest.mark.parametrize('response, channel', [(Text(), 'behavior:text'), + (Choice(['cat', 'dog'], confidence=Scale(1, 5)), + 'behavior:confidence')]) +def test_reported_answers_from_a_readout_only_model_are_not_supported(response, channel): + model, _ = _model(readout=True) + assert channel not in model.out_channels + with pytest.raises(NotImplementedError): + _ask(model, TaskContext(response=response, fitting_stimuli=STIMULI), [channel]) + + +def test_plain_behavior_request_returns_the_primary_answer(): + model, _ = _model(replies={s: 'dog' for s in 'abcd'}) + task = TaskContext(instruction='Which?', response=Choice(['cat', 'dog'])) + behavior = _ask(model, task, ['behavior']).collect('behavior') + np.testing.assert_array_equal(behavior.values.argmax(1), [1, 1, 1, 1]) + + +def test_task_type_and_response_must_agree(): + with pytest.raises(ValueError, match='contradicts'): + TaskContext(task_type='label', response=Scale(1, 7)) + with pytest.raises(ValueError, match='passive'): + TaskContext() + assert TaskContext(response=Choice(['a', 'b'])).label_set == ['a', 'b'] From 2631423c8dde66a80fc8056c6ad7d348c3af5f19 Mon Sep 17 00:00:00 2001 From: Kartik Pradeepan Date: Wed, 4 Nov 2026 07:57:40 -0500 Subject: [PATCH 5/8] ci: run unified tests from the task responses branch --- .github/workflows/umi-v2.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/umi-v2.yml b/.github/workflows/umi-v2.yml index 7248b236..2a72d192 100644 --- a/.github/workflows/umi-v2.yml +++ b/.github/workflows/umi-v2.yml @@ -11,9 +11,9 @@ jobs: coordinated: name: UMI v2 # Update the workflow and runner revision together. - uses: KartikP/brainscore-unified/.github/workflows/offline-tests.yml@af5309500c48b9cc27a6a468ebd2db93f7ec883f + uses: KartikP/brainscore-unified/.github/workflows/offline-tests.yml@5bf978697cd312116ac9aadbf297f8b0da61a364 with: candidate: core candidate_repository: ${{ github.event.pull_request.head.repo.full_name || github.repository }} candidate_revision: ${{ github.event.pull_request.head.sha || github.sha }} - unified_revision: af5309500c48b9cc27a6a468ebd2db93f7ec883f + unified_revision: 5bf978697cd312116ac9aadbf297f8b0da61a364 From e39fee1ed110b0cd2da50ac43ff30f903a1c12b1 Mon Sep 17 00:00:00 2001 From: Kartik Pradeepan Date: Wed, 4 Nov 2026 09:16:12 -0500 Subject: [PATCH 6/8] feat(channels): declare the visual forms a subject takes natively vision is now addressable by form: a vision-keyed extractor declares vision:still and a video-keyed one vision:clip. A subject with both takes each stimulus in its own form: clips go to the video extractor, stills to the vision one. --- brainscore_core/_model/dispatch.py | 18 +++++++++-- brainscore_core/brainscore_model.py | 7 ++++- brainscore_core/io_catalog.py | 3 ++ tests/test_channel_compatibility.py | 4 +-- tests/test_channel_registry.py | 4 +-- tests/test_subject_channels.py | 2 +- tests/test_visual_forms.py | 49 +++++++++++++++++++++++++++++ 7 files changed, 77 insertions(+), 10 deletions(-) create mode 100644 tests/test_visual_forms.py diff --git a/brainscore_core/_model/dispatch.py b/brainscore_core/_model/dispatch.py index 74ef64b0..d312077b 100644 --- a/brainscore_core/_model/dispatch.py +++ b/brainscore_core/_model/dispatch.py @@ -61,6 +61,14 @@ def register_input_handler(cls, event_type: type, handler_name: str, cls._INPUT_HANDLERS.append(entry) +CLIP_COLUMNS = ('video_path', 'video_file_name') + + +def _is_clip(stimuli) -> bool: + """Whether a stimulus table presents clips (video) rather than stills.""" + return any(column in getattr(stimuli, 'columns', ()) for column in CLIP_COLUMNS) + + def normalize_extractors(owner, preprocessors, activations_model): """Adapt registration forms once to ``extractor(stimuli, layers=...)``. @@ -249,9 +257,13 @@ def extract_for_modality(self, stimuli, modality: str, layers: List[str]): def _extract_for_modality_inner(self, stimuli, modality: str, layers: List[str]): owner = self.owner modality = canonical_modality(modality) - for key, extractor in owner._extractors.items(): - if canonical_modality(key) == modality: - return extractor(stimuli, layers=layers) + matching = [key for key in owner._extractors if canonical_modality(key) == modality] + if len(matching) > 1 and modality == 'vision': + # A still extractor and a clip extractor: the stimulus says which to use. + wanted = 'video' if _is_clip(stimuli) else 'vision' + matching = [key for key in matching if key == wanted] or matching + if matching: + return owner._extractors[matching[0]](stimuli, layers=layers) raise NotImplementedError( f"Model '{owner.identifier}' declares {modality!r} input but has no " 'activation extractor for it. Pass extractor={modality: extractor} to ' diff --git a/brainscore_core/brainscore_model.py b/brainscore_core/brainscore_model.py index 5f715284..cd8f6541 100644 --- a/brainscore_core/brainscore_model.py +++ b/brainscore_core/brainscore_model.py @@ -280,9 +280,14 @@ def _declared_capabilities(self): @property def in_channels(self): - return super().in_channels | set().union( + return super().in_channels | self._visual_forms() | set().union( *(c.declared_input_channels(self) for c in self._declared_capabilities())) + def _visual_forms(self): + """The visual forms taken natively: a ``vision`` key takes stills, a ``video`` key clips.""" + forms = {'vision': 'vision:still', 'video': 'vision:clip'} + return {forms[key] for key in self._preprocessors if key in forms} + @property def out_channels(self): # Every output comes from a registered part; nothing is found by name. diff --git a/brainscore_core/io_catalog.py b/brainscore_core/io_catalog.py index ea7859e8..e758f122 100644 --- a/brainscore_core/io_catalog.py +++ b/brainscore_core/io_catalog.py @@ -296,6 +296,9 @@ def validate( "(H, W, 3) uint8 image or (T, H, W, 3) uint8 clip, or a path to one", "image/video preprocessor + a vision model harness (wrapper)", meta_keys=("stimulus_id",), expected_ndim=(3, 4), expected_dtype="uint8", + addressing=("address is the form a subject takes natively: 'still' " + "(images) or 'clip' (video); a subject declares each form it " + "has, and a benchmark adapts a stimulus only when a form is missing"), materializers=("StimulusSet",)), CatalogEntry("text", INPUT, "StimulusSet column", "a Python str", diff --git a/tests/test_channel_compatibility.py b/tests/test_channel_compatibility.py index 89ab7090..540fe01f 100644 --- a/tests/test_channel_compatibility.py +++ b/tests/test_channel_compatibility.py @@ -223,11 +223,11 @@ def test_wrong_direction_channel_declaration_raises(): def test_invalid_addressed_channel_declaration_raises(): subject = _Subject( - in_channels={"vision"}, + in_channels={"text"}, out_channels={"neural:IT"}, ) benchmark = _Benchmark( - required_input_channels={"vision:left"}, + required_input_channels={"text:left"}, requested_output_channels={"neural:IT"}, ) diff --git a/tests/test_channel_registry.py b/tests/test_channel_registry.py index 337f5f58..ded8d199 100644 --- a/tests/test_channel_registry.py +++ b/tests/test_channel_registry.py @@ -92,8 +92,6 @@ def test_validate_reports_invalid_address_grammar(): def test_validate_reports_unaddressable_channel_family(): - image = np.zeros((224, 224, 3), dtype=np.uint8) - - failures = io_catalog.validate("vision:left", image, direction=INPUT) + failures = io_catalog.validate("text:left", "a sentence", direction=INPUT) assert any("not addressable" in failure for failure in failures) diff --git a/tests/test_subject_channels.py b/tests/test_subject_channels.py index 37053837..a7dd7ff0 100644 --- a/tests/test_subject_channels.py +++ b/tests/test_subject_channels.py @@ -98,7 +98,7 @@ def test_brainscore_model_gets_channel_identity_from_preprocessors(): behavioral_readout_layer="layer4", ) - assert model.in_channels == {"vision"} + assert model.in_channels == {"vision", "vision:still"} # a vision key takes stills assert model.out_channels == {"neural:IT", "behavior", "behavior:choice", "behavior:probabilities", "behavior:rating"} assert model.required_channels == set() diff --git a/tests/test_visual_forms.py b/tests/test_visual_forms.py new file mode 100644 index 00000000..120ac705 --- /dev/null +++ b/tests/test_visual_forms.py @@ -0,0 +1,49 @@ +"""A subject declares the visual forms it takes natively: stills, clips, or both.""" +import numpy as np +import pytest + +from brainscore_core.brainscore_model import BrainScoreModel +from brainscore_core.selection import LayerSelector +from brainscore_core.streaming_helpers import neural_response +from brainscore_core.supported_data_standards.brainio.assemblies import NeuroidAssembly +from brainscore_core.supported_data_standards.brainio.stimuli import StimulusSet + +pytestmark = pytest.mark.unit + + +def _extractor(value): + def extract(stimuli, layers, **kwargs): + ids = list(stimuli['stimulus_id']) + return NeuroidAssembly(np.full((len(ids), 2), value), dims=['presentation', 'neuroid'], + coords={'stimulus_id': ('presentation', ids), + 'neuroid_id': ('neuroid', ['n0', 'n1']), + 'layer': ('neuroid', [layers[0]] * 2)}) + return extract + + +def _stimuli(column): + stimuli = StimulusSet({'stimulus_id': ['s0', 's1'], column: ['a', 'b']}) + stimuli.identifier = f'forms-{column}' + return stimuli + + +@pytest.mark.parametrize('keys, forms', [(('vision',), {'vision:still'}), + (('video',), {'vision:clip'}), + (('vision', 'video'), {'vision:still', 'vision:clip'})]) +def test_forms_follow_the_extractor_keys(keys, forms): + model = BrainScoreModel('m', extractor={key: _extractor(1.0) for key in keys}) + assert 'vision' in model.in_channels + assert {c for c in model.in_channels if c.startswith('vision:')} == forms + + +def test_a_subject_with_both_forms_takes_each_stimulus_natively(): + model = BrainScoreModel('m', extractor={'vision': _extractor(1.0), 'video': _extractor(2.0)}, + region_layer_map={'IT': LayerSelector('blocks.1', modality='vision')}) + assert neural_response(model, _stimuli('image_path'), record='IT').values.max() == 1.0 + assert neural_response(model, _stimuli('video_path'), record='IT').values.max() == 2.0 + + +def test_models_without_an_extractor_declare_the_inputs_they_accept(): + model = BrainScoreModel('judge', supported_modalities={'vision'}, + answer=lambda stimulus_row, instruction, label_set: label_set[0]) + assert 'vision:still' in model.in_channels From 21f88d2f131c76535733f28cfd3fb5918dc59edf Mon Sep 17 00:00:00 2001 From: Kartik Pradeepan Date: Wed, 4 Nov 2026 09:33:53 -0500 Subject: [PATCH 7/8] feat(tasks): show each trial in order and keep clip time steps when recording reply and likelihood now receive a Prompt: the instruction, the stimulus parts in TaskContext.present order, then the response's directive. A stimulus with several parts and no declared order fails before any compute. Recording during a task picks the extractor that matches the stimulus's form and keeps a clip's time steps. --- brainscore_core/_model/behavioral.py | 64 +++++++++++++++++----------- brainscore_core/_model/dispatch.py | 21 ++++++--- brainscore_core/brainscore_model.py | 11 ++--- brainscore_core/contract.py | 10 +++++ brainscore_core/responses.py | 64 +++++++++++++++++++++++++++- tests/test_task_responses.py | 40 ++++++++++++++--- tests/test_visual_forms.py | 51 ++++++++++++++++++++++ 7 files changed, 217 insertions(+), 44 deletions(-) diff --git a/brainscore_core/_model/behavioral.py b/brainscore_core/_model/behavioral.py index 7ecff63a..ad5dde8e 100644 --- a/brainscore_core/_model/behavioral.py +++ b/brainscore_core/_model/behavioral.py @@ -320,7 +320,7 @@ def respond(self, stimuli, channels, regions=()): recorded during that presentation: the prompt pass of a reply, or the readout's own extraction. """ - from ..responses import Choice, Scale, Text, labels_of, prompt_for + from ..responses import Choice, Scale, Text, labels_of, trial_prompt owner, task = self.owner, self.task_context response, path = task.response, self.response_path layers = {region: self._region_layer(region) for region in regions} @@ -332,7 +332,8 @@ def respond(self, stimuli, channels, regions=()): 'cannot be recorded in one presentation with a generated ' "answer; use prefer_path='readout' or drop probabilities.") labels = labels_of(response) - prompt = prompt_for(task.instruction, response) + # Every trial's prompt is built first, so a bad layout fails before any compute. + prompts = [trial_prompt(task, row) for _, row in stimuli.iterrows()] neural, outputs = {}, {} if path == 'readout': @@ -347,17 +348,16 @@ def respond(self, stimuli, channels, regions=()): outputs['behavior:probabilities'] = _implied(values, 'readout') else: if path == 'likelihood': - def call(row): - return owner._likelihood_fn(stimulus_row=row, prompt=prompt, - candidates=list(labels)) + def call(row, prompt): + return owner._likelihood_fn(prompt, list(labels)) elif owner._reply_fn is not None: - def call(row): - return owner._reply_fn(stimulus_row=row, prompt=prompt) + def call(row, prompt): + return owner._reply_fn(prompt) else: # a registration's answer: it phrases and parses a Choice itself - def call(row): + def call(row, prompt): return owner._generation_fn(stimulus_row=row, instruction=task.instruction, label_set=list(labels)) - replies, neural = self._per_stimulus(stimuli, layers, call) + replies, neural = self._per_stimulus(stimuli, layers, call, prompts) if path == 'likelihood': probabilities = _softmax_rows(replies, labels) chosen = [labels[i] for i in probabilities.argmax(1)] @@ -384,25 +384,24 @@ def call(row): [response.parse_confidence(reply) for reply in replies], stimuli, 'reply', raw=replies) if self.probability_source is not None: - extra = self._probabilities_after(stimuli, prompt, labels) + extra = self._probabilities_after(stimuli, prompts, labels) outputs['behavior:probabilities'] = extra outputs['behavior'] = outputs[response.channels[0]] stamp = self._stamp(path, 'prompt' if path == 'generation' else 'readout extraction') return ({c: outputs[c] for c in channels}, {r: a.assign_attrs(recording=stamp) for r, a in neural.items()}) - def _probabilities_after(self, stimuli, prompt, labels): + def _probabilities_after(self, stimuli, prompts, labels): """Implied probabilities alongside a generated answer (extra passes, not recorded).""" owner = self.owner if self.probability_source == 'likelihood': - scores = [owner._likelihood_fn(stimulus_row=row, prompt=prompt, candidates=list(labels)) - for _, row in stimuli.iterrows()] + scores = [owner._likelihood_fn(prompt, list(labels)) for prompt in prompts] return _implied_array(_softmax_rows(scores, labels), stimuli, labels, 'likelihood') return _implied(self.readout_classifier.predict_proba( self.owner._extract_behavioral_features(stimuli)), 'readout') - def _per_stimulus(self, stimuli, layers, call): - """Call ``call(row)`` once per stimulus, recording ``layers`` on each first pass.""" + def _per_stimulus(self, stimuli, layers, call, prompts=None): + """Call ``call(row, prompt)`` once per stimulus, recording ``layers`` on each first pass.""" from ..io_catalog import canonical_modality owner = self.owner by_modality: Dict[str, list] = {} @@ -410,8 +409,8 @@ def _per_stimulus(self, stimuli, layers, call): by_modality.setdefault(self._region_modality(region), []).append(layer) recordings = {} for modality, modality_layers in by_modality.items(): - extractor = next((e for key, e in owner._extractors.items() - if canonical_modality(key) == modality), None) + from .dispatch import extractor_for + extractor = extractor_for(owner, modality, stimuli) if extractor is None: raise NotImplementedError(f'No {modality!r} extractor to record from.') try: @@ -425,8 +424,8 @@ def _per_stimulus(self, stimuli, layers, call): replies, trials = [], [] with ExitStack() as stack: live = {m: stack.enter_context(r) for m, r in recordings.items()} - for _, row in stimuli.iterrows(): - replies.append(call(row)) + for index, (_, row) in enumerate(stimuli.iterrows()): + replies.append(call(row, None if prompts is None else prompts[index])) trials.append({m: recording.take() for m, recording in live.items()}) neural = {region: _trials_to_assembly([t[self._region_modality(region)] for t in trials], stimuli, layer, region) @@ -494,8 +493,8 @@ def _generate_and_record(self, stimuli, layers: Dict[str, str]): by_modality.setdefault(self._region_modality(region), []).append(layer) recordings = {} for modality, modality_layers in by_modality.items(): - extractor = next((e for key, e in owner._extractors.items() - if canonical_modality(key) == modality), None) + from .dispatch import extractor_for + extractor = extractor_for(owner, modality, stimuli) if extractor is None: raise NotImplementedError(f'No {modality!r} extractor to record from.') try: @@ -703,23 +702,40 @@ def _trials_to_assembly(trials, stimuli, layer: str, region: str): from brainscore_core.supported_data_standards.brainio.assemblies import NeuroidAssembly rows = [] + timed = 'time_bin' in trials[0].dims for index, trial in enumerate(trials): selected = trial.isel(neuroid=trial['layer'].values == layer) if selected.sizes['neuroid'] == 0: raise ValueError( f'No activity reached layer {layer!r} while answering stimulus ' f'{stimuli["stimulus_id"].iloc[index]!r}; that input may not use it.') - rows.append(np.asarray(selected.values).reshape(-1)) + if timed: + # A clip keeps its time steps: (time_bin, neuroid) per presentation. + if selected.sizes['time_bin'] != trials[0].sizes['time_bin']: + raise ValueError('Clips of different lengths gave different numbers of time steps; ' + 'present clips of one length to record them together.') + rows.append(np.asarray(selected.transpose('presentation', 'time_bin', 'neuroid').values)[0]) + else: + rows.append(np.asarray(selected.values).reshape(-1)) first = trials[0].isel(neuroid=trials[0]['layer'].values == layer) index = first.indexes.get('neuroid') names = list(index.names) if isinstance(index, pd.MultiIndex) else [] names += [name for name, coord in first.coords.items() if coord.dims == ('neuroid',) and name != 'neuroid' and name not in names] coords = {name: ('neuroid', np.asarray(first[name].values)) for name in names} - coords['region'] = ('neuroid', [region] * len(rows[0])) + coords['region'] = ('neuroid', [region] * first.sizes['neuroid']) for column in stimuli.columns: coords[column] = ('presentation', list(stimuli[column].values)) - return NeuroidAssembly(np.stack(rows), coords=coords, dims=['presentation', 'neuroid']) + if not timed: + return NeuroidAssembly(np.stack(rows), coords=coords, dims=['presentation', 'neuroid']) + time_index = first.indexes.get('time_bin') + time_names = list(time_index.names) if isinstance(time_index, pd.MultiIndex) else [] + time_names += [name for name, coord in first.coords.items() + if coord.dims == ('time_bin',) and name != 'time_bin' and name not in time_names] + for name in time_names: + coords[name] = ('time_bin', np.asarray(first[name].values)) + return NeuroidAssembly(np.stack(rows), coords=coords, + dims=['presentation', 'time_bin', 'neuroid']) # Keep existing pickle references and public type names stable. diff --git a/brainscore_core/_model/dispatch.py b/brainscore_core/_model/dispatch.py index d312077b..8895c4b3 100644 --- a/brainscore_core/_model/dispatch.py +++ b/brainscore_core/_model/dispatch.py @@ -69,6 +69,17 @@ def _is_clip(stimuli) -> bool: return any(column in getattr(stimuli, 'columns', ()) for column in CLIP_COLUMNS) +def extractor_for(owner, modality: str, stimuli): + """The extractor for ``modality``; with a still and a clip extractor, the one + matching the stimulus's form. None when the model has no such extractor.""" + modality = canonical_modality(modality) + matching = [key for key in owner._extractors if canonical_modality(key) == modality] + if len(matching) > 1 and modality == 'vision': + wanted = 'video' if _is_clip(stimuli) else 'vision' + matching = [key for key in matching if key == wanted] or matching + return owner._extractors[matching[0]] if matching else None + + def normalize_extractors(owner, preprocessors, activations_model): """Adapt registration forms once to ``extractor(stimuli, layers=...)``. @@ -256,14 +267,10 @@ def extract_for_modality(self, stimuli, modality: str, layers: List[str]): def _extract_for_modality_inner(self, stimuli, modality: str, layers: List[str]): owner = self.owner + extractor = extractor_for(owner, modality, stimuli) + if extractor is not None: + return extractor(stimuli, layers=layers) modality = canonical_modality(modality) - matching = [key for key in owner._extractors if canonical_modality(key) == modality] - if len(matching) > 1 and modality == 'vision': - # A still extractor and a clip extractor: the stimulus says which to use. - wanted = 'video' if _is_clip(stimuli) else 'vision' - matching = [key for key in matching if key == wanted] or matching - if matching: - return owner._extractors[matching[0]](stimuli, layers=layers) raise NotImplementedError( f"Model '{owner.identifier}' declares {modality!r} input but has no " 'activation extractor for it. Pass extractor={modality: extractor} to ' diff --git a/brainscore_core/brainscore_model.py b/brainscore_core/brainscore_model.py index cd8f6541..a41646c2 100644 --- a/brainscore_core/brainscore_model.py +++ b/brainscore_core/brainscore_model.py @@ -132,13 +132,14 @@ def __init__( :param answer: ``(stimulus_row, instruction, label_set) -> label``; generates and parses a behavioral answer for instruction-following tasks. - :param reply: ``(stimulus_row, prompt) -> str``; the model's raw reply to - a prompt the task builds (instruction plus the response's directive). + :param reply: ``(prompt) -> str``; the model's raw reply to one trial. + ``prompt`` is a ``brainscore_core.responses.Prompt``: the trial's parts + in order (instruction, stimulus parts, the response's directive). Parsing is the task's. Gives choices, ratings, reported confidence and written answers. - :param likelihood: ``(stimulus_row, prompt, candidates) -> {candidate: - log probability}``; scores candidate replies without generating. - Gives choices and their implied probabilities. + :param likelihood: ``(prompt, candidates) -> {candidate: log + probability}``; scores candidate replies without generating. Gives + choices and their implied probabilities. :param policy: ``EnvironmentStep -> EnvironmentResponse``; chooses an action for closed-loop / embodied evaluation. :param intervention: ``StateChange -> (PerturbationApplied, cleanup)``; diff --git a/brainscore_core/contract.py b/brainscore_core/contract.py index aaf8889c..f4ae3c9d 100644 --- a/brainscore_core/contract.py +++ b/brainscore_core/contract.py @@ -45,6 +45,10 @@ class TaskContext: The older ``task_type`` values ``'label'``, ``'classification'`` and ``'probabilities'`` keep their meaning on the plain ``behavior`` channel. + ``present`` gives the order of a trial's parts when a stimulus has several + (e.g. ``('text', 'video')``: a sentence, then a clip). A replying model is + shown the instruction, those parts in order, then the response's directive. + ``readout_fit`` says which network a readout is fitted on while a change (a lesion, a stimulation) is active: ``'current'`` (default) fits on the changed network; ``'intact'`` fits on the network with every change undone, @@ -59,6 +63,7 @@ class TaskContext: metadata: Dict[str, Any] = field(default_factory=dict) readout_fit: str = 'current' # 'current' | 'intact' response: Any = None # a brainscore_core.responses Choice, Scale or Text + present: Optional[Tuple[str, ...]] = None # order of a trial's parts: 'text', 'image', 'video', 'audio' def __post_init__(self): # The response, when given, decides the task type; label lists follow it. @@ -75,6 +80,11 @@ def __post_init__(self): self.label_set = list(self.response.labels) elif self.task_type is None: raise ValueError("Give a response (Choice, Scale, Text) or task_type='passive'.") + if self.present is not None: + unknown = set(self.present) - {'text', 'image', 'video', 'audio'} + if unknown or isinstance(self.present, str): + raise ValueError(f"present lists parts from 'text', 'image', 'video', 'audio'; got {self.present!r}") + self.present = tuple(self.present) diff --git a/brainscore_core/responses.py b/brainscore_core/responses.py index f2e62c42..427948a1 100644 --- a/brainscore_core/responses.py +++ b/brainscore_core/responses.py @@ -13,7 +13,7 @@ """ import re from dataclasses import dataclass -from typing import Optional, Sequence, Tuple +from typing import Any, Optional, Sequence, Tuple @dataclass(frozen=True) @@ -106,6 +106,68 @@ def parse(self, reply: str) -> Tuple[Optional[str], Optional[str]]: return (reply, None) if isinstance(reply, str) else (None, 'the reply is not text') +# The stimulus columns that carry each kind of trial part. +PART_COLUMNS = { + 'text': ('sentence', 'text'), + 'image': ('image_file_name', 'image_path', 'filename'), + 'video': ('video_path', 'video_file_name'), + 'audio': ('audio_path', 'audio_file_name', 'audio_file'), +} + + +@dataclass(frozen=True) +class Prompt: + """What one trial shows a replying model, in order. + + ``parts`` are ``(kind, content)`` pairs: ``('text', str)``, or ``('image' | + 'video' | 'audio', path)``. The order is the trial's: the instruction, the + stimulus parts in ``TaskContext.present`` order, then the response's + directive. A ``reply`` or ``likelihood`` part shows them in this order. + """ + parts: Tuple[Tuple[str, Any], ...] + stimulus_id: Optional[str] = None + + def text(self) -> str: + """The text parts joined, for a model that reads text only.""" + return '\n\n'.join(content for kind, content in self.parts if kind == 'text') + + +def trial_prompt(task, stimulus_row) -> Prompt: + """Build one trial's Prompt from the task and a stimulus row (fails before any compute).""" + present = {} + for kind, columns in PART_COLUMNS.items(): + for column in columns: + value = _cell(stimulus_row, column) + if value is not None: + present[kind] = value + break + order = tuple(task.present) if task.present is not None else tuple(present) + if task.present is None and len(present) > 1: + raise ValueError(f'Each stimulus has several parts {sorted(present)}; say the order people ' + 'saw them in with TaskContext(present=(...)), e.g. present=("text", "video").') + missing = [kind for kind in order if kind not in present] + if missing: + raise ValueError(f'TaskContext.present names {missing}, but the stimulus has only {sorted(present)}.') + parts = [] + if (task.instruction or '').strip(): + parts.append(('text', task.instruction.strip())) + parts += [(kind, present[kind]) for kind in order] + directive = task.response.directive() if task.response is not None else '' + if directive: + parts.append(('text', directive)) + return Prompt(tuple(parts), stimulus_id=_cell(stimulus_row, 'stimulus_id')) + + +def _cell(row, column): + try: + value = row[column] + except (KeyError, TypeError, IndexError): + return None + if value is None or (isinstance(value, float) and value != value): + return None + return value + + def response_channels(response) -> Tuple[str, ...]: return () if response is None else tuple(response.channels) diff --git a/tests/test_task_responses.py b/tests/test_task_responses.py index 99b880fe..aef0639d 100644 --- a/tests/test_task_responses.py +++ b/tests/test_task_responses.py @@ -31,12 +31,12 @@ def _features(stimuli, layers, **kwargs): def _model(replies=None, scores=None, readout=False): prompts = [] - def reply(stimulus_row, prompt): - prompts.append(prompt) - return replies[stimulus_row['stimulus_id']] + def reply(prompt): + prompts.append(prompt.text()) + return replies[prompt.stimulus_id] - def likelihood(stimulus_row, prompt, candidates): - return {c: scores[stimulus_row['stimulus_id']][c] for c in candidates} + def likelihood(prompt, candidates): + return {c: scores[prompt.stimulus_id][c] for c in candidates} model = BrainScoreModel(identifier='toy', extractor={'text': _features}, reply=reply if replies else None, @@ -61,7 +61,7 @@ def test_choice_reply_is_parsed_and_invalid_replies_are_kept_as_invalid(): assert list(choice['answer_valid'].values) == [True, True, False, False] assert list(choice['raw_answer'].values) == ['Cat', 'a dog.', 'cat or dog', 'bird'] assert choice.attrs['measure'] == {'kind': 'reported', 'source': 'reply'} - assert prompts[0] == 'Which animal?\n\nAnswer with exactly one of: cat, dog.' + assert prompts[0] == 'Which animal?\n\nw\n\nAnswer with exactly one of: cat, dog.' def test_reported_confidence_comes_from_the_same_reply(): @@ -94,7 +94,7 @@ def test_written_answers_are_kept_as_written(): text = _ask(model, TaskContext(instruction='Describe it.', response=Text()), ['behavior:text']).collect('behavior:text') assert list(text.values) == [f'A story about {s}.' for s in 'abcd'] - assert prompts[0] == 'Describe it.' + assert prompts[0] == 'Describe it.\n\nw' def test_likelihood_gives_a_choice_and_its_implied_probabilities(): @@ -141,3 +141,29 @@ def test_task_type_and_response_must_agree(): with pytest.raises(ValueError, match='passive'): TaskContext() assert TaskContext(response=Choice(['a', 'b'])).label_set == ['a', 'b'] + + +def test_a_trial_is_shown_in_the_declared_order(): + from brainscore_core.responses import trial_prompt + row = {'stimulus_id': 's', 'sentence': 'Watch the ball.', 'video_path': 'clip.mp4'} + task = TaskContext(instruction='Which way?', response=Choice(['left', 'right']), + present=('text', 'video')) + assert trial_prompt(task, row).parts == ( + ('text', 'Which way?'), ('text', 'Watch the ball.'), ('video', 'clip.mp4'), + ('text', 'Answer with exactly one of: left, right.')) + with pytest.raises(ValueError, match='present'): + trial_prompt(TaskContext(instruction='Which way?', response=Choice(['left', 'right'])), row) + + +def test_an_unordered_multi_part_stimulus_fails_before_any_reply(): + calls = [] + model = BrainScoreModel(identifier='m', supported_modalities={'text', 'vision'}, + reply=lambda prompt: calls.append(prompt) or 'left') + both = StimulusSet({'stimulus_id': ['a'], 'sentence': ['x'], 'video_path': ['v.mp4']}) + both.identifier = 'both-parts' + session = BufferedSession(stimulus_set=both, task_context=TaskContext( + instruction='Which way?', response=Choice(['left', 'right'])), + requested_output_channels=['behavior:choice']) + with pytest.raises(ValueError, match='present'): + model.participate(session) + assert calls == [] diff --git a/tests/test_visual_forms.py b/tests/test_visual_forms.py index 120ac705..5d15b6b5 100644 --- a/tests/test_visual_forms.py +++ b/tests/test_visual_forms.py @@ -47,3 +47,54 @@ def test_models_without_an_extractor_declare_the_inputs_they_accept(): model = BrainScoreModel('judge', supported_modalities={'vision'}, answer=lambda stimulus_row, instruction, label_set: label_set[0]) assert 'vision:still' in model.in_channels + + +class _Recorder: + """record(layers) stand-in; a clip recording has two time steps.""" + + def __init__(self, value, timed): + self.value, self.timed, self.hooks = value, timed, [] + + def __call__(self, stimuli, layers, **kwargs): + raise AssertionError('a task presents through reply, not the extractor') + + def record(self, layers): + from contextlib import contextmanager + + @contextmanager + def recording(): + owner = self + + class Take: + @staticmethod + def take(): + coords = {'stimulus_id': ('presentation', ['?']), + 'neuroid_id': ('neuroid', ['n0', 'n1']), + 'layer': ('neuroid', [layers[0]] * 2)} + if not owner.timed: + return NeuroidAssembly(np.full((1, 2), owner.value), coords=coords, + dims=['presentation', 'neuroid']) + coords.update(time_bin_start_ms=('time_bin', [0.0, 1000.0]), + time_bin_end_ms=('time_bin', [1000.0, 2000.0])) + values = np.stack([np.full(2, owner.value), np.full(2, owner.value + 1)])[None] + return NeuroidAssembly(values, coords=coords, + dims=['presentation', 'time_bin', 'neuroid']) + yield Take() + return recording() + + +def test_a_clip_recorded_during_a_task_uses_the_clip_extractor_and_keeps_time(): + from brainscore_core.model_interface import TaskContext + from brainscore_core.responses import Choice + from brainscore_core.streaming import BufferedSession + model = BrainScoreModel('m', extractor={'vision': _Recorder(1.0, False), 'video': _Recorder(5.0, True)}, + region_layer_map={'IT': LayerSelector('blocks.1', modality='vision')}, + reply=lambda prompt: 'left') + task = TaskContext(instruction='Which way?', response=Choice(['left', 'right'])) + session = BufferedSession(stimulus_set=_stimuli('video_path'), task_context=task, + requested_output_channels=['behavior:choice', 'neural:IT']) + model.participate(session) + recorded = session.collect('neural:IT') + assert recorded.dims == ('presentation', 'time_bin', 'neuroid') + np.testing.assert_array_equal(recorded.values[0, :, 0], [5.0, 6.0]) + assert list(recorded['time_bin_start_ms'].values) == [0.0, 1000.0] From a08c209a60e729f7919afc0064a26b85d44f264b Mon Sep 17 00:00:00 2001 From: Kartik Pradeepan Date: Wed, 4 Nov 2026 09:48:33 -0500 Subject: [PATCH 8/8] ci: run unified tests from the folded experiment API branch --- .github/workflows/umi-v2.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/umi-v2.yml b/.github/workflows/umi-v2.yml index 2a72d192..438329b9 100644 --- a/.github/workflows/umi-v2.yml +++ b/.github/workflows/umi-v2.yml @@ -11,9 +11,9 @@ jobs: coordinated: name: UMI v2 # Update the workflow and runner revision together. - uses: KartikP/brainscore-unified/.github/workflows/offline-tests.yml@5bf978697cd312116ac9aadbf297f8b0da61a364 + uses: KartikP/brainscore-unified/.github/workflows/offline-tests.yml@fd3bec0605731b9c72900dfd7e85137a3aa56312 with: candidate: core candidate_repository: ${{ github.event.pull_request.head.repo.full_name || github.repository }} candidate_revision: ${{ github.event.pull_request.head.sha || github.sha }} - unified_revision: 5bf978697cd312116ac9aadbf297f8b0da61a364 + unified_revision: fd3bec0605731b9c72900dfd7e85137a3aa56312