From 202e079fa9cd04e842dafb1b693e11d68feeb42b Mon Sep 17 00:00:00 2001 From: Sebastian Walter Date: Wed, 11 Dec 2024 10:14:47 +0100 Subject: [PATCH] fix number of candidates checked --- python/text_utils/inference/__init__.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/python/text_utils/inference/__init__.py b/python/text_utils/inference/__init__.py index 1c7c09b..f9eeb37 100644 --- 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])