update logits in place for GuideLogitsProcessor

lapp0 · Oct 5, 2024 · 7514aff · 7514aff
1 parent 0da7039
commit 7514aff
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/outlines/processors/structured.py b/outlines/processors/structured.py
@@ -102,12 +102,13 @@ def process_logits(
 
             sequence_states.append(self._guide_states[curr_state_key])
 
-        mask = torch.full_like(logits, -math.inf)
+        mask = torch.ones_like(logits, dtype=torch.bool)
         for i, guide_state in enumerate(sequence_states):
             allowed_tokens = self.guide.get_next_instruction(guide_state).tokens
-            mask[i, allowed_tokens] = logits[i, allowed_tokens]
+            mask[i, allowed_tokens] = False
+        logits.masked_fill_(mask, float("-inf"))
 
-        return mask
+        return logits
 
     def copy(self) -> "GuideLogitsProcessor":
         """Return a copy of the logits processor."""