wip

Author: joecummings

torchtext/prototype/generate.py

@@ -136,6 +136,7 @@ def greedy_search(
             model_inputs = self.model.prepare_inputs_for_generation(input_ids, **model_kwargs)
             if self.is_huggingface_model:
                 model_inputs["return_dict"] = True
+                model_inputs["use_cache"] = True
                 model_inputs["output_hidden_states"] = True
 
             # Get model output
@@ -261,11 +262,15 @@ def update_func(emissions, N, T, prev_step_token_idxs, prev_step_model_states, t
 
                 # Forward pass
                 model_inputs = self.model.prepare_inputs_for_generation(state_and_tokens, **new_model_kwargs)
+                print(model_inputs.get("use_cache"), model_inputs.get("past_key_values"))
 
                 if self.is_huggingface_model:
                     model_inputs["return_dict"] = True
+                    model_inputs["use_cache"] = True
                     model_inputs["output_hidden_states"] = True
 
+                print(model_inputs.get("use_cache"), model_inputs.get("past_key_values"))
+
                 outputs = self.model(**model_inputs)
                 output_key = "logits" if self.is_huggingface_model else "decoder_output"
                 lm_scores = outputs[output_key]