Fix batch beam search

modeling_glm.py

@@ -873,6 +873,16 @@ class GLMForConditionalGeneration(GLMPreTrainedModel):
                 position_ids = position_ids[:, :, :seq_length]
             if attention_mask is not None:
                 attention_mask = attention_mask[:, :, :seq_length, :seq_length]
+        if position_ids is not None and input_ids.size(0) > position_ids.size(0):
+            batch_size = position_ids.size(0)
+            num_beams = input_ids.size(0) // batch_size
+            position_ids = position_ids.unsqueeze(1).expand(-1, num_beams, -1, -1)
+            position_ids = position_ids.reshape(batch_size * num_beams, *position_ids.shape[-2:])
+        if attention_mask is not None and input_ids.size(0) > attention_mask.size(0):
+            batch_size = attention_mask.size(0)
+            num_beams = input_ids.size(0) // batch_size
+            attention_mask = attention_mask.unsqueeze(1).expand(-1, num_beams, -1, -1, -1)
+            attention_mask = attention_mask.reshape(batch_size * num_beams, *attention_mask.shape[-3:])
         return {
             "input_ids": input_ids,
             "position_ids": position_ids,