Refine the code of Universal Assisted Generation (#34823)

* removed the useless attritbutes * add configs for window size * fixed the wrong kwargs * added docstring
2024-11-28 22:04:24 +08:00
parent 2b053fdf1a
commit 44af935ec5
3 changed files with 18 additions and 10 deletions
--- a/tests/generation/test_utils.py
+++ b/tests/generation/test_utils.py
@@ -3812,7 +3812,7 @@ class GenerationIntegrationTests(unittest.TestCase, GenerationIntegrationTestsMi
            do_sample=False,
            max_new_tokens=max_new_tokens_item,
            assistant_model=draft_model,
-            target_tokenizer=target_tokenizer,
+            tokenizer=target_tokenizer,
            assistant_tokenizer=assistant_tokenizer,
        )