wip feb 8

Author: joecummings

test/torchtext_unittest/prototype/test_generate.py

@@ -46,7 +46,7 @@ def test_generate_errors_with_incorrect_beams(self) -> None:
         with self.assertRaises(ValueError):
             generation_model.generate(self.inputs, num_beams=0)
 
-    @patch("logging.Logger.warning")
+    @patch("warnings.warn")
     def test_warns_when_no_max_len_provided(self, mock) -> None:
         generation_model = GenerationUtil(self.model)
         generation_model.generate(self.inputs)
@@ -91,12 +91,17 @@ def test_hf_DELETE(self) -> None:
             max_len=100,
             pad_idx=t5.config.pad_token_id,
             num_beams=10,
+            beam_size_token=t5.config.vocab_size,
         )
         end = time.time() - start
         print(t5_tokenizer.batch_decode(tokens, skip_special_tokens=True), end)
         exit()
 
     def test_jit_generate(self) -> None:
+        # jitted_model = torch.jit.script(self.model)
+        # encoder = jitted_model.get_encoder()
+        
+        
         generation_model = GenerationUtil(self.model)
         torch.jit.script(generation_model)
 

torchtext/prototype/generate.py

@@ -12,10 +12,10 @@
     get_obj_from_emitting_model_state,
 )
 
-import logging
 import warnings
 
-logger = logging.getLogger(__name__)
+
+MODEL_KWARGS_TYPE = Dict[str, Dict[str, Union[torch.Tensor, List[Optional[torch.Tensor]], List[torch.Tensor], None]]]
 
 
 @dataclass
@@ -56,9 +56,7 @@ def __init__(self, model: nn.Module, **kwargs) -> None:
         self.is_encoder_decoder = kwargs.pop("is_encoder_decoder", True)
         self.is_huggingface_model = kwargs.pop("is_huggingface_model", False)
 
-    def _prepare_encoder_decoder_kwargs_for_generation(
-        self, inputs: torch.Tensor, model_kwargs: Dict[str, Any]
-    ) -> Dict[str, Any]:
+    def _prepare_encoder_decoder_kwargs_for_generation(self, inputs: torch.Tensor) -> MODEL_KWARGS_TYPE:
         """Runs encoder and adds to model_kwargs for decoding. Modified from https://github.com/huggingface/transformers/blob/67d074874d285e616393c65a0e670088e1b6b74a/src/transformers/generation/utils.py#L592.
 
         Args:
@@ -72,40 +70,36 @@ def _prepare_encoder_decoder_kwargs_for_generation(
         encoder = self.model.get_encoder()
 
         # Create copy of encoder kwargs
-        encoder_kwargs = model_kwargs.copy()
+        encoder_kwargs: Dict[str, bool] = {}
 
-        # Forward pass
         if self.is_huggingface_model:
             encoder_kwargs["return_dict"] = True
 
-        # import pdb
-        # pdb.set_trace()
-        # print(encoder_kwargs.keys())
-
-        # assert torch.jit.isinstance(encoder_kwargs, Optional[Dict[str, bool]])
-
-        model_kwargs["encoder_outputs"] = encoder(inputs, **encoder_kwargs)
-
+        # Forward pass
+        # Explicitly call forward method to assert to assert this is a ScriptModule if JITted
+        model_kwargs = {"encoder_outputs": encoder.forward(inputs)}  # , **encoder_kwargs)
         return model_kwargs
 
     def _prepare_decoder_ids_for_generation(
         self,
         batch_size: int,
         pad_idx: int = 0,
         device: Optional[torch.device] = None,
-        model_kwargs: Optional[Dict[str, Any]] = None,
-    ):
+        model_kwargs: Optional[MODEL_KWARGS_TYPE] = None,
+    ) -> torch.Tensor:
         """Prepare decoder IDs for generation."""
         if model_kwargs is not None and "decoder_input_ids" in model_kwargs:
-            return model_kwargs.pop("decoder_input_ids")
+            decoder_input_ids = model_kwargs.pop("decoder_input_ids")
+            assert torch.jit.isinstance(decoder_input_ids, torch.Tensor)
+            return decoder_input_ids
         else:
             return torch.ones((batch_size, 1), dtype=torch.long, device=device) * pad_idx
 
     def _update_model_kwargs_for_generation(
         self,
         outputs: Dict[str, Any],
         model_kwargs: Dict[str, Any],
-    ) -> Dict[str, Any]:
+    ) -> MODEL_KWARGS_TYPE:
         """After a forward pass, update model_kwargs for faster decoding. Modified from https://github.com/huggingface/transformers/blob/67d074874d285e616393c65a0e670088e1b6b74a/src/transformers/generation/utils.py#L692.
 
         Args:
