huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-11 20:34:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

Check if allowed tokens is None

Browse Source
This commit is contained in:
Alex Weston 2024-10-25 11:21:17 -04:00
parent 803d697d3d
commit 3f79225472
1 changed files with 5 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -498,9 +498,10 @@ class GrammarLogitProcessor(LogitsProcessor):
): ):
if fsm_grammar_state == -1 or self.fsm is None: if fsm_grammar_state == -1 or self.fsm is None:
return logits 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 = torch.full_like(logits, -math.inf)
mask[:, allowed_tokens] = 0 if allowed_tokens is not None:
mask[:, allowed_tokens] = 0
biased_scores = logits + mask biased_scores = logits + mask
return biased_scores return biased_scores
@ -589,7 +590,8 @@ class HeterogeneousGrammarLogitProcessor(LogitsProcessor):
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.get_next_instruction(fsm_grammar_states[i]).tokens 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] logits[i] += mask[i]
return logits return logits