fix number of candidates checked

ad-freiburg · Dec 11, 2024 · 202e079 · 202e079
1 parent 0ddb602
commit 202e079
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/python/text_utils/inference/__init__.py b/python/text_utils/inference/__init__.py
@@ -218,16 +218,16 @@ def get_outputs(intermediate: bool) -> list[list[Beam]]:
 
             # reset current beams and fill with best candidates
             current_beams[idx] = []
-            for candidate in sorted(
-                candidates, key=lambda b: score_fn(b), reverse=True
-            )[:beam_width]:
+            for candidate in sorted(candidates, key=score_fn, reverse=True):
                 # update candidates
                 candidate = update_fn(candidate)
                 if candidate is None:
                     # skip invalid candidates
                     continue
-                else:
+                elif len(current_beams[idx]) < beam_width:
                     current_beams[idx].append(candidate)
+                else:
+                    break
 
             update_info[idx] = len(current_beams[idx])