Renormalize LMNextToken.sample() probs to fix floating point errors

probcomp · Jan 3, 2025 · f40ee1b · f40ee1b
1 parent f172d8b
commit f40ee1b
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/hfppl/distributions/lmcontext.py b/hfppl/distributions/lmcontext.py
@@ -27,6 +27,7 @@ async def log_prob(self, x):
 
     async def sample(self):
         probs = np.exp(self.ctx.next_token_logprobs)
+        probs /= np.sum(probs)  # Renormalize to fix floating point errors
         token_id = np.random.choice(len(probs), p=(probs))
         self.ctx.tokens.append(token_id)
         logprob = self.ctx.next_token_logprobs[token_id]