Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
chenpangpang
transformers
Commits
2d61823f
Unverified
Commit
2d61823f
authored
Apr 23, 2024
by
Fanli Lin
Committed by
GitHub
Apr 23, 2024
Browse files
[tests] add `require_torch_sdpa` for test that needs sdpa support (#30408)
* add cuda flag * check for sdpa * add bitsandbytes
parent
04ac3245
Changes
1
Show whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
2 additions
and
0 deletions
+2
-0
tests/models/cohere/test_modeling_cohere.py
tests/models/cohere/test_modeling_cohere.py
+2
-0
No files found.
tests/models/cohere/test_modeling_cohere.py
View file @
2d61823f
...
...
@@ -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.
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment