File size: 3,782 Bytes
b6ff324 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | # ------------------------------------------------------------------------
# Copyright (c) 2024-present, BAAI. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, esither express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ------------------------------------------------------------------------
"""Omegaconf utilities."""
import importlib
import json
from typing import List
import omegaconf
class OmegaConfEncoder(json.JSONEncoder):
"""Custom JSON encoder for omegaconf objects."""
def default(self, obj):
if isinstance(obj, (omegaconf.ListConfig, omegaconf.DictConfig)):
return omegaconf.OmegaConf.to_container(obj, resolve=True)
return super().default(obj)
def get_config() -> omegaconf.DictConfig:
"""Return omega configurations from CLI."""
cli_conf = omegaconf.OmegaConf.from_cli()
omegaconf.OmegaConf.register_new_resolver("eval", eval) # Register ``eval`` func.
return omegaconf.OmegaConf.merge(omegaconf.OmegaConf.load(cli_conf.config), cli_conf)
def save_config(config: omegaconf.DictConfig, f):
"""Save config to YAML format string."""
omegaconf.OmegaConf.save(config, f)
def config_to_yaml(config: omegaconf.DictConfig) -> str:
"""Dump config to YAML format string."""
return omegaconf.OmegaConf.to_yaml(config)
def config_to_class(config: omegaconf.DictConfig) -> object:
"""Return the class object from config."""
def get_obj_from_str(string, reload=False):
module, cls = string.rsplit(".", 1)
if reload:
module_imp = importlib.import_module(module)
importlib.reload(module_imp)
return getattr(importlib.import_module(module, package=None), cls)
if not config:
return None
if "target" not in config:
raise KeyError("Expected key `target` to instantiate.")
return get_obj_from_str(config["target"])
def config_to_object(config: omegaconf.DictConfig, **kwargs) -> object:
"""Instantiate an object from config."""
if not config:
return None
kwargs.update(config.get("params", dict()))
return config_to_class(config)(**kwargs)
def flatten_omega_conf(cfg, resolve=True) -> List:
"""Flatten omega configurations."""
ret = []
def handle_dict(key, value, resolve):
return [(f"{key}.{k1}", v1) for k1, v1 in flatten_omega_conf(value, resolve=resolve)]
def handle_list(key, value, resolve):
return [(f"{key}.{idx}", v1) for idx, v1 in flatten_omega_conf(value, resolve=resolve)]
if isinstance(cfg, omegaconf.DictConfig):
for k, v in cfg.items_ex(resolve=resolve):
if isinstance(v, omegaconf.DictConfig):
ret.extend(handle_dict(k, v, resolve=resolve))
elif isinstance(v, omegaconf.ListConfig):
ret.extend(handle_list(k, v, resolve=resolve))
else:
ret.append((str(k), v))
elif isinstance(cfg, omegaconf.ListConfig):
for idx, v in enumerate(cfg._iter_ex(resolve=resolve)):
if isinstance(v, omegaconf.DictConfig):
ret.extend(handle_dict(idx, v, resolve=resolve))
elif isinstance(v, omegaconf.ListConfig):
ret.extend(handle_list(idx, v, resolve=resolve))
else:
ret.append((str(idx), v))
return ret
|