Add sampling support to group beam search

src/transformers/generation/configuration_utils.py

@@ -735,8 +735,6 @@ def validate(self, strict=False):
                     "`diversity_penalty` is not 0.0 or `num_beam_groups` is not 1, triggering group beam search. In "
                     "this generation mode, "
                 )
-                if self.do_sample is True:
-                    raise ValueError(group_error_prefix + "`do_sample` must be set to `False`")
                 if self.num_beams % self.num_beam_groups != 0:
                     raise ValueError(group_error_prefix + "`num_beams` should be divisible by `num_beam_groups`")
                 if self.diversity_penalty == 0.0:

src/transformers/generation/utils.py

@@ -4408,12 +4408,36 @@ def _group_beam_search(
 
                 # reshape for beam search
                 next_token_scores = next_token_scores.view(batch_size, group_size * vocab_size)
-
-                # Sample 1 + len(eos_token_id) next tokens for each beam so we have at least 1 non eos token per beam.
-                n_eos_tokens = eos_token_id.shape[0] if eos_token_id is not None else 0
-                next_token_scores, next_tokens = torch.topk(
-                    next_token_scores, max(2, 1 + n_eos_tokens) * group_size, dim=1, largest=True, sorted=True
-                )
+                if generation_config.do_sample:
+                    warpers = [p for p in logits_processor if getattr(p,'is_warper',False)]
+                    for warper in warpers:
+                        next_token_scores = warper(group_input_ids, next_token_scores)
+                    num_possible_tokens = torch.sum( next_token_scores > -float('inf'), dim = -1).item()
+                    n_eos_tokens = eos_token_id.shape[0] if eos_token_id is not None else 0
+                    next_candidate = max(2,1 + n_eos_tokens)* group_size
+                    mini_possible_tokens = min(
+                        num_possible_tokens,
+                        next_candidate
+                    )
+                    probs = torch.nn.functional.softmax(
+                        next_token_scores,
+                        dim = -1
+                    )
+                    next_tokens = torch.multinomial(
+                        input = probs,
+                        num_samples = int(mini_possible_tokens)
+                    )
+                    next_token_scores = torch.gather(
+                        input = probs,
+                        dim= -1,
+                        index = next_tokens
+                    )
+                else :
+                    # Sample 1 + len(eos_token_id) next tokens for each beam so we have at least 1 non eos token per beam.
+                    n_eos_tokens = eos_token_id.shape[0] if eos_token_id is not None else 0
+                    next_token_scores, next_tokens = torch.topk(
+                        next_token_scores, max(2, 1 + n_eos_tokens) * group_size, dim=1, largest=True, sorted=True
+                    )
 
                 next_indices = torch.div(next_tokens, vocab_size, rounding_mode="floor")
                 next_tokens = next_tokens % vocab_size