Merge pull request #20 from probcomp/gg/renormalize

Renormalize LMNextToken.sample() probs to fix floating point errors
probcomp · Jan 4, 2025 · a191fca · a191fca
2 parents f172d8b + f40ee1b
commit a191fca
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]