config parser + yaml config

configs/default.yaml

@@ -0,0 +1,6 @@
+output_dir: train_outputs/
+bf16: true
+target_modules:
+  - down_proj
+  - up_proj
+  - gate_proj

hyperlora/configs.py

@@ -1,6 +1,7 @@
 import dataclasses
 import os
 import sys
+import yaml
 from dataclasses import dataclass, field
 from enum import Enum, auto
 from typing import Any, Dict, List, Literal, NewType, Optional, Tuple
@@ -8,6 +9,112 @@ from typing import Any, Dict, List, Literal, NewType, Optional, Tuple
 from transformers import MODEL_FOR_CAUSAL_LM_MAPPING, HfArgumentParser
 
 
+MODEL_CONFIG_CLASSES = list(MODEL_FOR_CAUSAL_LM_MAPPING.keys())
+MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
+
+
+DataClassType = NewType("DataClassType", Any)
+
+
+class ArgumentParser(HfArgumentParser):
+    def parse_yaml_and_args(
+        self, yaml_arg: str, other_args: Optional[List[str]] = None
+    ) -> List[dataclass]:
+        """
+        Parse a YAML file and overwrite the default/loaded values with the values provided to the command line.
+
+        Args:
+            yaml_arg (`str`):
+                The path to the config file used
+            other_args (`List[str]`, *optional`):
+                A list of strings to parse as command line arguments, e.g. ['--arg=val', '--arg2=val2'].
+
+        Returns:
+            [`List[dataclass]`]: a list of dataclasses with the values from the YAML file and the command line
+        """
+        arg_list = self.parse_yaml_file(os.path.abspath(yaml_arg))
+
+        outputs = []
+        # strip other args list into dict of key-value pairs
+        other_args = {
+            arg.split("=")[0].strip("-"): arg.split("=")[1] for arg in other_args
+        }
+        used_args = {}
+
+        # overwrite the default/loaded value with the value provided to the command line
+        # adapted from https://github.com/huggingface/transformers/blob/d0b5002378daabf62769159add3e7d66d3f83c3b/src/transformers/hf_argparser.py#L327
+        print(arg_list)
+        for data_yaml, data_class in zip(arg_list, self.dataclass_types):
+            keys = {f.name for f in dataclasses.fields(data_yaml) if f.init}
+            inputs = {k: v for k, v in vars(data_yaml).items() if k in keys}
+            for arg, val in other_args.items():
+                # add only if in keys
+                if arg in keys:
+                    base_type = data_yaml.__dataclass_fields__[arg].type
+                    inputs[arg] = val
+
+                    # cast type for ints, floats (default to strings)
+                    if base_type in [int, float]:
+                        inputs[arg] = base_type(val)
+
+                    if base_type == list[str]:
+                        inputs[arg] = [str(v) for v in val.split(",")]
+
+                    # bool of a non-empty string is True, so we manually check for bools
+                    if base_type == bool:
+                        if val in ["true", "True"]:
+                            inputs[arg] = True
+                        else:
+                            inputs[arg] = False
+
+                    if base_type == dict:
+                        inputs[arg] = yaml.load(val, Loader=yaml.FullLoader)
+
+                    # add to used-args so we can check if double add
+                    if arg not in used_args:
+                        used_args[arg] = val
+                    else:
+                        raise ValueError(
+                            f"Duplicate argument provided: {arg}, may cause unexpected behavior"
+                        )
+                # else:
+                #     raise ValueError(f"Argument provided not found in dataclass: {arg}")
+
+            obj = data_class(**inputs)
+            outputs.append(obj)
+
+        return outputs
+
+    def parse(self) -> DataClassType | Tuple[DataClassType]:
+        if len(sys.argv) == 2 and sys.argv[1].endswith(".yaml"):
+            # If we pass only one argument to the script and it's the path to a YAML file,
+            # let's parse it to get our arguments.
+            output = self.parse_yaml_file(os.path.abspath(sys.argv[1].split("=")[-1]))
+        # parse command line args and yaml file
+        elif len(sys.argv) > 2 and sys.argv[1].endswith(".yaml"):
+            output = self.parse_yaml_and_args(
+                os.path.abspath(sys.argv[1].split("=")[-1]), sys.argv[2:]
+            )
+        # parse --config for the yaml path and other command line args
+        elif any([arg.startswith("--config") for arg in sys.argv]):
+            yaml_arg = [
+                arg
+                for arg in sys.argv[1:]
+                if arg.startswith("--config") and arg.endswith(".yaml")
+            ][0]
+            other_args = [arg for arg in sys.argv[1:] if arg != yaml_arg]
+            output = self.parse_yaml_and_args(
+                os.path.abspath(yaml_arg.split("=")[-1]), other_args
+            )
+        # parse command line args only
+        else:
+            output = self.parse_args_into_dataclasses()
+
+        if len(output) == 1:
+            output = output[0]
+        return output
+
+
 class ExperimentSetup(Enum):
     LORA = "lora"
     HYPER_LORA = "hyper_lora"
@@ -20,7 +127,7 @@ class ModelArguments:
     Arguments for the base model.
     """
 
-    model_name_or_path: Optional[str] = field(
+    model_name_or_path: str = field(
         default=None,
         metadata={"help": ("Base model name or path.")},
     )

hyperlora/intx_sft.py

@@ -10,6 +10,8 @@ from data_utils import (
     tokenize_chat_messages,
 )
 from datasets import load_dataset
+
+from configs import ArgumentParser
 from model_loading import get_lora_config, get_model_and_tokenizer
 from modeling_utils import ModulatedPretrainedModel
 from training_utils import TRAINING_TASK, train_model
@@ -48,10 +50,10 @@ def main():
     # Set logging verbosity to INFO
     logging.basicConfig(level=logging.INFO)
 
-    parser = HfArgumentParser(
+    parser = ArgumentParser(
         (CtxTrainingArguments, ModelArguments, LoRAArguments, TrainingArguments)
     )
-    ctx_args, model_args, lora_args, training_args = parser.parse_args_into_dataclasses()
+    ctx_args, model_args, lora_args, training_args = parser.parse()
 
     training_args.label_names = ["labels"]
     training_args.eval_on_start = True
@@ -85,13 +87,14 @@ def main():
         # activate LoRA
         model.set_adapter("default")
 
+    print(model)
     log_num_train_params(model)
 
     # max_seq_len = 1024
 
     print("Loading dataset...")
-    train_file = "../data/raw_datasets/context_numbers/train.jsonl"
-    eval_file = "../data/raw_datasets/context_numbers/val.jsonl"
+    train_file = "data/raw_datasets/context_numbers/train.jsonl"
+    eval_file = "data/raw_datasets/context_numbers/val.jsonl"
     ds = load_dataset("json", data_files={"train": train_file, "eval": eval_file})
     # preprocessing
     ds = ds.map(get_preprocessing_fn("context_numbers"))