Only decode new tokens

Author: joecummings

torchtext/prototype/generate.py

@@ -77,7 +77,7 @@ def _prepare_encoder_decoder_kwargs_for_generation(self, inputs: torch.Tensor) -
 
         # 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)
+        model_kwargs = {"encoder_outputs": encoder.forward(inputs, **encoder_kwargs)}
         return model_kwargs
 
     def _prepare_decoder_ids_for_generation(
@@ -281,7 +281,7 @@ def update_func(emissions, N, T, prev_step_token_idxs, prev_step_hyp_idxs, prev_
                     ), f"state_and_tokens has shape {state_and_tokens.shape} = expected {(num_samples, timestep + 1)}"
                 else:
                     assert len(prev_model_state_sequences) == 1
-                    state_and_tokens = prev_model_state_sequences[0]  # dims: [1, 1]
+                    state_and_tokens = token_indices = prev_model_state_sequences[0]  # dims: [1, 1]
 
                 # Cleanup -- combine this with the above
                 if self.is_encoder_decoder:
@@ -292,13 +292,13 @@ def update_func(emissions, N, T, prev_step_token_idxs, prev_step_hyp_idxs, prev_
                     )
 
                 # Preprocess inputs for generation
-                model_inputs = self.model.prepare_inputs_for_generation(state_and_tokens, **new_model_kwargs)
+                model_inputs = self.model.prepare_inputs_for_generation(token_indices, **new_model_kwargs)
                 if self.is_huggingface_model:
                     model_inputs.update(self._huggingface_model_input_values)
                     if len(prev_step_hyp_idxs) > 1 and model_inputs["past_key_values"] is not None:
                         model_inputs["past_key_values"] = self.model._reorder_cache(
                             model_inputs["past_key_values"],
-                            torch.Tensor(prev_step_hyp_idxs).to(dtype=torch.int32),  # I think this is correct?
+                            torch.Tensor(prev_step_hyp_idxs).to(dtype=torch.int32),
                         )
 
                 # Forward pass