fix: validate layer names before extraction and name the missing input column

.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@89366a83eb82e41e7bb0dcefe6401210383942de
+    uses: KartikP/brainscore-unified/.github/workflows/offline-tests.yml@f03ca786311c6f56d2c8bfa501127b563d182717
     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: 89366a83eb82e41e7bb0dcefe6401210383942de
+      unified_revision: f03ca786311c6f56d2c8bfa501127b563d182717

brainscore_core/_model/recording.py

@@ -5,6 +5,63 @@
 from ..io_catalog import canonical_modality
 from ..selection import CompositeSelector, LayerSelector
 
+INPUT_SITE = ':input'
+
+
+def selector_layers(selector) -> Tuple[str, ...]:
+    return (tuple(selector.layer_paths) if isinstance(selector, CompositeSelector)
+            else (selector.layer_path,))
+
+
+def declared_layers(owner, modality: Optional[str] = None):
+    """Layer paths the extractors for ``modality`` (all, when None) declare.
+
+    An extractor declares its layer namespace through an optional public
+    ``layer_names()`` (paths relative to the module it wraps). Returns None when
+    there is no extractor or any candidate declares nothing (e.g. a plain
+    callable), so validation never rejects a layer it cannot see.
+    """
+    wanted = canonical_modality(modality) if modality else None
+    candidates = {id(e): e for key, e in owner._extractors.items()
+                  if wanted is None or canonical_modality(key) == wanted}
+    if not candidates:
+        return None
+    names = set()
+    for extractor in candidates.values():
+        try:
+            layer_names = extractor.layer_names
+        except AttributeError:
+            return None
+        names.update(layer_names())
+    return names
+
+
+def unknown_layer_message(owner, path: str, modality: Optional[str] = None) -> Optional[str]:
+    """Why ``path`` is not a layer of the routed extractor, or None if it is
+    (or cannot be checked). ``':input'`` sites check their module path."""
+    names = declared_layers(owner, modality)
+    base = path[:-len(INPUT_SITE)] if path.endswith(INPUT_SITE) else path
+    if names is None or base in names:
+        return None
+    import difflib
+    # A full path given where the extractor wraps a sub-module ends with a valid one.
+    suffixes = sorted((n for n in names if base.endswith('.' + n)), key=len, reverse=True)
+    close = list(dict.fromkeys(suffixes + difflib.get_close_matches(base, names, n=5, cutoff=0.6)))[:5]
+    tower = f'the {modality!r} extractor' if modality else 'the extractor'
+    hint = (f'Close matches: {close}.' if close else
+            f'Some valid layers: {sorted(names)[:10]} ({len(names)} in all; '
+            'see extractor.layer_names()).')
+    return f'{path!r} is not a layer of {tower} of {owner.identifier!r}. {hint}'
+
+
+def validate_region_layers(owner, selectors) -> None:
+    """Fail at registration when a region maps to a layer its extractor lacks."""
+    for region, selector in selectors.items():
+        for path in selector_layers(selector):
+            message = unknown_layer_message(owner, path, selector.modality)
+            if message:
+                raise ValueError(f'region_layer_map[{region!r}]: {message}')
+
 
 class Recorder:
     """Owns recording target state for a BrainScoreModel instance."""
