Sampling after LogitsProcessor ignores scores if previous token was `eos`
Author: freckletonj
Pull request: no
Labels: name: Generation

### System Info

# Environment Details

- `transformers` version: 4.41.0.dev0
- Platform: Linux-6.5.0-28-generic-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.19.4
- Safetensors version: 0.4.2
- Accelerate version: 0.25.0
- Accelerate config: 	- compute_environment: LOCAL_MACHINE
	- distributed_type: MULTI_GPU
	- mixed_precision: bf16
	- use_cpu: False
	- debug: False
	- num_processes: 2
	- machine_rank: 0
	- num_machines: 1
	- gpu_ids: all
	- rdzv_backend: static
	- same_network: True
	- main_training_function: main
	- downcast_bf16: no
	- tpu_use_cluster: False
	- tpu_use_sudo: False
	- tpu_env: []
- PyTorch version (GPU?): 2.2.0.dev20230926+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed

### Who can help?

@ArthurZucker , @younesbelkad

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

# Problem Details

I've created a little script that allows a `LogitsProcessor` to inject a `sentinel_token` after a `stop_token` has been seen.  This helps me filter out things beyond the `sentinel_token`.

Those `stop_token`s can be any valid token I want to filter beyond. But `
`, or, `tokenizer.eos_token_id` do not work. If I `breakpoint()` the scores in the `LogitsProcessor`, it clearly is setting every score to `-inf` besides my `sentinel_token`, which post sampling should mean that my `sentinel_token` is returned. This is not the case, it ignores the scores that I set in `LogitsProcessor`.

Works: results for batch of 2 if `stop_tokens = ['.', '."']` and `sentinel_token = '%'`

```
Echo.% ...
"J."% ...
```

Broken: results for batch of 2 if `stop_tokens = ['
']` and `sentinel_token = '%'`. Expected results would look like the above working example.

```
Echo.

"J."

```
This problem is expressly manifest in bigger batches with longer new token limits, where it will generate garbage that I need the `sentinel_token` to help filter out, eg it will go on with `




made up stuff etc.`. 

Is `model.generate` handling `eos` in some special undisclosed way? 

# Reproducible Example

```python
import torch
import os
from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor
from typing import List

# load model
model_path = os.path.expanduser("~/models/mistral/zephyr-7b-beta/")
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_path,
                                             torch_dtype=torch.float32,
                                             device_map="auto",
                                             trust_remote_code=False).eval()
tokenizer.pad_token = tokenizer.eos_token
tokenizer.pad_token_id = tokenizer.eos_token_id
tokenizer.padding_side = 'left'

class StopAfterTokenIsGenerated(LogitsProcessor):
    def __init__(self, base_len: int, stop_token_ids: List[int], sentinel_id: int):
        super().__init__()
        self.base_len = base_len
        self.stop_token_ids = torch.tensor(stop_token_ids)
        self.sentinel_id = sentinel_id  # insert after a stop_token is encountered

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        self.stop_token_ids = self.stop_token_ids.to(device=input_ids.device)
        if input_ids.size(1) > self.base_len:
            stop_mask = torch.any(input_ids[:, -1:] == self.stop_token_ids, dim=1)  # shape: [batch]
            scores[stop_mask] = torch.full_like(scores[stop_mask], -float("inf"))
            scores[stop_mask, self.sentinel_id] = 0
        return scores

prompts = [
    '

\n

\nsay a single word, then stop talking. no explanations.
\n
\n',
    '

\n

\nwrite the worlds shortest sentence, no explanations.
\n
\n',
]

# DOES NOT WORK
stop_tokens = [tokenizer.convert_tokens_to_ids('
')]

# WORKS
stop_tokens = [tokenizer.convert_tokens_to_ids('.'), tokenizer.convert_tokens_to_ids('."')]

sentinel = tokenizer.convert_tokens_to_ids('%')

tokens = tokenizer(prompts, return_tensors="pt", padding=True, add_special_tokens=False).to(model.device)

# logits processor
stop = StopAfterTokenIsGenerated(tokens['input_ids'].size(1), stop_tokens, sentinel)

output = model.generate(
    **tokens,
    max_new_tokens=32,
    do_sample=False,
    pad_token_id=tokenizer.pad_token_id,
    logits_processor=[stop, ],
)

for x in output:
    print('----------')
    print(tokenizer.decode(x))
```

### Expected behavior

Even if an `eos` token was seen, I should be able to manipulate the `scores` via a `LogitsProcessor` for the tokens that follow.

Comments:
@freckletonj hi! 

Yes, the `eos` token is handled differently by a stoppingCriteria and the generation runs in a loop until `eos` is generated. You can overcome it by removing the `eos_token_id` from model's config. But in that case you need another stopping criteria, for example (max_new_tokens=30), otherwise the loop will forever.

So a short snippet like this should work and the generation stops exactly after 32nd new token 

```python
# set eos tokens to None so that the generation goes on even after eos
model.config.eos_token_id = None
model.generation_config.eos_token_id = None
model.config.forced_eos_token_id = None

stop = StopAfterTokenIsGenerated(tokens['input_ids'].size(1), stop_tokens, sentinel)

output = model.generate(
    **tokens,
    max_new_tokens=32,
    do_sample=False,
    pad_token_id=tokenizer.pad_token_id,
    logits_processor=[stop, ],
)
```
 

Ah, that makes sense thank you!