[generate] ✨ vectorized beam search ✨ (#35802)

2025-03-18 18:39:36 +00:00
parent 12f2ebef63
commit 179d02ffb8
8 changed files with 426 additions and 271 deletions
--- a/tests/models/cohere2/test_modeling_cohere2.py
+++ b/tests/models/cohere2/test_modeling_cohere2.py
@@ -104,10 +104,6 @@ class Cohere2ModelTest(CohereModelTest, unittest.TestCase):
    def test_generate_continue_from_past_key_values(self):
        pass

-    @unittest.skip("Cohere2 has HybridCache and doesn't support low_memory generation")
-    def test_beam_search_low_memory(self):
-        pass
-
    @unittest.skip("Cohere2 has HybridCache and doesn't support contrastive generation")
    def test_contrastive_generate(self):
        pass