| import argparse |
| import os |
| import sys |
| from typing import Dict, Union |
|
|
| import yaml |
| from ml_collections import ConfigDict |
| from omegaconf import OmegaConf |
|
|
|
|
| def save_config(config: Union[ConfigDict, OmegaConf], save_path: str): |
| """ |
| Save a configuration object (ConfigDict or OmegaConf) to a file. |
| |
| Parameters: |
| ---------- |
| config : Union[ConfigDict, OmegaConf] |
| The configuration object to save. |
| save_path : str |
| The path where the configuration file will be saved. |
| |
| Raises: |
| ------ |
| ValueError: |
| If the configuration type is not supported. |
| """ |
|
|
| os.makedirs(os.path.dirname(save_path), exist_ok=True) |
|
|
| try: |
| with open(save_path, "w") as f: |
| if isinstance(config, ConfigDict): |
| yaml.dump( |
| config.to_dict(), |
| f, |
| default_flow_style=False, |
| sort_keys=False, |
| allow_unicode=True, |
| ) |
| elif isinstance(config, OmegaConf): |
| OmegaConf.save(config, save_path) |
| else: |
| OmegaConf.save(config, save_path) |
| except Exception as e: |
| raise ValueError( |
| f"Error saving configuration: {e}." |
| f"Unsupported configuration type. Supported types: ConfigDict, OmegaConf." |
| f"Config type is {type(config)}" |
| ) |
|
|
|
|
| def create_test_config( |
| original_config_path: str, new_config_path: str, model_type: str |
| ): |
| """ |
| Create a test configuration file based on an existing configuration. |
| |
| Parameters: |
| ---------- |
| original_config_path : str |
| Path to the original configuration file. |
| new_config_path : str |
| Path where the new configuration file will be saved. |
| model_type : str |
| The type of model (e.g., 'scnet', 'htdemucs'). |
| |
| Returns: |
| ------- |
| None |
| """ |
| from utils.settings import load_config |
|
|
| config = load_config(model_type=model_type, config_path=original_config_path) |
|
|
| config["inference"]["batch_size"] = 1 |
| config["training"]["batch_size"] = 1 |
| config["training"]["gradient_accumulation_steps"] = 1 |
| config["training"]["num_epochs"] = 2 |
| config["training"]["num_steps"] = 3 |
|
|
| save_config(config, new_config_path) |
| print(f"Test config created at: {new_config_path}") |
|
|
|
|
| def parse_args(dict_args: Union[Dict, None]) -> argparse.Namespace: |
| """ |
| Parse command-line arguments for configuring the model, dataset, and training parameters. |
| |
| Args: |
| dict_args: Dict of command-line arguments. If None, arguments will be parsed from sys.argv. |
| |
| Returns: |
| Namespace object containing parsed arguments and their values. |
| """ |
|
|
| parser = argparse.ArgumentParser() |
| parser.add_argument( |
| "--orig_config", type=str, default="", help="Path to the original config file." |
| ) |
| parser.add_argument("--model_type", type=str, default="", help="Model type") |
| parser.add_argument( |
| "--new_config", |
| type=str, |
| default="", |
| help="Path to save the new test configuration file.", |
| ) |
|
|
| if dict_args is not None: |
| args = parser.parse_args([]) |
| args_dict = vars(args) |
| args_dict.update(dict_args) |
| args = argparse.Namespace(**args_dict) |
| else: |
| args = parser.parse_args() |
|
|
| |
| if not args.new_config: |
| original_dir = os.path.dirname(args.orig_config) |
| tests_dir = os.path.join("tests_cache", original_dir) |
| os.makedirs(tests_dir, exist_ok=True) |
| args.new_config = os.path.join(tests_dir, os.path.basename(args.orig_config)) |
|
|
| return args |
|
|
|
|
| def redact_config(args): |
| |
| sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) |
|
|
| args = parse_args(args) |
|
|
| |
| create_test_config(args.orig_config, args.new_config, args.model_type) |
| return args.new_config |
|
|
|
|
| if __name__ == "__main__": |
| redact_config(None) |
|
|
