Refactor FSM idx tracking during batch generation #449
Conversation
| def reset(self) -> None: | ||
| """Reset the FSM to its initial state. Here this only resets the token counter.""" | ||
| self.num_tokens_generated = 0 | ||
| def copy(self) -> "StopAtTokenFSM": |
There was a problem hiding this comment.
Is the copy mechanism only necessary because we enforce max_token here? Same question for the other FSM. #451 is removing this mechanism from the FSMs.
There was a problem hiding this comment.
Yes, correct! If max_tokens is handled elsewhere, as suggested in #451, then the copy method for both StopAtTokenFSM and RegexFSM can just be return self, which removes all overhead. This is a great suggestion
There was a problem hiding this comment.
Great! This comforts me in the idea that max_tokens should not be handled at the FSM level :)
rlouf
force-pushed
the
refactor-fsm-idx-tracking
branch
from
dd47100 to
f206c5e
Compare
|
I rebased on main and added a small commit to format docstrings, don't forget to pull locally :) |
|
I think it makes sense to wait until #451 is merged so we can simplify the |
|
Does this mean we don't need to pass the seq idx as the first argument in #481? |
|
We pass the |
|
Should be up to date now post #451 merge. I changed the |
| num_sequences = len(prompts) | ||
| fsms = [self.fsm.copy() for _ in prompts] |
There was a problem hiding this comment.
Do we still need an array of FSMs now that they’re stateless?
There was a problem hiding this comment.
Yes, because of CFGFSM (which needs to track incremental parser state for each). That still needs to deepcopy new objects, but the others (RegexFSM and StopAtFSM) are stateless and can just return pointers to original object.
|
Great! I’ll review shortly. Could you rebase the branch on |
…ng has been generated
…erator instead of when initializing it
benlipkin
force-pushed
the
refactor-fsm-idx-tracking
branch
from
a887385 to
0c79c3e
Compare
Sorry about that. Just rebased now, but had made a bit of a mess before when I pulled main into this branch since a lot had happened in between. If this is still too messy, I can open a new PR on a new branch and just apply the diff with main to make the history cleaner. Let me know preferences. |
Yes, I am sorry for the extra work but this would be preferable. In the future, use |
|
Moved and updated to #510 |
As discussed in #433, here is a proposal for removing the
idxargument from theFSMinterface, as well as theresetmethod, at the cost of adding acopymethod. An FSM is initialized, then copied for each prompt during aSequenceGenerator.__call__()on a new batch. This solves themax_tokensissue, and cleans up state tracking forCFGFSM. The cost is overhead associated with building a newFSMobject for each sequence in a batch. There are ways to reduce this cost by improving the newcopymethod, but wanted to get thoughts on the overall interface first.