Fix test

Author: jackzhxng

extension/llm/export/test/test_export_llm.py

@@ -88,23 +88,19 @@ def test_with_config(self, mock_export_llama: MagicMock) -> None:
             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"
+                called_config.base.tokenizer_path, "/path/to/tokenizer.json"
             )
-            self.assertEqual(called_config["base"]["model_class"], "llama2")
-            self.assertEqual(called_config["base"]["preq_mode"].value, "8da4w")
-            self.assertEqual(called_config["model"]["dtype_override"].value, "fp16")
-            self.assertEqual(called_config["export"]["max_seq_length"], 256)
+            self.assertEqual(called_config.base.model_class, "llama2")
+            self.assertEqual(called_config.base.preq_mode.value, "8da4w")
+            self.assertEqual(called_config.model.dtype_override.value, "fp16")
+            self.assertEqual(called_config.export.max_seq_length, 256)
             self.assertEqual(
-                called_config["quantization"]["pt2e_quantize"].value, "xnnpack_dynamic"
+                called_config.quantization.pt2e_quantize.value, "xnnpack_dynamic"
             )
+            self.assertEqual(called_config.quantization.use_spin_quant.value, "cuda")
+            self.assertEqual(called_config.backend.coreml.quantize.value, "c4w")
             self.assertEqual(
-                called_config["quantization"]["use_spin_quant"].value, "cuda"
-            )
-            self.assertEqual(
-                called_config["backend"]["coreml"]["quantize"].value, "c4w"
-            )
-            self.assertEqual(
-                called_config["backend"]["coreml"]["compute_units"].value, "cpu_and_gpu"
+                called_config.backend.coreml.compute_units.value, "cpu_and_gpu"
             )
         finally:
             os.unlink(config_file)
@@ -142,13 +138,13 @@ def test_with_config_and_cli(self, mock_export_llama: MagicMock) -> None:
             mock_export_llama.assert_called_once()
             called_config = mock_export_llama.call_args[0][0]
             self.assertEqual(
-                called_config["base"]["model_class"], "stories110m"
+                called_config.base.model_class, "stories110m"
             )  # Override from CLI.
             self.assertEqual(
-                called_config["model"]["dtype_override"].value, "fp16"
+                called_config.model.dtype_override.value, "fp16"
             )  # From yaml.
             self.assertEqual(
-                called_config["backend"]["xnnpack"]["enabled"],
+                called_config.backend.xnnpack.enabled,
                 True,  # Override from CLI.
             )
         finally: