lit_config.json -> model_config.yaml (#1096)

Author: carmocca

eval/lm_eval_harness.py

@@ -157,7 +157,7 @@ def run_eval_harness(
     check_valid_checkpoint_dir(checkpoint_dir)
     tokenizer = Tokenizer(checkpoint_dir)
 
-    config = Config.from_json(checkpoint_dir / "lit_config.json")
+    config = Config.from_file(checkpoint_dir / "model_config.yaml")
 
     checkpoint_path = checkpoint_dir / "lit_model.pth"
 

litgpt/chat/base.py

@@ -132,7 +132,7 @@ def main(
 
     check_valid_checkpoint_dir(checkpoint_dir)
 
-    config = Config.from_json(checkpoint_dir / "lit_config.json")
+    config = Config.from_file(checkpoint_dir / "model_config.yaml")
 
     checkpoint_path = checkpoint_dir / "lit_model.pth"
 

litgpt/config.py

@@ -1,6 +1,6 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
 
-import json
+import yaml
 from copy import deepcopy
 from dataclasses import dataclass, field
 from pathlib import Path
@@ -120,8 +120,9 @@ def from_name(cls, name: str, **kwargs: Any) -> Self:
         return cls(**conf_dict)
 
     @classmethod
-    def from_json(cls, path: Union[str, Path], **kwargs: Any) -> Self:
+    def from_file(cls, path: Union[str, Path], **kwargs: Any) -> Self:
         with open(path, encoding="utf-8") as fp:
+<<<<<<< HEAD
             json_kwargs = json.load(fp)
 
         if "condense_ratio" in json_kwargs:  # legacy name
@@ -139,24 +140,29 @@ def from_json(cls, path: Union[str, Path], **kwargs: Any) -> Self:
 
         json_kwargs.update(kwargs)
         return cls(**json_kwargs)
+=======
+            file_kwargs = yaml.safe_load(fp)
+        file_kwargs.update(kwargs)
+        return cls(**file_kwargs)
+>>>>>>> 674f315 (lit_config.json -> model_config.yaml (#1096))
 
     @classmethod
     def from_checkpoint(cls, path: Path, **kwargs: Any) -> Self:
-        """Automatically load `lit_config.json` and if it doesn't exist - a matching config from `litgpt/config.py`."""
-        if (config_path := path / "lit_config.json").is_file():
-            return cls.from_json(config_path, **kwargs)
+        """Automatically load `model_config.yaml` and if it doesn't exist - a matching config from `litgpt/config.py`."""
+        if (config_path := path / "model_config.yaml").is_file():
+            return cls.from_file(config_path, **kwargs)
         if (model_name := path.name) in name_to_config:
             return cls.from_name(model_name, **kwargs)
-        raise FileNotFoundError(f"For {str(path)!r} neither 'lit_config.json' nor matching config exists.")
+        raise FileNotFoundError(f"For {str(path)!r} neither 'model_config.yaml' nor matching config exists.")
 
     @property
     def mlp_class(self) -> Type:
-        # `self.mlp_class_name` cannot be the type to keep the config json serializable
+        # `self.mlp_class_name` cannot be the type to keep the config serializable
         return getattr(litgpt.model, self.mlp_class_name)
 
     @property
     def norm_class(self) -> Type:
-        # `self.norm_class_name` cannot be the type to keep the config json serializable
+        # `self.norm_class_name` cannot be the type to keep the config serializable
         if self.norm_class_name == "RMSNorm":
             from functools import partial
 

litgpt/generate/adapter.py

@@ -61,7 +61,7 @@ def main(
 
     check_valid_checkpoint_dir(checkpoint_dir)
 
-    config = Config.from_json(checkpoint_dir / "lit_config.json")
+    config = Config.from_file(checkpoint_dir / "model_config.yaml")
 
     checkpoint_path = checkpoint_dir / "lit_model.pth"
 

litgpt/generate/adapter_v2.py

@@ -61,7 +61,7 @@ def main(
 
     check_valid_checkpoint_dir(checkpoint_dir)
 
-    config = Config.from_json(checkpoint_dir / "lit_config.json")
+    config = Config.from_file(checkpoint_dir / "model_config.yaml")
 
     checkpoint_path = checkpoint_dir / "lit_model.pth"
 

litgpt/generate/base.py

@@ -134,7 +134,7 @@ def main(
 
     check_valid_checkpoint_dir(checkpoint_dir)
 
-    config = Config.from_json(checkpoint_dir / "lit_config.json")
+    config = Config.from_file(checkpoint_dir / "model_config.yaml")
 
     checkpoint_path = checkpoint_dir / "lit_model.pth"
 

litgpt/generate/full.py

@@ -60,7 +60,7 @@ def main(
 
     check_valid_checkpoint_dir(checkpoint_dir)
 
-    config = Config.from_json(checkpoint_dir / "lit_config.json")
+    config = Config.from_file(checkpoint_dir / "model_config.yaml")
 
     checkpoint_path = finetuned_path
 

litgpt/generate/lora.py

@@ -70,8 +70,8 @@ def main(
 
     check_valid_checkpoint_dir(checkpoint_dir)
 
-    config = Config.from_json(
-        checkpoint_dir / "lit_config.json",
+    config = Config.from_file(
+        checkpoint_dir / "model_config.yaml",
         lora_r=lora_r,
         lora_alpha=lora_alpha,
         lora_dropout=lora_dropout,

litgpt/generate/sequentially.py

@@ -159,7 +159,7 @@ def main(
 
     check_valid_checkpoint_dir(checkpoint_dir)
 
-    config = Config.from_json(checkpoint_dir / "lit_config.json")
+    config = Config.from_file(checkpoint_dir / "model_config.yaml")
 
     checkpoint_path = checkpoint_dir / "lit_model.pth"
 

litgpt/generate/tp.py

@@ -138,7 +138,7 @@ def main(
 
     check_valid_checkpoint_dir(checkpoint_dir)
 
-    config = Config.from_json(checkpoint_dir / "lit_config.json")
+    config = Config.from_file(checkpoint_dir / "model_config.yaml")
 
     model_file = "lit_model.pth"
     checkpoint_path = checkpoint_dir / model_file