code: def extract_submodules_state_dict(
      state_dict: Dict[str, torch.Tensor], submodule_names: List[str]
  ):
      """
      Extract the sub state-dict corresponding to a list of given submodules.
  
      Args:
          state_dict (`Dict[str, torch.Tensor]`): The state dict to extract from.
          submodule_names (`List[str]`): The list of submodule names we want to extract.
      """
      result = {}
      for module_name in submodule_names:
          # We want to catch module_name parameter (module_name.xxx) or potentially module_name, but not any of the
          # submodules that could being like module_name (transformers.h.1 and transformers.h.10 for instance)
          result.update(
              {
                  key: param
                  for key, param in state_dict.items()
                  if key == module_name
                  or key.startswith(module_name + ".")
              }
          )
      return result
  
imports: import json
  import os
  from collections.abc import Mapping
  from typing import Dict, List, Optional, Union
  import numpy
  import torch
  from safetensors import safe_open
label_desc: Use image node_id 5 for calling a global function with example usage: extract_submodules_state_dict(state_dict, submodule_names) and returns: result