Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Perform additional check for owner string in is_<library>_available functions #2859

Merged
merged 3 commits into from
Aug 23, 2024
Merged
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 12 additions & 9 deletions sentence_transformers/util.py
Original file line number Diff line number Diff line change
Expand Up @@ -1304,22 +1304,25 @@ def get_device_name() -> Literal["mps", "cuda", "npu", "hpu", "cpu"]:
return "cpu"


def is_accelerate_available() -> bool:
def check_package_availability(package_name: str, owner: str) -> bool:
"""
Returns True if the accelerate library is available.
Checks if a package is available from the correct owner.
"""
return importlib.util.find_spec("accelerate") is not None
spec = importlib.util.find_spec(package_name)
if spec and owner in spec.origin:
return True
return False
leblancfg marked this conversation as resolved.
Show resolved Hide resolved


def is_datasets_available() -> bool:
def is_accelerate_available() -> bool:
"""
Returns True if the datasets library is available.
Returns True if the Huggingface accelerate library is available.
"""
return importlib.util.find_spec("datasets") is not None
return check_package_availability("accelerate", "huggingface")


def is_training_available() -> bool:
def is_datasets_available() -> bool:
"""
Returns True if we have the required dependencies for training Sentence Transformer models
Returns True if the Huggingface datasets library is available.
"""
return is_accelerate_available() and is_datasets_available()
return check_package_availability("datasets", "huggingface")
Loading