transformers-community
/

custom_generate_example

custom_generate

Model card Files Files and versions

RaushanTurganbay HF Staff commited on Feb 23

Commit

3217931

·

verified ·

1 Parent(s): ec99db5

Fix for the test

Files changed (1) hide show

custom_generate/generate.py +1 -1

custom_generate/generate.py CHANGED Viewed

@@ -6,7 +6,7 @@ def generate(model, input_ids, generation_config=None, left_padding=None, **kwar
     generation_config = generation_config or model.generation_config  # default to the model generation config
     cur_length = input_ids.shape[1]
-    max_length = generation_config.max_length or cur_length + generation_config.max_new_tokens
     # Example of custom argument: add `left_padding` (integer) pad tokens before the prompt
     if left_padding is not None:

     generation_config = generation_config or model.generation_config  # default to the model generation config
     cur_length = input_ids.shape[1]
+    max_length = generation_config.max_length or cur_length + kwargs.get("max_new_tokens")
     # Example of custom argument: add `left_padding` (integer) pad tokens before the prompt
     if left_padding is not None: