code: def language_model(
      *,
      X,
      sequence_lengths,
      embed_weights,
      hidden,
      config,
      reuse=None,
      train=False,
  ):
      """
      A language model output and loss for the language modelling objective described in the original finetune paper.
      This language model uses weights that are tied to the input embedding.
      :param X: The raw token ids fed to the featurizer.
      :param M: A loss mask, with 1's where losses should be counted and 0's elsewhere.
      :param embed_weights: The word embedding matrix, normally the one returned by the featurizer.
      :param hidden: Output of the featurizer.
      :param config: A config object.
      :param reuse: A Flag passed through to the tf.variable_scope context manager.
      :return: A dict containing:
          logits: The un-normalised log-probabilities over each word in the vocabulary.
          loss: The masked language modelling loss.
  
      """
      X = merge_leading_dims(X, 2)
      M = tf.sequence_mask(sequence_lengths, dtype=tf.float32)
      hidden = merge_leading_dims(hidden, 3)
  
      batch, seq = shape_list(X)
      vocab_size, hidden_dim = shape_list(embed_weights)
  
      with tf.compat.v1.variable_scope(
          "model/language-model", reuse=reuse
      ):
          # language model ignores last hidden state because we don't have a target
          lm_h = tf.reshape(
              hidden, [-1, config.n_embed]
          )  # [batch, seq_len, embed] --> [batch * seq_len, embed]
          lm_logits = tf.matmul(
              lm_h, embed_weights, transpose_b=True
          )  # tied weights
          lm_logits = tf.cast(lm_logits, tf.float32)
          hidden_shape = tf.shape(input=hidden)
          logits = tf.reshape(
              lm_logits,
              shape=tf.concat(
                  [hidden_shape[:-1], [vocab_size]], axis=0
              ),
          )
          lm_logits_offset = tf.reshape(
              logits[:, :-1], [-1, vocab_size]
          )
  
          lm_losses = tf.compat.v1.losses.sparse_softmax_cross_entropy(
              logits=lm_logits_offset,
              labels=tf.reshape(X[:, 1:], [-1]),
              weights=tf.reshape(M[:, 1:], [-1]),
          )
  
          perplexity = tf.math.divide_no_nan(
              tf.reduce_sum(
                  input_tensor=tf.exp(lm_losses) * M[:, 1:], axis=1
              ),
              tf.reduce_sum(input_tensor=M[:, 1:], axis=1),
          )
  
          return {
              "logits": logits,
              "losses": lm_losses,
              "perplexity": perplexity,
          }
  
imports: import functools
  import tensorflow
  from tensorflow_addons.text.crf import crf_log_likelihood
  from finetune.base_models.gpt.featurizer import attn, dropout, norm
  from finetune.util.shapes import shape_list, merge_leading_dims
  from finetune.optimizers.recompute_grads import recompute_grad
  from finetune.errors import FinetuneError
  from finetune.nn.activations import act_fns
  from finetune.nn.nn_utils import norm
  from tensorflow.python.framework import function
label_desc: Use image node_id 3 for calling a global function with example usage: language_model() and returns: dict