Implement prompt/generation alignment

[RobinPicard]

Opening this PR to discuss the implementation of token alignment (#161)

This is not intended to be merged, I was just wondering whether you think this is a promising direction to look into

The idea is to modify the self.states_to_token_maps of the fsm associated to each prompt to create a new state to accommodate the fact that a token has been removed compared to when this map was created during the initialization of the fsm

Disadvantage:

• The case in which the text after the "boundary" of a token matching the end of the prompt does not exist in the vocabulary by itself is not covered

Advantage:

• It takes very little time to update the fsm in that way

[rlouf]

This is not intended to be merged, I was just wondering whether you think this is a promising direction to look into

I think this is the right general direction.

• The case in which the text after the "boundary" of a token matching the end of the prompt does not exist in the vocabulary by itself is not covered

Could you illustrate this? I had a PR opened (can't find it right now) where I iterated once over the vocabulary to find the overlapping tokens.

[RobinPicard]

Could you illustrate this? I had a PR opened (can't find it right now) where I iterated once over the vocabulary to find the overlapping tokens.

Making up a fake example. My prompt is "Good mor". Let's say there's a token for "mor" and it's the last one of the prompt. We would want token alignment to replace "mor" with "morning". However, if the token "ning" by itself does not exist, then there's nothing in the states_to_token_maps that correspond to it as, at this point, the character-based FSM that would allow to generate "ning" has already been turned into a token-based mapping.

I was looking at creating states_to_token_maps only after the call is made (and the FSM is updated) but that would add too much overhead.

I was then thinking that a solution could be to create at initialization a mapping that contains both information about characters and about tokens (so we would have some states with no tokens leading to them that would be used for the token alignement)

[rlouf]

How about looping over the entire vocabulary and store the tokens that accept mor as a prefix. Then, in the unconstrained case the first state of the FSM would have transitions to the overlapping tokens only?

Haven't taken the time to think about the constrained case yet.

[RobinPicard]

I had not realized that I could walk the states_to_token_maps character by character for the postfix part of the crossing tokens in the constrained case. I think it works with almost no additional overhead like that. Let me know if you think it's fine and I'll update the tests afterward

[rlouf]

Yes I think that's the right approach. There's some stuff to figure out in terms of design, but otherwise looks good.

[RobinPicard]

I'll write unit tests next if you think having those separate functions is the right design

[rlouf]

I have made several comments on the overall design, but nothing that would dramatically affect your implementation. You can start implementing tests.

[shawnz]

Hi there, I am just a user here who is looking forward to this change. However, I noticed that there is an error if the model is running on a GPU. I think it could be fixed by passing the device in these two statements here (at least, this fixes it for me)

[rlouf]

We're getting really close. There are a few design changes remaining, and mostly we should have comprehensive tests before merging.

[rlouf]

I rebased your branch on main after a big refactor of the FSM interface. I will take a closer look this week.

[RobinPicard]

Is this still something we want to work on?

[rlouf]

Yes! I'm currently thinking about how we could integrate that to the logits processors since most integration are going to use this :)