Overhaul Conversation class and prompt templating
Author: Rocketknight1
Pull request: yes

Current status: Ready for review!

## Background
We currently have limited support for chat formatting for models. The main tool for this task is the private method `_build_conversation_input_ids`, which tokenizes a `Conversation`. This method is currently defined at the model class level, which causes problems because sometimes we have multiple models in the same class (such as different fine-tunes of `LLaMA`) which use different chat tokens. As a result, I suspect a lot of users are using an incorrect prompt format with their model and are totally unaware of this, which will result in serious silent performance degradation that’s almost impossible to debug.

Additionally, `_build_conversation_input_ids()` is a private method that is only used in `ConversationPipeline`. It is not documented and I suspect most users don't know about it at all. In many cases users are just documenting their model's chat format in the model card or something like that.

## Solution
- ~We add a `PromptConfig` class to store information about the prompt and chat control tokens the model was trained with.~
- ~`PromptConfig` is attached to the `tokenizer`, and is saved/loaded when `save_pretrained` / `load_pretrained` is called. It saves to `prompt_config.json`.~
- `apply_chat_template()` is now a documented public method on `PreTrainedTokenizerBase`. Deprecate or delete all of the private `_build_conversation_input_ids()` methods.
- We add fallback `default_chat_template` properties to the classes that used to have private `_build_conversation_input_ids` methods, for backward compatibility.
- `apply_chat_template()` reads a `chat_template` attribute, and if missing it will read `default_chat_template`.
- `apply_chat_template()` is no longer locked into the `Conversation` format - it can also accept lists of dicts, similar to the `OpenAI` API.
- The `Conversation` object is also overhauled to use an internal list of dicts representation.

## Notes
Several existing classes already have class-specific `_build_conversation_input_ids()` methods. We need some way to preserve the class-specific behaviour for backward compatibility when moving to the public `apply_chat_template()` method, and this is achieved via the `default_chat_template` properties on those classes. It's a bit inelegant, but it works!

## Needs discussion

- [X] ~Is string substitutions the right way to do things, or should we allow templating with either a full template engine or regexes?~
- [X] ~Is this the right location for the classes?~
- [X] ~Is `PromptConfig` the right name? Maybe `ChatConfig` or `DialogConfig`?~
- [X] ~Is this the right way to break up functionality between `Tokenizer` and `PromptConfig`?~
- [x] What defaults should we set for models where we don't have to ensure backward compatibility? Should we pick 'standard' chat tokens and encourage people to use them going forward?
- [ ] Will we need to change anything for chats to support audio or image messages? Can we make our lives easier by preparing for that in advance?

## TODO

- [x] Convert to a Jinja template system
- [x] Massively lock down the template system and disallow callables, etc.
- [x] Ensure compatibility with `ConversationPipeline`
- [x] Ensure we get the same results with `Conversation` objects and lists of dicts in `ChatML` format
- [x] ~Ensure `prompt_config` is saved/loaded correctly~
- [x] Ensure backward compatibility on any model that had a private conversation method
- [x] Clear out any other TODOs in the code
- [x] Add optional dependency boilerplate for `jinja2`
- [x] Add tests
- [x] Make triple-sure the LLaMA conversion matches the old function

## Final checklist before merging

- [x] Add proper docstring and documentation for `build_conversation_input_ids()`
- [x] ~Add proper docstring and documentation for `PromptConfig`~
- [x] Add proper docstring and documentation for `Conversation`
- [x] Add usage example for the templates
- [x] Figure out how to make this whole thing discoverable
- [x] Remove old private `_build_conversation_from_input_ids` methods - right now they're staying so I can do equivalence testing

Comments:
_The documentation is not available anymore as the PR was closed or merged._

100% Agree on the objective.

I feel like the current proposed config is quite confusing to understand/reason about (therefore I doubt users will too).

I though we had in mind something more like and templating engine.
```
{
 "template": "[[SYS]]{{ prefix }}[[/SYS]] [EOS] {% for persona, message in conversation %}[{persona}]: {message}[EOS]{ % endfor %}"
}
```
I'm roughly following jinja's code for the templating system.

