# AUTOGENERATED! DO NOT EDIT! File to edit: ../src/config_loader.ipynb.

# %% auto 0
__all__ = ['class_to_str', 'load_config', 'config_to_dict', 'save_dataclass_yaml', 'save_dict_yaml', 'get_obj_from_str',
           'instantiate_from_config', 'load_model_from_config']

# %% ../src/config_loader.ipynb 3
from .imports import *
from omegaconf import OmegaConf

# %% ../src/config_loader.ipynb 5
def class_to_str(cls):
    return str(cls)[8:-2]

# %% ../src/config_loader.ipynb 6
def load_config(file_path):
    return OmegaConf.load(f"{file_path}")

# %% ../src/config_loader.ipynb 7
def config_to_dict(config):
    return OmegaConf.to_container(config)

# %% ../src/config_loader.ipynb 8
def save_dataclass_yaml(data_obj, file_path):
    conf = OmegaConf.structured(data_obj)
    with open(file_path, 'w') as f:
        OmegaConf.save(config=conf, f=f)

# %% ../src/config_loader.ipynb 9
def save_dict_yaml(dict_obj, file_path):
    conf = OmegaConf.create(dict_obj)
    with open(file_path, 'w') as f:
        OmegaConf.save(config=conf, f=f)

# %% ../src/config_loader.ipynb 14
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)

# %% ../src/config_loader.ipynb 15
def instantiate_from_config(config):
    if not "target" in config: raise KeyError("Expected key `target` to instantiate.")
    if not "params" in config: print(f"[WARNING] Expected key `params` to instantiate.")
    return get_obj_from_str(config["target"])(**config.get("params", dict()))

# %% ../src/config_loader.ipynb 16
def load_model_from_config(config, ckpt, device):
    
    print(f"Loading model from {ckpt}")
    pl_sd = torch.load(ckpt, map_location=torch.device(device).type, weights_only=True)
          
    model = instantiate_from_config(config.model)
    
    sd = pl_sd["state_dict"]
    m, u = model.load_state_dict(sd, strict=True)
    
    return model.to(device)
