From 272283e97ddbf4532e44cc11eef153c93f015387 Mon Sep 17 00:00:00 2001 From: Kartik Pradeepan Date: Wed, 4 Nov 2026 07:07:51 -0500 Subject: [PATCH 1/2] 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 2/2] 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