code: def gpt_featurizer(
      X,
      encoder,
      config,
      train=False,
      reuse=None,
      explain=False,
      **kwargs,
  ):
      """
      The transformer element of the finetuning model. Maps from tokens ids to a dense, embedding of the sequence.
  
      :param X: A tensor of token indexes with shape [batch_size, sequence_length, token_idx]
      :param encoder: A TextEncoder object.
      :param config: A config object, containing all parameters for the featurizer.
      :param train: If this flag is true, dropout and losses are added to the graph.
      :param reuse: Should reuse be set within this scope.
      :return: A dict containing;
          embed_weights: the word embedding matrix.
          features: The output of the featurizer_final state.
          sequence_features: The output of the featurizer at each timestep.
      """
      initial_shape = tf.shape(input=X)
      X = tf.reshape(X, shape=tf.concat(([-1], initial_shape[-1:]), 0))
      x_shape = tf.shape(input=X)
      sequence_length = x_shape[1]
      pos_values = get_pos_values(sequence_length, encoder.vocab_size)
      X = tf.stack((X, tf.tile(pos_values, [x_shape[0], 1])), 2)
  
      with tf.compat.v1.variable_scope("model/featurizer", reuse=reuse):
          embed_weights = tf.compat.v1.get_variable(
              name="we",
              shape=[
                  encoder.vocab_size + config.max_length,
                  config.n_embed,
              ],
              initializer=tf.compat.v1.random_normal_initializer(
                  stddev=config.weight_stddev
              ),
          )
          if config.train_embeddings:
              embed_weights = dropout(
                  embed_weights, config.embed_p_drop, train
              )
          else:
              embed_weights = tf.stop_gradient(embed_weights)
  
          clf_token = encoder.end_token
          pool_idx = tf.cast(
              tf.argmax(
                  input=tf.cast(
                      tf.equal(X[:, :, 0], clf_token), tf.float32
                  ),
                  axis=1,
              ),
              tf.int32,
          )
  
          if explain:
              X = add_explain_tokens(X, sequence_length, pool_idx)
  
          h = embed(X, embed_weights)
          for layer in range(config.n_layer):
              if (
                  (config.n_layer - layer) == config.num_layers_trained
                  and config.num_layers_trained != config.n_layer
              ):
                  h = tf.stop_gradient(h)
                  train_layer = False
              else:
                  train_layer = train
  
              with tf.compat.v1.variable_scope("h%d_" % layer):
                  block_fn = functools.partial(
                      block,
                      n_head=config.n_heads,
                      act_fn=config.act_fn,
                      resid_pdrop=config.resid_p_drop,
                      attn_pdrop=config.attn_p_drop,
                      scope="h%d" % layer,
                      train=train_layer,
                      scale=True,
                      explain=explain,
                  )
                  if config.low_memory_mode and train_layer:
                      block_fn = recompute_grad(
                          block_fn, use_entire_scope=True
                      )
                  if layer < config.n_layer - 1:
                      h = block_fn(h)
                  else:
                      h_out = block_fn(h)
  
              # get the attention weights from the last layer
              if layer == config.n_layer - 1:
                  with tf.compat.v1.variable_scope(
                      "h%d_/h%d/attn" % (layer, layer), reuse=True
                  ):
                      q, k, v = multihead_qkv(
                          h,
                          n_state=shape_list(h)[-1],
                          n_head=config.n_heads,
                          train=train,
                      )
                      w = attn_weights(q, k, v, scale=True)
  
          if explain:
              explain_out = h_out[:, initial_shape[1] :]
              explain_out = tf.reshape(
                  explain_out,
                  shape=tf.concat((initial_shape, [config.n_embed]), 0),
              )
              h_out = h_out[:, : initial_shape[1]]
  
          # Use hidden state at classifier token as input to final proj. + softmax
          clf_h = tf.reshape(
              h_out, [-1, config.n_embed]
          )  # [batch * seq_len, embed]
          clf_h = tf.gather(
              clf_h,
              tf.range(shape_list(X)[0], dtype=tf.int32)
              * sequence_length
              + pool_idx,
          )
          clf_h = tf.reshape(
              clf_h,
              shape=tf.concat(
                  (initial_shape[:-1], [config.n_embed]), 0
              ),
          )
          seq_feats = tf.reshape(
              h_out,
              shape=tf.concat((initial_shape, [config.n_embed]), 0),
          )
  
          lengths = lengths_from_eos_idx(
              eos_idx=pool_idx, max_length=sequence_length
          )
  
          out = {
              "embed_weights": embed_weights,
              "features": clf_h,
              "sequence_features": seq_feats,
              "eos_idx": pool_idx,
              "lengths": lengths,
              "attention_weights": w,  # [n_heads, seq_len, seq_len]
          }
          if explain:
              out["explain_out"] = explain_out
          return out
  
imports: import functools
  import numpy
  import tensorflow
  from finetune.optimizers.recompute_grads import recompute_grad
  from finetune.util.shapes import shape_list, lengths_from_eos_idx
  from finetune.nn.activations import act_fns
  from finetune.nn.nn_utils import dropout, norm
label_desc: Use image node_id 18 for calling a global function with example usage: gpt_featurizer(X, encoder, config, train, reuse, explain) and returns: out