Skip to content

Commit

Permalink
add test marker to run all tests with @require_bitsandbytes
Browse files Browse the repository at this point in the history
  • Loading branch information
Titus-von-Koeller committed Feb 1, 2024
1 parent 7bc6d76 commit 38a1989
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 2 deletions.
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -32,4 +32,5 @@ doctest_optionflags="NUMBER NORMALIZE_WHITESPACE ELLIPSIS"
doctest_glob="**/*.md"
markers = [
"flash_attn_test: marks tests related to flash attention (deselect with '-m \"not flash_attn_test\"')",
"bitsandbytes: select (or deselect with `not`) bitsandbytes integration tests",
]
12 changes: 10 additions & 2 deletions src/transformers/testing_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -957,9 +957,17 @@ def require_apex(test_case):

def require_bitsandbytes(test_case):
"""
Decorator for bits and bytes (bnb) dependency
Decorator marking a test that requires the bitsandbytes library. Will be skipped when the library or its hard dependency torch is not installed.
"""
return unittest.skipUnless(is_bitsandbytes_available(), "test requires bnb")(test_case)
if is_bitsandbytes_available() and is_torch_available():
try:
import pytest

return pytest.mark.bitsandbytes(test_case)
except ImportError:
return test_case
else:
return unittest.skip("test requires bitsandbytes and torch")(test_case)


def require_optimum(test_case):
Expand Down

0 comments on commit 38a1989

Please sign in to comment.