def model_state_dict_private_to_public( private_state_dict: ModuleStateDict, ) -> Dict[str, ModelParameter]: public_state_dict: Dict[str, ModelParameter] = {} for private_name, tensor in private_state_dict.items(): if not isinstance(tensor, torch.Tensor): raise RuntimeError("Isn't the state dict supposed to be " "a shallow key-to-tensor mapping?!") for mapping in MODEL_STATE_DICT_MAPPINGS: try: public_name = mapping.private_to_public.map(private_name) except ValueError: continue else: break else: raise RuntimeError( f"Couldn't find a match for state dict key: {private_name}") public_state_dict[public_name] = ModelParameter(private_name, tensor) return public_state_dict
def save_model_state_dict(hf: h5py.File, state_dict: ModuleStateDict) -> None: g = hf.create_group(MODEL_STATE_DICT_GROUP, track_order=True) for private_key, tensor in state_dict.items(): if not isinstance(tensor, torch.Tensor): raise RuntimeError("Isn't the state dict supposed to be " "a shallow key-to-tensor mapping?!") for mapping in MODEL_STATE_DICT_MAPPINGS: try: public_key = mapping.private_to_public.map(private_key) except ValueError: continue else: break else: raise RuntimeError("Couldn't find a match for state dict key: %s" % private_key) dataset = g.create_dataset(public_key, data=tensor.numpy()) dataset.attrs[STATE_DICT_KEY_ATTR] = private_key