Repository: oumi-ai/oumi
File: tests/unit/core/configs/test_config.py (Python)

Code before:


def test_config_equality():
    config_a = TrainingConfig()
    config_b = TrainingConfig()
    assert config_a == config_b

    config_a.model.model_name = "test_model"
    assert config_a != config_b


def test_config_override():
    low_priority_config = TrainingConfig()
    low_priority_config.model.model_name = "model_low_priority"

    high_priority_config = TrainingConfig()
    high_priority_config.model.model_name = "model_high_priority"

    # Override with CLI arguments if provided
    merged_config = OmegaConf.merge(low_priority_config, high_priority_config)
    assert merged_config.model.model_name == "model_high_priority"
    assert merged_config == high_priority_config
    assert merged_config != low_priority_config


def test_config_from_yaml_and_arg_list(tmp_path):
    task = EvaluationTaskParams(
        evaluation_platform="lm_harness",
        task_name="mmlu",
        num_samples=5,
        eval_kwargs={"num_fewshot": 5},
    )
    config = EvaluationConfig(tasks=[task])
    config.model.model_name = "foo"
    config_path = tmp_path / "eval.yaml"
    config.to_yaml(config_path)

    new_config_override_vals = EvaluationConfig.from_yaml_and_arg_list(
        config_path,
        [
            "tasks[0].num_samples=1",  # override field
            "tasks[0].eval_kwargs.num_fewshot=1",  # override nested dict field
            "tasks[0].eval_kwargs.foo=bar",  # add new field
        ],
    )
    assert new_config_override_vals.tasks == [
        EvaluationTaskParams(
            evaluation_platform="lm_harness",
            task_name="mmlu",
            num_samples=1,
            eval_kwargs={"num_fewshot": 1, "foo": "bar"},

Review comment (bug):
A few additional tests you should add:

1) Modifying an element in a list when the list has more than 1 element.

2) Trying to modify an element in a list (ex, list[0]) when the list is empty (I'd expect this to result in an error)

3) Trying to modify an element at an incorrect index in the list (say list[5] for a list of 2 elements).

Code after:
        loaded_config = TrainingConfig.from_yaml(Path(filename))
        assert original_config == loaded_config


def test_config_equality():
    config_a = TrainingConfig()
    config_b = TrainingConfig()
    assert config_a == config_b

    config_a.model.model_name = "test_model"
    assert config_a != config_b


def test_config_override():
    low_priority_config = TrainingConfig()
    low_priority_config.model.model_name = "model_low_priority"

    high_priority_config = TrainingConfig()
    high_priority_config.model.model_name = "model_high_priority"

    # Override with CLI arguments if provided
    merged_config = OmegaConf.merge(low_priority_config, high_priority_config)
    assert merged_config.model.model_name == "model_high_priority"
    assert merged_config == high_priority_config
    assert merged_config != low_priority_config


def test_config_from_yaml_and_arg_list_override_values(tmp_path):
    task1 = EvaluationTaskParams(
        evaluation_platform="lm_harness",
        task_name="mmlu",
        num_samples=5,
        eval_kwargs={"num_fewshot": 5},
    )
    task2 = EvaluationTaskParams(
        evaluation_platform="lm_harness",
        task_name="hellaswag",
        num_samples=5,
        eval_kwargs={"num_fewshot": 5},
    )
    config = EvaluationConfig(tasks=[task1, task2])
    config.model.model_name = "foo"
    config_path = tmp_path / "eval.yaml"
    config.to_yaml(config_path)

    new_config = EvaluationConfig.from_yaml_and_arg_list(
        config_path,
        [
            "tasks[1].num_samples=1",  # override field
            "tasks[1].eval_kwargs.num_fewshot=1",  # override nested dict field
            "tasks[1].eval_kwargs.foo=bar",  # add new field