From 7bc2c97bd9209c19c0f8f675f74313e4ed4a4abd Mon Sep 17 00:00:00 2001 From: Alex Weston <43505988+aW3st@users.noreply.github.com> Date: Mon, 28 Oct 2024 00:10:55 -0400 Subject: [PATCH] Check if allowed tokens is None (#2694) * Upgrade outlines to 0.1.1 * Update for new API * Check if allowed tokens is None --------- Co-authored-by: Nicolas Patry --- server/poetry.lock | 2 +- server/text_generation_server/utils/logits_process.py | 8 +++++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/server/poetry.lock b/server/poetry.lock index 4060842d..89231d21 100644 --- a/server/poetry.lock +++ b/server/poetry.lock @@ -3976,4 +3976,4 @@ torch = ["torch"] [metadata] lock-version = "2.0" python-versions = ">=3.9,<3.13" -content-hash = "5b2536579dd1d4013da43c5260666cab1c6575c8c1dd6ebb1ccfd5dc0f2874fd" +content-hash = "40be820ced080c2457b0794ed61fdd5340615f0fe75420985eaaca7b2b6c3968" \ No newline at end of file diff --git a/server/text_generation_server/utils/logits_process.py b/server/text_generation_server/utils/logits_process.py index e79a4bb0..ec2813a1 100644 --- a/server/text_generation_server/utils/logits_process.py +++ b/server/text_generation_server/utils/logits_process.py @@ -498,9 +498,10 @@ class GrammarLogitProcessor(LogitsProcessor): ): if fsm_grammar_state == -1 or self.fsm is None: return logits - allowed_tokens = self.fsm.allowed_token_ids(fsm_grammar_state) + allowed_tokens = self.fsm.get_next_instruction(fsm_grammar_state).tokens mask = torch.full_like(logits, -math.inf) - mask[:, allowed_tokens] = 0 + if allowed_tokens is not None: + mask[:, allowed_tokens] = 0 biased_scores = logits + mask return biased_scores @@ -589,7 +590,8 @@ class HeterogeneousGrammarLogitProcessor(LogitsProcessor): if fsm_grammar_states[i] == -1 or fsm is None: continue allowed_tokens = fsm.get_next_instruction(fsm_grammar_states[i]).tokens - mask[i, allowed_tokens] = 0 + if allowed_tokens is not None: + mask[i, allowed_tokens] = 0 logits[i] += mask[i] return logits