diff --git a/brainscore_language/model_helpers/huggingface.py b/brainscore_language/model_helpers/huggingface.py index 6a7fd855..9c1c8091 100644 --- a/brainscore_language/model_helpers/huggingface.py +++ b/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() diff --git a/brainscore_language/model_helpers/localize.py b/brainscore_language/model_helpers/localize.py index a56eea55..95b09897 100644 --- a/brainscore_language/model_helpers/localize.py +++ b/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) diff --git a/brainscore_language/models/qwen3_0.6b/__init__.py b/brainscore_language/models/qwen3_0.6b/__init__.py new file mode 100644 index 00000000..31fb606c --- /dev/null +++ b/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, + } +) \ No newline at end of file diff --git a/brainscore_language/models/qwen3_1.7b/__init__.py b/brainscore_language/models/qwen3_1.7b/__init__.py new file mode 100644 index 00000000..e9fbdeb6 --- /dev/null +++ b/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, + } +) \ No newline at end of file diff --git a/brainscore_language/models/qwen3_14b/__init__.py b/brainscore_language/models/qwen3_14b/__init__.py new file mode 100644 index 00000000..6eeb7b2e --- /dev/null +++ b/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, + } +) \ No newline at end of file diff --git a/brainscore_language/models/qwen3_32b/__init__.py b/brainscore_language/models/qwen3_32b/__init__.py new file mode 100644 index 00000000..f3878e3d --- /dev/null +++ b/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, + } +) \ No newline at end of file diff --git a/brainscore_language/models/qwen3_4b/__init__.py b/brainscore_language/models/qwen3_4b/__init__.py new file mode 100644 index 00000000..1b1959f9 --- /dev/null +++ b/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, + } +) \ No newline at end of file diff --git a/brainscore_language/models/qwen3_8b/__init__.py b/brainscore_language/models/qwen3_8b/__init__.py new file mode 100644 index 00000000..1c39fc54 --- /dev/null +++ b/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, + } +) \ No newline at end of file diff --git a/examples/score_localization.py b/examples/score_localization.py index 781f3939..6acb06f3 100644 --- a/examples/score_localization.py +++ b/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) \ No newline at end of file + 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) \ No newline at end of file diff --git a/tests/test_model_helpers/test_huggingface.py b/tests/test_model_helpers/test_huggingface.py index 3e54c71f..3f080b7d 100644 --- a/tests/test_model_helpers/test_huggingface.py +++ b/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): """ diff --git a/tests/test_model_helpers/test_localize.py b/tests/test_model_helpers/test_localize.py new file mode 100644 index 00000000..f2eaecdc --- /dev/null +++ b/tests/test_model_helpers/test_localize.py @@ -0,0 +1,47 @@ +import numpy as np + +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)