Token healing
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 
- @ArthurZucker 

Comments:
@ArthurZucker Follow-up on https://github.com/huggingface/transformers/pull/29081 :
* No more pygtrie dependency, I've added functionality to HF's Trie instead
* A -hopefully- prefix-robust space tokenization.

Tests failing on CI don't seem to be related to my code.

Please let me know if you see anything I can improve.

@LeonardoEmili Follow-up on https://github.com/huggingface/transformers/pull/29081

I have an example script on `examples/research_projects/token-healing/run_token_healing.py` let me know if that helps! Otherwise, please share the string and input_ids that caused the error.

@gante @ArthurZucker CI is green :)

cc @ArthurZucker (to check the tokenization side)

Failing test fixed by https://github.com/huggingface/transformers/pull/30389

Reviewing asap

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

feel free to ping me again for approval! 

> feel free to ping me again for approval!

@ArthurZucker It should be ok now :)

Also make sure to rebase! 

@ArthurZucker Thank you for your review!
I could be wrong but the failing CI tests seem unrelated to my code. Is the convention to periodically fetch & rebase until it is fixed?

CI passed :) @ArthurZucker @gante 

Awesome! 🤗 

I'll do a last pass now! 🤗 

Thanks for you work @ahmed-moubtahij and congrats on the merge 🔥 

> Thanks for you work @ahmed-moubtahij and congrats on the merge 🔥

🥳🥳 Thank you and @gante for the assistance and follow-ups! I learned a lot :)