token healing impl
Author: ahmed-moubtahij
Pull request: yes

# What does this PR do?

Token healing rectifies the token boundary bias in greedy tokenization. It does this by trimming and regrowing the prompt to better align with the model's tokenizer, thus enhancing generation quality. The improvement is clearest with completion models.

Token boundary bias is a silent performance killer that doesn't seem very well known. It has clear impact on completion quality.

A more thorough explanation of the problem: [The Art of Prompt Design: Prompt Boundaries and Token Healing | by Scott Lundberg](https://towardsdatascience.com/the-art-of-prompt-design-prompt-boundaries-and-token-healing-3b2448b0be38).

### Motivation

Given a completion prompt with a partial url ending with `:`, the model might have seen the expected completion `://` as a _single_ token in training. However, the prompt's tail token `:` tells it that the next token is not `//`, and so it generates a wrong completion. Such errors compound in auto-regressive language models.


Fixes [#28346](https://github.com/huggingface/transformers/issues/28346)


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?


## Who can review?

- @gante 

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_29081). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

CI is failing due to an automatic update in the pytest package, we are tracking it. Will let you know when it is sorted -- it will need a rebase 

> CI is failing due to an automatic update in the pytest package, we are tracking it. Will let you know when it is sorted -- it will need a rebase

Thanks for the follow-up!

@Ayenem `main` is fixed, rebasing should make CI green except if there are PR-specific issues :)

![image](https://github.com/huggingface/transformers/assets/50707385/96063985-5688-4225-8e88-f6a8d55a84c7)
In case it's relevant, here are (some) listed remotes with `git branch -r`:
```
  origin/HEAD -> origin/main
  origin/heal_tokens
  origin/main
  origin/token_healing
  upstream/'delete-delete-doc'
  upstream/BritneyMuller-housekeeping-patch
  upstream/_dummy_fix_weight_only_usage
  upstream/_dummy_fix_weight_only_usage_2
  upstream/add-chat-glm
  upstream/add-deci-lm
  upstream/add-encode-special-tokens
  upstream/add-flash-decoding
  upstream/add-mamba
  upstream/add-prefix-space
  upstream/add-quantization-workflow
  ```

(@Ayenem we're trying to fix the merge conflicts for you, and we're experimenting with a few GH permissions on our side. You may see a few test commits 🤗 )

Now rebased after #29320 was merged, which was causing the last set of errors seen here. If everything went well, we should see a green CI here 🤞 

@Ayenem FYI, I've reverted the tokenizer input to your original suggestion (tokenizer passed to `generate`), after a discussion I had with @Rocketknight1. That way, the input is standardized and matches another incoming PR (https://github.com/huggingface/transformers/pull/28932) 🤗 

> @Ayenem FYI, I've reverted the tokenizer input to your original suggestion (tokenizer passed to `generate`), after a discussion I had with @Rocketknight1. That way, the input is standardized and matches another incoming PR (#28932) 🤗

It does feel better to offload the tokenizer choice and loading to the caller. Thanks again for following up on this 🙏

CI is green! It was possible :')

ping @ArthurZucker :)

Sorry for the late review on it!

Hi @Ayenem , thanks for this feature. I was curious to look into this as an early feature to see how this works on my domain data but had some issues with some generation using the example data provided (stacktrace attached below). Could you share some example script how to test it?
```
Traceback (most recent call last):
  File "token_heal.py", line 33, in <module>
    output = model.generate(
  File "/home/leonardo/.local/lib/python3.8/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/leonardo/projects/transformers/src/transformers/generation/utils.py", line 1439, in generate
    input_ids = self.heal_tokens(input_ids, tokenizer)
  File "/home/leonardo/projects/transformers/src/transformers/generation/utils.py", line 1861, in heal_tokens
    seq_bias[(tail_id,)] += 1.0
KeyError: (518,)
```

Environment used:
- **transformers**: I'm checked out at the head of your fork, this specific [commit](https://github.com/Ayenem/transformers/commit/249a8f7ffec06a0387d7148114a7fb9e23b7085e)
- **pygtrie version**: 2.5.0
- **Python version**: 3.8.10
- **Model**: AutoModelForCausalLM("meta-llama/Llama-2-7b-hf")