@@ -152,7 +146,7 @@ def greedy_search(
         max_len: int,
         eos_idx: int,
         pad_idx: Optional[int] = None,
-        model_kwargs: Optional[Dict[str, Any]] = {},
+        model_kwargs: Optional[MODEL_KWARGS_TYPE] = {},
     ) -> torch.Tensor:
         """Greedy search decoding for text generation. Takes the most likely next token every time.
 
@@ -184,9 +178,8 @@ def greedy_search(
             _, next_tokens = torch.topk(probs, 1)
 
             # For any finished sequences, padding idx should be the last token
-            if eos_idx is not None:
-                if pad_idx is not None:
-                    next_tokens = next_tokens * unfinished_sequences + pad_idx * (1 - unfinished_sequences)
+            if eos_idx is not None and pad_idx is not None:
+                next_tokens = next_tokens * unfinished_sequences + pad_idx * (1 - unfinished_sequences)
 
             # Append the next tokens to the previous tokens
             input_ids = torch.cat([input_ids, next_tokens], dim=-1)
@@ -233,7 +226,7 @@ def beam_search(
             encoder_output_key = "last_hidden_state" if self.is_huggingface_model else "encoder_output"
             encoder_output = model_kwargs["encoder_outputs"][encoder_output_key]
 
-        def update_func(emissions, N, T, prev_step_token_idxs, prev_step_model_states, timestep):
+        def update_func(emissions, N, T, prev_step_token_idxs, prev_step_hyp_idxs, prev_step_model_states, timestep):
             # `emissions` and `N` are unused in this current implementation
 
             i = T  # Hacky access to the current seq in inputs
@@ -269,7 +262,7 @@ def update_func(emissions, N, T, prev_step_token_idxs, prev_step_model_states, t
                 if end > curr_beam_size:
                     end = curr_beam_size
 
-                num_samples = end - start  # Is this always just gunna be equal to curr_beam_size?
+                num_samples = end - start
 
                 if prev_step_token_idxs != [-1]:
                     state_sequences = torch.cat(prev_model_state_sequences[start:end], dim=0)
@@ -303,9 +296,6 @@ def update_func(emissions, N, T, prev_step_token_idxs, prev_step_model_states, t
                 if self.is_huggingface_model:
                     model_inputs.update(self._huggingface_model_input_values)
 
-                from typing import MappingProxyType
-
-                model_inputs = MappingProxyType(model_inputs)
                 # Forward pass
                 outputs = self.model(**model_inputs)
 
@@ -315,17 +305,14 @@ def update_func(emissions, N, T, prev_step_token_idxs, prev_step_model_states, t
 
                 # HF optimizations to reduce overhead in future `forward` calls
                 if self.is_huggingface_model:
-                    new_model_kwargs = self._update_model_kwargs_for_generation(
-                        outputs, new_model_kwargs, is_encoder_decoder=self.is_encoder_decoder
-                    )
-                    if new_model_kwargs["past"] is not None:
-                        import pdb
-
-                        pdb.set_trace()
-                        beam_indices += [start for _ in range(num_samples)]
+                    new_model_kwargs = self._update_model_kwargs_for_generation(outputs, new_model_kwargs)
+                    if new_model_kwargs["past"] is not None and len(prev_step_hyp_idxs) > 1:
+                        if len(prev_step_hyp_idxs) == 9:
+                            import pdb
+                            pdb.set_trace()
                         new_model_kwargs["past"] = self.model._reorder_cache(
                             new_model_kwargs["past"],
-                            torch.Tensor(beam_indices).to(dtype=torch.int32),  # I think this is correct?
+                            torch.Tensor(prev_step_hyp_idxs).to(dtype=torch.int32),  # I think this is correct?
                         )
 
                 # Keep track of probabilities over vocab for this pairing
@@ -404,7 +391,7 @@ def is_not_neg_one(elem: int) -> bool:
             return final_tokens_as_tensors
 
         if num_python_workers > 1:
-            logger.warning("Multiprocessing has not yet been implemented.")
+            warnings.warn("Multiprocessing has not yet been implemented.")
 
         all_final_tokens = [beam_decode_step(i) for i in range(len(input_ids))]
 
@@ -473,28 +460,28 @@ def generate(
             1. `num_beams` == 1 or `num_beams` is None -> greedy search
             2. `num_beams` > 1 -> beam search
         """
-        model_kwargs = {}
+        model_kwargs: MODEL_KWARGS_TYPE = {}
 
         if self.is_encoder_decoder:
-            model_kwargs = self._prepare_encoder_decoder_kwargs_for_generation(inputs, model_kwargs)
+            assert torch.jit.isinstance(inputs, torch.Tensor)
+            model_kwargs = self._prepare_encoder_decoder_kwargs_for_generation(inputs)
             inputs = self._prepare_decoder_ids_for_generation(
                 len(inputs), device=inputs.device, model_kwargs=model_kwargs
             )
 
         if max_len is None:
             # Too hard to try to figure out the exact max_seq_length for each model
-            logger.warning("`max_len` was not specified. Defaulting to 256 tokens.")
+            warnings.warn("`max_len` was not specified. Defaulting to 256 tokens.")
             max_len = 256
 
-        if num_beams == 1 or num_beams is None:
+        if num_beams is None or num_beams == 1:
             if num_python_workers > 1:
-                logger.warning(f"Multiprocessing is not implemented for greedy search.")
+                warnings.warn(f"Multiprocessing is not implemented for greedy search.")
             return self.greedy_search(inputs, max_len, eos_idx, pad_idx=pad_idx, model_kwargs=model_kwargs)
         elif num_beams > 1:
             if beam_size_token is None:
                 raise ValueError(
-                    "`beam_size_token` must be specified for beam search. \
-                    If confused about what to put, you can default to the vocab size of the model you are using."
+                    "`beam_size_token` must be specified for beam search. If confused about what to put, you can default to the vocab size of the model you are using."
                 )
             return self.beam_search(
                 inputs,

torchtext/prototype/models/t5/model.py

@@ -134,13 +134,15 @@ def __init__(
             for p in self.parameters():
                 p.requires_grad = False
 
-    def prepare_inputs_for_generation(self, input_ids, encoder_outputs):
+    @torch.jit.export
+    def prepare_inputs_for_generation(self, input_ids: torch.Tensor, encoder_outputs: torch.Tensor) -> Dict[str, torch.Tensor]:
         return {"decoder_tokens": input_ids, "encoder_outputs": encoder_outputs}
 
+    @torch.jit.export
     def get_encoder(self) -> T5Encoder:
         return self.encoder
 
-    @torch.jit.ignore
+    @torch.jit.export
     def get_decoder(self) -> Optional[T5Decoder]:
         if self.decoder is None:
             warnings.warn("Decoder is not set on this model.")