diff --git a/.github/workflows/umi-v2.yml b/.github/workflows/umi-v2.yml index 438329b9..6bdfe354 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@fd3bec0605731b9c72900dfd7e85137a3aa56312 + uses: KartikP/brainscore-unified/.github/workflows/offline-tests.yml@346b90ab5093d520d4c49c10752acce0afa6fb39 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: fd3bec0605731b9c72900dfd7e85137a3aa56312 + unified_revision: 346b90ab5093d520d4c49c10752acce0afa6fb39 diff --git a/brainscore_core/_evaluation/compatibility.py b/brainscore_core/_evaluation/compatibility.py index 222c52eb..6ce24ffb 100644 --- a/brainscore_core/_evaluation/compatibility.py +++ b/brainscore_core/_evaluation/compatibility.py @@ -160,6 +160,14 @@ def check_channel_compatibility(subject: Subject, benchmark) -> CompatibilityMat f"'{benchmark_id}' sends {_format_channels(bench_inputs)}." ) + required_traits = set(getattr(benchmark, 'required_traits', None) or ()) + missing_traits = required_traits - set(subject.traits) if required_traits else set() + if missing_traits: + raise CompatibilityError( + f"Benchmark '{benchmark_id}' requires subject traits {sorted(missing_traits)}; " + f"subject '{subject_id}' declares {sorted(subject.traits) or 'none'}." + ) + if getattr(benchmark, "uses_session", False): probe = getattr(subject, "check_session_support", None) if callable(probe): diff --git a/brainscore_core/_model/behavioral.py b/brainscore_core/_model/behavioral.py index ad5dde8e..7ad05e1a 100644 --- a/brainscore_core/_model/behavioral.py +++ b/brainscore_core/_model/behavioral.py @@ -1,7 +1,8 @@ """Behavioral readout and generation support for BrainScoreModel.""" from contextlib import ExitStack, contextmanager, nullcontext -from typing import Any, Callable, Dict, Optional, Sequence, Tuple +from dataclasses import replace +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple from ..contract import TaskContext @@ -25,6 +26,8 @@ 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 + # With memory='session': the session's earlier trials as (prompt, reply). + self.history: list = [] # Response tasks (TaskContext.response): how the answer and probabilities are produced. self.response_path: Optional[str] = None self.probability_source: Optional[str] = None @@ -181,6 +184,7 @@ def start_task(self, task_context_or_task, fitting_stimuli=None, ) def reset(self) -> None: + self.history = [] self.response_path = self.probability_source = None self.task_context = None self.readout_classifier = None @@ -298,7 +302,7 @@ def _start_response_task(self, task_context) -> None: + (' 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 isinstance(response, Choice) and (response.probabilities or response.reaction_time): if path != 'readout' and owner._likelihood_fn is not None: source = 'likelihood' elif readout_wired: @@ -316,21 +320,16 @@ def _start_response_task(self, task_context) -> None: 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. + Each stimulus is presented once on the answer's path, and regions are + recorded during that presentation: the prompt pass of a reply or of + likelihood scoring, or the readout's own extraction. Implied + probabilities next to a reply are scored first, from the same prompt, + so the recording is still the prompt pass. """ 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} - 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.") + layers = {region: self._region_layers(region) for region in regions} labels = labels_of(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()] @@ -347,17 +346,39 @@ def respond(self, stimuli, channels, regions=()): if self.probability_source == 'readout': outputs['behavior:probabilities'] = _implied(values, 'readout') else: + score_first = path == 'generation' and self.probability_source == 'likelihood' + remember = owner._memory == 'session' + if remember and owner._reply_fn is None and owner._likelihood_fn is None: + raise ValueError("memory='session' needs reply= or likelihood=; answer= cannot see " + 'earlier trials.') + + def with_history(prompt): + return replace(prompt, history=tuple(self.history)) if remember else prompt + if path == 'likelihood': def call(row, prompt): - return owner._likelihood_fn(prompt, list(labels)) + scores = owner._likelihood_fn(with_history(prompt), list(labels)) + if remember: + self.history.append((prompt, max(scores, key=scores.get))) + return scores elif owner._reply_fn is not None: def call(row, prompt): - return owner._reply_fn(prompt) + shown = with_history(prompt) + # Scoring first keeps the recorded first pass the prompt alone. + scores = owner._likelihood_fn(shown, list(labels)) if score_first else None + reply = owner._reply_fn(shown) + if remember: + self.history.append((prompt, str(reply))) + return reply, scores else: # a registration's answer: it phrases and parses a Choice itself 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, prompts) + scores = None + if path == 'generation' and owner._reply_fn is not None: + scores = [s for _, s in replies] + replies = [r for r, _ in replies] if path == 'likelihood': probabilities = _softmax_rows(replies, labels) chosen = [labels[i] for i in probabilities.argmax(1)] @@ -383,54 +404,92 @@ def call(row, prompt): 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, prompts, labels) - outputs['behavior:probabilities'] = extra + if score_first: + outputs['behavior:probabilities'] = _implied_array( + _softmax_rows(scores, labels), stimuli, labels, 'likelihood') + elif self.probability_source == 'readout': + outputs['behavior:probabilities'] = _implied(self.readout_classifier.predict_proba( + owner._extract_behavioral_features(stimuli)), 'readout') + if isinstance(response, Choice) and response.reaction_time: + outputs['behavior:rt'] = _reaction_time(outputs['behavior:probabilities'], + response.reaction_time, stimuli) outputs['behavior'] = outputs[response.channels[0]] - stamp = self._stamp(path, 'prompt' if path == 'generation' else 'readout extraction') + if prompts: + # A multi-image trial records one step per image: name each by its part. + image_names = [name for (kind, _), name in zip(prompts[0].parts, prompts[0].names) + if kind == 'image'] + neural = {region: (assembly.assign_coords(part_name=('part', image_names)) + if 'part' in assembly.dims and assembly.sizes['part'] == len(image_names) + else assembly) + for region, assembly in neural.items()} + stamp = self._stamp(path, 'readout extraction' if path == 'readout' else 'prompt') 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, prompts, labels): - """Implied probabilities alongside a generated answer (extra passes, not recorded).""" - owner = self.owner - if self.probability_source == 'likelihood': - 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 _open_recordings(self, stimuli, layers): + """One live recording per input type, before any compute. - 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 + An extractor's own ``record(layers)`` when it has one; otherwise forward + hooks on the model passed as ``BrainScoreModel(model=...)``, with layer + paths read from that model. + """ owner = self.owner + from .dispatch import extractor_for by_modality: Dict[str, list] = {} - for region, layer in layers.items(): - by_modality.setdefault(self._region_modality(region), []).append(layer) + for region, region_layers in layers.items(): + by_modality.setdefault(self._region_modality(region), []).extend(region_layers) recordings = {} for modality, modality_layers in by_modality.items(): - from .dispatch import extractor_for + modality_layers = list(dict.fromkeys(modality_layers)) extractor = extractor_for(owner, modality, stimuli) if extractor is None: raise NotImplementedError(f'No {modality!r} extractor to record from.') try: record = extractor.record except AttributeError: + record = None + if record is not None: + recordings[modality] = record(modality_layers) + elif owner._model is not None: + recordings[modality] = _hook_recording(owner._model, modality_layers, owner.identifier) + else: 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))) + f'The {modality!r} extractor cannot record while the model answers: give it ' + 'a record(layers) method, or pass the network as BrainScoreModel(model=...).') + return recordings + + def _per_stimulus(self, stimuli, layers, call, prompts=None): + """Call ``call(row, prompt)`` once per stimulus, recording ``layers`` on each first pass.""" + recordings = self._open_recordings(stimuli, layers) replies, trials = [], [] with ExitStack() as stack: live = {m: stack.enter_context(r) for m, r in recordings.items()} 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) - for region, layer in layers.items()} - return replies, neural + return replies, self._trial_regions(trials, stimuli, layers) + + def _trial_regions(self, trials, stimuli, layers): + """Per-trial recordings -> one NeuroidAssembly per region (composite regions gathered).""" + import xarray as xr + neural = {} + for region, region_layers in layers.items(): + modality = self._region_modality(region) + per_layer = [_trials_to_assembly([trial[modality] for trial in trials], stimuli, layer, region) + for layer in region_layers] + neural[region] = self._gather(region, per_layer) + return neural + + def _gather(self, region, per_layer): + """A region's units: one layer as is; a composite region's units gathered from its layers.""" + import xarray as xr + from ..selection import CompositeSelector + selector = self.owner._region_layer_selectors.get(region) + if not isinstance(selector, CompositeSelector): + return per_layer[0] + recorder = self.owner._recorder + combined = xr.concat([recorder.ensure_unit_indices(a) for a in per_layer], dim='neuroid') + return recorder.gather_composite_units(combined, selector) @contextmanager def suspended(self): @@ -451,27 +510,22 @@ def respond_and_record(self, stimuli, regions: Sequence[str]): 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. + prompt, before any generated token). With a readout, the region layers + are extracted in the same call as the readout layer. """ - layers = {region: self._region_layer(region) for region in regions} + layers = {region: self._region_layers(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: + def _region_layers(self, region: str) -> List[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 + return list(selector.layer_paths) + return [selector.layer_path] def _region_modality(self, region: str) -> str: from ..io_catalog import canonical_modality @@ -479,45 +533,25 @@ def _region_modality(self, region: str) -> str: if modality is not None: return canonical_modality(modality) modalities = {canonical_modality(key) for key in self.owner._extractors} + if len({id(e) for e in self.owner._extractors.values()}) == 1: + return sorted(modalities)[0] # one (fused) extractor records every region 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(): - 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: - 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))) + def _generate_and_record(self, stimuli, layers): + recordings = self._open_recordings(stimuli, 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 - ).assign_attrs(recording=self._stamp('answer', 'prompt')) - return behavior, neural + stamp = self._stamp('answer', 'prompt') + return behavior, {region: assembly.assign_attrs(recording=stamp) + for region, assembly in self._trial_regions(trials, stimuli, layers).items()} def _stamp(self, path, forward_pass): """How a recording was made, so task and passive recordings stay distinct.""" @@ -526,7 +560,7 @@ def _stamp(self, path, forward_pass): 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], raw: bool = False): + def _read_out_and_record(self, stimuli, layers, raw: bool = False): from ..responses import Scale owner = self.owner readout_layer = owner._behavioral_readout_layer @@ -537,7 +571,7 @@ def _read_out_and_record(self, stimuli, layers: Dict[str, str], raw: bool = Fals 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()])) + all_layers = list(dict.fromkeys([readout_layer, *(l for ls in layers.values() for l in ls)])) extracted = owner._recorder.ensure_layer_provenance( owner._extract_for_modality(stimuli, modality, all_layers), all_layers, description='Readout extraction') @@ -550,11 +584,11 @@ def _read_out_and_record(self, stimuli, layers: Dict[str, str], raw: bool = Fals else: behavior = self.readout_classifier.predict_proba(features) neural = {} - for region, layer in layers.items(): - selected = extracted.isel(neuroid=layer_of == layer) + for region, region_layers in layers.items(): + per_layer = [extracted.isel(neuroid=layer_of == layer) for layer in region_layers] + selected = self._gather(region, per_layer) neural[region] = selected.assign_coords( - region=('neuroid', [region] * selected.sizes['neuroid']) - ) + region=('neuroid', [region] * selected.sizes['neuroid'])) if not raw: neural[region] = neural[region].assign_attrs( recording=self._stamp('readout', 'readout extraction')) @@ -630,6 +664,58 @@ def generate_predictions(self, stimuli, after_each: Optional[Callable[[], None]] ) +@contextmanager +def _hook_recording(model, layers, identifier): + """Forward hooks on ``model`` for extractors without ``record()``. + + ``take()`` returns the first pass since the previous take, each layer's + output flattened per presentation, as a still-image extractor packages it. + """ + from brainscore_core.supported_data_standards.brainio.assemblies import NeuroidAssembly + + def module_at(path): + module = model + for part in path.split('.'): + module = module[int(part)] if part.isdigit() else getattr(module, part, None) + if module is None: + raise NotImplementedError(f'Layer {path!r} is not a module of the model passed as model=; ' + 'give the extractor a record(layers) method.') + return module + + modules = {layer: module_at(layer) for layer in layers} # unknown paths fail before compute + captured = {} + + def keep(layer): + def hook(_module, _inputs, output): + if layer not in captured: + output = output[0] if isinstance(output, (tuple, list)) else output + captured[layer] = np.asarray(output.detach().float().cpu().numpy()) + return hook + + handles = [module.register_forward_hook(keep(layer)) for layer, module in modules.items()] + + class Recording: + @staticmethod + def take(): + missing = [layer for layer in layers if layer not in captured] + if missing: + raise ValueError(f'No forward pass reached {missing} since the last presentation.') + rows = {layer: captured[layer].reshape(1, -1) for layer in layers} + captured.clear() + values = np.concatenate([rows[layer] for layer in layers], axis=1) + neuroid_layers = [layer for layer in layers for _ in range(rows[layer].shape[1])] + numbers = [i for layer in layers for i in range(rows[layer].shape[1])] + return NeuroidAssembly(values, dims=['presentation', 'neuroid'], coords={ + 'stimulus_id': ('presentation', ['?']), + 'neuroid_id': ('neuroid', [f'{identifier}.{l}.{i}' for l, i in zip(neuroid_layers, numbers)]), + 'neuroid_num': ('neuroid', numbers), 'layer': ('neuroid', neuroid_layers)}) + try: + yield Recording() + finally: + for handle in handles: + handle.remove() + + def _presentation(stimuli): return {column: ('presentation', list(stimuli[column].values)) for column in stimuli.columns} @@ -676,6 +762,20 @@ def _texts(parsed, replies, stimuli): attrs={'measure': {'kind': 'reported', 'source': 'reply'}}) +def _reaction_time(probabilities, link, stimuli): + """An implied reaction time per stimulus from label probabilities (larger = slower).""" + from brainscore_core.supported_data_standards.brainio.assemblies import BehavioralAssembly + values = np.asarray(probabilities.values, dtype=float) + ordered = np.sort(values, axis=1)[:, ::-1] + if link == 'evidence_margin': + rt = 1.0 - (ordered[:, 0] - ordered[:, 1]) + else: # surprisal of the chosen (most probable) label + rt = -np.log2(np.clip(ordered[:, 0], 1e-12, 1.0)) + return BehavioralAssembly(rt, dims=['presentation'], coords=_presentation(stimuli), + attrs={'measure': {'kind': 'implied', 'link': link, + 'source': probabilities.attrs.get('measure', {}).get('source')}}) + + 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) @@ -702,7 +802,12 @@ 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 + # A clip keeps its time steps and a multi-image trial its parts: one extra dimension. + extra = [dim for dim in trials[0].dims if dim not in ('presentation', 'neuroid')] + if len(extra) > 1: + raise ValueError(f'A recording during a task may have one extra dimension; got {extra}.') + timed = bool(extra) + extra = extra[0] if extra else None for index, trial in enumerate(trials): selected = trial.isel(neuroid=trial['layer'].values == layer) if selected.sizes['neuroid'] == 0: @@ -710,11 +815,10 @@ def _trials_to_assembly(trials, stimuli, layer: str, region: str): f'No activity reached layer {layer!r} while answering stimulus ' f'{stimuli["stimulus_id"].iloc[index]!r}; that input may not use it.') 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]) + if selected.sizes[extra] != trials[0].sizes[extra]: + raise ValueError(f'Trials gave different numbers of {extra} steps; present trials of one ' + 'length (clips of one duration, sequences of one length) together.') + rows.append(np.asarray(selected.transpose('presentation', extra, 'neuroid').values)[0]) else: rows.append(np.asarray(selected.values).reshape(-1)) first = trials[0].isel(neuroid=trials[0]['layer'].values == layer) @@ -728,14 +832,14 @@ def _trials_to_assembly(trials, stimuli, layer: str, region: str): coords[column] = ('presentation', list(stimuli[column].values)) 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)) + extra_index = first.indexes.get(extra) + extra_names = list(extra_index.names) if isinstance(extra_index, pd.MultiIndex) else [] + extra_names += [name for name, coord in first.coords.items() + if coord.dims == (extra,) and name != extra and name not in extra_names] + for name in extra_names: + coords[name] = (extra, np.asarray(first[name].values)) return NeuroidAssembly(np.stack(rows), coords=coords, - dims=['presentation', 'time_bin', 'neuroid']) + dims=['presentation', extra, '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 8895c4b3..416e7c78 100644 --- a/brainscore_core/_model/dispatch.py +++ b/brainscore_core/_model/dispatch.py @@ -69,6 +69,16 @@ def _is_clip(stimuli) -> bool: return any(column in getattr(stimuli, 'columns', ()) for column in CLIP_COLUMNS) +def fused_for(owner, detected): + """The fused extractor taking every detected input type, if one is registered.""" + if len(detected) < 2: + return None + for types, extractor in owner._fused_extractors: + if set(detected) <= types: + return extractor + return None + + 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.""" @@ -93,7 +103,18 @@ def normalize_extractors(owner, preprocessors, activations_model): import inspect if isinstance(activations_model, Mapping): - extractors = dict(activations_model) + # A fused extractor is registered under every input type it takes at + # once, {('vision', 'text'): fused}. It also serves each type alone + # unless a single-type extractor is registered for it. + fused = [(frozenset(canonical_modality(m) for m in key), value) + for key, value in activations_model.items() if isinstance(key, tuple)] + extractors = {} + for types, value in fused: + for modality in types: + extractors.setdefault(modality, value) + extractors.update({key: value for key, value in activations_model.items() + if not isinstance(key, tuple)}) + owner._fused_extractors = fused elif activations_model is not None: if not preprocessors: raise ValueError( @@ -283,6 +304,13 @@ def process_multi_modality(self, stimuli, detected: Set[str], import xarray as xr owner = self.owner + fused = fused_for(owner, detected) + if fused is not None: + # One fused extractor takes every input type: one forward pass. + assembly = fused(stimuli, layers=layers) + if owner._recorder.is_multi_region: + assembly = owner._tag_neuroids_with_regions(assembly) + return assembly ordered = sorted( detected, key=lambda m: ( diff --git a/brainscore_core/_sessions/drivers.py b/brainscore_core/_sessions/drivers.py index c2b1c141..16a2ec96 100644 --- a/brainscore_core/_sessions/drivers.py +++ b/brainscore_core/_sessions/drivers.py @@ -394,10 +394,6 @@ def drive_behavior_and_neural_session(model, session, driver: str = "interact") 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( @@ -407,24 +403,30 @@ def drive_behavior_and_neural_session(model, session, driver: str = "interact") answers = [c for c in requested if parse_channel(c)[0] == "behavior"] _refuse_variance(model, session) model.start_task(task_context) + model._behavioral.history = [] # memory spans one session, not the next stimuli = _behavior_stimuli_from_session(session, task_context) if stimuli is None: raise ValueError("Recording during a task needs stimuli to present.") - 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, - meta={"driver": driver, "record": region, - "during": "behavior"})) + if task_context.response is None and set(answers) - {"behavior"}: + raise ValueError(f"{answers} need TaskContext(response=...); task_type " + f"{task_context.task_type!r} answers on 'behavior'.") + # A streaming session gets each trial's outputs as soon as that trial ends; + # otherwise the whole set is answered and sent at once. + streaming = getattr(session, "streaming", False) + batches = ([(index, _preserve_stimulus_set_type(stimuli, stimuli.iloc[[index]])) + for index in range(len(stimuli))] if streaming else [(0, stimuli)]) + for index, batch in batches: + if task_context.response is not None: + behavior, neural = model._behavioral.respond(batch, answers, regions) + else: + plain, neural = model._behavioral.respond_and_record(batch, regions) + behavior = {"behavior": plain} if answers else {} + for channel in answers: + session.emit(StreamEvent(channel=channel, payload=behavior[channel], t_ms=float(index), + 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=float(index), + meta={"driver": driver, "record": region, "during": "behavior"})) _drive_neural_session_via_process = drive_neural_session diff --git a/brainscore_core/brainscore_model.py b/brainscore_core/brainscore_model.py index a41646c2..fe93b875 100644 --- a/brainscore_core/brainscore_model.py +++ b/brainscore_core/brainscore_model.py @@ -96,6 +96,7 @@ def __init__( intervention: Optional[Callable] = None, reply: Optional[Callable] = None, likelihood: Optional[Callable] = None, + memory: Optional[str] = None, behavioral_readout_layer: Optional[str] = None, generation_fn: Optional[Callable] = None, action_fn: Optional[Callable] = None, @@ -132,6 +133,10 @@ def __init__( :param answer: ``(stimulus_row, instruction, label_set) -> label``; generates and parses a behavioral answer for instruction-following tasks. + :param extractor: also accepts a fused extractor under a tuple of input + types, ``{('vision', 'text'): fused}``: one forward pass over all of a + stimulus's parts. It declares each type, and takes any stimulus that + brings several of them at once. :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). @@ -139,7 +144,14 @@ def __init__( and written answers. :param likelihood: ``(prompt, candidates) -> {candidate: log probability}``; scores candidate replies without generating. Gives - choices and their implied probabilities. + choices and their implied probabilities. Run the prompt on its own + before any candidate, so a recording during scoring is the prompt pass. + :param memory: ``'session'`` keeps what the model saw and said across the + trials of a session, until ``reset()``: each trial's prompt and reply + reach ``reply`` and ``likelihood`` as ``prompt.history``. It declares + the ``memory:session`` trait. A stateful extractor (a recurrent model) + keeps its own state the same way. ``None`` (default): every trial is + independent. :param policy: ``EnvironmentStep -> EnvironmentResponse``; chooses an action for closed-loop / embodied evaluation. :param intervention: ``StateChange -> (PerturbationApplied, cleanup)``; @@ -220,6 +232,7 @@ def __init__( from ._model.dispatch import normalize_extractors self._preprocessors: Dict[str, Optional[Callable]] = dict(preprocessors) self._activations_model = activations_model + self._fused_extractors = [] # (input types, extractor) registered under a tuple self._extractors, self._layer_extractors = normalize_extractors( self, preprocessors, activations_model) for modality in supported_modalities or (): @@ -234,6 +247,9 @@ def __init__( 'behavioral_readout_layer') self._generation_fn = capability_values.get('generation_fn') self._reply_fn, self._likelihood_fn = reply, likelihood + if memory not in (None, 'session'): + raise ValueError("memory must be None or 'session'") + self._memory = memory self._action_fn = capability_values.get('action_fn') self._state_change_fn = capability_values.get('state_change_fn') @@ -289,6 +305,10 @@ def _visual_forms(self): forms = {'vision': 'vision:still', 'video': 'vision:clip'} return {forms[key] for key in self._preprocessors if key in forms} + @property + def traits(self): + return {'memory:session'} if self._memory == 'session' else set() + @property def out_channels(self): # Every output comes from a registered part; nothing is found by name. diff --git a/brainscore_core/capabilities/neural.py b/brainscore_core/capabilities/neural.py index 423f3b95..6a8dc7fc 100644 --- a/brainscore_core/capabilities/neural.py +++ b/brainscore_core/capabilities/neural.py @@ -65,7 +65,8 @@ def process(self, model, stimuli, **kwargs): f"{sorted(model._region_layer_map_dict.keys())}." ) - if multi_modality and len(detected) > 1: + from .._model.dispatch import fused_for + if (multi_modality or fused_for(model, detected) is not None) and len(detected) > 1: return model._process_multi_modality(stimuli, detected, layers) if not multi_modality and len(detected) > 1: diff --git a/brainscore_core/capabilities/sessions.py b/brainscore_core/capabilities/sessions.py index a0ae6265..97c0b855 100644 --- a/brainscore_core/capabilities/sessions.py +++ b/brainscore_core/capabilities/sessions.py @@ -71,14 +71,14 @@ def declared_output_channels(self, model): # 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'} + channels |= {'behavior:choice', 'behavior:probabilities', 'behavior:rating', 'behavior:rt'} 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'} + channels |= {'behavior:choice', 'behavior:probabilities', 'behavior:rt'} return channels def supports_session(self, model, channels) -> bool: diff --git a/brainscore_core/contract.py b/brainscore_core/contract.py index f4ae3c9d..ff3fbd91 100644 --- a/brainscore_core/contract.py +++ b/brainscore_core/contract.py @@ -46,7 +46,9 @@ class TaskContext: ``'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 + (e.g. ``('text', 'video')``: a sentence, then a clip). Name columns to show + several parts of one kind in sequence: ``('image:instruction', 'image:word_1', + 'image:word_2')``. 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 @@ -81,9 +83,11 @@ def __post_init__(self): 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'} + kinds = {'text', 'image', 'video', 'audio'} + unknown = [entry for entry in self.present if entry.split(':', 1)[0] not in kinds] if unknown or isinstance(self.present, str): - raise ValueError(f"present lists parts from 'text', 'image', 'video', 'audio'; got {self.present!r}") + raise ValueError("present lists parts as a kind ('text', 'image', 'video', 'audio') " + f"or 'kind:column' (e.g. 'image:word_1'); got {self.present!r}") self.present = tuple(self.present) @@ -123,6 +127,16 @@ def required_channels(self) -> Set[str]: """Inputs the subject cannot run without; none by default.""" return set() + @property + def traits(self) -> Set[str]: + """Properties of the subject a benchmark may require; none by default. + + ``'memory:session'``: the subject keeps what it saw and said across the + trials of one session, until ``reset()``. A benchmark that needs a trait + lists it in ``required_traits``; a subject without it is N/A. + """ + return set() + @abstractmethod def participate(self, session: "Session") -> None: """Consume session inputs and emit responses until the session ends.""" diff --git a/brainscore_core/io_catalog.py b/brainscore_core/io_catalog.py index e758f122..a56591a4 100644 --- a/brainscore_core/io_catalog.py +++ b/brainscore_core/io_catalog.py @@ -363,7 +363,7 @@ def validate( "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"), + "choice, probabilities, confidence, rating, rt 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 index 427948a1..8b9433d7 100644 --- a/brainscore_core/responses.py +++ b/brainscore_core/responses.py @@ -51,21 +51,32 @@ class Choice: ``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. + ``reaction_time`` also returns an implied reaction time from those + probabilities, scored by correlation with human RTs as Futrell2018 scores + surprisal: ``'evidence_margin'`` (1 minus the gap between the two most + probable labels; close calls are slow) or ``'surprisal'`` (-log2 of the + chosen label's probability). """ labels: Tuple[str, ...] probabilities: bool = False confidence: Optional[Scale] = None + reaction_time: Optional[str] = None # 'evidence_margin' or 'surprisal': an implied RT + + REACTION_TIMES = ('evidence_margin', 'surprisal') def __post_init__(self): if isinstance(self.labels, str) or len(self.labels) < 2: raise ValueError('Choice needs at least two labels') + if self.reaction_time not in (None, *self.REACTION_TIMES): + raise ValueError(f'reaction_time is one of {self.REACTION_TIMES}, not {self.reaction_time!r}') 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 ())) + + (('behavior:confidence',) if self.confidence is not None else ()) + + (('behavior:rt',) if self.reaction_time is not None else ())) def directive(self) -> str: text = f'Answer with exactly one of: {", ".join(self.labels)}.' @@ -126,6 +137,8 @@ class Prompt: """ parts: Tuple[Tuple[str, Any], ...] stimulus_id: Optional[str] = None + names: Tuple[str, ...] = () # one per part: 'instruction', a column name, or 'directive' + history: Tuple[Tuple['Prompt', str], ...] = () # earlier trials of the session: (prompt, reply) def text(self) -> str: """The text parts joined, for a model that reads text only.""" @@ -145,17 +158,28 @@ def trial_prompt(task, stimulus_row) -> Prompt: 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 = [] + shown = [] # (kind, value, name) + for entry in order: + if ':' in entry: + # A named part: 'image:word_1' shows the image in column word_1. + kind, column = entry.split(':', 1) + value = _cell(stimulus_row, column) + if value is None: + raise ValueError(f'TaskContext.present names column {column!r}, which this stimulus lacks.') + shown.append((kind, value, column)) + elif entry in present: + shown.append((entry, present[entry], entry)) + else: + raise ValueError(f'TaskContext.present names {entry!r}, but the stimulus has only {sorted(present)}.') + parts, names = [], [] if (task.instruction or '').strip(): - parts.append(('text', task.instruction.strip())) - parts += [(kind, present[kind]) for kind in order] + parts.append(('text', task.instruction.strip())); names.append('instruction') + for kind, value, name in shown: + parts.append((kind, value)); names.append(name) 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')) + parts.append(('text', directive)); names.append('directive') + return Prompt(tuple(parts), stimulus_id=_cell(stimulus_row, 'stimulus_id'), names=tuple(names)) def _cell(row, column): diff --git a/tests/test_compatibility.py b/tests/test_compatibility.py index e586bd77..572c07c3 100644 --- a/tests/test_compatibility.py +++ b/tests/test_compatibility.py @@ -395,7 +395,7 @@ def _model(self): def test_every_part_declares_registered_channels(self): model = self._model() assert model.out_channels == {'neural:IT', 'behavior', 'behavior:choice', - 'behavior:probabilities', 'behavior:rating', + 'behavior:probabilities', 'behavior:rating', 'behavior:rt', 'motor', 'perturbation'} assert {'vision', 'observation', 'lesion', 'stimulation'} <= model.in_channels for output in sorted(model.out_channels): diff --git a/tests/test_fused_extractor.py b/tests/test_fused_extractor.py new file mode 100644 index 00000000..18105ed6 --- /dev/null +++ b/tests/test_fused_extractor.py @@ -0,0 +1,62 @@ +"""A fused extractor: one forward pass over several input types, split by input.""" +import numpy as np +import pytest + +from brainscore_core.brainscore_model import BrainScoreModel +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 + + +class Fused: + """Sees the image and the sentence of each stimulus together; returns one step per input.""" + identifier = backbone_id = 'fused' + + def __init__(self): + self.calls = [] + + def __call__(self, stimuli, layers, **kwargs): + self.calls.append(sorted(c for c in stimuli.columns if c != 'stimulus_id')) + ids = list(stimuli['stimulus_id']) + values = np.stack([[[1.0, 2.0], [3.0, 4.0]] for _ in ids]) # parts: image, text + return NeuroidAssembly(values, dims=['presentation', 'part', 'neuroid'], coords={ + 'stimulus_id': ('presentation', ids), 'part_name': ('part', ['image', 'text']), + 'part_index': ('part', [0, 1]), + 'neuroid_id': ('neuroid', ['n0', 'n1']), 'layer': ('neuroid', [layers[0]] * 2)}) + + +def test_a_fused_extractor_declares_each_input_and_runs_once(): + fused = Fused() + model = BrainScoreModel('fused', extractor={('vision', 'text'): fused}, + region_layer_map={'A': 'layers.3'}) + assert {'vision', 'text', 'vision:still'} <= model.in_channels + both = StimulusSet({'stimulus_id': ['a', 'b'], 'image_path': ['a.png', 'b.png'], + 'sentence': ['one', 'two']}) + both.identifier = 'fused-both' + session = BufferedSession(stimulus_set=both, requested_output_channels=['neural:A']) + model.participate(session) + recorded = session.collect('neural:A') + assert len(fused.calls) == 1 and fused.calls[0] == ['image_path', 'sentence'] + assert recorded.dims == ('presentation', 'part', 'neuroid') + assert list(recorded['part_name'].values) == ['image', 'text'] + + +def test_single_input_extractors_still_take_single_input_stimuli(): + fused, vision_calls = Fused(), [] + + def vision(stimuli, layers, **kwargs): + vision_calls.append(1) + ids = list(stimuli['stimulus_id']) + return NeuroidAssembly(np.ones((len(ids), 2)), dims=['presentation', 'neuroid'], coords={ + 'stimulus_id': ('presentation', ids), 'neuroid_id': ('neuroid', ['v0', 'v1']), + 'layer': ('neuroid', [layers[0]] * 2)}) + model = BrainScoreModel('mixed', extractor={'vision': vision, ('vision', 'text'): fused}, + region_layer_map={'A': LayerSelector('layers.3', modality='vision')}) + images = StimulusSet({'stimulus_id': ['a'], 'image_path': ['a.png']}) + images.identifier = 'images-only' + session = BufferedSession(stimulus_set=images, requested_output_channels=['neural:A']) + model.participate(session) + assert session.collect('neural:A').dims == ('presentation', 'neuroid') diff --git a/tests/test_recording_everywhere.py b/tests/test_recording_everywhere.py new file mode 100644 index 00000000..4dd3df27 --- /dev/null +++ b/tests/test_recording_everywhere.py @@ -0,0 +1,147 @@ +"""Regions record on every task path: likelihood scoring, a reply with implied +probabilities, composite regions, streaming sessions, and extractors without +record() (through forward hooks on the model).""" +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.responses import Choice +from brainscore_core.selection import CompositeSelector, 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 + +STIMULI = StimulusSet({'stimulus_id': ['a', 'b', 'c'], 'sentence': ['x', 'yy', 'zzz']}) +STIMULI.identifier = 'everywhere' + + +class Net: + """Each pass sees a text; 'early' has 3 units, 'late' 2. Passes are logged.""" + + def __init__(self): + self.passes, self.hooks = [], [] + + def forward(self, text): + self.passes.append(text) + out = {'early': np.array([len(text), 2.0 * len(text), 3.0 * len(text)]), + 'late': np.array([-len(text), 10.0])} + for hook in self.hooks: + hook(out) + return out + + +class Recordable: + identifier = backbone_id = 'toy' + + def __init__(self, net): + self.net = net + + def __call__(self, stimuli, layers, **kwargs): + raise AssertionError('recorded during the task') + + @contextmanager + def record(self, layers): + seen = [] + + class Take: + @staticmethod + def take(): + first = seen[0] + seen.clear() + values = np.concatenate([first[layer] for layer in layers])[None] + units = [(layer, i) for layer in layers for i in range(len(first[layer]))] + return NeuroidAssembly(values, dims=['presentation', 'neuroid'], coords={ + 'stimulus_id': ('presentation', ['?']), + 'neuroid_id': ('neuroid', [f'{l}.{i}' for l, i in units]), + 'neuroid_num': ('neuroid', [i for _, i in units]), + 'layer': ('neuroid', [l for l, _ in units])}) + self.net.hooks.append(seen.append) + try: + yield Take() + finally: + self.net.hooks.remove(seen.append) + + +def _model(net, region_map=None): + def reply(prompt): + net.forward(prompt.text()) + return 'cat' + + def likelihood(prompt, candidates): + net.forward(prompt.text()) # the prompt alone first + return {c: -float(len(net.forward(prompt.text() + c)['late'])) - i + for i, c in enumerate(candidates)} + return BrainScoreModel('toy', extractor={'text': Recordable(net)}, + region_layer_map=region_map or {'A': LayerSelector('early', modality='text')}, + reply=reply, likelihood=likelihood) + + +def _run(model, task, channels, streaming=False): + session = BufferedSession(stimulus_set=STIMULI, task_context=task, + requested_output_channels=channels, streaming=streaming) + model.participate(session) + return session + + +def _prompt_text(stimulus_text): + return f'Which?\n\n{stimulus_text}\n\nAnswer with exactly one of: cat, dog.' + + +def test_recording_during_likelihood_scoring_is_the_prompt_pass(): + net = Net() + task = TaskContext(instruction='Which?', response=Choice(['cat', 'dog'], probabilities=True), + prefer_path='likelihood') + recorded = _run(_model(net), task, ['behavior:probabilities', 'neural:A']).collect('neural:A') + expected = [Net().forward(_prompt_text(s))['early'] for s in ('x', 'yy', 'zzz')] + np.testing.assert_array_equal(recorded.values, expected) + + +def test_a_reply_with_implied_probabilities_records_the_prompt_pass(): + net = Net() + task = TaskContext(instruction='Which?', response=Choice(['cat', 'dog'], probabilities=True), + prefer_path='generation') + session = _run(_model(net), task, ['behavior:choice', 'behavior:probabilities', 'neural:A']) + expected = [Net().forward(_prompt_text(s))['early'] for s in ('x', 'yy', 'zzz')] + np.testing.assert_array_equal(session.collect('neural:A').values, expected) + assert session.collect('behavior:probabilities').attrs['measure']['source'] == 'likelihood' + + +def test_a_composite_region_records_its_chosen_units(): + net = Net() + region = CompositeSelector(layers=(('early', (0, 2)), ('late', None)), modality='text') + task = TaskContext(instruction='Which?', response=Choice(['cat', 'dog'])) + recorded = _run(_model(net, {'R': region}), task, ['behavior:choice', 'neural:R']).collect('neural:R') + first = Net().forward(_prompt_text('x')) + np.testing.assert_array_equal(recorded.values[0], [first['early'][0], first['early'][2], *first['late']]) + + +def test_a_streaming_session_gets_each_trial_as_it_ends(): + net = Net() + task = TaskContext(instruction='Which?', response=Choice(['cat', 'dog'])) + session = _run(_model(net), task, ['behavior:choice', 'neural:A'], streaming=True) + assert [e.t_ms for e in session.emitted if e.channel == 'neural:A'] == [0.0, 1.0, 2.0] + assert session.collect('neural:A').sizes['presentation'] == 3 + + +def test_an_extractor_without_record_is_recorded_through_the_model(): + torch = pytest.importorskip('torch') + network = torch.nn.Sequential(torch.nn.Linear(2, 3), torch.nn.ReLU(), torch.nn.Linear(3, 1)) + + def extract(stimuli, layers, **kwargs): + raise AssertionError('recorded during the task') + + def reply(prompt): + network(torch.tensor([[float(len(prompt.text())), 1.0]])) + return 'cat' + model = BrainScoreModel('plain', model=network, extractor={'text': extract}, + region_layer_map={'A': LayerSelector('0', modality='text')}, reply=reply) + task = TaskContext(instruction='Which?', response=Choice(['cat', 'dog'])) + recorded = _run(model, task, ['behavior:choice', 'neural:A']).collect('neural:A') + expected = network[0](torch.tensor([[float(len(_prompt_text('x'))), 1.0]])).detach().numpy()[0] + np.testing.assert_allclose(recorded.values[0], expected, rtol=1e-6) + assert not network[0]._forward_hooks diff --git a/tests/test_session_memory.py b/tests/test_session_memory.py new file mode 100644 index 00000000..9497a113 --- /dev/null +++ b/tests/test_session_memory.py @@ -0,0 +1,79 @@ +"""Memory across trials: a declared trait, required by benchmarks, kept through a session.""" +import pytest + +from brainscore_core.brainscore_model import BrainScoreModel +from brainscore_core.compatibility import CompatibilityError +from brainscore_core.model_interface import TaskContext +from brainscore_core.responses import Choice +from brainscore_core.streaming import BufferedSession +from brainscore_core.supported_data_standards.brainio.stimuli import StimulusSet +from brainscore_core._evaluation.compatibility import check_channel_compatibility + +pytestmark = pytest.mark.unit + +WORDS = ['cat', 'cat', 'dog', 'sun', 'sun', 'dog'] +ONE_BACK = [None, 'yes', 'no', 'no', 'yes', 'no'] # same as the previous trial? +STIMULI = StimulusSet({'stimulus_id': [f's{i}' for i in range(6)], 'sentence': WORDS}) +STIMULI.identifier = 'one-back' +TASK = TaskContext(instruction='Is this word the same as the previous one?', + response=Choice(['yes', 'no'])) + + +def _word(prompt): + return [content for kind, content in prompt.parts if kind == 'text'][1] + + +def one_back(prompt): + """Answers from the session's history; with no history it can only say no.""" + previous = _word(prompt.history[-1][0]) if prompt.history else None + return 'yes' if previous == _word(prompt) else 'no' + + +def _answers(model): + session = BufferedSession(stimulus_set=STIMULI, task_context=TASK, + requested_output_channels=['behavior:choice']) + model.participate(session) + return list(session.collect('behavior:choice')['raw_answer'].values) + + +def test_a_model_with_memory_sees_earlier_trials(): + model = BrainScoreModel('remembers', supported_modalities={'text'}, reply=one_back, memory='session') + assert model.traits == {'memory:session'} + assert _answers(model)[1:] == ONE_BACK[1:] + + +def test_without_memory_every_trial_is_alone(): + model = BrainScoreModel('forgets', supported_modalities={'text'}, reply=one_back) + assert model.traits == set() + assert _answers(model) == ['no'] * 6 + + +def test_memory_lasts_one_session(): + model = BrainScoreModel('remembers', supported_modalities={'text'}, reply=one_back, memory='session') + _answers(model) + assert len(model._behavioral.history) == 6 + model.reset() + assert model._behavioral.history == [] + assert _answers(model)[1:] == ONE_BACK[1:] # a new session starts fresh + + +class _MemoryBenchmark: + identifier = 'one-back' + required_input_channels = {'text'} + requested_output_channels = {'behavior:choice'} + required_traits = {'memory:session'} + + +def test_a_memory_benchmark_marks_a_model_without_memory_na(): + forgets = BrainScoreModel('forgets', supported_modalities={'text'}, reply=one_back) + with pytest.raises(CompatibilityError, match='memory:session'): + check_channel_compatibility(forgets, _MemoryBenchmark()) + remembers = BrainScoreModel('remembers', supported_modalities={'text'}, reply=one_back, memory='session') + check_channel_compatibility(remembers, _MemoryBenchmark()) + + +def test_memory_needs_a_part_that_can_see_history(): + model = BrainScoreModel('answers', supported_modalities={'text'}, memory='session', + answer=lambda stimulus_row, instruction, label_set: 'no') + with pytest.raises(ValueError, match='reply'): + _answers(model) diff --git a/tests/test_subject_channels.py b/tests/test_subject_channels.py index a7dd7ff0..bba6b0d2 100644 --- a/tests/test_subject_channels.py +++ b/tests/test_subject_channels.py @@ -100,7 +100,7 @@ def test_brainscore_model_gets_channel_identity_from_preprocessors(): 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"} + "behavior:probabilities", "behavior:rating", "behavior:rt"} assert model.required_channels == set() diff --git a/tests/test_task_responses.py b/tests/test_task_responses.py index aef0639d..6db13331 100644 --- a/tests/test_task_responses.py +++ b/tests/test_task_responses.py @@ -167,3 +167,34 @@ def test_an_unordered_multi_part_stimulus_fails_before_any_reply(): with pytest.raises(ValueError, match='present'): model.participate(session) assert calls == [] + + +def test_named_parts_show_a_sequence_in_order(): + from brainscore_core.responses import trial_prompt + row = {'stimulus_id': 't1', 'instruction': 'rules.png', 'word_1': 'w1.png', 'word_2': 'w2.png'} + task = TaskContext(response=Choice(['yes', 'no']), + present=('image:instruction', 'image:word_1', 'image:word_2')) + prompt = trial_prompt(task, row) + assert prompt.parts[:3] == (('image', 'rules.png'), ('image', 'w1.png'), ('image', 'w2.png')) + assert prompt.names == ('instruction', 'word_1', 'word_2', 'directive') + with pytest.raises(ValueError, match='word_3'): + trial_prompt(TaskContext(response=Choice(['yes', 'no']), present=('image:word_3',)), row) + with pytest.raises(ValueError, match='kind'): + TaskContext(response=Choice(['yes', 'no']), present=('picture:word_1',)) + + +@pytest.mark.parametrize('link', ['evidence_margin', 'surprisal']) +def test_implied_reaction_times_rise_for_close_calls(link): + scores = {'a': {'cat': -0.1, 'dog': -5.0}, 'b': {'cat': -1.0, 'dog': -1.2}, + 'c': {'cat': -3.0, 'dog': -0.05}, 'd': {'cat': -0.7, 'dog': -0.69}} + model, _ = _model(scores=scores) + task = TaskContext(instruction='Which?', response=Choice(['cat', 'dog'], reaction_time=link)) + rt = _ask(model, task, ['behavior:rt']).collect('behavior:rt') + assert rt.attrs['measure'] == {'kind': 'implied', 'link': link, 'source': 'likelihood'} + fast, slow = rt.values[0], rt.values[3] # a clear call, a near tie + assert slow > fast + + +def test_a_reply_only_model_cannot_give_reaction_times(): + model, _ = _model(replies={s: 'cat' for s in 'abcd'}) + assert 'behavior:rt' not in model.out_channels diff --git a/tests/test_visual_forms.py b/tests/test_visual_forms.py index 5d15b6b5..48332a94 100644 --- a/tests/test_visual_forms.py +++ b/tests/test_visual_forms.py @@ -98,3 +98,38 @@ def test_a_clip_recorded_during_a_task_uses_the_clip_extractor_and_keeps_time(): 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] + + +def test_a_multi_image_trial_is_recorded_per_image(): + from contextlib import contextmanager + from brainscore_core.model_interface import TaskContext + from brainscore_core.responses import Choice + from brainscore_core.streaming import BufferedSession + + class PerImage: + def __call__(self, stimuli, layers, **kwargs): + raise AssertionError('recorded during the reply') + + @contextmanager + def record(self, layers): + class Take: + @staticmethod + def take(): + values = np.arange(6, dtype=float).reshape(1, 3, 2) # 3 images, 2 units + return NeuroidAssembly(values, dims=['presentation', 'part', 'neuroid'], coords={ + 'stimulus_id': ('presentation', ['?']), 'part_index': ('part', [0, 1, 2]), + 'neuroid_id': ('neuroid', ['n0', 'n1']), 'layer': ('neuroid', [layers[0]] * 2)}) + yield Take() + + model = BrainScoreModel('m', extractor={'vision': PerImage()}, + region_layer_map={'IT': LayerSelector('blocks.1', modality='vision')}, + reply=lambda prompt: 'yes') + trials = StimulusSet({'stimulus_id': ['t1'], 'rules': ['r.png'], 'word_1': ['a.png'], 'word_2': ['b.png']}) + trials.identifier = 'sequence' + task = TaskContext(response=Choice(['yes', 'no']), present=('image:rules', 'image:word_1', 'image:word_2')) + session = BufferedSession(stimulus_set=trials, task_context=task, + requested_output_channels=['behavior:choice', 'neural:IT']) + model.participate(session) + recorded = session.collect('neural:IT') + assert recorded.dims == ('presentation', 'part', 'neuroid') + assert list(recorded['part_name'].values) == ['rules', 'word_1', 'word_2']