add option to provide config file via kwargs

Author: A9isha

src/maxtext/configs/pyconfig.py

@@ -81,6 +81,11 @@ def _resolve_or_infer_config(argv: list[str] | None = None, **kwargs) -> tuple[s
   """Resolves or infers config file path from module."""
   if argv is None:
     argv = [""]
+
+  if kwargs.get("base_config"):
+    logger.info(f"Using config : {kwargs['base_config']}")
+    return resolve_config_path(kwargs["base_config"]), argv[1:] if len(argv) > 1 else []
+
   if len(argv) >= 2 and argv[1].endswith(".yml"):
     return resolve_config_path(argv[1]), argv[2:]
   module = _module_from_path(argv[0]) if len(argv) > 0 else None

src/maxtext/utils/model_creation_utils.py

@@ -252,7 +252,7 @@ def from_pretrained(config, original_mesh=None, devices=None, model_mode=MODEL_M
   """Creates a NNX model with sharded parameters, possibly loading from a checkpoint."""
   mesh = original_mesh
   if config.convert_checkpoint_if_possible:
-    if not not (epath.Path(config.base_output_directory) / "0" / "items").exists():
+    if not (epath.Path(config.base_output_directory) / "0" / "items").exists():
       # Try to convert checkpoint on the fly
       if not config.hf_access_token:
         raise ValueError("hf_access_token must be provided when not providing a pre-existing checkpoint")