The main advantage to use a real templating engine is that users can "write code" to create the final string, which avoids the big chunk of code with all the options you suggest and is probably more future proof.

One big con is that the string might be hard to read and parse because spaces are *crucial* to get the correct input_ids so it is relatively error prone. We could find a way to make testing super easy (maybe add the test values directly in the config ?) but I don't t see how we can make those always *readable*.

We don't need to pull in all of `jinja` right now, since we mostly need the loop, (maybe an if) to evaluate the template.
If what I'm saying is actually true, we can probably write our own engine in a few lines of regexp (which is fine since templates are supposed to be simple).
Since we own the templating there's no risk of executing arbitrary code either.

It seems roughly what you are trying to do with all the config options, however, I think we can evolve it much more easily by adding more syntax to our template engine if required.

In order to use the template engine user need to :

- Understand the template syntax (Let's steal jinja's syntax I think it's the most popular)
- Understand which variables are available. Here's I'm supposing a `Conversation` object which is exactly :
```
{"messages": [("persona", 'message")]`, "system": "xxx"}
```

I'm using non custom objects so it's simpler to understand. I'm using a `tuple` to simplify my template showcase (I think a `dict` is a bit more appropriate if we ever are to add new fields to the messages)

**Whatever struct we decide here initially, it can never ever be modified to remove/modify existing fields**.

We can always add more later, but removing is forbidden.
I would spend some time thinking about the `system` thing for instance, again I chose here what made my example pretty, I think you went to treat `system` as any other participant which has lots of merits (but can the system prompt really change during the conversation ?).

Another thing to not, is that as having implemented those things for llama, and trying to make sure things align, having a small tool to be able to modify the conversations live and see the actual sent string and/or input_ids generated, would be nice to make sure that the template we're choosing is actually producing the desired inputs for the model.


@Narsil That's a cool idea! I considered templating as one of the options in the proposal doc, but I thought adding `jinja` as a dependency and opening potential attack surfaces there might be a bad idea. If there's a simple and secure way to do it that doesn't add a heavy dependency, then that's probably better than the string substitutions I'm using right now.

> adding jinja as a dependency and opening potential attack surfaces there might be a bad idea. 

Very good thought ! I agree completely. I had in mind to tackle only a minimal subset I think, which would avoid both
I have never written those, so I expect it to be moderately hard to custom implement with very limited syntax (just bail on anything out of the ordinary).

I find this: https://aosabook.org/en/500L/a-template-engine.html. If we can't have a simple template engine I agree with your current proposal.

Update: Just realized `jinja2` is in PyTorch's `requirements.txt` already, so it shouldn't be a problem for us to have it as a requirement either since the vast majority of users will already have it installed.

> Update: Just realized jinja2 is in PyTorch's requirements.txt already, so it shouldn't be a problem for us to have it as a requirement either since the vast majority of users will already have it installed.

Careful about limiting syntax if things like `eval` exist.

Templating is in!

I think this should be ready for actual review now! cc @ArthurZucker @Narsil @gante 

The main thing it's missing is tests, but I'm working on that now, and I'm not sure I want to add them until I'm certain about where the classes will live - I'm still not totally sure about the separate `PromptConfig`, especially after the template refactor. The number of attributes in the config has shrunk down a lot, so we might just be able to make these tokenizer attributes and remove the `PromptConfig` class entirely.

A question that I did not get clear from the PR: How can a user store/load the template along its fine-tuned model? Through the tokenizer `chat_template` field? (If so, documenting a usage example would be very helpful :D)

@gante Correct! The original plan was to use a separate `PromptConfig`, but now it's just a single attribute that gets saved and loaded with the tokenizer. That could definitely be hard for people to find, though, and some people might not be familiar with Jinja templating. I'm considering adding a helper function like `build_template()` to construct a template from a series of arguments, like the old `role_prefix` arguments that `PromptConfig` had before the refactor.

I'll think about where I could add usage examples and documentation to make sure it's visible!

It's late on Friday, but I think this is ready for final review now! There's quite a few people in this thread, so pinging them @lewtun @ArthurZucker @Narsil @gante - you don't have to review the whole thing again if you've already taken a look, but let me know if you think there's anything missing before I merge! Optionally @amyeroberts since she's on core maintainer watch next week too (but maybe Arthur is enough?)

In particular, let me know if anyone has any strong opinions about two things:

1) Should we pick a preferred Hugging Face format for chat tokens, and set it as the default on the base class, so we can standardize this in future?
2) Is there anything in this implementation that will require breaking changes if we want to support images/audio in chat later? Can we make the breaking changes now to save the pain later?

I think all the major issues have been dealt with at this point, cc @amyeroberts for core maintainer review!

For Amy and anyone else reviewing it, I'd suggest starting with the doc [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_25323/en/chat_templating) to understand what's going on with the PR.

After that, the only major changes are the new `apply_chat_template` function in `tokenization_utils_base` and the rewrite of the `Conversation` class in `pipelines/conversational.py`. Almost everything else in the PR is just model-specific tests or default templates to retain backward compatibility for those classes, so hopefully it should be less intimidating to review than it looks at first!

@LysandreJik I thought about your comment about inverting the chat templates to turn a processed string into a list of conversation dicts. This is something that @lewtun has requested as well, but I'm not sure it's possible! The problem is that because templates are very flexible and general, I don't think there's a process that can invert an arbitrary template (and also information like roles may be discarded by some templates, which is then impossible for us to reconstruct).

We could write 'inverse' template functions for existing templates, but I don't think a general solution is possible, except to ask users to also include the inverse map in a new attribute when defining `tokenizer.chat_template`, but I can see multiple issues here, especially since the inverse function would probably (?) have to be Python instead of jinja, which opens us up to arbitrary code execution. I think this is probably something we should not try to include in this PR for now! 

All comments have been addressed now (except the `tokenizer_kwargs` - looking at that now to see which ones I should make explicit arguments!)

There's one change I didn't make though - @ArthurZucker suggested removing the `tokenize` argument to `apply_chat_template`. This makes a lot of sense to me, but I'm a bit worried about it, because we generally want to tokenize with `add_special_tokens=False`. The reason for this is that the special tokens, if needed, should be contained in the template. I think if we expect users to tokenize the outputs themselves then they will naturally use the default options, which include `add_special_tokens=True`, and then get silent performance degradation!

I think leaving a `tokenize` argument is the right choice for now, because at least then we have a correct 'reference' tokenization that users can compare to, and disable if they want. I'm open to suggestions, though!

Finished breaking out the tokenizer kwargs and updating the doc - this should be ready for re-review now! cc @ArthurZucker @LysandreJik (but no rush - I know you're both busy!)

There is one failing test in an unrelated pipeline - I can fix this with a rebase, but rebasing this PR is quite annoying, so I'll leave that to the very end.

One other comment for reviewers - the slowest part of this process is Jinja compiling the template. Right now, we do the compilation every time `apply_chat_template` is called. I could cache compiled templates easily to make this function about 5X faster. However, it's probably not that significant because chat datasets are much smaller than LLM pretraining datasets, so I don't think preprocessing will be a major bottleneck, and so it might be easier to just skip cacheing for now to reduce code complexity and avoid adding the extra cache attributes to the tokenizer. WDYT?

Quick update: I addressed all of @ArthurZucker's comments, but I realized that my simple template for LLaMA missed some edge cases. I had to do a proper rewrite to match the behaviour of the original code which made it a lot bigger, but at least it's fully correct now! Tests were updated to match.

Since the LLaMA template is referenced in the docs I'll need to update those as well, which I'll handle tomorrow!

I've also added caching using `@lru_cache` on a tokenizer method. 

Everything should be ready now - Jinja templates now support exception raising and caching is implemented.

Awesome 🔥 Should I do a last review?

Sure! Also, there's one function-pickling error that I can quickly fix tomorrow, probably by moving `raise_exception` out into a utilities file or something or something.

Rebased and fixed the last issues!