From cd5bd8db595c4fce2b46759d14ecf4b5e2dedff9 Mon Sep 17 00:00:00 2001 From: Sebastian Walter Date: Thu, 12 Dec 2024 22:45:07 +0100 Subject: [PATCH] make estimated score default --- python/text_utils/inference/__init__.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/python/text_utils/inference/__init__.py b/python/text_utils/inference/__init__.py index 530b6d8..ac70b42 100644 --- a/python/text_utils/inference/__init__.py +++ b/python/text_utils/inference/__init__.py @@ -34,7 +34,7 @@ def beam_search( logit_fns: list[LogitFn] | None = None, kwargs_select_fn: MaskSelectFn | None = None, kwargs_update_fn: MaskUpdateFn | None = None, - stop_condition: str | None = None, + stop_condition: str = "estimated_score", max_new_tokens: int | None = None, return_incomplete: bool = False, yield_intermediate: bool = False, @@ -43,8 +43,6 @@ def beam_search( assert ( max_new_tokens is None or max_new_tokens > 0 ), "max_new_tokens must be None or positive" - if stop_condition is None: - stop_condition = "max_score" assert stop_condition in { "max_score", "estimated_score",