code: def upgrade_state_dict_with_xlm_weights(
      state_dict: Dict[str, Any], pretrained_xlm_checkpoint: str
  ) -> Dict[str, Any]:
      """
      Load XLM weights into a Transformer encoder or decoder model.
  
      Args:
          state_dict: state dict for either TransformerEncoder or
              TransformerDecoder
          pretrained_xlm_checkpoint: checkpoint to load XLM weights from
  
      Raises:
          AssertionError: If architecture (num layers, attention heads, etc.)
              does not match between the current Transformer encoder or
              decoder and the pretrained_xlm_checkpoint
      """
      if not os.path.exists(pretrained_xlm_checkpoint):
          raise IOError(
              "Model file not found: {}".format(
                  pretrained_xlm_checkpoint
              )
          )
  
      state = checkpoint_utils.load_checkpoint_to_cpu(
          pretrained_xlm_checkpoint
      )
      xlm_state_dict = state["model"]
      for key in xlm_state_dict.keys():
          for search_key in [
              "embed_tokens",
              "embed_positions",
              "layers",
          ]:
              if search_key in key:
                  subkey = key[key.find(search_key) :]
                  assert subkey in state_dict, (
                      "{} Transformer encoder / decoder "
                      "state_dict does not contain {}. Cannot "
                      "load {} from pretrained XLM checkpoint "
                      "{} into Transformer.".format(
                          str(state_dict.keys()),
                          subkey,
                          key,
                          pretrained_xlm_checkpoint,
                      )
                  )
  
                  state_dict[subkey] = xlm_state_dict[key]
      return state_dict
  
file_path: fairseq/fairseq/models/transformer_from_pretrained_xlm.py
imports: import os
  from typing import Any, Dict
  from fairseq import checkpoint_utils
  from fairseq.data.legacy.masked_lm_dictionary import MaskedLMDictionary
  from fairseq.models import register_model, register_model_architecture
  from fairseq.models.transformer import TransformerDecoder, TransformerEncoder, TransformerModel, base_architecture
label_desc: Use image node_id 1 for calling a global function with example usage: upgrade_state_dict_with_xlm_weights(state_dict, pretrained_xlm_checkpoint) and returns: state_dict