fix chatglmv1/2 beam search bug (#7017)

* fix chatglmv1/2 beam search bug

* group beam search fix bug

Author: wtmlon

paddlenlp/generation/utils.py

@@ -1497,6 +1497,7 @@ def _post_process_(outputs, input_ids, cur_len, origin_len, scores, unfinished_f
 
     def reorder_cache(self, cache, beam_idx):
         cache = map_structure(lambda x: paddle.index_select(x, beam_idx), cache)
+        return cache
 
     def beam_search(
         self,
@@ -1626,7 +1627,7 @@ def beam_search(
             cache_name = "cache" if "cache" in model_kwargs else "past_key_values"
             if model_kwargs[cache_name] is not None:
                 # reorder the cache
-                self.reorder_cache(model_kwargs[cache_name], beam_idx)
+                model_kwargs[cache_name] = self.reorder_cache(model_kwargs[cache_name], beam_idx)
 
         pred_ids, scores = beam_scorer.finalize(
             input_ids,
@@ -1774,7 +1775,7 @@ def group_beam_search(
             cache_name = "cache" if "cache" in model_kwargs else "past_key_values"
             if model_kwargs[cache_name] is not None:
                 # reorder the cache
-                self.reorder_cache(model_kwargs[cache_name], beam_idx)
+                model_kwargs[cache_name] = self.reorder_cache(model_kwargs[cache_name], reordering_indices)
 
         pred_ids, scores = beam_scorer.finalize(
             input_ids,

paddlenlp/transformers/chatglm/modeling.py

@@ -25,6 +25,7 @@
 from paddle import Tensor
 from paddle.distributed import fleet
 from paddle.distributed.fleet.utils import recompute
+from paddle.utils import map_structure
 
 from ...utils.env import CONFIG_NAME
 from ...utils.log import logger
@@ -842,6 +843,10 @@ def prepare_inputs_for_generation(
                 "attention_mask": attention_mask,
             }
 
+    def reorder_cache(self, cache: paddle.Tensor, beam_idx):
+        cache = map_structure(lambda x: paddle.index_select(x, beam_idx, axis=1), cache)
+        return cache
+
     def update_model_kwargs_for_generation(
         self,
         outputs,

paddlenlp/transformers/chatglm_v2/modeling.py

@@ -768,6 +768,7 @@ def __init__(self, config: ChatGLMv2Config):
 
     def reorder_cache(self, cache: paddle.Tensor, beam_idx):
         cache = map_structure(lambda x: paddle.index_select(x, beam_idx, axis=1), cache)
+        return cache
 
     def update_model_kwargs_for_generation(
         self,