code: def get_nll(
      self,
      seq_logits,
      doc_scores,
      target,
      reduce_loss=False,
      epsilon=0.0,
      exclude_bos_score=False,
      n_docs=None,
  ):
      # shift tokens left
      target = tf.concat(
          [
              target[:, 1:],
              tf.fill(
                  [target.shape[0], 1],
                  tf.cast(
                      self.config.generator.pad_token_id, target.dtype
                  ),
              ),
          ],
          axis=1,
      )
  
      # bos_token_id is None for T5
      bos_token_id = (
          self.config.bos_token_id or self.config.generator.bos_token_id
      )
      n_docs = n_docs if n_docs is not None else self.config.n_docs
      equal_bos_token_id_all = tf.reduce_all(
          tf.equal(target[:, 0], bos_token_id)
      )
      use_bos = bos_token_id is not None and equal_bos_token_id_all
  
      def _mask_pads(ll, smooth_obj):
          pad_mask = tf.equal(
              target,
              tf.cast(self.config.generator.pad_token_id, target.dtype),
          )
          if tf.reduce_any(pad_mask):
              ll = tf.where(pad_mask, 0.0, ll)
              smooth_obj = tf.where(pad_mask, 0.0, smooth_obj)
          return tf.squeeze(ll, axis=-1), tf.squeeze(
              smooth_obj, axis=-1
          )
  
      # seq_logits.shape = (batch*n_docs, tgt_len , vocabs)
      seq_logprobs = tf.nn.log_softmax(seq_logits, axis=-1)
      seq_logprobs = tf.reshape(
          seq_logprobs,
          (
              seq_logits.shape[0] // n_docs,
              n_docs,
              -1,
              seq_logits.shape[-1],
          ),
      )  # (batch_size, n_docs, tgt_len, vocabs)
      doc_logprobs = tf.nn.log_softmax(doc_scores, axis=1)
      doc_logprobs = tf.expand_dims(doc_logprobs, axis=-1)
      doc_logprobs = tf.expand_dims(
          doc_logprobs, axis=-1
      )  # done twice to get 4-D
  
      # RAG-sequence marginalization
      first_token_scores = seq_logprobs[:, :, :1, :]
      second_token_scores = seq_logprobs[:, :, 1:2, :]
      remainder = seq_logprobs[:, :, 2:, :]
      rag_logprobs = tf.concat(
          [
              first_token_scores,
              second_token_scores + doc_logprobs,
              remainder,
          ],
          axis=2,
      )
  
      # calculate loss
      target = tf.expand_dims(target, axis=1)  # n_docs dimension
      target = tf.expand_dims(target, axis=-1)  # logits dimension
      target = tf.repeat(target, n_docs, axis=1)
      assert len(target.shape) == len(rag_logprobs.shape)
  
      # last-axis gathering only - use 2D-reshape-trick for Torch's style nD gathering
      def torch_gather(param, id_tensor):
          # 2d-gather torch equivalent: https://stackoverflow.com/questions/52129909/tensorflow-equivalent-of-torch-gather
          def gather2d(target, id_tensor):
              idx = tf.stack(
                  [
                      tf.range(
                          tf.shape(id_tensor)[0], dtype=id_tensor.dtype
                      ),
                      id_tensor[:, 0],
                  ],
                  axis=-1,
              )
              result = tf.gather_nd(target, idx)
              return tf.expand_dims(result, axis=-1)
  
          target = tf.reshape(
              param, (-1, param.shape[-1])
          )  # reshape 2D
          target_shape = id_tensor.shape
  
          id_tensor = tf.reshape(id_tensor, (-1, 1))  # also 2D-index
          result = gather2d(target, id_tensor)
          return tf.reshape(result, target_shape)
  
      ll = torch_gather(rag_logprobs, id_tensor=target)
      smooth_obj = tf.reduce_sum(
          rag_logprobs, axis=-1, keepdims=True
      )  # total sum of all (normalised) logits
  
      ll, smooth_obj = _mask_pads(ll, smooth_obj)
  
      # sum over tokens, exclude bos while scoring
      if exclude_bos_score and use_bos:
          ll = tf.reduce_sum(ll[:, :, 1:], axis=2)
      else:
          ll = tf.reduce_sum(ll, axis=2)
  
      smooth_obj = tf.reduce_sum(smooth_obj, axis=2)
      ll = tf.math.reduce_logsumexp(ll, axis=1)  # logsumexp over docs
      smooth_obj = tf.math.reduce_logsumexp(smooth_obj, axis=1)
  
      nll_loss = -ll
      smooth_loss = -smooth_obj
  
      if reduce_loss:
          nll_loss = tf.reduce_sum(nll_loss)
          smooth_loss = tf.reduce_sum(smooth_loss)
  
      eps_i = epsilon / rag_logprobs.shape[-1]
      loss = (1.0 - epsilon) * nll_loss + eps_i * smooth_loss
      return loss
  
imports: from __future__ import annotations
  import copy
  from dataclasses import dataclass
  from typing import List, Optional, Tuple, Union
  import numpy
  import tensorflow
  from ...configuration_utils import PretrainedConfig
  from ...generation import TFLogitsProcessorList
  from ...modeling_tf_utils import TFCausalLanguageModelingLoss, TFModelInputType, TFPreTrainedModel, shape_list, unpack_inputs
  from ...utils import ModelOutput, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
  from .configuration_rag import RagConfig
  from .retrieval_rag import RagRetriever
label_desc: Use image node_id 7 for calling the TFRagSequenceForGeneration obj's underlying member method code with example usage: obj.get_nll(seq_logits, doc_scores, target, reduce_loss, epsilon, exclude_bos_score, n_docs) and returns: loss, tf,  tf, tf, tf