Skip to content

Commit

Permalink
[tests] add require_bitsandbytes marker (#30116)
Browse files Browse the repository at this point in the history
* add bnb flag

* move maker

* add accelerator maker
  • Loading branch information
faaany authored Apr 8, 2024
1 parent 5e673ed commit d16f0ab
Showing 1 changed file with 4 additions and 0 deletions.
4 changes: 4 additions & 0 deletions tests/models/persimmon/test_modeling_persimmon.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
from transformers import PersimmonConfig, is_torch_available, set_seed
from transformers.testing_utils import (
backend_empty_cache,
require_bitsandbytes,
require_torch,
require_torch_accelerator,
require_torch_fp16,
Expand Down Expand Up @@ -465,6 +466,8 @@ def test_model_rope_scaling(self):
@require_torch
class PersimmonIntegrationTest(unittest.TestCase):
@slow
@require_torch_accelerator
@require_bitsandbytes
def test_model_8b_chat_logits(self):
input_ids = [1, 306, 4658, 278, 6593, 310, 2834, 338]
model = PersimmonForCausalLM.from_pretrained(
Expand Down Expand Up @@ -492,6 +495,7 @@ def test_model_8b_chat_logits(self):
@slow
@require_torch_accelerator
@require_torch_fp16
@require_bitsandbytes
def test_model_8b_chat_greedy_generation(self):
EXPECTED_TEXT_COMPLETION = """human: Simply put, the theory of relativity states that?\n\nadept: The theory of relativity states that the laws of physics are the same for all observers, regardless of their relative motion."""
prompt = "human: Simply put, the theory of relativity states that?\n\nadept:"
Expand Down

0 comments on commit d16f0ab

Please sign in to comment.