From 1fff2bc25508747b9f1d1a321822da17b5cb39cb Mon Sep 17 00:00:00 2001 From: Stefano Fiorucci <44616784+anakin87@users.noreply.github.com> Date: Tue, 21 Nov 2023 11:15:37 +0100 Subject: [PATCH] merge lazy import blocks (#6358) --- .../preview/components/generators/hugging_face_local.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/haystack/preview/components/generators/hugging_face_local.py b/haystack/preview/components/generators/hugging_face_local.py index 6bd5021e2..025028913 100644 --- a/haystack/preview/components/generators/hugging_face_local.py +++ b/haystack/preview/components/generators/hugging_face_local.py @@ -9,12 +9,8 @@ logger = logging.getLogger(__name__) SUPPORTED_TASKS = ["text-generation", "text2text-generation"] -with LazyImport( - message="PyTorch is needed to run this component. Please install it by following the instructions at https://pytorch.org/" -) as torch_import: +with LazyImport(message="Run 'pip install transformers[torch]'") as torch_and_transformers_import: import torch - -with LazyImport(message="Run 'pip install transformers'") as transformers_import: from huggingface_hub import model_info from transformers import ( pipeline, @@ -127,8 +123,7 @@ class HuggingFaceLocalGenerator: For some chat models, the output includes both the new text and the original prompt. In these cases, it's important to make sure your prompt has no stop words. """ - transformers_import.check() - torch_import.check() + torch_and_transformers_import.check() pipeline_kwargs = pipeline_kwargs or {} generation_kwargs = generation_kwargs or {}