Skip to content

Commit

Permalink
Added unit tests for bidirectional huggingface models
Browse files Browse the repository at this point in the history
  • Loading branch information
shehadak committed Nov 9, 2023
1 parent a4c52f4 commit bc44421
Showing 1 changed file with 33 additions and 12 deletions.
45 changes: 33 additions & 12 deletions tests/test_model_helpers/test_huggingface.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,36 +11,38 @@


class TestNextWord:
@pytest.mark.parametrize('model_identifier, expected_next_word', [
pytest.param('bert-base-uncased', '.', marks=pytest.mark.memory_intense),
pytest.param('gpt2-xl', 'jumps', marks=pytest.mark.memory_intense),
('distilgpt2', 'es'),
@pytest.mark.parametrize('model_identifier, expected_next_word, bidirectional', [
pytest.param('bert-base-uncased', 'and', True, marks=pytest.mark.memory_intense),
pytest.param('bert-base-uncased', '.', False, marks=pytest.mark.memory_intense),
pytest.param('gpt2-xl', 'jumps', False, marks=pytest.mark.memory_intense),
('distilgpt2', 'es', False),
])
def test_single_string(self, model_identifier, expected_next_word):
def test_single_string(self, model_identifier, expected_next_word, bidirectional):
"""
This is a simple test that takes in text = 'the quick brown fox', and tests the next word.
This test is a stand-in prototype to check if our model definitions are correct.
"""

model = HuggingfaceSubject(model_id=model_identifier, region_layer_mapping={})
model = HuggingfaceSubject(model_id=model_identifier, region_layer_mapping={}, bidirectional=bidirectional)
text = 'the quick brown fox'
_logger.info(f'Running {model.identifier()} with text "{text}"')
model.start_behavioral_task(task=ArtificialSubject.Task.next_word)
next_word = model.digest_text(text)['behavior'].values
assert next_word == expected_next_word

@pytest.mark.parametrize('model_identifier, expected_next_words', [
pytest.param('bert-base-uncased', ['.', '.', '.'], marks=pytest.mark.memory_intense),
pytest.param('gpt2-xl', ['jumps', 'the', 'dog'], marks=pytest.mark.memory_intense),
('distilgpt2', ['es', 'the', 'fox']),
@pytest.mark.parametrize('model_identifier, expected_next_words, bidirectional', [
pytest.param('bert-base-uncased', [';', 'the', 'water'], True, marks=pytest.mark.memory_intense),
pytest.param('bert-base-uncased', ['.', '.', '.'], False, marks=pytest.mark.memory_intense),
pytest.param('gpt2-xl', ['jumps', 'the', 'dog'], False, marks=pytest.mark.memory_intense),
('distilgpt2', ['es', 'the', 'fox'], False),
])
def test_list_input(self, model_identifier, expected_next_words):
def test_list_input(self, model_identifier, expected_next_words, bidirectional):
"""
This is a simple test that takes in text = ['the quick brown fox', 'jumps over', 'the lazy'], and tests the
next word for each text part in the list.
This test is a stand-in prototype to check if our model definitions are correct.
"""
model = HuggingfaceSubject(model_id=model_identifier, region_layer_mapping={})
model = HuggingfaceSubject(model_id=model_identifier, region_layer_mapping={}, bidirectional=bidirectional)
text = ['the quick brown fox', 'jumps over', 'the lazy']
_logger.info(f'Running {model.identifier()} with text "{text}"')
model.start_behavioral_task(task=ArtificialSubject.Task.next_word)
Expand Down Expand Up @@ -173,6 +175,25 @@ def test_one_text_single_target(self):
assert len(representations['neuroid']) == 768
_logger.info(f'representation shape is correct: {representations.shape}')

@pytest.mark.memory_intense
def test_one_text_single_target_bidirectional(self):
"""
This is a simple test that takes in text = 'the quick brown fox', and asserts that a bidirectiona BERT model
layer indexed by `representation_layer` has 1 text presentation and 768 neurons. This test is a stand-in prototype
to check if our model definitions are correct.
"""
model = HuggingfaceSubject(model_id='bert-base-uncased', region_layer_mapping={
ArtificialSubject.RecordingTarget.language_system: 'bert.encoder.layer.4'})
text = 'the quick brown fox'
_logger.info(f'Running {model.identifier()} with text "{text}"')
model.start_neural_recording(recording_target=ArtificialSubject.RecordingTarget.language_system,
recording_type=ArtificialSubject.RecordingType.fMRI)
representations = model.digest_text(text)['neural']
assert len(representations['presentation']) == 1
assert representations['stimulus'].squeeze() == text
assert len(representations['neuroid']) == 768
_logger.info(f'representation shape is correct: {representations.shape}')

@pytest.mark.memory_intense
def test_one_text_two_targets(self):
model = HuggingfaceSubject(model_id='distilgpt2', region_layer_mapping={
Expand Down

0 comments on commit bc44421

Please sign in to comment.