checkpoint_utils.py

def adapt_state_dict_for_model(model, state_dict):
    model_keys = model.state_dict().keys()
    model_wants_prefix = any(k.startswith("_orig_mod.") for k in model_keys)
    state_has_prefix = any(k.startswith("_orig_mod.") for k in state_dict.keys())

    if model_wants_prefix == state_has_prefix:
        return state_dict

    if model_wants_prefix:
        return {
            k if k.startswith("_orig_mod.") else f"_orig_mod.{k}": v
            for k, v in state_dict.items()
        }

    prefix_len = len("_orig_mod.")
    return {
        k[prefix_len:] if k.startswith("_orig_mod.") else k: v
        for k, v in state_dict.items()
    }


def load_model_state(model, checkpoint, key="model_state_dict", strict=True):
    state_dict = checkpoint.get(key, checkpoint.get("model"))
    if state_dict is None:
        raise KeyError(f"Checkpoint does not contain '{key}' or 'model'")
    return model.load_state_dict(adapt_state_dict_for_model(model, state_dict), strict=strict)