@@ -48,7 +105,15 @@ def start_recording(self,
             raise ValueError('Recording targets must be nonempty region or layer names.')
         import warnings
         for target in targets:
-            if target not in owner._region_layer_selectors:
+            if target in owner._region_layer_selectors:
+                continue
+            message = unknown_layer_message(owner, target)
+            if message:
+                raise ValueError(
+                    f"Recording target {target!r} is not a region in region_layer_map "
+                    f"(known regions: {sorted(owner._region_layer_selectors)}). As a layer "
+                    f"path: {message}")
+            if declared_layers(owner) is None:
                 warnings.warn(
                     f"Recording target {target!r} is not a region in region_layer_map; "
                     "treating it as a raw layer path. If you meant a region, check the spelling.",
@@ -92,9 +157,8 @@ def regions_for_modality(self, modality):
                 or canonical_modality(routing[region]) == canonical_modality(modality)]
 
     def layers_for_region(self, region):
-        selector = self.owner._region_layer_selectors.get(region, LayerSelector(region))
-        return (list(selector.layer_paths) if isinstance(selector, CompositeSelector)
-                else [selector.layer_path])
+        return list(selector_layers(
+            self.owner._region_layer_selectors.get(region, LayerSelector(region))))
 
     def current_layers(self) -> List[str]:
         if self.recording_layers:

brainscore_core/brainscore_model.py

@@ -366,6 +366,8 @@ def _install_region_layer_map(self, region_layer_map, routing=None) -> None:
                     routed = copy(selector)
                     routed.modality = modality
                     selectors[region] = routed
+        from ._model.recording import validate_region_layers
+        validate_region_layers(self, selectors)
         self._region_layer_selectors = selectors
         self._region_layer_map_dict = {r: sel.layer_path for r, sel in selectors.items()}
         self._region_modality_map = {

brainscore_core/capabilities/neural.py

@@ -27,11 +27,18 @@ def process(self, model, stimuli, **kwargs):
         detected = model._detect_modalities(stimuli)
 
         if not detected:
+            from ..io_catalog import stimulus_columns
+            mapping = {**model.COLUMN_TO_MODALITY, **stimulus_columns()}
+            accepted = {m: sorted(c for c, cm in mapping.items()
+                                  if canonical_modality(cm) == canonical_modality(m))
+                        for m in model.supported_modalities}
             raise ValueError(
-                f"No recognized modality columns in stimulus set. "
-                f"Columns present: {list(stimuli.columns)}. "
-                f"Known column mappings: {model.COLUMN_TO_MODALITY}. "
-                f"Model supports: {model.supported_modalities}."
+                f"No column in this stimulus set says what kind of input it holds "
+                f"(columns present: {list(stimuli.columns)}). Add one column per input "
+                f"this model takes; each row holds the input itself or its file path, "
+                f"for example stimuli['image_path'] = [...] for images. Columns this "
+                f"model reads: {accepted}. A stimulus_paths dict alone is not enough: "
+                f"the column names the input type."
             )
 
         layers = model._recorder.current_layers()

brainscore_core/extractor.py

@@ -16,6 +16,12 @@ class Extractor(Protocol):
     ``identifier`` and ``backbone_id`` key the activation cache (registrations
     that share weights share ``backbone_id``). An extractor that keeps state
     between calls also defines ``reset()``.
+
+    An extractor may also define ``layer_names()``, returning every path it
+    can record, relative to the module it wraps. ``BrainScoreModel`` then
+    rejects an unknown ``region_layer_map`` path or raw recording target
+    before any forward pass; a ``':input'`` site is checked by its module
+    path. An extractor without ``layer_names()`` skips this early check.
     """
 
     identifier: str

tests/test_layer_validation.py

@@ -0,0 +1,89 @@
+"""Layer paths are checked against the extractor's declared layer names before any forward pass."""
+import warnings
+
+import pytest
+
+from brainscore_core.brainscore_model import BrainScoreModel
+from brainscore_core.selection import CompositeSelector, LayerSelector
+
+
+class Declared:
+    """An extractor that declares its layer namespace; it must never be called here."""
+
+    def __init__(self, names):
+        self.names = names
+        self.identifier = 'declared'
+
+    def layer_names(self):
+        return list(self.names)
+
+    def __call__(self, stimuli, layers):
+        raise AssertionError('validation must not run a forward pass')
+
+
+VISION = Declared(['conv1', 'layer1', 'layer2', 'layer4', 'logits'])
+TEXT = Declared(['encoder.layers.0', 'encoder.layers.1'])
+
+
+def _model(region_layer_map, **kwargs):
+    return BrainScoreModel('toy', extractor={'vision': VISION}, region_layer_map=region_layer_map, **kwargs)
+
+
+def test_typo_in_region_layer_map_fails_at_construction_with_close_matches():
+    with pytest.raises(ValueError, match=r"region_layer_map\['IT'\].*'layr4'.*Close matches: \['layer4'"):
+        _model({'IT': 'layr4'})
+
+
+def test_region_layer_map_setters_are_validated():
+    model = _model({'IT': 'layer4'})
+    with pytest.raises(ValueError, match='layr2'):
+        model.region_layer_map = {'V1': 'layr2'}
+    with pytest.raises(ValueError, match='layr2'):
+        model.set_region_layer('V1', 'layr2')
+    assert dict(model.region_layer_map) == {'IT': 'layer4'}
+
+
+def test_unmapped_region_that_is_not_a_layer_fails_in_start_recording():
+    model = _model({'V1': 'conv1', 'V4': 'layer2', 'IT': 'layer4'})
+    with pytest.raises(ValueError, match=r"'V2' is not a region.*\['IT', 'V1', 'V4'\]"):
+        model.start_recording('V2')
+    with pytest.raises(ValueError, match="'V2'"):
+        model.start_recording(['IT', 'V2'])
+
+
+def test_existing_raw_layer_records_without_a_warning():
+    model = _model({'IT': 'layer4'})
+    with warnings.catch_warnings():
+        warnings.simplefilter('error')
+        model.start_recording('layer1')
+        model.start_recording('layer1:input')
+    assert model._recorder.current_layers() == ['layer1:input']
+
+
+def test_composite_selector_layers_are_each_checked():
+    with pytest.raises(ValueError, match="'layr2'"):
+        _model({'IT': CompositeSelector(layers=(('layer4', None), ('layr2', (0,))))})
+
+
+def test_routed_region_is_checked_against_its_own_tower():
+    extractor = {'vision': VISION, 'text': TEXT}
+    model = BrainScoreModel('two', extractor=extractor, region_layer_map={
+        'IT': LayerSelector('layer4', modality='vision'),
+        'language': LayerSelector('encoder.layers.1', modality='text')})
+    assert model.region_layer_map['language'].modality == 'text'
+    with pytest.raises(ValueError, match="'text' extractor"):
+        BrainScoreModel('two', extractor=extractor,
+                        region_layer_map={'language': LayerSelector('layer4', modality='text')})
+    with pytest.raises(ValueError, match="'vision' extractor"):
+        BrainScoreModel('two', extractor=extractor, region_layer_map={'IT': 'encoder.layers.1'},
+                        region_modality_map={'IT': 'vision'})
+
+
+def test_callable_extractor_without_layer_names_is_not_validated():
+    def extract(stimuli, layers):
+        raise AssertionError('not called')
+
+    model = BrainScoreModel('plain', extractor={'vision': extract}, region_layer_map={'IT': 'anything'})
+    with pytest.warns(UserWarning, match='raw layer path'):
+        model.start_recording('V2')
+

tests/test_model_interface.py

@@ -493,7 +493,16 @@ def test_no_modality_columns_raises(self):
             preprocessors={'vision': make_stub_preprocessor()},
         )
         stimuli = StubStimulusSet(columns=['unknown_column'])
-        with pytest.raises(ValueError, match="No recognized modality columns"):
+        with pytest.raises(ValueError, match="No column in this stimulus set"):
+            m.process(stimuli)
+
+    def test_paths_only_stimulus_set_error_names_the_column_to_add(self):
+        import pandas as pd
+        from brainscore_core.supported_data_standards.brainio.stimuli import StimulusSet
+        m = BrainScoreModel(identifier='test', extractor={'vision': lambda stimuli, layers: None})
+        stimuli = StimulusSet(pd.DataFrame({'stimulus_id': ['a']}))
+        stimuli.stimulus_paths = {'a': '/tmp/a.png'}
+        with pytest.raises(ValueError, match=r"No column.*image_path.*stimulus_paths dict alone"):
             m.process(stimuli)
 
     def test_single_modality_detected_from_multi_preprocessor_model(self):