docstr: Returns:
  
  Example:
  
  ```python
  >>> from transformers import AutoTokenizer, RagRetriever, TFRagModel
  >>> import torch
  
  >>> tokenizer = AutoTokenizer.from_pretrained("facebook/rag-token-base")
  >>> retriever = RagRetriever.from_pretrained(
  ...     "facebook/rag-token-base", index_name="exact", use_dummy_dataset=True
  ... )
  >>> # initialize with RagRetriever to do everything in one forward call
  >>> model = TFRagModel.from_pretrained("facebook/rag-token-base", 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"
  ... )
  >>> input_ids = input_dict["input_ids"]
  >>> outputs = model(input_ids)
  ```
docstr_tok: ["Returns",":","Example",":","``","`","python",">",">",">","from","transformers","import","AutoTokenizer",",","RagRetriever",",","TFRagModel",">",">",">","import","torch",">",">",">","tokenizer","=","AutoTokenizer.from_pretrained","(","``","facebook\/rag-token-base","''",")",">",">",">","retriever","=","RagRetriever.from_pretrained","(","...","``","facebook\/rag-token-base","''",",","index_name=","''","exact","''",",","use_dummy_dataset=True","...",")",">",">",">","#","initialize","with","RagRetriever","to","do","everything","in","one","forward","call",">",">",">","model","=","TFRagModel.from_pretrained","(","``","facebook\/rag-token-base","''",",","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","''","...",")",">",">",">","input_ids","=","input_dict","[","``","input_ids","''","]",">",">",">","outputs","=","model","(","input_ids",")","``","`"]
code: def call(
      self,
      input_ids: TFModelInputType | None = None,
      attention_mask: np.ndarray | tf.Tensor | None = None,
      encoder_outputs: np.ndarray | tf.Tensor | None = None,
      decoder_input_ids: np.ndarray | tf.Tensor | None = None,
      decoder_attention_mask: 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,
      return_dict: bool | None = None,
      training: bool = False,
      **kwargs,
  ) -> TFRetrievAugLMOutput:
      r"""
      Returns:
  
      Example:
  
      ```python
      >>> from transformers import AutoTokenizer, RagRetriever, TFRagModel
      >>> import torch
  
      >>> tokenizer = AutoTokenizer.from_pretrained("facebook/rag-token-base")
      >>> retriever = RagRetriever.from_pretrained(
      ...     "facebook/rag-token-base", index_name="exact", use_dummy_dataset=True
      ... )
      >>> # initialize with RagRetriever to do everything in one forward call
      >>> model = TFRagModel.from_pretrained("facebook/rag-token-base", 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"
      ... )
      >>> input_ids = input_dict["input_ids"]
      >>> outputs = model(input_ids)
      ```"""
      assert (
          "decoder_cached_states" not in kwargs
      ), "Please use past_key_values to cache intermediate outputs"  # from modeling_tf_bart.py
  
      # aliasing to minimize code changing
      n_docs = n_docs if n_docs is not None else self.config.n_docs
  
      # whether retriever has to be used
      has_to_retrieve = (
          self.retriever is not None
          and (
              context_input_ids is None
              or context_attention_mask is None
              or doc_scores is None
          )
          and encoder_outputs is None
      )
  
      # encoder_outputs are pre-computed during RAG-token generation
      if encoder_outputs is None:
          if has_to_retrieve:
              question_enc_outputs = self.question_encoder(
                  input_ids,
                  attention_mask=attention_mask,
                  return_dict=True,
                  training=training,
              )
              # see https://github.com/huggingface/transformers/blob/main/src/transformers/models/dpr/modeling_tf_dpr.py#L91
              question_encoder_last_hidden_state = question_enc_outputs[
                  0
              ]  # hidden states of question encoder => pooler_output
  
              retriever_outputs = self.retriever(
                  input_ids,
                  question_encoder_last_hidden_state.numpy(),
                  prefix=self.generator.config.prefix,
                  n_docs=n_docs,
                  return_tensors="tf",
              )
              (
                  context_input_ids,
                  context_attention_mask,
                  retrieved_doc_embeds,
                  retrieved_doc_ids,
              ) = (
                  retriever_outputs["context_input_ids"],
                  retriever_outputs["context_attention_mask"],
                  retriever_outputs["retrieved_doc_embeds"],
                  retriever_outputs["doc_ids"],
              )
  
              context_input_ids = tf.cast(context_input_ids, tf.int32)
              context_attention_mask = tf.cast(
                  context_attention_mask, tf.int32
              )
              retrieved_doc_embeds = tf.cast(
                  retrieved_doc_embeds, tf.float32
              )
              retrieved_doc_ids = tf.cast(retrieved_doc_ids, tf.int32)
  
              # compute doc_scores
              doc_scores = tf.squeeze(
                  tf.matmul(
                      tf.expand_dims(
                          question_encoder_last_hidden_state, axis=1
                      ),
                      retrieved_doc_embeds,
                      transpose_b=True,
                  ),
                  axis=1,
              )
  
          else:
              assert context_input_ids is not None, (
                  "Make sure that `context_input_ids` are passed, if no `retriever` is set. Alternatively, you can"
                  " set a retriever using the `set_retriever(...)` function."
              )
              assert context_attention_mask is not None, (
                  "Make sure that `context_attention_mask` are passed, if no `retriever` is set. Alternatively, you"
                  " can set a retriever using the `set_retriever(...)` function."
              )
              assert doc_scores is not None, (
                  "Make sure that `doc_scores` are passed, if no `retriever` is set. Alternatively, you can set a"
                  " retriever using the `set_retriever(...)` function."
              )
  
      assert (
          doc_scores is not None
      ), "Make sure that `doc_scores` are passed when passing `encoder_outputs` to the forward function."
  
      assert (doc_scores.shape[1] % n_docs) == 0, (
          f" The first dimension of `context_input_ids` should be a multiple of `n_docs`={n_docs}, but is"
          f" {context_input_ids.shape[0]}."
      )
  
      # Decoder input without context documents
      if decoder_input_ids is not None:
          decoder_input_ids = tf.repeat(
              decoder_input_ids, n_docs, axis=0
          )
  
      if decoder_attention_mask is not None:
          decoder_attention_mask = tf.repeat(
              decoder_attention_mask, n_docs, axis=0
          )
  
      gen_outputs = self.generator(
          context_input_ids,
          attention_mask=context_attention_mask,
          encoder_outputs=encoder_outputs,
          decoder_input_ids=decoder_input_ids,
          decoder_attention_mask=decoder_attention_mask,
          past_key_values=past_key_values,
          use_cache=use_cache,
          return_dict=True,
          training=training,
      )
  
      if not has_to_retrieve:
          question_encoder_last_hidden_state = None
          question_enc_hidden_states = None
          question_enc_attentions = None
          retrieved_doc_embeds = None
          retrieved_doc_ids = None
      else:
          question_enc_hidden_states = (
              question_enc_outputs.hidden_states
          )
          question_enc_attentions = question_enc_outputs.attentions
  
      if not has_to_retrieve or not output_retrieved:
          # don't output retrieved docs
          context_input_ids = (None,)
          context_attention_mask = None
          retrieved_doc_embeds = None
          retrieved_doc_ids = None
  
      return TFRetrievAugLMOutput(
          logits=gen_outputs.logits,
          doc_scores=doc_scores,
          past_key_values=gen_outputs.past_key_values,
          context_input_ids=context_input_ids,
          context_attention_mask=context_attention_mask,
          retrieved_doc_embeds=retrieved_doc_embeds,
          retrieved_doc_ids=retrieved_doc_ids,
          question_encoder_last_hidden_state=question_encoder_last_hidden_state,
          question_enc_hidden_states=question_enc_hidden_states,
          question_enc_attentions=question_enc_attentions,
          generator_enc_last_hidden_state=gen_outputs.encoder_last_hidden_state,
          generator_enc_hidden_states=gen_outputs.encoder_hidden_states,
          generator_enc_attentions=gen_outputs.encoder_attentions,
          generator_dec_hidden_states=gen_outputs.decoder_hidden_states,
          generator_dec_attentions=gen_outputs.decoder_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 3 for calling the TFRagModel obj's underlying member method code with example usage: obj.call(input_ids, attention_mask, encoder_outputs, decoder_input_ids, decoder_attention_mask, past_key_values, doc_scores, context_input_ids, context_attention_mask, use_cache, output_attentions, output_hidden_states, output_retrieved, n_docs, return_dict, training) and returns: TFRetrievAugLMOutput