From f3ee6052db0239a7dfcd4d2efc1f746c150c7e53 Mon Sep 17 00:00:00 2001 From: Matthias Freiberger Date: Thu, 13 Feb 2025 12:03:44 +0100 Subject: [PATCH] Adding default parameter for max_new_tokens in TransformersModel (#604) --- src/smolagents/models.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/smolagents/models.py b/src/smolagents/models.py index 19d9aa4..75dda78 100644 --- a/src/smolagents/models.py +++ b/src/smolagents/models.py @@ -599,7 +599,16 @@ class TransformersModel(Model): model_id = default_model_id logger.warning(f"`model_id`not provided, using this default tokenizer for token counts: '{model_id}'") self.model_id = model_id + + default_max_tokens = 5000 + max_new_tokens = kwargs.get("max_new_tokens") or kwargs.get("max_tokens") + if not max_new_tokens: + kwargs["max_new_tokens"] = default_max_tokens + logger.warning( + f"`max_new_tokens` not provided, using this default value for `max_new_tokens`: {default_max_tokens}" + ) self.kwargs = kwargs + if device_map is None: device_map = "cuda" if torch.cuda.is_available() else "cpu" logger.info(f"Using device: {device_map}")