class_docstr: A mixin class to support loading pre-trained checkpoints for modules
  that implement the GPT2 model.
  
  The GPT2 model was proposed in
  `Language Models are Unsupervised Multitask Learners`_
  by `Radford et al.` from OpenAI. It is a unidirectional Transformer model
  pre-trained using the vanilla language modeling objective on a large corpus.
  
  The available GPT2 models are as follows:
  
    * ``gpt2-small``: Small version of GPT-2, 124M parameters.
    * ``gpt2-medium``: Medium version of GPT-2, 355M parameters.
    * ``gpt2-large``: Large version of GPT-2, 774M parameters.
    * ``gpt2-xl``: XL version of GPT-2, 1558M parameters.
  
  We provide the following GPT2 classes:
  
    * :class:`~texar.tf.modules.GPT2Encoder` for text encoding.
    * :class:`~texar.tf.modules.GPT2Decoder` for text generation and
      decoding.
    * :class:`~texar.tf.modules.GPT2Classifier` for text classification and
      sequence tagging.
  
  .. _`Language Models are Unsupervised Multitask Learners`:
      https://openai.com/blog/better-language-models/
class_docstr_tok: ["A","mixin","class","to","support","loading","pre-trained","checkpoints","for","modules","that","implement","the","GPT2","model",".","The","GPT2","model","was","proposed","in","`","Language","Models","are","Unsupervised","Multitask","Learners","`","_","by","`","Radford","et","al.","`","from","OpenAI",".","It","is","a","unidirectional","Transformer","model","pre-trained","using","the","vanilla","language","modeling","objective","on","a","large","corpus",".","The","available","GPT2","models","are","as","follows",":","*","``","gpt2-small","``",":","Small","version","of","GPT-2",",","124M","parameters",".","*","``","gpt2-medium","``",":","Medium","version","of","GPT-2",",","355M","parameters",".","*","``","gpt2-large","``",":","Large","version","of","GPT-2",",","774M","parameters",".","*","``","gpt2-xl","``",":","XL","version","of","GPT-2",",","1558M","parameters",".","We","provide","the","following","GPT2","classes",":","*",":","class",":","`","~texar.tf.modules.GPT2Encoder","`","for","text","encoding",".","*",":","class",":","`","~texar.tf.modules.GPT2Decoder","`","for","text","generation","and","decoding",".","*",":","class",":","`","~texar.tf.modules.GPT2Classifier","`","for","text","classification","and","sequence","tagging",".","..","_","`","Language","Models","are","Unsupervised","Multitask","Learners","`",":","https",":","\/\/openai.com\/blog\/better-language-models\/"]
docstr: Initialize model parameters from weights stored in the pre-trained
  checkpoint.
  
  Args:
      pretrained_model_name (str): Name of the pre-trained model.
      cache_dir (str): Path to the cache directory.
      scope_name (str): Scope name of the model.
      load_output_layer (bool): If `False`, will not load weights of the
          output layer. Set this argument to `False` when loading weights
          into a GPT2 encoder. Defaults to `True`.
