Token healing (under 40 LOC)
Author: ahmed-moubtahij
Pull request: no
Labels: name: Feature request

### Feature request

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, though I'm not sure where it would fit as a transformers feature.

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.

### Your contribution

My implementation (under 40 LOC): https://github.com/Ayenem/TokenHealer/tree/main

Comments:
FYI @gante 

@Ayenem that's a really cool use of `sequence_bias`! I had read about token healing in the past, and never thought about it when I added `sequence_bias`.

Would you be up for adding it to `transformers`? I think we could include it inside `generate` (e.g. [here](https://github.com/huggingface/transformers/blob/6c78bbcb8320d316434262ef003251ca997db0d1/src/transformers/generation/utils.py#L1638)), passing a simple `token_healing=True` flag. What do you think?

I'm glad you appreciate this use of your feature! Hyrum's law at play :)

Sounds good, I'll get started on a dev branch. Thanks for the pointer!

Hi @Ayenem 👋 

I noticed you closed your latest PR. Is there anything we can help you with? 🤗 

Hi @gante thanks for the follow-up!

I deleted my branch because it was becoming a mess and it automatically closed this PR. I didn't think about it but it makes sense. I think I'll just open a new PR when my changes are done.

I've been working on removing the dependency to pygtrie by adding functionality to HF's [Trie](https://github.com/huggingface/transformers/blob/main/src/transformers/tokenization_utils.py#L52). I've implemented it in [my repo](https://github.com/Ayenem/TokenHealer/blob/hf_trie/tokenhealing/trie.py) and it works.

I added my changes to `tokenization_utils::Trie` and my [Trie tests](https://github.com/Ayenem/TokenHealer/blob/hf_trie/test/test_trie.py) to `test_tokenization_utils::TrieTest` but running the latter throws Attribute Error on my added Trie functions, as if `test_tokenization_utils.py` doesn't see them.

I commented out the entire Trie in `tokenization_utils.py` and successfully ran `test_tokenization_utils`, which is weird. I don't know if `test_tokenization_utils` imported Trie from some cache or some config  or pathing was in play but it ran successfuly with the entire Trie class commented out.

I'll retry from scratch with a new PR 😅 

EDIT:
I made a new fork. Same issue:
![image](https://github.com/huggingface/transformers/assets/50707385/bf654d33-b36b-45c5-b656-9e74a18dc21b)
Left: `TrieTest` in `test_tokenization_utils.py`
Middle: `Trie` class in `tokenization_utils.py`, commented out.
Right: tests passing for `test_tokenization_utils.py` even though the Trie class shouldn't be accessible.

So whenever I add new tests to `TrieTest`, it complains about Attribute Error in the Trie class even when the attributes are actually there.

EDIT 2:
I think I found the issue...`test_tokenization_utils.py` is probably importing Trie from the transformers packages installed in my environment, instead of the Trie from source code. Investigating...

EDIT 3:
Problem solved, re-opened in https://github.com/huggingface/transformers/pull/30081