fix: cover corner case of rescoring completed sample in tree search

benlipkin · Oct 22, 2024 · 31c3d83 · 31c3d83
1 parent 0d73820
commit 31c3d83
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/decoding/generators.py b/decoding/generators.py
@@ -277,12 +277,15 @@ def _TreeSearch(
         for sample, passed, failed in zip(beam, stop_pass, stop_fail, strict=True):
             if passed:
                 finished.add(sample)
+                continue
             if failed:
                 continue
             prompts.append(sample.item)
         if len(finished) >= search_params.n:
             return sort_samples(finished)[: search_params.n]
-        beam = _BestOfN(prompts, llm, scorer, sampling_params)[: search_params.width]
+        beam = _BestOfN(prompts, llm, scorer, sampling_params)
+        if len(beam) > search_params.width:
+            beam = beam[: search_params.width]
     return _handle_maxsteps(finished)