Filter out input ids not being used in inference. Fixes issue dottxt-…

…ai#1294
priyank9320 · Dec 2, 2024 · ac54946 · ac54946
1 parent 36f1bf2
commit ac54946
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/outlines/processors/structured.py b/outlines/processors/structured.py
@@ -110,6 +110,7 @@ def process_logits(
             allowed_tokens = self.guide.get_next_instruction(guide_state).tokens.to(
                 mask.device, non_blocking=True
             )
+            allowed_tokens = allowed_tokens[allowed_tokens < mask.shape[-1]] # filter out input ids exceeding the mask length
             allowed_tokens_batch.append(allowed_tokens)
             batch_indices.append(
                 torch.full_like(allowed_tokens, i)