| 1234567891011121314151617181920212223242526272829303132333435 |
- from pathlib import Path
- import torch
- import yaml
- class NoAliasSafeDumper(yaml.SafeDumper):
- # Disable anchor/alias in yaml because looks ugly
- def ignore_aliases(self, data):
- return True
- def yaml_no_alias_safe_dump(data, stream=None, **kwargs):
- """Safe-dump in yaml with no anchor/alias"""
- return yaml.dump(
- data, stream, allow_unicode=True, Dumper=NoAliasSafeDumper, **kwargs
- )
- def gen_conf(file, out_dir):
- conf = torch.load(file)["config"]
- conf["oss_bucket"] = "null"
- print(conf)
- output_dir = Path(out_dir)
- output_dir.mkdir(parents=True, exist_ok=True)
- with (output_dir / "config.yaml").open("w", encoding="utf-8") as f:
- yaml_no_alias_safe_dump(conf, f, indent=4, sort_keys=False)
- if __name__ == "__main__":
- import sys
- in_f = sys.argv[1]
- out_f = sys.argv[2]
- gen_conf(in_f, out_f)
|
