saymrwulf/transformers
1
0
Fork 0
mirror of https://github.com/saymrwulf/transformers.git synced 2026-05-14 20:58:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix test_generated_length_assisted_generation (#34935)

Browse source 
fix test_generated_length_assisted_generation
This commit is contained in:
Nadav Timor 2025-01-29 07:03:45 -05:00 committed by GitHub
parent ec7afad609
commit 42c8ccfd4c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 8 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
tests/generation/test_utils.py
View file

@ -3405,7 +3405,14 @@ class GenerationIntegrationTests(unittest.TestCase, GenerationIntegrationTestsMi
assistant_model=assistant,
min_new_tokens=10,
)
self.assertTrue((input_length + 10) <= out.shape[-1] <= 20)
self.assertTrue((input_length + 10) <= out.shape[-1])
out = model.generate(
input_ids,
assistant_model=assistant,
max_new_tokens=7,
)
self.assertTrue(out.shape[-1] <= (input_length + 7))
def test_model_kwarg_assisted_decoding_decoder_only(self):
# PT-only test: TF doesn't support assisted decoding yet.