Mirrors
/
hf_text-generation-inference
mirror of https://github.com/huggingface/text-generation-inference.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
server/text_generation_server/utils/logits_process.py
View File

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