From e53aabcbfc32a289ca5f1d774ad94fa735041ca0 Mon Sep 17 00:00:00 2001 From: elronbandel Date: Thu, 27 Feb 2025 10:14:58 +0200 Subject: [PATCH] Add batch size control to Huggingface pipeline based inference engine Signed-off-by: elronbandel --- src/unitxt/inference.py | 5 +++-- utils/.secrets.baseline | 6 +++--- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/src/unitxt/inference.py b/src/unitxt/inference.py index 8f1a0afeb9..9ac19c97f5 100644 --- a/src/unitxt/inference.py +++ b/src/unitxt/inference.py @@ -849,6 +849,8 @@ class HFPipelineBasedInferenceEngine( model_name: str label: str = "hf_pipeline_inference_engine" + batch_size: int = 30 + use_fast_tokenizer: bool = True use_fp16: bool = True load_in_8bit: bool = False @@ -857,7 +859,6 @@ class HFPipelineBasedInferenceEngine( device_map: Any = None - pipe: Any = InternalField(default=None) _requirements_list = { "transformers": "Install huggingface package using 'pip install --upgrade transformers", @@ -965,7 +966,7 @@ def _infer( if not self._is_loaded(): self._prepare_engine() - outputs = self.model([instance["source"] for instance in dataset]) + outputs = self.model([instance["source"] for instance in dataset], batch_size=self.batch_size) return [ self.get_return_object(output[0], instance["source"], return_meta_data) diff --git a/utils/.secrets.baseline b/utils/.secrets.baseline index 639c121312..deed4ed7d6 100644 --- a/utils/.secrets.baseline +++ b/utils/.secrets.baseline @@ -133,14 +133,14 @@ "filename": "src/unitxt/inference.py", "hashed_secret": "aa6cd2a77de22303be80e1f632195d62d211a729", "is_verified": false, - "line_number": 1305 + "line_number": 1306 }, { "type": "Secret Keyword", "filename": "src/unitxt/inference.py", "hashed_secret": "c8f16a194efc59559549c7bd69f7bea038742e79", "is_verified": false, - "line_number": 1808 + "line_number": 1809 } ], "src/unitxt/loaders.py": [ @@ -178,5 +178,5 @@ } ] }, - "generated_at": "2025-02-25T18:27:39Z" + "generated_at": "2025-02-27T08:14:19Z" }