big test refactor

huggingface · May 4, 2020 · c113844 · c113844
1 parent d84253e
commit c113844
Show file tree

Hide file tree

Showing 2 changed files with 456 additions and 623 deletions.
diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py
@@ -141,6 +141,7 @@ def test_attention_outputs(self):
             self.assertEqual(model.config.output_attentions, True)
             self.assertEqual(model.config.output_hidden_states, False)
             self.assertEqual(len(attentions), self.model_tester.num_hidden_layers)
+
             if chunk_length is not None:
                 self.assertListEqual(
                     list(attentions[0].shape[-4:]),
@@ -648,8 +649,8 @@ def test_lm_head_model_random_no_beam_search_generate(self):
         config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
         input_ids = inputs_dict["input_ids"] if "input_ids" in inputs_dict else inputs_dict["inputs"]
 
-        # max length of input_ids should be < max_length
-        input_ids = input_ids[..., :10]
+        # make sure that input_ids is at most of size 15
+        input_ids = input_ids[..., :15]
 
         # iterate over all generative models
         for model_class in self.all_generative_model_classes:
@@ -693,8 +694,8 @@ def test_lm_head_model_random_beam_search_generate(self):
             torch_device
         )
 
-        # max length of input_ids should be < max_length
-        input_ids = input_ids[..., :10]
+        # make sure that input_ids is at most of size 15
+        input_ids = input_ids[..., :15]
 
         for model_class in self.all_generative_model_classes:
             model = model_class(config).to(torch_device)