docstr: Implements TFRAG token decoding.
  
  Args:
      input_ids (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
          The sequence used as a prompt for the generation. If `input_ids` is not passed, then
          `context_input_ids` has to be provided.
      attention_mask (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
          Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
  
          - 1 for tokens that are **not masked**,
          - 0 for tokens that are **masked**.
  
          [What are attention masks?](../glossary#attention-mask)
      context_input_ids (`tf.Tensor` of shape `(batch_size * config.n_docs, config.max_combined_length)`, *optional*, returned when *output_retrieved=True*):
          Input IDs post-processed from the retrieved documents and the question encoder `input_ids` by the
          retriever.
  
          If the model has is not initialized with a `retriever`, `context_input_ids` has to be provided to the
          forward pass. `context_input_ids` are returned by [`~RagRetriever.__call__`].
      context_attention_mask (`tf.Tensor` of shape `(batch_size * config.n_docs, config.max_combined_length)`, *optional*, returned when *output_retrieved=True*):
          Attention mask post-processed from the retrieved documents and the question encoder `input_ids` by the
          retriever.
  
          If the model has is not initialized with a `retriever`, `context_input_ids` has to be provided to the
          forward pass. `context_input_ids` are returned by [`~RagRetriever.__call__`].
      doc_scores (`tf.Tensor` of shape `(batch_size, config.n_docs)`):
          Score between each retrieved document embeddings (see `retrieved_doc_embeds`) and
          `question_encoder_last_hidden_state`.
  
          If the model has is not initialized with a `retriever`, `context_input_ids` has to be provided to the
          forward pass. `context_input_ids` are returned by [`~RagRetriever.__call__`].
      n_docs (`int`, *optional*, defaults to `config.n_docs`)
          Number of documents to retrieve and/or number of documents for which to generate an answer.
      generation_config (`~generation.GenerationConfig`, *optional*):
          The generation configuration to be used as base parametrization for the generation call. `**kwargs`
          passed to generate matching the attributes of `generation_config` will override them. If
          `generation_config` is not provided, the default will be used, which had the following loading
          priority: 1) from the `generation_config.json` model file, if it exists; 2) from the model
          configuration. Please note that unspecified parameters will inherit [`~generation.GenerationConfig`]'s
          default values, whose documentation should be checked to parameterize generation.
      logits_processor (`TFLogitsProcessorList`, *optional*):
          Custom logits processors that complement the default logits processors built from arguments and a
          model's config. If a logit processor is passed that is already created with the arguments or a model's
          config an error is thrown.
      kwargs (`Dict[str, Any]`, *optional*):
          Ad hoc parametrization of `generate_config` and/or additional model-specific kwargs that will be
          forwarded to the `forward` function of the model.
  
  Return:
      `tf.Tensor` of shape `(batch_size * num_return_sequences, sequence_length)`: The generated sequences. The
      second dimension (sequence_length) is either equal to `max_length` or shorter if all batches finished early
      due to the `eos_token_id`.
docstr_tok: ["Implements","TFRAG","token","decoding",".","Args",":","input_ids","(","`","tf.Tensor","`","of","shape","`","(","batch_size",",","sequence_length",")","`",",","*","optional","*",")",":","The","sequence","used","as","a","prompt","for","the","generation",".","If","`","input_ids","`","is","not","passed",",","then","`","context_input_ids","`","has","to","be","provided",".","attention_mask","(","`","tf.Tensor","`","of","shape","`","(","batch_size",",","sequence_length",")","`",",","*","optional","*",")",":","Mask","to","avoid","performing","attention","on","padding","token","indices",".","Mask","values","selected","in","`","[","0",",","1","]","`",":","-","1","for","tokens","that","are","*","*","not","masked","*","*",",","-","0","for","tokens","that","are","*","*","masked","*","*",".","[","What","are","attention","masks","?","]","(","..","\/glossary","#","attention-mask",")","context_input_ids","(","`","tf.Tensor","`","of","shape","`","(","batch_size","*","config.n_docs",",","config.max_combined_length",")","`",",","*","optional","*",",","returned","when","*","output_retrieved=True","*",")",":","Input","IDs","post-processed","from","the","retrieved","documents","and","the","question","encoder","`","input_ids","`","by","the","retriever",".","If","the","model","has","is","not","initialized","with","a","`","retriever","`",",","`","context_input_ids","`","has","to","be","provided","to","the","forward","pass",".","`","context_input_ids","`","are","returned","by","[","`","~RagRetriever.__call__","`","]",".","context_attention_mask","(","`","tf.Tensor","`","of","shape","`","(","batch_size","*","config.n_docs",",","config.max_combined_length",")","`",",","*","optional","*",",","returned","when","*","output_retrieved=True","*",")",":","Attention","mask","post-processed","from","the","retrieved","documents","and","the","question","encoder","`","input_ids","`","by","the","retriever",".","If","the","model","has","is","not","initialized","with","a","`","retriever","`",",","`","context_input_ids","`","has","to","be","provided","to","the","forward","pass",".","`","context_input_ids","`","are","returned","by","[","`","~RagRetriever.__call__","`","]",".","doc_scores","(","`","tf.Tensor","`","of","shape","`","(","batch_size",",","config.n_docs",")","`",")",":","Score","between","each","retrieved","document","embeddings","(","see","`","retrieved_doc_embeds","`",")","and","`","question_encoder_last_hidden_state","`",".","If","the","model","has","is","not","initialized","with","a","`","retriever","`",",","`","context_input_ids","`","has","to","be","provided","to","the","forward","pass",".","`","context_input_ids","`","are","returned","by","[","`","~RagRetriever.__call__","`","]",".","n_docs","(","`","int","`",",","*","optional","*",",","defaults","to","`","config.n_docs","`",")","Number","of","documents","to","retrieve","and\/or","number","of","documents","for","which","to","generate","an","answer",".","generation_config","(","`","~generation.GenerationConfig","`",",","*","optional","*",")",":","The","generation","configuration","to","be","used","as","base","parametrization","for","the","generation","call",".","`","*","*","kwargs","`","passed","to","generate","matching","the","attributes","of","`","generation_config","`","will","override","them",".","If","`","generation_config","`","is","not","provided",",","the","default","will","be","used",",","which","had","the","following","loading","priority",":","1",")","from","the","`","generation_config.json","`","model","file",",","if","it","exists",";","2",")","from","the","model","configuration",".","Please","note","that","unspecified","parameters","will","inherit","[","`","~generation.GenerationConfig","`","]","'s","default","values",",","whose","documentation","should","be","checked","to","parameterize","generation",".","logits_processor","(","`","TFLogitsProcessorList","`",",","*","optional","*",")",":","Custom","logits","processors","that","complement","the","default","logits","processors","built","from","arguments","and","a","model","'s","config",".","If","a","logit","processor","is","passed","that","is","already","created","with","the","arguments","or","a","model's","config","an","error","is","thrown",".","kwargs","(","`","Dict","[","str",",","Any","]","`",",","*","optional","*",")",":","Ad","hoc","parametrization","of","`","generate_config","`","and\/or","additional","model-specific","kwargs","that","will","be","forwarded","to","the","`","forward","`","function","of","the","model",".","Return",":","`","tf.Tensor","`","of","shape","`","(","batch_size","*","num_return_sequences",",","sequence_length",")","`",":","The","generated","sequences",".","The","second","dimension","(","sequence_length",")","is","either","equal","to","`","max_length","`","or","shorter","if","all","batches","finished","early","due","to","the","`","eos_token_id","`","."]
code: def generate(
      self,
      input_ids: TFModelInputType | None = None,
      attention_mask: tf.Tensor | None = None,
      context_input_ids=None,
      context_attention_mask=None,
      doc_scores=None,
      n_docs=None,
      generation_config=None,
      logits_processor=TFLogitsProcessorList(),
      **kwargs,
  ):
      """
      Implements TFRAG token decoding.
  
      Args:
          input_ids (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
              The sequence used as a prompt for the generation. If `input_ids` is not passed, then
              `context_input_ids` has to be provided.
          attention_mask (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
              Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
  
              - 1 for tokens that are **not masked**,
              - 0 for tokens that are **masked**.
  
              [What are attention masks?](../glossary#attention-mask)
          context_input_ids (`tf.Tensor` of shape `(batch_size * config.n_docs, config.max_combined_length)`, *optional*, returned when *output_retrieved=True*):
              Input IDs post-processed from the retrieved documents and the question encoder `input_ids` by the
              retriever.
  
              If the model has is not initialized with a `retriever`, `context_input_ids` has to be provided to the
              forward pass. `context_input_ids` are returned by [`~RagRetriever.__call__`].
          context_attention_mask (`tf.Tensor` of shape `(batch_size * config.n_docs, config.max_combined_length)`, *optional*, returned when *output_retrieved=True*):
              Attention mask post-processed from the retrieved documents and the question encoder `input_ids` by the
              retriever.
  
              If the model has is not initialized with a `retriever`, `context_input_ids` has to be provided to the
              forward pass. `context_input_ids` are returned by [`~RagRetriever.__call__`].
          doc_scores (`tf.Tensor` of shape `(batch_size, config.n_docs)`):
              Score between each retrieved document embeddings (see `retrieved_doc_embeds`) and
              `question_encoder_last_hidden_state`.
  
              If the model has is not initialized with a `retriever`, `context_input_ids` has to be provided to the
              forward pass. `context_input_ids` are returned by [`~RagRetriever.__call__`].
          n_docs (`int`, *optional*, defaults to `config.n_docs`)
              Number of documents to retrieve and/or number of documents for which to generate an answer.
          generation_config (`~generation.GenerationConfig`, *optional*):
              The generation configuration to be used as base parametrization for the generation call. `**kwargs`
              passed to generate matching the attributes of `generation_config` will override them. If
              `generation_config` is not provided, the default will be used, which had the following loading
              priority: 1) from the `generation_config.json` model file, if it exists; 2) from the model
              configuration. Please note that unspecified parameters will inherit [`~generation.GenerationConfig`]'s
              default values, whose documentation should be checked to parameterize generation.
          logits_processor (`TFLogitsProcessorList`, *optional*):
              Custom logits processors that complement the default logits processors built from arguments and a
              model's config. If a logit processor is passed that is already created with the arguments or a model's
              config an error is thrown.
          kwargs (`Dict[str, Any]`, *optional*):
              Ad hoc parametrization of `generate_config` and/or additional model-specific kwargs that will be
              forwarded to the `forward` function of the model.
  
      Return:
          `tf.Tensor` of shape `(batch_size * num_return_sequences, sequence_length)`: The generated sequences. The
          second dimension (sequence_length) is either equal to `max_length` or shorter if all batches finished early
          due to the `eos_token_id`.
      """
      # Handle `generation_config` and kwargs that might update it
      if generation_config is None:
          generation_config = self.generation_config
      generation_config = copy.deepcopy(generation_config)
      model_kwargs = generation_config.update(
          **kwargs
      )  # All unused kwargs must be model kwargs
  
      # set default parameters
      n_docs = n_docs if n_docs is not None else self.config.n_docs
  
      # retrieve docs
      if self.retriever is not None and context_input_ids is None:
          question_hidden_states = self.question_encoder(
              input_ids, attention_mask=attention_mask
          )[0]
          out = self.retriever(
              input_ids,
              question_hidden_states.numpy().astype(np.float32),
              prefix=self.generator.config.prefix,
              n_docs=n_docs,
              return_tensors="tf",
          )
          (
              context_input_ids,
              context_attention_mask,
              retrieved_doc_embeds,
          ) = (
              out["context_input_ids"],
              out["context_attention_mask"],
              out["retrieved_doc_embeds"],
          )
  
          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
          )
  
          # compute doc_scores
          doc_scores = tf.matmul(
              tf.expand_dims(question_hidden_states, axis=1),
              retrieved_doc_embeds,
              transpose_b=True,
          )
          doc_scores = tf.squeeze(doc_scores, axis=1)
  
      assert (context_input_ids.shape[0] % 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]}."
      )
  
      batch_size = context_input_ids.shape[0] // n_docs
  
      encoder = self.rag.generator.get_encoder()
      encoder_outputs = encoder(
          input_ids=context_input_ids,
          attention_mask=context_attention_mask,
          output_attentions=generation_config.output_attentions,
          output_hidden_states=generation_config.output_hidden_states,
          return_dict=True,
      )
  
      decoder_input_ids = tf.fill(
          (batch_size * generation_config.num_beams, 1),
          tf.cast(generation_config.decoder_start_token_id, tf.int32),
      )
      last_hidden_state = encoder_outputs["last_hidden_state"]
  
      def extend_enc_output(tensor, num_beams=None):
          """
          Broadcast tensor with `num_beams` replica, with correct order Input: tensor of shape (batch_size*n_docs ,
          d) Output: tensor of shape (batch_size*num_beams*n_docs , d)
          """
  
          # expand batch_size & num_beam dimensions
          d_shape_list = tensor.shape[1:]
  
          # split n_docs dimensions
          new_shape = (batch_size, 1, n_docs) + d_shape_list
          tensor = tf.reshape(tensor, new_shape)
  
          # repeat same last hidden states over `num_beams` dimension
          new_shape = (batch_size, num_beams, n_docs) + d_shape_list
          tensor = tf.broadcast_to(tensor, new_shape)
  
          # merge `batch_size`, `num_beams`, `num_docs` dims again
          new_shape = (batch_size * num_beams * n_docs,) + d_shape_list
          return tf.reshape(tensor, new_shape)
  
      # correctly extend last_hidden_state and attention mask
      context_attention_mask = extend_enc_output(
          context_attention_mask, num_beams=generation_config.num_beams
      )
      encoder_outputs["last_hidden_state"] = extend_enc_output(
          last_hidden_state, num_beams=generation_config.num_beams
      )
  
      doc_scores = tf.repeat(
          doc_scores, generation_config.num_beams, axis=0
      )
  
      # define start_len & additional parameters
      model_kwargs["doc_scores"] = doc_scores
      model_kwargs["encoder_outputs"] = encoder_outputs
      model_kwargs["attention_mask"] = context_attention_mask
      model_kwargs["n_docs"] = n_docs
  
      pre_processor = self._get_logits_processor(
          generation_config=generation_config,
          input_ids_seq_length=tf.shape(decoder_input_ids)[-1],
          logits_processor=logits_processor,
      )
  
      if generation_config.num_beams == 1:
          return self.greedy_search(
              input_ids=decoder_input_ids,
              max_length=generation_config.max_length,
              pad_token_id=generation_config.pad_token_id,
              eos_token_id=generation_config.eos_token_id,
              logits_processor=pre_processor,
              output_attentions=generation_config.output_attentions,
              output_hidden_states=generation_config.output_hidden_states,
              output_scores=generation_config.output_scores,
              return_dict_in_generate=generation_config.return_dict_in_generate,
              **model_kwargs,
          )
      elif generation_config.num_beams > 1:
          if (
              generation_config.num_beams
              < generation_config.num_return_sequences
          ):
              raise ValueError(
                  "Beam search decoding cannot return more sequences than it has beams. Please set num_beams >="
                  f" num_return_sequences, got {generation_config.num_beams} and"
                  f" {generation_config.num_return_sequences} (respectivelly)"
              )
  
          def unflatten_beam_dim(tensor):
              """Unflattens the first, flat batch*beam dimension of a non-scalar array."""
              shape = shape_list(tensor)
              return tf.reshape(
                  tensor, [-1, generation_config.num_beams] + shape[1:]
              )
  
          decoder_input_ids = unflatten_beam_dim(decoder_input_ids)
          model_kwargs["attention_mask"] = unflatten_beam_dim(
              model_kwargs["attention_mask"]
          )
          model_kwargs["encoder_outputs"][
              "last_hidden_state"
          ] = unflatten_beam_dim(
              model_kwargs["encoder_outputs"]["last_hidden_state"]
          )
  
          return self.beam_search(
              input_ids=decoder_input_ids,
              max_length=generation_config.max_length,
              pad_token_id=generation_config.pad_token_id,
              eos_token_id=generation_config.eos_token_id,
              logits_processor=pre_processor,
              output_attentions=generation_config.output_attentions,
              output_hidden_states=generation_config.output_hidden_states,
              output_scores=generation_config.output_scores,
              return_dict_in_generate=generation_config.return_dict_in_generate,
              **model_kwargs,
          )
      else:
          raise ValueError(
              f"`num_beams` has to be an integer strictly superior to 0 (≥ 1), but is {generation_config.num_beams}"
          )
  
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 10 for calling the TFRagTokenForGeneration obj's underlying member method code with example usage: obj.generate(input_ids, attention_mask, context_input_ids, context_attention_mask, doc_scores, n_docs, generation_config, logits_processor) and returns: tf, self, self, tf