Initializes the configuration by loading YAML files, and applying CLI, env, and kwarg overrides. Returns pydantic MaxTextConfig class whereas `initialize` returns the og `HyperParameters`
(argv: list[str], **kwargs)
| 211 | |
| 212 | |
| 213 | def initialize_pydantic(argv: list[str], **kwargs) -> MaxTextConfig: |
| 214 | """Initializes the configuration by loading YAML files, and applying CLI, env, and kwarg overrides. |
| 215 | Returns pydantic MaxTextConfig class whereas `initialize` returns the og `HyperParameters` |
| 216 | """ |
| 217 | # 1. Load base and inherited configs from file(s) |
| 218 | config_path = resolve_config_path(argv[1]) |
| 219 | base_yml_config = _load_config(config_path) |
| 220 | |
| 221 | # 2. Get overrides from CLI and kwargs |
| 222 | cli_cfg = omegaconf.OmegaConf.from_cli(argv[2:]) |
| 223 | kwargs_cfg = omegaconf.OmegaConf.create(kwargs) |
| 224 | overrides_cfg = omegaconf.OmegaConf.merge(cli_cfg, kwargs_cfg) |
| 225 | |
| 226 | # 3. Handle model-specific config |
| 227 | temp_cfg = omegaconf.OmegaConf.merge(base_yml_config, overrides_cfg) |
| 228 | model_name = temp_cfg.get("model_name", "default") |
| 229 | model_cfg = {} |
| 230 | if model_name != "default": |
| 231 | # First try relative to base config path |
| 232 | model_config_path = os.path.join(os.path.dirname(config_path), "models", f"{model_name}.yml") |
| 233 | if not os.path.isfile(model_config_path): |
| 234 | # Fallback to default location within package |
| 235 | dir_path = os.path.dirname(os.path.realpath(__file__)) |
| 236 | model_config_path = os.path.join(dir_path, "configs", "models", f"{model_name}.yml") |
| 237 | |
| 238 | if os.path.exists(model_config_path): |
| 239 | model_loaded_cfg = omegaconf.OmegaConf.load(model_config_path) |
| 240 | # if override_model_config=True, only apply model configs for keys not present in overrides. |
| 241 | if temp_cfg.get("override_model_config"): |
| 242 | model_cfg = {k: v for k, v in model_loaded_cfg.items() if k not in overrides_cfg} |
| 243 | else: |
| 244 | model_cfg = model_loaded_cfg |
| 245 | else: |
| 246 | logger.warning("Model config for '%s' not found at %s", model_name, model_config_path) |
| 247 | |
| 248 | # 4. Final merge (base, model, then overrides) |
| 249 | model_cfg_oc = omegaconf.OmegaConf.create(model_cfg) |
| 250 | |
| 251 | # 4. Manually merge logical_axis_rules to avoid OmegaConf's list replacement behavior. |
| 252 | base_rules_oc = base_yml_config.get("logical_axis_rules", []) |
| 253 | model_rules_oc = model_cfg_oc.get("logical_axis_rules", []) |
| 254 | overrides_rules_oc = overrides_cfg.get("logical_axis_rules", []) |
| 255 | |
| 256 | base_rules = omegaconf.OmegaConf.to_container(base_rules_oc, resolve=True) if base_rules_oc else [] |
| 257 | model_rules = omegaconf.OmegaConf.to_container(model_rules_oc, resolve=True) if model_rules_oc else [] |
| 258 | overrides_rules = omegaconf.OmegaConf.to_container(overrides_rules_oc, resolve=True) if overrides_rules_oc else [] |
| 259 | |
| 260 | merged_rules = _apply_rules(base_rules, model_rules, model_cfg_oc) |
| 261 | merged_rules = _apply_rules(merged_rules, overrides_rules, overrides_cfg) |
| 262 | |
| 263 | # Remove the rules from the original configs before the main merge |
| 264 | if "logical_axis_rules" in base_yml_config: |
| 265 | del base_yml_config["logical_axis_rules"] |
| 266 | if "logical_axis_rules" in model_cfg_oc: |
| 267 | del model_cfg_oc["logical_axis_rules"] |
| 268 | if "logical_axis_rules" in overrides_cfg: |
| 269 | del overrides_cfg["logical_axis_rules"] |
| 270 |