Repository: huggingface/transformers (Python) — Bug Fix

Issue:
RuntimeError with prefix_allowed_tokens_fn and do_sample=True When Allowed Tokens List is Empty
### System Info

- `transformers` version: 4.36.0.dev0
- Platform: macOS-13.4.1-arm64-arm-64bit
- Python version: 3.10.13
- Huggingface_hub version: 0.19.3
- Safetensors version: 0.3.3
- Accelerate version: 0.23.0
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.0.1 (False)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@gante 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```python
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

def main():
    model_id = "gpt2"
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    tokenizer.pad_token = tokenizer.eos_token
    model = AutoModelForCausalLM.from_pretrained(model_id)
    prefix = "Hello"

    input_ids = tokenizer([prefix], add_special_tokens=False, return_tensors="pt", padding=True)["input_ids"]

    def empty_prefix_allowed_tokens_fn(batch_id, sent):
        return []

    try:
        output = model.generate(
            input_ids,
            do_sample=False,
            max_length=10,
            num_beams=1,
            prefix_allowed_tokens_fn=empty_prefix_allowed_tokens_fn,
            num_return_sequences=1
        )
        generations = tokenizer.batch_decode(output, skip_special_tokens=True)
        print(generations)
    except RuntimeError as e:
        print("RuntimeError encountered:", e)

if __name__ == '__main__':
    main()
```
While the example above may seems to be very artificial, but actually this is a very common problem. Suppose you want to use LLM to generate a json object, c.f. #27557, and the constraints will return an empty set of allowed tokens once the json object is complete such as `{"ip": "127.0.0.1"}`, then this issue will occur.

### Expected behavior

I expect the sampling should behave like in greedy search, i.e. when the allowed token list is empty, the model will return PAD instead of throwing an error. 


Related issues
- [model.generate with prefix_allowed_tokens_fn throws RuntimeError: probability tensor contains either inf, nan or element < 0 #15169](https://github.com/huggingface/transformers/issues/15169)
- [RunTimeError when using prefix_allowed_tokens_fn and top-k/top-p sampling in model.generate #13707](https://github.com/huggingface/transformers/issues/13707)

The above two issues are about the same problem as ours, but no fix has been done so far. 





Discussion:
Hey! Thanks for reporting, this would require use to check that the output of `self._prefix_allowed_tokens_fn(batch_id, sent)` on each token is not `[]` before applying the mask. It does make sense because we don't specify that the list cannot be empty. 
Would you like to open a PR for a fix? (meaning something like:
```diff 

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        mask = torch.full_like(scores, -math.inf)
        for batch_id, beam_sent in enumerate(input_ids.view(-1, self._num_beams, input_ids.shape[-1])):
            for beam_id, sent in enumerate(beam_sent):
-                mask[batch_id * self._num_beams + beam_id, self._prefix_allowed_tokens_fn(batch_id, sent)] = 0
+                allowed_tokens = self._prefix_allowed_tokens_fn(batch_id, sent)
+                if len(allowed_tokens) > 0:
+                mask[batch_id * self._num_beams + beam_id, allowed_tokens] = 0

      return scores + mask

```
Ok, I will raise a fix for it. 
Hi there @Saibo-creator! 👋 (cc @ArthurZucker)

I echo Patrick's comment in a related issue [here](https://github.com/huggingface/transformers/issues/15169#issuecomment-1018617055): this is not an issue of `generate`, but rather an issue of the user-defined `prefix_allowed_tokens_fn`. As such, we shouldn't attempt to fix the problem for the user, as it might result in unexpected behavior. 

Instead, we should raise an informative exception: the user has set an unfeasible set of constraints in `prefix_allowed_tokens_fn` :)
yep raising an exception sounds good as well! Less silent changes on our side 🤗 
Hey, thanks for your responses. 
I agree with you about rasing an execption. 

I would like to summarize and double check  that our motivations align.

As the example I gave above, a widely demanded usage of LLM is to generate json object reliably.

Force LLM to generate structured objects like json means we need to constrain the LLM's generation process, including let it stop when necessary, e.g. when the json object is complete. 

There we can have two ways to handle it:
1. the constraints such as `prefix_allowed_tokens_fn` should return an allowed set of `{EOS}` instead of `{}` to let the generation stop
2. the constraints such as `prefix_allowed_tokens_fn` should return  an empty set of tokens `{}`

From discussion above, it seems everyone agree that the option 1 is better and option 2 should be considered an exception.

Do you agree ?





@Saibo-creator agreed. 

I would even go beyond the case of constrained generation here: if, at any part of text generation, the set of possible tokens is `{}`, then there is a logical error. We have been adopting this idea in other parts of text generation :)

Resolving patch:
diff --git a/src/transformers/generation/logits_process.py b/src/transformers/generation/logits_process.py
--- a/src/transformers/generation/logits_process.py
+++ b/src/transformers/generation/logits_process.py
@@ -1071,7 +1071,14 @@ def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> to
         mask = torch.full_like(scores, -math.inf)
         for batch_id, beam_sent in enumerate(input_ids.view(-1, self._num_beams, input_ids.shape[-1])):
             for beam_id, sent in enumerate(beam_sent):
-                mask[batch_id * self._num_beams + beam_id, self._prefix_allowed_tokens_fn(batch_id, sent)] = 0
+                prefix_allowed_tokens = self._prefix_allowed_tokens_fn(batch_id, sent)
+                if len(prefix_allowed_tokens) == 0:
+                    raise ValueError(
+                        f"`prefix_allowed_tokens_fn` returned an empty list for batch ID {batch_id}."
+                        f"This means that the constraint is unsatisfiable. Please check your implementation"
+                        f"of `prefix_allowed_tokens_fn` "
+                    )
+                mask[batch_id * self._num_beams + beam_id, prefix_allowed_tokens] = 0
 
         return scores + mask
 


Test patch:
diff --git a/tests/generation/test_logits_process.py b/tests/generation/test_logits_process.py
--- a/tests/generation/test_logits_process.py
+++ b/tests/generation/test_logits_process.py
@@ -610,6 +610,13 @@ def prefix_allowed_tokens_fn(batch_id, inputs_ids):
             torch.isinf(filtered_scores).tolist(), [[False, False, True, True, True], [True, True, False, False, True]]
         )
 
+        def empty_prefix_allowed_tokens_fn(batch_id, inputs_ids):
+            return []
+
+        prefix_constrained_logits_proc = PrefixConstrainedLogitsProcessor(empty_prefix_allowed_tokens_fn, 1)
+
+        self.assertRaises(ValueError, prefix_constrained_logits_proc, input_ids, scores.clone())
+
     def test_hamming_diversity(self):
         vocab_size = 4
         num_beams = 2