code: def _get_mapped_state_dict(
      module: torch.nn.Module,
      state_dict: StateDictType,
      mapping: Optional[Dict[str, str]] = None,
  ) -> StateDictType:
      # First fix all top-level keys according to `combined_mapping`.
      combined_mapping = (
          module._get_mapping(mapping)
          if isinstance(module, TransformerModule)
          else {}
      )
      for hf_key, cls_key in sorted(
          # Sort by most specific key first.
          combined_mapping.items(),
          key=lambda x: x[0].count("."),
          reverse=True,
      ):
          relevant_keys = set(
              [
                  key
                  for key in state_dict.keys()
                  if (key == hf_key or key.startswith(hf_key + "."))
              ]
          )
          for key in relevant_keys:
              new_key = key.replace(hf_key, cls_key, 1)
              # We have to be careful not to overwrite an entry that we might have updated
              # on a previous iteration of this loop due to having a more specific key.
              if new_key not in state_dict:
                  state_dict[new_key] = state_dict.pop(key)
  
      # Now loop through the submodules, calling this function on each submodule.
      for name, submodule in module.named_children():
          # Submodule might be a wrapped, sharded module.
          if isinstance(submodule, ShardedModuleMixin):
              submodule = submodule.get_original_module()
  
          # Pull-out the part of the state_dict corresponding to just this submodule.
          relevant_keys = set(
              [
                  key
                  for key in state_dict.keys()
                  if key.startswith(name + ".")
              ]
          )
          module_state_dict = {
              key.replace(name + ".", "", 1): state_dict.pop(key)
              for key in relevant_keys
          }
          # Recursively call this function from the submodule to map this part
          # of the state_dict.
          module_state_dict = _get_mapped_state_dict(
              submodule, module_state_dict
          )
          # And then update the full state_dict.
          for key, value in module_state_dict.items():
              state_dict[name + "." + key] = value
  
      return state_dict
  
imports: import logging
  import os
  from os import PathLike
  from typing import TYPE_CHECKING, Optional, Dict, Union, List, Any, TypeVar, Type
  import re
  import warnings
  import torch
  import torch.distributed
  from allennlp.common.util import is_distributed, is_global_primary
  from allennlp.nn.parallel import ShardedModuleMixin
  from allennlp.nn.module import Module
  from allennlp.nn.util import StateDictType, read_state_dict, _check_incompatible_keys
label_desc: Use image node_id 1 for calling a global function with example usage: _get_mapped_state_dict(module, state_dict, mapping) and returns: state_dict