From 7ab0580734dd1377bccc848f2049e91fb146682b Mon Sep 17 00:00:00 2001 From: Piotr Bielak Date: Thu, 17 Jul 2025 15:45:00 +0300 Subject: [PATCH] Use profiler in text-generation-pipeline The examples in `text-generation-pipeline` fail due to profiler args still being passed separately instead of a single `profiler` argument. This commit updates the creation of the profiler and fixes the arguments passed to `self.model.generate(...)`. --- .../text-generation-pipeline/pipeline.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/examples/text-generation/text-generation-pipeline/pipeline.py b/examples/text-generation/text-generation-pipeline/pipeline.py index 15cb96a3d4..6361e25fa2 100644 --- a/examples/text-generation/text-generation-pipeline/pipeline.py +++ b/examples/text-generation/text-generation-pipeline/pipeline.py @@ -58,14 +58,21 @@ def __call__(self, prompt): if torch.is_tensor(model_inputs[t]): model_inputs[t] = model_inputs[t].to(self.device) + from optimum.habana.utils import HabanaProfile + + profiler = HabanaProfile( + warmup=self.profiling_warmup_steps, + active=self.profiling_steps, + record_shapes=self.profiling_record_shapes, + name="generate", + ) + output = self.model.generate( **model_inputs, generation_config=self.generation_config, lazy_mode=True, hpu_graphs=self.use_hpu_graphs, - profiling_steps=self.profiling_steps, - profiling_warmup_steps=self.profiling_warmup_steps, - profiling_record_shapes=self.profiling_record_shapes, + profiler=profiler, ).cpu() if use_batch: