AI changes the rest of the call sites

Author: jackzhxng

tests/models/cohere2/test_modeling_cohere2.py

@@ -275,7 +275,9 @@ def test_export_static_cache(self):
         max_new_tokens = 30 - prompt_token_ids.shape[-1]
 
         # Static Cache + export
-        exported_program = convert_and_export_with_cache(model)
+        exported_program = convert_and_export_with_cache(
+            model, config=model.config, generation_config=model.generation_config
+        )
         ep_generated_ids = TorchExportableModuleWithStaticCache.generate(
             exported_program=exported_program, prompt_token_ids=prompt_token_ids, max_new_tokens=max_new_tokens
         )

tests/models/exaone4/test_modeling_exaone4.py

@@ -400,7 +400,9 @@ def test_export_static_cache(self):
         max_new_tokens = max_generation_length - prompt_token_ids.shape[-1]
 
         # Static Cache + export
-        exported_program = convert_and_export_with_cache(model)
+        exported_program = convert_and_export_with_cache(
+            model, config=model.config, generation_config=model.generation_config
+        )
         ep_generated_ids = TorchExportableModuleWithStaticCache.generate(
             exported_program=exported_program, prompt_token_ids=prompt_token_ids, max_new_tokens=max_new_tokens
         )

tests/models/gemma/test_modeling_gemma.py

@@ -459,8 +459,12 @@ def test_export_static_cache(self):
         # Static Cache + export
         from transformers.integrations.executorch import TorchExportableModuleForDecoderOnlyLM
 
-        exportable_module = TorchExportableModuleForDecoderOnlyLM(model)
-        exported_program = exportable_module.export()
+        exportable_module = TorchExportableModuleForDecoderOnlyLM(
+            model, config=model.config, generation_config=model.generation_config
+        )
+        exported_program = exportable_module.export(
+            input_ids=prompt_token_ids, cache_position=torch.arange(prompt_token_ids.shape[-1], dtype=torch.long, device=model.device)
+        )
         ep_generated_ids = TorchExportableModuleWithStaticCache.generate(
             exported_program=exported_program, prompt_token_ids=prompt_token_ids, max_new_tokens=max_new_tokens
         )

tests/models/gemma2/test_modeling_gemma2.py

@@ -364,8 +364,12 @@ def test_export_static_cache(self):
         # Static Cache + export
         from transformers.integrations.executorch import TorchExportableModuleForDecoderOnlyLM
 
-        exportable_module = TorchExportableModuleForDecoderOnlyLM(model)
-        exported_program = exportable_module.export()
+        exportable_module = TorchExportableModuleForDecoderOnlyLM(
+            model, config=model.config, generation_config=model.generation_config
+        )
+        exported_program = exportable_module.export(
+            input_ids=prompt_token_ids, cache_position=torch.arange(prompt_token_ids.shape[-1], dtype=torch.long, device=model.device)
+        )
         ep_generated_ids = TorchExportableModuleWithStaticCache.generate(
             exported_program=exported_program, prompt_token_ids=prompt_token_ids, max_new_tokens=max_new_tokens
         )
@@ -388,8 +392,13 @@ def test_export_hybrid_cache(self):
 
         # Export + HybridCache
         model.eval()
-        exportable_module = TorchExportableModuleForDecoderOnlyLM(model)
-        exported_program = exportable_module.export()
+        exportable_module = TorchExportableModuleForDecoderOnlyLM(
+            model, config=model.config, generation_config=model.generation_config
+        )
+        exported_program = exportable_module.export(
+            input_ids=torch.tensor([[1]], dtype=torch.long, device=model.device), 
+            cache_position=torch.tensor([0], dtype=torch.long, device=model.device)
+        )
 
         # Test generation with the exported model
         prompt = "What is the capital of France?"

tests/models/gemma3/test_modeling_gemma3.py

@@ -808,8 +808,13 @@ def test_export_text_only_with_hybrid_cache(self):
 
         # Export + HybridCache
         model.eval()
-        exportable_module = TorchExportableModuleForDecoderOnlyLM(model)
-        exported_program = exportable_module.export()
+        exportable_module = TorchExportableModuleForDecoderOnlyLM(
+            model, config=model.config, generation_config=model.generation_config
+        )
+        exported_program = exportable_module.export(
+            input_ids=torch.tensor([[1]], dtype=torch.long, device=model.device), 
+            cache_position=torch.tensor([0], dtype=torch.long, device=model.device)
+        )
         logging.info(f"\nExported program: {exported_program}")
 
         # Test generation with the exported model

