add gpt2 model and tests

brain-score · Oct 31, 2023 · 79048fc · 79048fc
1 parent 3db0ac2
commit 79048fc
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 4 deletions.
diff --git a/brainscore_language/models/gpt/__init__.py b/brainscore_language/models/gpt/__init__.py
@@ -16,3 +16,6 @@
 
 model_registry['gpt-neo-1.3B'] = lambda: HuggingfaceSubject(model_id='EleutherAI/gpt-neo-1.3B', region_layer_mapping={
     ArtificialSubject.RecordingTarget.language_system: 'transformer.h.18'})
+
+model_registry['gpt2'] = lambda: HuggingfaceSubject(model_id='gpt2', region_layer_mapping={
+    ArtificialSubject.RecordingTarget.language_system: 'transformer.h.11'})
diff --git a/brainscore_language/models/gpt/test.py b/brainscore_language/models/gpt/test.py
@@ -14,7 +14,9 @@
     ('gpt-neo-2.7B', [np.nan, 15.07522869,  3.6358602 ,  0.04999408,  1.42219079,
                       0.0399301 ,  0.02614061,  0.02547451]),
     ('gpt-neo-1.3B', [np.nan, 15.36009979,  5.54412651,  0.11744193,  0.60116327,
-                      0.04266951,  0.08952015,  0.09213546])
+                      0.04266951,  0.08952015,  0.09213546]),
+    ('gpt2', [np.nan, 13.00629139, 12.12623215,  9.79956627,  8.60373306,
+        3.44214535,  1.31260252, 10.23834896])
 ])
 def test_reading_times(model_identifier, expected_reading_times):
     model = load_model(model_identifier)
@@ -29,7 +31,8 @@ def test_reading_times(model_identifier, expected_reading_times):
     ('distilgpt2', ['es', 'the', 'fox']),
     ('gpt2-xl', ['jumps', 'the', 'dog']),
     ('gpt-neo-2.7B', ['jumps', 'the', 'dog']),
-    ('gpt-neo-1.3B', ['jumps', 'the', 'dog'])
+    ('gpt-neo-1.3B', ['jumps', 'the', 'dog']),
+    ('gpt2', ['es', 'the', ','])
 ])
 def test_next_word(model_identifier, expected_next_words):
     model = load_model(model_identifier)
@@ -44,7 +47,8 @@ def test_next_word(model_identifier, expected_next_words):
     ('distilgpt2', 768),
     ('gpt2-xl', 1600),
     ('gpt-neo-1.3B', 2048),
-    ('gpt-neo-2.7B', 2560)
+    ('gpt-neo-2.7B', 2560),
+    ('gpt2', 768)
 ])
 def test_neural(model_identifier, feature_size):
     model = load_model(model_identifier)

diff --git a/tests/test_integration.py b/tests/test_integration.py
@@ -45,7 +45,8 @@
         ('distilgpt2', 'syntaxgym-subordination', approx(0.2173913, abs=.0005)),
         ('distilgpt2', 'syntaxgym-subordination_orc-orc', approx(0.95652174, abs=.0005)),
         ('distilgpt2', 'syntaxgym-subordination_pp-pp', approx(0.47826087, abs=.0005)),
-        ('distilgpt2', 'syntaxgym-subordination_src-src', approx(0.56521739, abs=.0005))
+        ('distilgpt2', 'syntaxgym-subordination_src-src', approx(0.56521739, abs=.0005)),
+        ('gpt2', 'Futrell2018-pearsonr', approx(0.31013720, abs=.0005))
     ]
 )
 def test_score(model_identifier, benchmark_identifier, expected_score):