format

Author: joecummings

test/torchtext_unittest/prototype/test_generate.py

@@ -57,27 +57,29 @@ def test_warns_when_mp_with_greedy(self, mock) -> None:
 
     def test_beam_search_with_t5_(self) -> None:
         generation_model = GenerationUtil(self.model)
-        tokens = generation_model.generate(self.inputs, num_beams=3, max_len=30, beam_size_token=self.model.config.vocab_size)
+        tokens = generation_model.generate(
+            self.inputs, num_beams=3, max_len=30, beam_size_token=self.model.config.vocab_size
+        )
         generated_text = self.transform.decode(tokens.tolist())
 
         expected_generated_text = [
-            'kate mccartney: a dog is good for you . she says studies have shown that dog ownership is good for',
-            'Das ist gut.',
-            'acceptable',
-            '4.0',
-            'a tornado ripped through a swath of a lake in st. louis . a s'
+            "kate mccartney: a dog is good for you . she says studies have shown that dog ownership is good for",
+            "Das ist gut.",
+            "acceptable",
+            "4.0",
+            "a tornado ripped through a swath of a lake in st. louis . a s",
         ]
 
         self.assertEqual(generated_text, expected_generated_text)
 
-
-
     def test_hf_DELETE(self) -> None:
         from transformers import T5ForConditionalGeneration, T5Tokenizer
         from torchtext.prototype.generate import GenerationUtil
 
         t5 = T5ForConditionalGeneration.from_pretrained("t5-base")
-        test_sequence = ["summarize: studies have shown that owning a dog is good for you"]#, "Q: what is the capital of Alaska?"]
+        test_sequence = [
+            "summarize: studies have shown that owning a dog is good for you"
+        ]  # , "Q: what is the capital of Alaska?"]
         generative_hf_t5 = GenerationUtil(t5, is_encoder_decoder=True, is_huggingface_model=True)
         t5_tokenizer = T5Tokenizer.from_pretrained("t5-base")
         test_sequence_tk = t5_tokenizer(test_sequence, padding=True, return_tensors="pt").input_ids
@@ -89,18 +91,14 @@ def test_hf_DELETE(self) -> None:
             max_len=100,
             pad_idx=t5.config.pad_token_id,
             num_beams=10,
-           
-           
         )
         end = time.time() - start
         print(t5_tokenizer.batch_decode(tokens, skip_special_tokens=True), end)
         exit()
-    
+
     def test_jit_generate(self) -> None:
         generation_model = GenerationUtil(self.model)
         torch.jit.script(generation_model)
-    
 
-    
     def test_beam_search_speed(self) -> None:
         pass

torchtext/prototype/generate.py

@@ -14,6 +14,7 @@
 
 import logging
 import warnings
+
 logger = logging.getLogger(__name__)
 
 
@@ -47,25 +48,23 @@ class GenerationUtil(nn.Module):
     More examples can be found in the `notebooks` directory of this repository.
     """
 
-    _huggingface_model_input_values = {
-        "return_dict": True,
-        "use_cache": True,
-        "output_hidden_states": True
-    }
+    _huggingface_model_input_values = {"return_dict": True, "use_cache": True, "output_hidden_states": True}
 
     def __init__(self, model: nn.Module, **kwargs) -> None:
         super().__init__()
         self.model = model
         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: Dict[str, Any]
+    ) -> Dict[str, Any]:
         """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:
             inputs: (Tensor): Tokenized startings sequence(s).
             model_kwargs (Dict[str, Any]): Model keyword arguments to be modified for decoding.
-        
+
         Returns:
             Modified model_kwargs with addition of encoded input sequence(s).
         """
@@ -78,19 +77,23 @@ def _prepare_encoder_decoder_kwargs_for_generation(self, inputs: torch.Tensor, m
         # 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)
 
         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
+        self,
+        batch_size: int,
+        pad_idx: int = 0,
+        device: Optional[torch.device] = None,
+        model_kwargs: Optional[Dict[str, Any]] = None,
     ):
         """Prepare decoder IDs for generation."""
         if model_kwargs is not None and "decoder_input_ids" in model_kwargs:
@@ -108,7 +111,7 @@ def _update_model_kwargs_for_generation(
         Args:
             outputs (Dict[str, Any]): LM output.
             model_kwargs (Dict[str, Any]): Model keyword args to be modified for future runs.
-        
+
         Returns:
             Modified model_kwargs w/ updated past, token_type_ids, and attention_mask.
         """
