Merge pull request #49 from huggingface/min_new_tokens

add min new tokens

LLM/language_model.py

@@ -75,7 +75,7 @@ class LanguageModelHandler(BaseHandler):
         dummy_input_text = "Write me a poem about Machine Learning."
         dummy_chat = [{"role": self.user_role, "content": dummy_input_text}]
         warmup_gen_kwargs = {
-            "min_new_tokens": self.gen_kwargs["max_new_tokens"],
+            "min_new_tokens": self.gen_kwargs["min_new_tokens"],
             "max_new_tokens": self.gen_kwargs["max_new_tokens"],
             **self.gen_kwargs,
         }

STT/whisper_stt_handler.py

@@ -71,7 +71,7 @@ class WhisperSTTHandler(BaseHandler):
             # generating more tokens than previously will trigger CUDA graphs capture
             # one should warmup with a number of generated tokens above max tokens targeted for subsequent generation
             warmup_gen_kwargs = {
-                "min_new_tokens": self.gen_kwargs["max_new_tokens"],
+                "min_new_tokens": self.gen_kwargs["min_new_tokens"],
                 "max_new_tokens": self.gen_kwargs["max_new_tokens"],
                 **self.gen_kwargs,
             }

arguments_classes/language_model_arguments.py

@@ -45,6 +45,12 @@ class LanguageModelHandlerArguments:
             "help": "Maximum number of new tokens to generate in a single completion. Default is 128."
         },
     )
+    lm_gen_min_new_tokens: int = field(
+        default=0,
+        metadata={
+            "help": "Minimum number of new tokens to generate in a single completion. Default is 0."
+        },
+    )
     lm_gen_temperature: float = field(
         default=0.0,
         metadata={

arguments_classes/whisper_stt_arguments.py

@@ -33,6 +33,12 @@ class WhisperSTTHandlerArguments:
             "help": "The maximum number of new tokens to generate. Default is 128."
         },
     )
+    stt_gen_min_new_tokens: int = field(
+        default=0,
+        metadata={
+            "help": "The minimum number of new tokens to generate. Default is 0."
+        },
+    )
     stt_gen_num_beams: int = field(
         default=1,
         metadata={