fix: Scale token probability by number of top logprobs in confidence …

…score calculation
okdshin · Jan 7, 2025 · c539445 · c539445
1 parent daac810
commit c539445
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/reasoning_llm_mcts/reasoning_state.py b/src/reasoning_llm_mcts/reasoning_state.py
@@ -139,7 +139,7 @@ def calc_confidence_score(
         # ci_sum += math.exp(token_lp) / sum([math.exp(top_lps) for top_lps in top_lps])
 
         max_lp = max(top_lps)
-        ci_sum += math.exp(token_lp - max_lp) / sum(math.exp(lp - max_lp) for lp in top_lps)
+        ci_sum += (len(top_lps) * math.exp(token_lp - max_lp)) / sum(math.exp(lp - max_lp) for lp in top_lps)
 
     confidence_score = ci_sum / len(token_logprobs)
     return confidence_score