bug in the token healing
Author: alpertunga-bile
Pull request: no
Labels: name: bug

name: Generation

### System Info

- `transformers` version: 4.45.1
- Platform: Windows-11-10.0.22631-SP0
- Python version: 3.12.5
- Huggingface_hub version: 0.24.6
- Safetensors version: 0.4.4
- Accelerate version: 0.34.2
- Accelerate config:    not found
- PyTorch version (GPU?): 2.4.0+cu124 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: no
- Using GPU in script?: yes
- GPU type: NVIDIA GeForce GTX 1660 Ti

### Who can help?

@ArthurZucker @gante 

### 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

I am using the below script to generate outputs. I am using a custom trained GPT2 model.

```python
inputs = self.tokenizer(input, padding=True, return_tensors="pt").to(self.dev)
generated_ids = self.model.generate(
    **inputs,
    **get_variable_dictionary(args),
    pad_token_id=self.tokenizer.eos_token_id,
    renormalize_logits=True,
    token_healing=True,
    tokenizer=self.tokenizer,
)
```

Below code block is filling the **GenerationConfig** parameters:
```python 
**get_variable_dictionary(args)
```

The script is run without issues when ```token_healing``` is disabled. When ```token_healing``` is enabled, this error is occured:

```bash
!!! Exception during processing !!! 'ExtensionsTrie' object has no attribute 'values'
Traceback (most recent call last):
  File "D:\sd\ComfyUI\execution.py", line 323, in execute
    output_data, output_ui, has_subgraph = get_output_data(obj, input_data_all, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb)
                                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\sd\ComfyUI\execution.py", line 198, in get_output_data
    return_values = _map_node_over_list(obj, input_data_all, obj.FUNCTION, allow_interrupt=True, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb) 
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 
  File "D:\sd\ComfyUI\custom_nodes\prompt-generator-comfyui\generator\generate.py", line 102, in generate_multiple_texts
    generated_ids = self.model.generate(
                    ^^^^^^^^^^^^^^^^^^^^
  File "D:\sd\ComfyUI\venv\Lib\site-packages\torch\utils\_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "D:\sd\ComfyUI\venv\Lib\site-packages\transformers\generation\utils.py", line 1882, in generate
    input_ids = self.heal_tokens(input_ids, tokenizer)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\sd\ComfyUI\venv\Lib\site-packages\transformers\generation\utils.py", line 2295, in heal_tokens
    seq_bias = {(alt_tok,): 10.0 for alt_tok in vocab_trie.values(prefix=tail_tok)}
                                                ^^^^^^^^^^^^^^^^^
AttributeError: 'ExtensionsTrie' object has no attribute 'values'
```
I did some changes to the code in [src/transformers/generation/utils.py](https://github.com/huggingface/transformers/blob/main/src/transformers/generation/utils.py#L2295) file at line 2295 and it is working after the below updates:

```python
seq_bias = {(tokenizer.convert_tokens_to_ids(alt_tok),): 10.0 for alt_tok in vocab_trie.extensions(prefix=tail_tok)}
```
As I understand from the exceptions; ```sequence_bias``` need the keys to be integer tuple values, but the older version is given the string tuple as the key. And **ExtensionsTrie** doesn't have values function but have extensions function.

I can't be sure if the error is a general one because I saw that there are already tests for ```token_healing``` and it passed those tests.

### Expected behavior

When using ```token_healing``` variable, the program has to be not terminated with error.

Comments:
I am trying to create PR for the bug but I have encountered with errors and have questions about them.

First of all, I think I figured it out why this bug is not occured in the tests. Because the test requires ```auto-gptq``` package and when it is not installed, it just skips the tests. As I checked from the [CircleCI link](https://app.circleci.com/pipelines/github/huggingface/transformers/106008/workflows/76361f4c-5548-4dde-a997-5a56fb255fbe/jobs/1407893), the environment is not installing the ```auto-gptq``` package.

After installing ```auto-gptq``` and ```optimum``` packages, I can finally run the tests. But I didn't understand why the test is using [TheBloke/deepseek-llm-7B-base-GPTQ](https://huggingface.co/TheBloke/deepseek-llm-7B-base-GPTQ) rather than using a small model like [distilbert/distilgpt2](https://huggingface.co/distilbert/distilgpt2) for testing tokenizer functionality. Assuming that the tests can be passed with that model but I tried it in Google Colab and it can't pass the tests so I didn't understand its purpose. As I understand from the [article](https://towardsdatascience.com/the-art-of-prompt-design-prompt-boundaries-and-token-healing-3b2448b0be38), token healing can be applied to any tokenizer. So I used smaller model to work easily and tried the tests.

I had to refactor some parts of the code to pass the specified token healing tests. After the final changes, it is failed only on one test which is ```An example ["like this"] and another example [```. I have question for this one too. Isn't it wrong to write a test about a specific case for a specific tokenizer model? Because some of the tokenizers might not include quote after square bracket word in their vocabulary. Of course the same question will be asked to the url test case which is ```The link is <a href="http:```. I am asking this question because the ExtensionsTree data structure is not checking if the character is in the node but assumes it exists already and this behaviour leads to failing 3 tests but when the check is applied, it is failing on 1 test.