@@ -144,7 +147,12 @@ def _update_model_kwargs_for_generation(
         return model_kwargs
 
     def greedy_search(
-        self, input_ids: torch.Tensor, max_len: int, eos_idx: int, pad_idx: Optional[int] = None, model_kwargs: Optional[Dict[str, Any]] = {}
+        self,
+        input_ids: torch.Tensor,
+        max_len: int,
+        eos_idx: int,
+        pad_idx: Optional[int] = None,
+        model_kwargs: Optional[Dict[str, Any]] = {},
     ) -> torch.Tensor:
         """Greedy search decoding for text generation. Takes the most likely next token every time.
 
@@ -217,7 +225,7 @@ def beam_search(
             eos_idx (int): End-of-sequence index.
             num_python_workers (int): Number of python workers to use for multiprocessing.
             model_kwargs
-        
+
         Returns:
             Tensor of the generated sequences.
         """
@@ -227,9 +235,9 @@ def beam_search(
 
         def update_func(emissions, N, T, prev_step_token_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
-            
+
             # Copy over the `model_kwargs` in order to modify
             new_model_kwargs = model_kwargs.copy()
 
@@ -254,18 +262,22 @@ def update_func(emissions, N, T, prev_step_token_idxs, prev_step_model_states, t
                 max_inference_batch_size, 1000 / (timestep + 1)
             )  # many hypotheses will EOS, so increase the batch size gradually
             curr_beam_size = len(prev_step_token_idxs)
-            
+
             # 2. Batched inference to get next tokens
             while start < curr_beam_size:  # catch the remainder
                 end = start + step
                 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  # Is this always just gunna be equal to curr_beam_size?
 
                 if prev_step_token_idxs != [-1]:
                     state_sequences = torch.cat(prev_model_state_sequences[start:end], dim=0)
-                    token_indices = torch.Tensor(prev_step_token_idxs[start:end]).to(dtype=torch.long, device=self.model.device).reshape(num_samples, 1)
+                    token_indices = (
+                        torch.Tensor(prev_step_token_idxs[start:end])
+                        .to(dtype=torch.long, device=self.model.device)
+                        .reshape(num_samples, 1)
+                    )
 
                     state_and_tokens = torch.cat(
                         [state_sequences, token_indices], dim=-1
@@ -303,14 +315,17 @@ 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)
+                    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["past"] = self.model._reorder_cache(
                             new_model_kwargs["past"],
-                            torch.Tensor(beam_indices).to(dtype=torch.int32) # I think this is correct?
+                            torch.Tensor(beam_indices).to(dtype=torch.int32),  # I think this is correct?
                         )
 
                 # Keep track of probabilities over vocab for this pairing
@@ -337,7 +352,7 @@ def update_func(emissions, N, T, prev_step_token_idxs, prev_step_model_states, t
                                 )
                             )
                         )
-                
+
                 start += step
 
             return out_probs, model_states
@@ -392,11 +407,10 @@ def is_not_neg_one(elem: int) -> bool:
             logger.warning("Multiprocessing has not yet been implemented.")
 
         all_final_tokens = [beam_decode_step(i) for i in range(len(input_ids))]
-        
+
         # 5. Return top hypotheses for all input sequences
         return torch.stack(all_final_tokens, dim=0)
 
-
     def forward(
         self,
         inputs: Optional[torch.Tensor] = None,
@@ -460,10 +474,12 @@ def generate(
             2. `num_beams` > 1 -> beam search
         """
         model_kwargs = {}
-        
+
         if self.is_encoder_decoder:
             model_kwargs = self._prepare_encoder_decoder_kwargs_for_generation(inputs, model_kwargs)
-            inputs = self._prepare_decoder_ids_for_generation(len(inputs), device=inputs.device, model_kwargs=model_kwargs)
+            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