Move test under integration b/c relies on T5

pytorch · Feb 17, 2023 · 2658bef · 2658bef
1 parent 4bc5321
commit 2658bef
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/...htext_unittest/prototype/test_generate.py → ...egration_tests/prototype/test_generate.py b/...htext_unittest/prototype/test_generate.py → ...egration_tests/prototype/test_generate.py
@@ -28,7 +28,7 @@ def setUp(self) -> None:
     def test_greedy_generate_with_t5(self) -> None:
         generation_model = GenerationUtils(self.model)
 
-        tokens = generation_model.generate(self.inputs, num_beams=1, max_len=30)
+        tokens = generation_model.generate(self.inputs, num_beams=1, max_length=30)
         generated_text = self.transform.decode(tokens.tolist())
 
         expected_generated_text = [

diff --git a/torchtext/prototype/generate.py b/torchtext/prototype/generate.py
@@ -36,8 +36,8 @@ class GenerationUtils:
 
     def __init__(self, model: nn.Module, **kwargs) -> None:
         self.model = model
-        self.is_encoder_decoder = kwargs.pop("is_encoder_decoder")
-        self.is_huggingface_model = kwargs.pop("is_huggingface_model")
+        self.is_encoder_decoder = kwargs.pop("is_encoder_decoder", True)
+        self.is_huggingface_model = kwargs.pop("is_huggingface_model", False)
 
     def _prepare_decoder_ids_for_generation(
         self, batch_size: int, pad_idx: int = 0, device: Optional[torch.device] = None, **model_kwargs