tests/models/llama/test_modeling_llama.py

@@ -352,8 +352,12 @@ def test_export_static_cache(self):
             # Static Cache + export
             from transformers.integrations.executorch import TorchExportableModuleForDecoderOnlyLM
 
-            exportable_module = TorchExportableModuleForDecoderOnlyLM(model)
-            exported_program = exportable_module.export()
+            exportable_module = TorchExportableModuleForDecoderOnlyLM(
+                model, config=model.config, generation_config=model.generation_config
+            )
+            exported_program = exportable_module.export(
+                input_ids=prompt_token_ids, cache_position=torch.arange(prompt_token_ids.shape[-1], dtype=torch.long, device=model.device)
+            )
             ep_generated_ids = TorchExportableModuleWithStaticCache.generate(
                 exported_program=exported_program, prompt_token_ids=prompt_token_ids, max_new_tokens=max_new_tokens
             )

tests/models/olmo/test_modeling_olmo.py

@@ -383,8 +383,12 @@ def test_export_static_cache(self):
         # Static Cache + export
         from transformers.integrations.executorch import TorchExportableModuleForDecoderOnlyLM
 
-        exportable_module = TorchExportableModuleForDecoderOnlyLM(model)
-        exported_program = exportable_module.export()
+        exportable_module = TorchExportableModuleForDecoderOnlyLM(
+            model, config=model.config, generation_config=model.generation_config
+        )
+        exported_program = exportable_module.export(
+            input_ids=prompt_token_ids, cache_position=torch.arange(prompt_token_ids.shape[-1], dtype=torch.long, device=model.device)
+        )
         ep_generated_ids = TorchExportableModuleWithStaticCache.generate(
             exported_program=exported_program, prompt_token_ids=prompt_token_ids, max_new_tokens=max_new_tokens
         )

tests/models/olmo2/test_modeling_olmo2.py

@@ -383,7 +383,9 @@ def test_export_static_cache(self):
         self.assertEqual(EXPECTED_TEXT_COMPLETION, eager_generated_text)
 
         # Static Cache + export
-        exported_program = convert_and_export_with_cache(model)
+        exported_program = convert_and_export_with_cache(
+            model, config=model.config, generation_config=model.generation_config
+        )
         ep_generated_ids = TorchExportableModuleWithStaticCache.generate(
             exported_program=exported_program, prompt_token_ids=prompt_token_ids, max_new_tokens=max_new_tokens
         )

tests/models/phi3/test_modeling_phi3.py

@@ -416,8 +416,12 @@ def test_export_static_cache(self):
         # Static Cache + export
         from transformers.integrations.executorch import TorchExportableModuleForDecoderOnlyLM
 
-        exportable_module = TorchExportableModuleForDecoderOnlyLM(model)
-        exported_program = exportable_module.export()
+        exportable_module = TorchExportableModuleForDecoderOnlyLM(
+            model, config=model.config, generation_config=model.generation_config
+        )
+        exported_program = exportable_module.export(
+            input_ids=prompt_token_ids, cache_position=torch.arange(prompt_token_ids.shape[-1], dtype=torch.long, device=model.device)
+        )
         ep_generated_ids = TorchExportableModuleWithStaticCache.generate(
             exported_program=exported_program, prompt_token_ids=prompt_token_ids, max_new_tokens=max_new_tokens
         )

tests/models/qwen2/test_modeling_qwen2.py

@@ -299,11 +299,15 @@ def test_export_static_cache(self):
         # Static Cache + export
         from transformers.integrations.executorch import TorchExportableModuleForDecoderOnlyLM
 
-        exportable_module = TorchExportableModuleForDecoderOnlyLM(model)
+        exportable_module = TorchExportableModuleForDecoderOnlyLM(
+            model, config=model.config, generation_config=model.generation_config
+        )
         strict = version.parse(torch.__version__) != version.parse(
             "2.7.0"
         )  # Due to https://github.com/pytorch/pytorch/issues/150994
-        exported_program = exportable_module.export(strict=strict)
+        exported_program = exportable_module.export(
+            input_ids=prompt_token_ids, cache_position=torch.arange(prompt_token_ids.shape[-1], dtype=torch.long, device=model.device), strict=strict
+        )
         ep_generated_ids = TorchExportableModuleWithStaticCache.generate(
             exported_program=exported_program, prompt_token_ids=prompt_token_ids, max_new_tokens=max_new_tokens
         )