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 {}