Skip to content

Commit

Permalink
merge lazy import blocks (#6358)
Browse files Browse the repository at this point in the history
  • Loading branch information
anakin87 authored Nov 21, 2023
1 parent 2943b83 commit 1fff2bc
Showing 1 changed file with 2 additions and 7 deletions.
9 changes: 2 additions & 7 deletions haystack/preview/components/generators/hugging_face_local.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,12 +9,8 @@

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,
Expand Down Expand Up @@ -127,8 +123,7 @@ def __init__(
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 {}
Expand Down

0 comments on commit 1fff2bc

Please sign in to comment.