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