Update

[ghstack-poisoned]

Author: jackzhxng

extension/llm/export/export_llm.py

@@ -34,12 +34,11 @@
 from typing import Any, List, Tuple
 
 import hydra
-import yaml
 
 from executorch.examples.models.llama.config.llm_config import LlmConfig
 from executorch.examples.models.llama.export_llama_lib import export_llama
 from hydra.core.config_store import ConfigStore
-from omegaconf import DictConfig, OmegaConf
+from omegaconf import OmegaConf
 
 cs = ConfigStore.instance()
 cs.store(name="llm_config", node=LlmConfig)
@@ -79,7 +78,7 @@ def main() -> None:
                 "Cannot specify additional CLI arguments when using --config. "
                 f"Found: {remaining_args}. Use either --config file or hydra CLI args, not both."
             )
-        
+
         config_file_path = pop_config_arg()
         default_llm_config = LlmConfig()
         llm_config_from_file = OmegaConf.load(config_file_path)

extension/llm/export/test/test_export_llm.py

@@ -10,8 +10,11 @@
 import unittest
 from unittest.mock import MagicMock, patch
 
-from executorch.examples.models.llama.config.llm_config import LlmConfig
-from executorch.extension.llm.export.export_llm import main, parse_config_arg, pop_config_arg
+from executorch.extension.llm.export.export_llm import (
+    main,
+    parse_config_arg,
+    pop_config_arg,
+)
 
 
 class TestExportLlm(unittest.TestCase):
@@ -45,12 +48,14 @@ def test_with_config(self, mock_export_llama: MagicMock) -> None:
         """Test main function with --config file and no hydra args."""
         # Create a temporary config file
         with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
-            f.write("""
+            f.write(
+                """
 base:
   tokenizer_path: /path/to/tokenizer.json
 export:
   max_seq_length: 256
-""")
+"""
+            )
             config_file = f.name
 
         try:
@@ -61,7 +66,9 @@ def test_with_config(self, mock_export_llama: MagicMock) -> None:
             # Verify export_llama was called with config
             mock_export_llama.assert_called_once()
             called_config = mock_export_llama.call_args[0][0]
-            self.assertEqual(called_config["base"]["tokenizer_path"], "/path/to/tokenizer.json")
+            self.assertEqual(
+                called_config["base"]["tokenizer_path"], "/path/to/tokenizer.json"
+            )
             self.assertEqual(called_config["export"]["max_seq_length"], 256)
         finally:
             os.unlink(config_file)
@@ -70,7 +77,9 @@ def test_with_cli_args(self) -> None:
         """Test main function with only hydra CLI args."""
         test_argv = ["script.py", "debug.verbose=True"]
         with patch.object(sys, "argv", test_argv):
-            with patch("executorch.extension.llm.export.export_llm.hydra_main") as mock_hydra:
+            with patch(
+                "executorch.extension.llm.export.export_llm.hydra_main"
+            ) as mock_hydra:
                 main()
                 mock_hydra.assert_called_once()
 
@@ -86,9 +95,12 @@ def test_config_with_cli_args_error(self) -> None:
             with patch.object(sys, "argv", test_argv):
                 with self.assertRaises(ValueError) as cm:
                     main()
-                
+
                 error_msg = str(cm.exception)
-                self.assertIn("Cannot specify additional CLI arguments when using --config", error_msg)
+                self.assertIn(
+                    "Cannot specify additional CLI arguments when using --config",
+                    error_msg,
+                )
         finally:
             os.unlink(config_file)
 
@@ -99,7 +111,13 @@ def test_config_rejects_multiple_cli_args(self) -> None:
             config_file = f.name
 
         try:
-            test_argv = ["script.py", "--config", config_file, "debug.verbose=True", "export.output_dir=/tmp"]
+            test_argv = [
+                "script.py",
+                "--config",
+                config_file,
+                "debug.verbose=True",
+                "export.output_dir=/tmp",
+            ]
             with patch.object(sys, "argv", test_argv):
                 with self.assertRaises(ValueError):
                     main()
@@ -109,4 +127,3 @@ def test_config_rejects_multiple_cli_args(self) -> None:
 
 if __name__ == "__main__":
     unittest.main()
-