docstr: do_marginalize (`bool`, *optional*):
      If `True`, the logits are marginalized over all documents by making use of
      `torch.nn.functional.log_softmax`.
  labels (`tf.Tensor` or `np.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
      Labels for computing the cross entropy classification loss according to Rag-Token model formulation See
      https://arxiv.org/pdf/2005.11401.pdf Section 2.1 for details about Rag-Token formulation. Indices should be
      in `[0, ..., config.vocab_size - 1]`.
  reduce_loss (`bool`, *optional*):
      Only relevant if `labels` is passed. If `True`, the NLL loss is reduced using the `tf.Tensor.sum`
      operation.
  kwargs (`Dict[str, any]`, optional, defaults to *{}*):
      Legacy dictionary, which is required so that model can use *generate()* function.
  
  Returns:
  
  Example:
  
  ```python
  >>> import tensorflow as tf
  >>> from transformers import AutoTokenizer, RagRetriever, TFRagTokenForGeneration
  
  >>> tokenizer = AutoTokenizer.from_pretrained("facebook/rag-token-nq")
  >>> retriever = RagRetriever.from_pretrained(
  ...     "facebook/rag-token-nq", index_name="exact", use_dummy_dataset=True
  ... )
  >>> # initialize with RagRetriever to do everything in one forward call
  >>> model = TFRagTokenForGeneration.from_pretrained("facebook/rag-token-nq", retriever=retriever, from_pt=True)
  
  >>> input_dict = tokenizer.prepare_seq2seq_batch(
  ...     "How many people live in Paris?", "In Paris, there are 10 million people.", return_tensors="tf"
  ... )
  >>> outputs = model(input_dict, output_retrieved=True)
  
  >>> # or use retriever separately
  >>> # 1. Encode
  >>> input_ids = input_dict["input_ids"]
  >>> question_hidden_states = model.question_encoder(input_ids)[0]
  >>> # 2. Retrieve
  >>> docs_dict = retriever(input_ids.numpy(), question_hidden_states.numpy(), return_tensors="tf")
  >>> doc_scores = tf.squeeze(
  ...     tf.matmul(
  ...         tf.expand_dims(question_hidden_states, axis=1), docs_dict["retrieved_doc_embeds"], transpose_b=True
  ...     ),
  ...     axis=1,
  ... )
  >>> # 3. Forward to generator
  >>> outputs = model(
  ...     inputs=None,
  ...     context_input_ids=docs_dict["context_input_ids"],
  ...     context_attention_mask=docs_dict["context_attention_mask"],
  ...     doc_scores=doc_scores,
  ...     decoder_input_ids=input_dict["labels"],
  ... )
  
  >>> # or directly generate
  >>> generated = model.generate(
  ...     context_input_ids=docs_dict["context_input_ids"],
  ...     context_attention_mask=docs_dict["context_attention_mask"],
  ...     doc_scores=doc_scores,
  ... )
  >>> generated_string = tokenizer.batch_decode(generated, skip_special_tokens=True)
  ```
docstr_tok: ["do_marginalize","(","`","bool","`",",","*","optional","*",")",":","If","`","True","`",",","the","logits","are","marginalized","over","all","documents","by","making","use","of","`","torch.nn.functional.log_softmax","`",".","labels","(","`","tf.Tensor","`","or","`","np.ndarray","`","of","shape","`","(","batch_size",",","sequence_length",")","`",",","*","optional","*",")",":","Labels","for","computing","the","cross","entropy","classification","loss","according","to","Rag-Token","model","formulation","See","https",":","\/\/arxiv.org\/pdf\/2005.11401.pdf","Section","2.1","for","details","about","Rag-Token","formulation",".","Indices","should","be","in","`","[","0",",","...",",","config.vocab_size","-","1","]","`",".","reduce_loss","(","`","bool","`",",","*","optional","*",")",":","Only","relevant","if","`","labels","`","is","passed",".","If","`","True","`",",","the","NLL","loss","is","reduced","using","the","`","tf.Tensor.sum","`","operation",".","kwargs","(","`","Dict","[","str",",","any","]","`",",","optional",",","defaults","to","*","{","}","*",")",":","Legacy","dictionary",",","which","is","required","so","that","model","can","use","*","generate","(",")","*","function",".","Returns",":","Example",":","``","`","python",">",">",">","import","tensorflow","as","tf",">",">",">","from","transformers","import","AutoTokenizer",",","RagRetriever",",","TFRagTokenForGeneration",">",">",">","tokenizer","=","AutoTokenizer.from_pretrained","(","``","facebook\/rag-token-nq","''",")",">",">",">","retriever","=","RagRetriever.from_pretrained","(","...","``","facebook\/rag-token-nq","''",",","index_name=","''","exact","''",",","use_dummy_dataset=True","...",")",">",">",">","#","initialize","with","RagRetriever","to","do","everything","in","one","forward","call",">",">",">","model","=","TFRagTokenForGeneration.from_pretrained","(","``","facebook\/rag-token-nq","''",",","retriever=retriever",",","from_pt=True",")",">",">",">","input_dict","=","tokenizer.prepare_seq2seq_batch","(","...","``","How","many","people","live","in","Paris","?","``",",","``","In","Paris",",","there","are","10","million","people",".","``",",","return_tensors=","''","tf","''","...",")",">",">",">","outputs","=","model","(","input_dict",",","output_retrieved=True",")",">",">",">","#","or","use","retriever","separately",">",">",">","#","1",".","Encode",">",">",">","input_ids","=","input_dict","[","``","input_ids","''","]",">",">",">","question_hidden_states","=","model.question_encoder","(","input_ids",")","[","0","]",">",">",">","#","2",".","Retrieve",">",">",">","docs_dict","=","retriever","(","input_ids.numpy","(",")",",","question_hidden_states.numpy","(",")",",","return_tensors=","''","tf","''",")",">",">",">","doc_scores","=","tf.squeeze","(","...","tf.matmul","(","...","tf.expand_dims","(","question_hidden_states",",","axis=1",")",",","docs_dict","[","``","retrieved_doc_embeds","''","]",",","transpose_b=True","...",")",",","...","axis=1",",","...",")",">",">",">","#","3",".","Forward","to","generator",">",">",">","outputs","=","model","(","...","inputs=None",",","...","context_input_ids=docs_dict","[","``","context_input_ids","''","]",",","...","context_attention_mask=docs_dict","[","``","context_attention_mask","''","]",",","...","doc_scores=doc_scores",",","...","decoder_input_ids=input_dict","[","``","labels","''","]",",","...",")",">",">",">","#","or","directly","generate",">",">",">","generated","=","model.generate","(","...","context_input_ids=docs_dict","[","``","context_input_ids","''","]",",","...","context_attention_mask=docs_dict","[","``","context_attention_mask","''","]",",","...","doc_scores=doc_scores",",","...",")",">",">",">","generated_string","=","tokenizer.batch_decode","(","generated",",","skip_special_tokens=True",")","``","`"]
code: def call(
      self,
      input_ids: TFModelInputType | None = None,
      attention_mask: np.ndarray | tf.Tensor | None = None,
      decoder_input_ids: np.ndarray | tf.Tensor | None = None,
      decoder_attention_mask: np.ndarray | tf.Tensor | None = None,
      encoder_outputs: np.ndarray | tf.Tensor | None = None,
      past_key_values: Tuple[Tuple[Union[np.ndarray, tf.Tensor]]]
      | None = None,
      doc_scores: np.ndarray | tf.Tensor | None = None,
      context_input_ids: np.ndarray | tf.Tensor | None = None,
      context_attention_mask: np.ndarray | tf.Tensor | None = None,
      use_cache: bool | None = None,
      output_attentions: bool | None = None,
      output_hidden_states: bool | None = None,
      output_retrieved: bool | None = None,
      n_docs: int | None = None,
      do_marginalize: bool | None = None,
      labels: np.ndarray | tf.Tensor | None = None,
      reduce_loss: bool | None = None,
      return_dict: bool | None = None,
      training: bool = False,
      **kwargs,  # needs kwargs for generation
  ) -> TFRetrievAugLMMarginOutput:
      r"""
      do_marginalize (`bool`, *optional*):
          If `True`, the logits are marginalized over all documents by making use of
          `torch.nn.functional.log_softmax`.
      labels (`tf.Tensor` or `np.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
          Labels for computing the cross entropy classification loss according to Rag-Token model formulation See
          https://arxiv.org/pdf/2005.11401.pdf Section 2.1 for details about Rag-Token formulation. Indices should be
          in `[0, ..., config.vocab_size - 1]`.
      reduce_loss (`bool`, *optional*):
          Only relevant if `labels` is passed. If `True`, the NLL loss is reduced using the `tf.Tensor.sum`
          operation.
      kwargs (`Dict[str, any]`, optional, defaults to *{}*):
          Legacy dictionary, which is required so that model can use *generate()* function.
  
      Returns:
  
      Example:
  
      ```python
      >>> import tensorflow as tf
      >>> from transformers import AutoTokenizer, RagRetriever, TFRagTokenForGeneration
  
      >>> tokenizer = AutoTokenizer.from_pretrained("facebook/rag-token-nq")
      >>> retriever = RagRetriever.from_pretrained(
      ...     "facebook/rag-token-nq", index_name="exact", use_dummy_dataset=True
      ... )
      >>> # initialize with RagRetriever to do everything in one forward call
      >>> model = TFRagTokenForGeneration.from_pretrained("facebook/rag-token-nq", retriever=retriever, from_pt=True)
  
      >>> input_dict = tokenizer.prepare_seq2seq_batch(
      ...     "How many people live in Paris?", "In Paris, there are 10 million people.", return_tensors="tf"
      ... )
      >>> outputs = model(input_dict, output_retrieved=True)
  
      >>> # or use retriever separately
      >>> # 1. Encode
      >>> input_ids = input_dict["input_ids"]
      >>> question_hidden_states = model.question_encoder(input_ids)[0]
      >>> # 2. Retrieve
      >>> docs_dict = retriever(input_ids.numpy(), question_hidden_states.numpy(), return_tensors="tf")
      >>> doc_scores = tf.squeeze(
      ...     tf.matmul(
      ...         tf.expand_dims(question_hidden_states, axis=1), docs_dict["retrieved_doc_embeds"], transpose_b=True
      ...     ),
      ...     axis=1,
      ... )
      >>> # 3. Forward to generator
      >>> outputs = model(
      ...     inputs=None,
      ...     context_input_ids=docs_dict["context_input_ids"],
      ...     context_attention_mask=docs_dict["context_attention_mask"],
      ...     doc_scores=doc_scores,
      ...     decoder_input_ids=input_dict["labels"],
      ... )
  
      >>> # or directly generate
      >>> generated = model.generate(
      ...     context_input_ids=docs_dict["context_input_ids"],
      ...     context_attention_mask=docs_dict["context_attention_mask"],
      ...     doc_scores=doc_scores,
      ... )
      >>> generated_string = tokenizer.batch_decode(generated, skip_special_tokens=True)
      ```"""
  
      assert (
          "decoder_cached_states" not in kwargs
      ), "Please use past_key_values to cache intermediate outputs"  # from modeling_tf_bart.py
  
      do_marginalize = (
          do_marginalize
          if do_marginalize
          else self.config.do_marginalize
      )
      reduce_loss = (
          reduce_loss if reduce_loss else self.config.reduce_loss
      )
  
      if labels is not None:
          if decoder_input_ids is None:
              decoder_input_ids = labels
          use_cache = False
  
      outputs = self.rag(
          input_ids,
          attention_mask=attention_mask,
          encoder_outputs=encoder_outputs,
          decoder_input_ids=decoder_input_ids,
          decoder_attention_mask=decoder_attention_mask,
          context_input_ids=context_input_ids,
          context_attention_mask=context_attention_mask,
          doc_scores=doc_scores,
          past_key_values=past_key_values,
          use_cache=use_cache,
          output_attentions=output_attentions,
          output_hidden_states=output_hidden_states,
          output_retrieved=output_retrieved,
          n_docs=n_docs,
          training=training,
      )
  
      loss = None
      logits = outputs.logits
      if labels is not None:
          assert decoder_input_ids is not None
          loss = self.get_nll(
              outputs.logits,
              outputs.doc_scores,
              labels,
              reduce_loss=reduce_loss,
              epsilon=self.config.label_smoothing,
              n_docs=n_docs,
          )
  
      if do_marginalize:
          logits = self.marginalize(logits, outputs.doc_scores, n_docs)
  
      return TFRetrievAugLMMarginOutput(
          loss=loss,
          logits=logits,
          past_key_values=outputs.past_key_values,
          doc_scores=outputs.doc_scores,
          context_input_ids=outputs.context_input_ids,
          context_attention_mask=outputs.context_attention_mask,
          retrieved_doc_embeds=outputs.retrieved_doc_embeds,
          retrieved_doc_ids=outputs.retrieved_doc_ids,
          question_encoder_last_hidden_state=outputs.question_encoder_last_hidden_state,
          question_enc_hidden_states=outputs.question_enc_hidden_states,
          question_enc_attentions=outputs.question_enc_attentions,
          generator_enc_last_hidden_state=outputs.generator_enc_last_hidden_state,
          generator_enc_hidden_states=outputs.generator_enc_hidden_states,
          generator_enc_attentions=outputs.generator_enc_attentions,
          generator_dec_hidden_states=outputs.generator_dec_hidden_states,
          generator_dec_attentions=outputs.generator_dec_attentions,
      )
  
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 9 for calling the TFRagTokenForGeneration obj's underlying member method code with example usage: obj.call(input_ids, attention_mask, decoder_input_ids, decoder_attention_mask, encoder_outputs, past_key_values, doc_scores, context_input_ids, context_attention_mask, use_cache, output_attentions, output_hidden_states, output_retrieved, n_docs, do_marginalize, labels, reduce_loss, return_dict, training) and returns: TFRetrievAugLMMarginOutput