Update for new API

This commit is contained in:
Nicolas Patry 2024-10-25 10:46:05 +02:00
parent 0721649fc6
commit b49cff3f07
No known key found for this signature in database
GPG Key ID: D2920555C90F704C
1 changed files with 2 additions and 2 deletions

View File

@ -513,7 +513,7 @@ class GrammarLogitProcessor(LogitsProcessor):
def _advance(next_token_id, fsm_grammar_state, fsm):
if fsm_grammar_state == -1:
return fsm_grammar_state
return fsm.next_state(fsm_grammar_state, next_token_id)
return fsm.get_next_state(fsm_grammar_state, next_token_id)
# TODO: move grammar compilation into the router
@staticmethod
@ -588,7 +588,7 @@ class HeterogeneousGrammarLogitProcessor(LogitsProcessor):
fsm = self.fsms[i]
if fsm_grammar_states[i] == -1 or fsm is None:
continue
allowed_tokens = fsm.allowed_token_ids(fsm_grammar_states[i])
allowed_tokens = fsm.get_next_instruction(fsm_grammar_states[i]).tokens
mask[i, allowed_tokens] = 0
logits[i] += mask[i]
return logits