Adding default parameter for max_new_tokens in TransformersModel (#604)

huggingface · Feb 13, 2025 · f3ee605 · f3ee605
1 parent a427c84
commit f3ee605
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/src/smolagents/models.py b/src/smolagents/models.py
@@ -599,7 +599,16 @@ def __init__(
             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}")