docstr_tok: ["Initialize","model","parameters","from","weights","stored","in","the","pre-trained","checkpoint",".","Args",":","pretrained_model_name","(","str",")",":","Name","of","the","pre-trained","model",".","cache_dir","(","str",")",":","Path","to","the","cache","directory",".","scope_name","(","str",")",":","Scope","name","of","the","model",".","load_output_layer","(","bool",")",":","If","`","False","`",",","will","not","load","weights","of","the","output","layer",".","Set","this","argument","to","`","False","`","when","loading","weights","into","a","GPT2","encoder",".","Defaults","to","`","True","`","."]
code: def _init_from_checkpoint(
      self,
      pretrained_model_name,
      cache_dir,
      scope_name,
      load_output_layer=True,
      **kwargs,
  ):
      r"""Initialize model parameters from weights stored in the pre-trained
      checkpoint.
  
      Args:
          pretrained_model_name (str): Name of the pre-trained model.
          cache_dir (str): Path to the cache directory.
          scope_name (str): Scope name of the model.
          load_output_layer (bool): If `False`, will not load weights of the
              output layer. Set this argument to `False` when loading weights
              into a GPT2 encoder. Defaults to `True`.
      """
      init_checkpoint = os.path.abspath(
          os.path.join(cache_dir, "model.ckpt")
      )
      ckpt = tf.train.load_checkpoint(init_checkpoint)
      ckpt_params = {
          key: ckpt.get_tensor(key)
          for key in ckpt.get_variable_to_shape_map().keys()
      }
  
      tvars = tf.trainable_variables()
      name_to_variable = collections.OrderedDict()
      for var in tvars:
          name = var.name
          m = re.match("^(.*):\\d+$", name)
          if m is not None:
              name = m.group(1)
          name_to_variable[name] = var
  
      if load_output_layer:
          global_tensor_map = {
              "model/wte": scope_name + "/word_embeddings/w",
              "model/wpe": scope_name + "/position_embeddings/w",
              "model/ln_f/b": scope_name + "/decoder/beta",
              "model/ln_f/g": scope_name + "/decoder/gamma",
          }
  
          layer_tensor_map = {
              "ln_1/b": scope_name + "/layer_{}/beta",
              "ln_1/g": scope_name + "/layer_{}/gamma",
              "ln_2/b": scope_name + "/layer_{}/past_poswise_ln/beta",
              "ln_2/g": scope_name + "/layer_{}/past_poswise_ln/gamma",
              "mlp/c_fc/b": scope_name + "/decoder/layer_{}"
              "/ffn/intermediate/bias",
              "mlp/c_fc/w": scope_name + "/decoder/layer_{}"
              "/ffn/intermediate/kernel",
              "mlp/c_proj/b": scope_name
              + "/decoder/layer_{}/ffn/output/"
              "bias",
              "mlp/c_proj/w": scope_name
              + "/decoder/layer_{}/ffn/output/"
              "kernel",
              "attn/c_attn/b": None,
              "attn/c_attn/w": None,
              "attn/c_proj/b": scope_name + "/decoder/layer_{}"
              "/self_attention/self/output/"
              "bias",
              "attn/c_proj/w": scope_name + "/decoder/layer_{}"
              "/self_attention/self/output/"
              "kernel",
          }
      else:
          global_tensor_map = {
              "model/wte": scope_name + "/word_embeddings/w",
              "model/wpe": scope_name + "/position_embeddings/w",
              "model/ln_f/b": scope_name + "/encoder/LayerNorm/beta",
              "model/ln_f/g": scope_name + "/encoder/LayerNorm/gamma",
          }
  
          layer_tensor_map = {
              "ln_1/b": scope_name + "/encoder/layer_{}/LayerNorm/beta",
              "ln_1/g": scope_name
              + "/encoder/layer_{}/LayerNorm/gamma",
              "ln_2/b": scope_name + "/encoder/layer_{}/output/"
              "LayerNorm/beta",
              "ln_2/g": scope_name + "/encoder/layer_{}/output/"
              "LayerNorm/gamma",
              "mlp/c_fc/b": scope_name + "/encoder/layer_{}"
              "/ffn/intermediate/bias",
              "mlp/c_fc/w": scope_name + "/encoder/layer_{}"
              "/ffn/intermediate/kernel",
              "mlp/c_proj/b": scope_name
              + "/encoder/layer_{}/ffn/output/"
              "bias",
              "mlp/c_proj/w": scope_name
              + "/encoder/layer_{}/ffn/output/"
              "kernel",
              "attn/c_attn/b": None,
              "attn/c_attn/w": None,
              "attn/c_proj/b": scope_name + "/encoder/layer_{}"
              "/attention/self/output/bias",
              "attn/c_proj/w": scope_name + "/encoder/layer_{}"
              "/attention/self/output/kernel",
          }
  
      for name, array in ckpt_params.items():
          if name in global_tensor_map:
              v_name = global_tensor_map[name]
              pointer = name_to_variable[v_name]
              pointer._initializer_op = tf.assign(
                  pointer._variable, array
              )
          else:
              name_tmp = name.split("/")
              layer_no = name_tmp[1][1:]
              name = "/".join(name_tmp[2:])
  
              if name in layer_tensor_map:
                  if name == "attn/c_attn/b":
                      if load_output_layer:
                          K = name_to_variable[
                              scope_name
                              + "/decoder/layer_"
                              + layer_no
                              + "/self_attention/self/key/bias"
                          ]
                          Q = name_to_variable[
                              scope_name
                              + "/decoder/layer_"
                              + layer_no
                              + "/self_attention/self/query/bias"
                          ]
                          V = name_to_variable[
                              scope_name
                              + "/decoder/layer_"
                              + layer_no
                              + "/self_attention/self/value/bias"
                          ]
                      else:
                          K = name_to_variable[
                              scope_name
                              + "/encoder/layer_"
                              + layer_no
                              + "/attention/self/key/bias"
                          ]
                          Q = name_to_variable[
                              scope_name
                              + "/encoder/layer_"
                              + layer_no
                              + "/attention/self/query/bias"
                          ]
                          V = name_to_variable[
                              scope_name
                              + "/encoder/layer_"
                              + layer_no
                              + "/attention/self/value/bias"
                          ]
  
                      index_d = array.shape[-1] // 3
  
                      Q_w = array[:index_d]
                      K_w = array[index_d : 2 * index_d]
                      V_w = array[2 * index_d :]
  
                      K._initializer_op = tf.assign(K._variable, K_w)
                      Q._initializer_op = tf.assign(Q._variable, Q_w)
                      V._initializer_op = tf.assign(V._variable, V_w)
                  elif name == "attn/c_attn/w":
                      if load_output_layer:
                          K = name_to_variable[
                              scope_name
                              + "/decoder/layer_"
                              + layer_no
                              + "/self_attention/self/key/kernel"
                          ]
                          Q = name_to_variable[
                              scope_name
                              + "/decoder/layer_"
                              + layer_no
                              + "/self_attention/self/query/kernel"
                          ]
                          V = name_to_variable[
                              scope_name
                              + "/decoder/layer_"
                              + layer_no
                              + "/self_attention/self/value/kernel"
                          ]
                      else:
                          K = name_to_variable[
                              scope_name
                              + "/encoder/layer_"
                              + layer_no
                              + "/attention/self/key/kernel"
                          ]
                          Q = name_to_variable[
                              scope_name
                              + "/encoder/layer_"
                              + layer_no
                              + "/attention/self/query/kernel"
                          ]
                          V = name_to_variable[
                              scope_name
                              + "/encoder/layer_"
                              + layer_no
                              + "/attention/self/value/kernel"
                          ]
  
                      index_d = array.shape[-1] // 3
  
                      Q_w = np.transpose(array[0, :, :index_d])
                      K_w = np.transpose(
                          array[0, :, index_d : 2 * index_d]
                      )
                      V_w = np.transpose(array[0, :, 2 * index_d :])
  
                      K._initializer_op = tf.assign(K._variable, K_w)
                      Q._initializer_op = tf.assign(Q._variable, Q_w)
                      V._initializer_op = tf.assign(V._variable, V_w)
                  elif (
                      name == "attn/c_proj/w"
                      or name == "mlp/c_fc/w"
                      or name == "mlp/c_proj/w"
                  ):
                      v_name = layer_tensor_map[name]
                      pointer = name_to_variable[
                          v_name.format(layer_no)
                      ]
                      pointer._initializer_op = tf.assign(
                          pointer._variable, array[0]
                      )
                  else:
                      v_name = layer_tensor_map[name]
                      pointer = name_to_variable[
                          v_name.format(layer_no)
                      ]
                      pointer._initializer_op = tf.assign(
                          pointer._variable, array
                      )
  
file_path: texar/texar/tf/modules/pretrained/gpt2.py
imports: import collections
  import json
  import os
  import re
  import warnings
  from abc import ABC
  from typing import Any, Dict
  import tensorflow
  import numpy
  from texar.tf.modules.pretrained.pretrained_base import PretrainedMixin
label_desc: Use image node_id 2 for calling the PretrainedGPT2Mixin obj's underlying member method code with example usage: obj._init_from_checkpoint(pretrained_model_name, cache_dir, scope_name, load_output_layer) without return types