Use generation_config

huggingface · Apr 29, 2024 · 8c8c7a5 · 8c8c7a5
1 parent 72a98c5
commit 8c8c7a5
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 2 deletions.
diff --git a/src/transformers/models/idefics2/fine_tune_idefics2.py b/src/transformers/models/idefics2/fine_tune_idefics2.py
@@ -275,6 +275,7 @@ def compute_metrics(eval_preds):
     report_to="none",
     eval_do_concat_batches=False,
     predict_with_generate=True,
+    generation_config={"max_new_tokens": 200},
 )
 
 # important: we need to disable caching during training

diff --git a/src/transformers/trainer_seq2seq.py b/src/transformers/trainer_seq2seq.py
@@ -308,8 +308,6 @@ def prediction_step(
                 k: v for k, v in inputs.items() if k not in ("decoder_input_ids", "decoder_attention_mask")
             }
 
-        # TODO fix this
-        gen_kwargs["max_new_tokens"] = 200
         generated_tokens = self.model.generate(**generation_inputs, **gen_kwargs)
 
         # Temporary hack to ensure the generation config is not initialized for each iteration of the evaluation loop