Set default generation config for TextGenerationPipeline and TextToAudioPipeline (#2339)

AKloniecki · web-flow · commit f03351d8ac86 · 2025-11-06T11:54:28.000+01:00
Signed-off-by: Artur Kloniecki &lt;arturx.kloniecki@intel.com&gt;
diff --git a/optimum/habana/transformers/modeling_utils.py b/optimum/habana/transformers/modeling_utils.py
@@ -425,6 +425,12 @@ def adapt_transformers_to_gaudi():
             max_new_tokens=256,
         )
     )
+    transformers.pipelines.text_generation.TextGenerationPipeline._default_generation_config = GaudiGenerationConfig(
+        max_new_tokens=256,
+    )
+    transformers.pipelines.text_to_audio.TextToAudioPipeline._default_generation_config = GaudiGenerationConfig(
+        max_new_tokens=256,
+    )
 
     # Optimization for BLOOM generation on Gaudi
     transformers.models.bloom.modeling_bloom.BloomAttention.forward = gaudi_bloom_attention_forward

Original file line number	Diff line number	Diff line change
`@@ -425,6 +425,12 @@ def adapt_transformers_to_gaudi():`
`425`	`425`	`max_new_tokens=256,`
`426`	`426`	`)`
`427`	`427`	`)`
	`428`	`+ transformers.pipelines.text_generation.TextGenerationPipeline._default_generation_config = GaudiGenerationConfig(`
	`429`	`+ max_new_tokens=256,`
	`430`	`+ )`
	`431`	`+ transformers.pipelines.text_to_audio.TextToAudioPipeline._default_generation_config = GaudiGenerationConfig(`
	`432`	`+ max_new_tokens=256,`
	`433`	`+ )`
`428`	`434`
`429`	`435`	`# Optimization for BLOOM generation on Gaudi`
`430`	`436`	`transformers.models.bloom.modeling_bloom.BloomAttention.forward = gaudi_bloom_attention_forward`