Improve localization: ratio-based top_k, BF16 fix, and rename localize_fed10

brainscore_language/model_helpers/huggingface.py

@@ -17,7 +17,7 @@
 from brainscore_language.artificial_subject import ArtificialSubject
 from brainscore_language.model_helpers.preprocessing import prepare_context
 from brainscore_language.utils import fullname
-from brainscore_language.model_helpers.localize import localize_fed10
+from brainscore_language.model_helpers.localize import localize_fedorenko2010
 
 
 class HuggingfaceSubject(ArtificialSubject):
@@ -92,7 +92,7 @@ def __init__(
 
         if self.use_localizer:
             layer_names = region_layer_mapping["language_system"]
-            self.language_mask = localize_fed10(model_id=self.model_id, 
+            self.language_mask = localize_fedorenko2010(model_id=self.model_id, 
                 model=self.basemodel, 
                 tokenizer=self.tokenizer, 
                 layer_names=layer_names,
@@ -308,7 +308,7 @@ def _setup_hooks(self):
     def output_to_representations(self, layer_representations: Dict[Tuple[str, str, str], np.ndarray], stimuli_coords):
         representation_values = np.concatenate([
             # Choose to use last token (-1) of values[batch, token, unit] to represent passage.
-            values[:, -1:, :].squeeze(0).cpu() for values in layer_representations.values()],
+            self._tensor_to_numpy(values[:, -1:, :].squeeze(0)) for values in layer_representations.values()],
             axis=-1)  # concatenate along neuron axis
         neuroid_coords = {
             'layer': ('neuroid', np.concatenate([[layer] * values.shape[-1]
@@ -404,4 +404,4 @@ def hook_function(_layer: torch.nn.Module, _input, output: torch.Tensor, key=key
         return hook
 
     def _tensor_to_numpy(self, tensor: torch.Tensor) -> np.ndarray:
-        return tensor.cpu().data.numpy()
+        return tensor.float().cpu().data.numpy()

brainscore_language/model_helpers/localize.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Tuple
 from collections import OrderedDict
 
 import os
@@ -124,14 +124,14 @@ def extract_representations(
         batch_rand_actv = extract_batch(model, non_words_tokens["input_ids"], non_words_tokens["attention_mask"], layer_names)
 
         for layer_name in layer_names:
-            final_layer_representations["sentences"][layer_name][batch_idx*batch_size:(batch_idx+1)*batch_size] = torch.stack(batch_real_actv[layer_name]).numpy()
-            final_layer_representations["non-words"][layer_name][batch_idx*batch_size:(batch_idx+1)*batch_size] = torch.stack(batch_rand_actv[layer_name]).numpy()
+            final_layer_representations["sentences"][layer_name][batch_idx*batch_size:(batch_idx+1)*batch_size] = torch.stack(batch_real_actv[layer_name]).float().numpy()
+            final_layer_representations["non-words"][layer_name][batch_idx*batch_size:(batch_idx+1)*batch_size] = torch.stack(batch_rand_actv[layer_name]).float().numpy()
 
     return final_layer_representations
 
-def localize_fed10(model_id: str,
+def localize_fedorenko2010(model_id: str,
     model: torch.nn.Module, 
-    top_k: int, 
+    top_k: int|float, 
     tokenizer: transformers.PreTrainedTokenizer, 
     hidden_dim: int, 
     layer_names: List[str], 
@@ -142,7 +142,11 @@ def localize_fed10(model_id: str,
     Localize the model by selecting the top `top_k` units.
     """
 
-    save_path = f"{BRAINIO_CACHE}/{model_id}_language_mask.npy"
+    model_id_path = model_id.replace("/", "_")
+    if not os.path.exists(BRAINIO_CACHE):
+        os.makedirs(BRAINIO_CACHE, exist_ok=True)
+
+    save_path = f"{BRAINIO_CACHE}/{model_id_path}_language_mask.npy"
 
     if os.path.exists(save_path):
         logger.debug(f"Loading language mask from {save_path}")
@@ -163,6 +167,11 @@ def localize_fed10(model_id: str,
     def is_topk(a, k=1):
         _, rix = np.unique(-a, return_inverse=True)
         return np.where(rix < k, 1, 0).reshape(a.shape)
+
+    if isinstance(top_k, float) and 0 < top_k < 1:
+        top_k = int(top_k * np.prod(t_values_matrix.shape))
+    else:
+        top_k = int(top_k)
 
     language_mask = is_topk(t_values_matrix, k=top_k)
 

brainscore_language/models/qwen3_0.6b/__init__.py

@@ -0,0 +1,25 @@
+from brainscore_language import model_registry
+from brainscore_language import ArtificialSubject
+from brainscore_language.model_helpers.huggingface import HuggingfaceSubject
+from transformers import AutoConfig
+
+# Qwen3-0.6B
+batch_size = 4
+top_k = 0.1 # take 10% of total number of units
+model_name = 'Qwen/Qwen3-0.6B'
+
+config = AutoConfig.from_pretrained(model_name)
+num_blocks = config.num_hidden_layers 
+hidden_size = config.hidden_size
+
+layer_names = [f'model.layers.{block}' for block in range(num_blocks)]
+
+model_registry['qwen3-0.6b-loc-10'] = lambda: HuggingfaceSubject(model_id=model_name, 
+    region_layer_mapping={ArtificialSubject.RecordingTarget.language_system: layer_names},
+    use_localizer=True,
+    localizer_kwargs={
+        'hidden_dim': hidden_size,
+        'batch_size': batch_size,
+        "top_k": top_k,
+    }
+)

brainscore_language/models/qwen3_1.7b/__init__.py

@@ -0,0 +1,25 @@
+from brainscore_language import model_registry
+from brainscore_language import ArtificialSubject
+from brainscore_language.model_helpers.huggingface import HuggingfaceSubject
+from transformers import AutoConfig
+
+# Qwen3-1.7B
+batch_size = 4
+top_k = 0.1 # take 10% of total number of units
+model_name = 'Qwen/Qwen3-1.7B'
+
+config = AutoConfig.from_pretrained(model_name)
+num_blocks = config.num_hidden_layers 
+hidden_size = config.hidden_size
+
+layer_names = [f'model.layers.{block}' for block in range(num_blocks)]
+
+model_registry['qwen3-1.7b-loc-10'] = lambda: HuggingfaceSubject(model_id=model_name, 
+    region_layer_mapping={ArtificialSubject.RecordingTarget.language_system: layer_names},
+    use_localizer=True,
+    localizer_kwargs={
+        'hidden_dim': hidden_size,
+        'batch_size': batch_size,
+        "top_k": top_k,
+    }
+)

brainscore_language/models/qwen3_14b/__init__.py

@@ -0,0 +1,25 @@
+from brainscore_language import model_registry
+from brainscore_language import ArtificialSubject
+from brainscore_language.model_helpers.huggingface import HuggingfaceSubject
+from transformers import AutoConfig
+
+# Qwen3-14B
+batch_size = 4
+top_k = 0.1 # take 10% of total number of units
+model_name = 'Qwen/Qwen3-14B'
+
+config = AutoConfig.from_pretrained(model_name)
+num_blocks = config.num_hidden_layers 
+hidden_size = config.hidden_size
+
+layer_names = [f'model.layers.{block}' for block in range(num_blocks)]
+
+model_registry['qwen3-14b-loc-10'] = lambda: HuggingfaceSubject(model_id=model_name, 
+    region_layer_mapping={ArtificialSubject.RecordingTarget.language_system: layer_names},
+    use_localizer=True,
+    localizer_kwargs={
+        'hidden_dim': hidden_size,
+        'batch_size': batch_size,
+        "top_k": top_k,
+    }
+)

brainscore_language/models/qwen3_32b/__init__.py

@@ -0,0 +1,25 @@
+from brainscore_language import model_registry
+from brainscore_language import ArtificialSubject
+from brainscore_language.model_helpers.huggingface import HuggingfaceSubject
+from transformers import AutoConfig
+
+# Qwen3-32B
+batch_size = 4
+top_k = 0.1 # take 10% of total number of units
+model_name = 'Qwen/Qwen3-32B'
+
+config = AutoConfig.from_pretrained(model_name)
+num_blocks = config.num_hidden_layers 
+hidden_size = config.hidden_size
+
+layer_names = [f'model.layers.{block}' for block in range(num_blocks)]
+
+model_registry['qwen3-32b-loc-10'] = lambda: HuggingfaceSubject(model_id=model_name, 
+    region_layer_mapping={ArtificialSubject.RecordingTarget.language_system: layer_names},
+    use_localizer=True,
+    localizer_kwargs={
+        'hidden_dim': hidden_size,
+        'batch_size': batch_size,
+        "top_k": top_k,
+    }
+)

brainscore_language/models/qwen3_4b/__init__.py

@@ -0,0 +1,25 @@
+from brainscore_language import model_registry
+from brainscore_language import ArtificialSubject
+from brainscore_language.model_helpers.huggingface import HuggingfaceSubject
+from transformers import AutoConfig
+
+# Qwen3-4B
+batch_size = 4
+top_k = 0.1 # take 10% of total number of units
+model_name = 'Qwen/Qwen3-4B'
+
+config = AutoConfig.from_pretrained(model_name)
+num_blocks = config.num_hidden_layers 
+hidden_size = config.hidden_size
+
+layer_names = [f'model.layers.{block}' for block in range(num_blocks)]
+
+model_registry['qwen3-4b-loc-10'] = lambda: HuggingfaceSubject(model_id=model_name, 
+    region_layer_mapping={ArtificialSubject.RecordingTarget.language_system: layer_names},
+    use_localizer=True,
+    localizer_kwargs={
+        'hidden_dim': hidden_size,
+        'batch_size': batch_size,
+        "top_k": top_k,
+    }
+)

brainscore_language/models/qwen3_8b/__init__.py

@@ -0,0 +1,25 @@
+from brainscore_language import model_registry
+from brainscore_language import ArtificialSubject
+from brainscore_language.model_helpers.huggingface import HuggingfaceSubject
+from transformers import AutoConfig
+
+# Qwen3-8B
+batch_size = 4
+top_k = 0.1 # take 10% of total number of units
+model_name = 'Qwen/Qwen3-8B'
+
+config = AutoConfig.from_pretrained(model_name)
+num_blocks = config.num_hidden_layers 
+hidden_size = config.hidden_size
+
+layer_names = [f'model.layers.{block}' for block in range(num_blocks)]
+
+model_registry['qwen3-8b-loc-10'] = lambda: HuggingfaceSubject(model_id=model_name, 
+    region_layer_mapping={ArtificialSubject.RecordingTarget.language_system: layer_names},
+    use_localizer=True,
+    localizer_kwargs={
+        'hidden_dim': hidden_size,
+        'batch_size': batch_size,
+        "top_k": top_k,
+    }
+)

examples/score_localization.py

@@ -1,26 +1,38 @@
-from tqdm import tqdm
 from brainscore_language import load_benchmark
 from brainscore_language.model_helpers.huggingface import HuggingfaceSubject
 from brainscore_language import ArtificialSubject
+from transformers import AutoConfig
 
-benchmark = load_benchmark('Pereira2018.243sentences-linear-shuffle')
+def score_hf_model(model_name: str, benchmark_name: str, batch_size: int=16, top_k: int|float=0.1):
+    config = AutoConfig.from_pretrained(model_name)
+    num_blocks = config.num_hidden_layers if hasattr(config, 'num_hidden_layers') else config.n_layer
+    hidden_size = config.hidden_size if hasattr(config, 'hidden_size') else config.d_model
 
-num_blocks = 12
-layer_names = [f'transformer.h.{block}.{layer_type}' 
-    for block in range(num_blocks) 
-    for layer_type in ['ln_1', 'attn', 'ln_2', 'mlp']
-]
+    layer_names = [f'model.layers.{block}' for block in range(num_blocks)]
 
-model = HuggingfaceSubject(model_id='gpt2', 
-    region_layer_mapping={ArtificialSubject.RecordingTarget.language_system: layer_names},
-    use_localizer=True,
-    localizer_kwargs={
-        'hidden_dim': 768,
-        'batch_size': 16,
-        "top_k": 4096,
-    }
-)
+    model = HuggingfaceSubject(model_id=model_name, 
+        region_layer_mapping={ArtificialSubject.RecordingTarget.language_system: layer_names},
+        use_localizer=True,
+        localizer_kwargs={
+            'hidden_dim': hidden_size,
+            'batch_size': batch_size,
+            "top_k": top_k,
+        }
+    )
 
-model_score = benchmark(model)
+    benchmark = load_benchmark(benchmark_name)
 
-print(model_score)
+    model_score = benchmark(model)
+
+    return model_score
+
+if __name__ == '__main__':
+
+    model_name = 'Qwen/Qwen3-4B'
+    benchmark_name = 'Pereira2018.243sentences-ridge'
+    top_k = 0.1 # take 10% of total number of units
+    batch_size = 16
+
+    model_score = score_hf_model(model_name, benchmark_name, batch_size, top_k)
+
+    print(model_score)

tests/test_model_helpers/test_huggingface.py

@@ -5,6 +5,8 @@
 import torch
 from pytest import approx
 
+from transformers import AutoModelForCausalLM
+
 from brainscore_language.artificial_subject import ArtificialSubject
 from brainscore_language.model_helpers.huggingface import HuggingfaceSubject
 
@@ -158,6 +160,17 @@ def test_list_input(self):
         np.testing.assert_array_equal(representations['stimulus'], text)
         assert len(representations['neuroid']) == 768
 
+    def test_bf16_model(self):
+        # bf16 activations must be cast to float32 before numpy conversion
+        basemodel = AutoModelForCausalLM.from_pretrained('distilgpt2', torch_dtype=torch.bfloat16)
+        model = HuggingfaceSubject(model_id='distilgpt2', model=basemodel, region_layer_mapping={
+            ArtificialSubject.RecordingTarget.language_system: 'transformer.h.0.ln_1'})
+        model.start_neural_recording(recording_target=ArtificialSubject.RecordingTarget.language_system,
+                                     recording_type=ArtificialSubject.RecordingType.fMRI)
+        representations = model.digest_text('the quick brown fox')['neural']
+        assert representations.values.dtype == np.float32
+        assert len(representations['neuroid']) == 768
+
     @pytest.mark.memory_intense
     def test_one_text_single_target(self):
         """

tests/test_model_helpers/test_localize.py

@@ -0,0 +1,64 @@
+import numpy as np
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from brainscore_language.model_helpers import localize
+
+
+def _fake_representations(layer_names, hidden_dim, n_samples=8, seed=0):
+    rng = np.random.RandomState(seed)
+    sentences = {name: rng.randn(n_samples, hidden_dim) + 1 for name in layer_names}
+    non_words = {name: rng.randn(n_samples, hidden_dim) for name in layer_names}
+    return {"sentences": sentences, "non-words": non_words}
+
+
+def _run_localize(monkeypatch, cache_dir, top_k, layer_names, hidden_dim,
+                  model_id="Qwen/Qwen3-4B"):
+    def fake_extract(*args, **kwargs):
+        return _fake_representations(layer_names, hidden_dim)
+
+    monkeypatch.setattr(localize, "BRAINIO_CACHE", str(cache_dir))
+    monkeypatch.setattr(localize, "extract_representations", fake_extract)
+    return localize.localize_fedorenko2010(
+        model_id=model_id, model=None, top_k=top_k, tokenizer=None,
+        hidden_dim=hidden_dim, layer_names=layer_names, batch_size=4, device="cpu")
+
+
+def test_top_k_ratio_selects_fraction_of_units(monkeypatch, tmp_path):
+    mask = _run_localize(monkeypatch, tmp_path, 0.1, ["layer.0", "layer.1"], 50)
+    assert mask.shape == (2, 50)
+    assert mask.sum() == 10
+
+
+def test_top_k_int_is_absolute_count(monkeypatch, tmp_path):
+    mask = _run_localize(monkeypatch, tmp_path, 7, ["layer.0", "layer.1"], 50)
+    assert mask.sum() == 7
+
+
+def test_mask_cached_to_slash_replaced_path(monkeypatch, tmp_path):
+    mask = _run_localize(monkeypatch, tmp_path, 7, ["layer.0", "layer.1"], 50)
+    assert (tmp_path / "Qwen_Qwen3-4B_language_mask.npy").exists()
+
+    def _fail(*args, **kwargs):
+        raise AssertionError("should load from cache")
+
+    monkeypatch.setattr(localize, "extract_representations", _fail)
+    reloaded = localize.localize_fedorenko2010(
+        model_id="Qwen/Qwen3-4B", model=None, top_k=7, tokenizer=None,
+        hidden_dim=50, layer_names=["layer.0", "layer.1"], batch_size=4, device="cpu")
+    np.testing.assert_array_equal(reloaded, mask)
+
+
+def test_localize_runs_on_bf16_model(monkeypatch, tmp_path):
+    # bf16 activations must be cast before numpy conversion; a completed run proves the fix
+    monkeypatch.setattr(localize, "BRAINIO_CACHE", str(tmp_path))
+    model = AutoModelForCausalLM.from_pretrained("distilgpt2", torch_dtype=torch.bfloat16)
+    tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
+    layer_names = ["transformer.h.0", "transformer.h.1"]
+    hidden_dim = model.config.n_embd
+    mask = localize.localize_fedorenko2010(
+        model_id="distilgpt2", model=model, top_k=0.1, tokenizer=tokenizer,
+        hidden_dim=hidden_dim, layer_names=layer_names, batch_size=8, device="cpu")
+    assert mask.shape == (len(layer_names), hidden_dim)
+    assert set(np.unique(mask)) <= {0, 1}
+    assert mask.sum() == int(0.1 * len(layer_names) * hidden_dim)