set self.ctx.model_mask to empty set when len(good_tokens) == 0

probcomp · Jan 6, 2025 · ec6b35a · ec6b35a
1 parent 582de08
commit ec6b35a
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/hfppl/distributions/lmcontext.py b/hfppl/distributions/lmcontext.py
@@ -80,9 +80,11 @@ async def log_prob(self, v):
         )
         if len(good_tokens) == 0:
             # If there are no good tokens, the log probability of v under the mask is -inf
-            return float("-inf")
+            logprob_good = float("-inf")
+        else:
+            logprob_good = logsumexp(self.ctx.next_token_logprobs[list(good_tokens)])
+
         bad_tokens = [i for i in self.ctx.model_mask if i not in good_tokens]
-        logprob_good = logsumexp(self.ctx.next_token_logprobs[list(good_tokens)])
         self.ctx.next_token_logprobs[bad_tokens] = float("-inf")
         self.ctx.next_token_logprobs -= logprob_good
         self.ctx.model_mask = good_tokens