riaz.somc/transformers
1
0
Fork 0
mirror of https://github.com/huggingface/transformers.git synced 2025-07-24 23:08:57 +06:00
Code Issues Actions 35 Packages Projects Releases Wiki Activity

[tests] add require_torch_sdpa for test that needs sdpa support (#30408)

Browse Source 
* add cuda flag

* check for sdpa

* add bitsandbytes
This commit is contained in:
Fanli Lin 2024-04-23 17:39:38 +08:00 committed by GitHub
parent 04ac3245e4
commit 2d61823fa2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tests/models/cohere/test_modeling_cohere.py
View File

@ -374,6 +374,7 @@ class CohereModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMix
@slow
class CohereIntegrationTest(unittest.TestCase):
@require_torch_multi_gpu
@require_bitsandbytes
def test_batched_4bit(self):
model_id = "CohereForAI/c4ai-command-r-v01-4bit"
@ -393,6 +394,7 @@ class CohereIntegrationTest(unittest.TestCase):
output = model.generate(**inputs, max_new_tokens=40, do_sample=False)
self.assertEqual(tokenizer.batch_decode(output, skip_special_tokens=True), EXPECTED_TEXT)
@require_torch_sdpa
def test_batched_small_model_logits(self):
# Since the model is very large, we created a random cohere model so that we can do a simple
# logits check on it.