code: def load_exported_model(
      layout: ModelLayout,
      model,
      model_H: Hyperparams,
      load_path: str,
      load_heads_map: Dict[str, str],
      use_cache: bool = False,
  ):
      """
      :param load_heads_map: maps name in model -> name to load from
      """
      if use_cache and blobs.is_blob_url(load_path):
          load_path = blobs.download_directory_cached(load_path)
  
      with bf.BlobFile(os.path.join(load_path, "info.json")) as f:
          info = json.load(f)
  
      old_model_H = Hyperparams(**info["model_hparams"])
      original_n_shards = old_model_H.n_shards
      if "n_shards" in info:
          assert info["n_shards"] == original_n_shards
      assert layout.n_shards == model_H.n_shards
      # print("orig n_shards", original_n_shards, "new n_shards", layout.n_shards)
  
      def fetch_single_piece(fine_piece_name):
          with bf.BlobFile(
              os.path.join(load_path, "checkpoint", fine_piece_name),
              "rb",
          ) as f:
              return torch.load(f, map_location=torch.device("cpu"))
  
      if original_n_shards % layout.n_shards == 0:
          n_chkpt_shards_per_rank = exact_div(
              original_n_shards, layout.n_shards
          )
          shard_idx_start = n_chkpt_shards_per_rank * layout.shard_idx
          load_shard_idxs = range(
              shard_idx_start, shard_idx_start + n_chkpt_shards_per_rank
          )
  
          def fetch(chkpt_prefix: str, module_name: str = ""):
              sharded_pieces = [
                  fetch_single_piece(
                      f"{chkpt_prefix}_shard_{shard_idx:03d}.pkl"
                  )
                  for shard_idx in load_shard_idxs
              ]
  
              model_piece = {}
              for k in sharded_pieces[0].keys():
                  parameter_name = (
                      ".".join([module_name, k]) if module_name else k
                  )
  
                  sharding_dim = dim_to_shard(parameter_name)
  
                  if sharding_dim is None:
                      val = sharded_pieces[0][k]
                  else:
                      val = torch.cat(
                          [piece[k] for piece in sharded_pieces],
                          dim=sharding_dim,
                      )
                  fix_factor = get_shard_fix_factor(
                      parameter_name, model_H, old_model_H
                  )
                  model_piece[k] = (val.float() * fix_factor).to(
                      val.dtype
                  )
  
              return model_piece
  
      elif layout.n_shards % original_n_shards == 0:
          n_ranks_per_chkpt_shard = exact_div(
              layout.n_shards, original_n_shards
          )
          shard_idx_to_load = (
              layout.shard_idx // n_ranks_per_chkpt_shard
          )
          shard_slice_idx = layout.shard_idx % n_ranks_per_chkpt_shard
  
          def fetch(chkpt_prefix: str, module_name: str = ""):
              unsharded_piece = fetch_single_piece(
                  f"{chkpt_prefix}_shard_{shard_idx_to_load:03d}.pkl"
              )
  
              model_piece = {}
              for k in unsharded_piece.keys():
                  parameter_name = (
                      ".".join([module_name, k]) if module_name else k
                  )
  
                  sharding_dim = dim_to_shard(parameter_name)
  
                  if sharding_dim is None:
                      val = unsharded_piece[k]
                  else:
                      split_size = exact_div(
                          unsharded_piece[k].size()[sharding_dim],
                          n_ranks_per_chkpt_shard,
                      )
                      val = torch.split(
                          unsharded_piece[k],
                          [split_size] * n_ranks_per_chkpt_shard,
                          dim=sharding_dim,
                      )[shard_slice_idx]
                  fix_factor = get_shard_fix_factor(
                      parameter_name, model_H, old_model_H
                  )
                  model_piece[k] = (val.float() * fix_factor).to(
                      val.dtype
                  )
  
              return model_piece
  
      else:
          raise NotImplementedError(
              f"Tried running a model that was originally created with "
              f"{original_n_shards} shards with {layout.n_shards} shards. The new number "
              f"of shards must evenly divide or be divisible by the original number of shards."
          )
  
      if model.include_input_embeddings:
          model.embedding.load_state_dict(
              fetch("input_embeddings", "embedding")
          )
  
          if model.include_pos_embeddings:
              model.position_embedding.load_state_dict(
                  fetch("position_embedding", "position_embedding")
              )
  
      # fetch the resblocks
      for resblock_idx, resblock in enumerate(model.torso.resblocks):
          d = fetch(
              f"resblock_{resblock_idx:04d}",
              f"torso.resblocks.{resblock_idx}",
          )
          if not model_H.get("key_bias"):
              d = {
                  k: v
                  for (k, v) in d.items()
                  if "attn.k_proj.bias" not in k
              }
          resblock.load_state_dict(d)
  
      # fetch the final_layer_norm
      if model.include_final_layer_norm:
          model.ln_f.load_state_dict(fetch("final_layer_norm", "ln_f"))
  
      # fetch the unembeddings
      if model.include_output_unembeddings:
          # Pull in the one piece
          model.load_state_dict(
              fetch("output_unembeddings"), strict=False
          )
  
      for model_head, save_head in load_heads_map.items():
          model.scalar_heads[model_head].load_state_dict(
              fetch(
                  f"output_head_{save_head}",
                  f"scalar_heads.{model_head}",
              )
          )
  
imports: import json
  import os
  import re
  import shutil
  import uuid
  from dataclasses import dataclass, field
  from typing import Callable, Optional, Set, List, Dict
  import blobfile
  import numpy
  import torch
  from mpi4py import MPI
  import summarize_from_feedback
  from summarize_from_feedback.model_layout import ModelLayout
  from summarize_from_feedback.models import sample_fns
  from summarize_from_feedback.models.loss_functions import softmax_xent_loss_fn
  from summarize_from_feedback.models.transformer import Hyperparams
  from summarize_from_feedback.models.transformer import build_with_random_weights
  from summarize_from_feedback.utils import exact_div, hyperparams
  from summarize_from_feedback.utils import blobs
  from summarize_from_feedback.utils.dist_utils import setup_cuda_device_and_dist, create_data_parallel_comm, create_within_replica_comm
  from summarize_from_feedback.utils.nested import map_nested
  from summarize_from_feedback.utils.torch_utils import nans, to_numpy
label_desc: Use image node_id 7 for calling a global function with example usage: load_exported_model(layout, model, model_H, load_path, load_heads_map, use_cache) and returns: torch, model_piece, model_piece