fix next-tokens after update_model_kwargs_for_generation (#6919)

PaddlePaddle · Sep 5, 2023 · 3a9d41b · 3a9d41b
1 parent 19a9207
commit 3a9d41b
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/paddlenlp/experimental/transformers/generation_utils.py b/paddlenlp/experimental/transformers/generation_utils.py
@@ -209,6 +209,8 @@ def update_model_kwargs_for_generation(self, cache, just_decoder, next_tokens, e
                 model_kwargs["seq_len_decoder"],
                 model_kwargs["seq_len_decoder"] + 1,
             )
+
+        model_kwargs["next_tokens"] = next_tokens
         return model_kwargs
 
     def sample(
@@ -271,6 +273,7 @@ def _post_process_(outputs, top_p, temperature, step_idx_ori, model_kwargs):
             model_kwargs = self.update_model_kwargs_for_generation(
                 cache, just_decoder, next_tokens, eos_token_id, model_kwargs
             )
+            next_tokens = model_kwargs["next_tokens"]
 
             save_with_output(
                 next_tokens,