docstr: Implements RAG token decoding.
  
  Args:
      input_ids (`torch.LongTensor` 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 (`torch.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 (`torch.LongTensor` 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 (`torch.LongTensor` 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 (`torch.FloatTensor` 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__`].
      max_length (`int`, *optional*, defaults to 20):
          The maximum length of the sequence to be generated.
      min_length (`int`, *optional*, defaults to 10):
          The minimum length of the sequence to be generated.
      early_stopping (`bool`, *optional*, defaults to `False`):
          Whether or not to stop the beam search when at least `num_beams` sentences are finished per batch or
          not.
      use_cache: (`bool`, *optional*, defaults to `True`):
          Whether or not the model should use the past last key/values attentions (if applicable to the model) to
          speed up decoding.
      pad_token_id (`int`, *optional*):
          The id of the *padding* token.
      bos_token_id (`int`, *optional*):
          The id of the *beginning-of-sequence* token.
      eos_token_id (`int`, *optional*):
          The id of the *end-of-sequence* token.
      length_penalty (`float`, *optional*, defaults to 1.0):
          Exponential penalty to the length that is used with beam-based generation. It is applied as an exponent
          to the sequence length, which in turn is used to divide the score of the sequence. Since the score is
          the log likelihood of the sequence (i.e. negative), `length_penalty` > 0.0 promotes longer sequences,
          while `length_penalty` < 0.0 encourages shorter sequences.
      no_repeat_ngram_size (`int`, *optional*, defaults to 0):
          If set to int > 0, all ngrams of that size can only occur once.
      encoder_no_repeat_ngram_size (`int`, *optional*, defaults to 0):
          If set to int > 0, all ngrams of that size that occur in the `encoder_input_ids` cannot occur in the
          `decoder_input_ids`.
      bad_words_ids(`List[int]`, *optional*):
          List of token ids that are not allowed to be generated. In order to get the tokens of the words that
          should not appear in the generated text, use `tokenizer.encode(bad_word, add_prefix_space=True)`.
      num_beams (`int`, *optional*, defaults to 1):
          Number of beams for beam search. 1 means no beam search.
      num_beam_groups (`int`, *optional*, defaults to 1):
          Number of groups to divide `num_beams` into in order to ensure diversity among different groups of
          beams. [this paper](https://arxiv.org/pdf/1610.02424.pdf) for more details.
      diversity_penalty (`float`, *optional*, defaults to 0.0):
          This value is subtracted from a beam's score if it generates a token same as any beam from other group
          at a particular time. Note that `diversity_penalty` is only effective if `group beam search` is
          enabled.
      num_return_sequences(`int`, *optional*, defaults to 1):
          The number of independently computed returned sequences for each element in the batch. Note that this
          is not the value we pass to the `generator`'s `[`~generation_utils.GenerationMixin.generate`] function,
          where we set `num_return_sequences` to `num_beams`. decoder_start_token_id (`int`, *optional*): If an
          encoder-decoder model starts decoding with a different token than *bos*, the id of that token.
      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.
      prefix_allowed_tokens_fn: (`Callable[[int, torch.Tensor], List[int]]`, *optional*):
          If provided, this function constraints the beam search to allowed tokens only at each step. If not
          provided no constraint is applied. This function takes 2 arguments `inputs_ids` and the batch ID
          `batch_id`. It has to return a list with the allowed tokens for the next generation step conditioned on
          the previously generated tokens `inputs_ids` and the batch ID `batch_id`. This argument is useful for
          constrained generation conditioned on the prefix, as described in [Autoregressive Entity
          Retrieval](https://arxiv.org/abs/2010.00904).
      logits_processor (`LogitsProcessorList`, *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.
      stopping_criteria (`StoppingCriteriaList`, *optional*):
           Custom stopping criteria that complement the default stopping criteria built from arguments and a
           model's config. If a stopping criteria is passed that is already created with the arguments or a
           model's config an error is thrown.
      forced_bos_token_id (`int`, *optional*):
          The id of the token to force as the first generated token after the `decoder_start_token_id`. Useful
          for multilingual models like [mBART](../model_doc/mbart) where the first generated token needs to be
          the target language token.
      forced_eos_token_id (`int`, *optional*):
          The id of the token to force as the last generated token when `max_length` is reached.
      remove_invalid_values (`bool`, *optional*):
          Whether to remove possible *nan* and *inf* outputs of the model to prevent the generation method to
          crash. Note that using `remove_invalid_values` can slow down generation.
  
  Return:
      `torch.LongTensor` 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: Optional[torch.LongTensor] = None,
      attention_mask: Optional[torch.LongTensor] = None,
      context_input_ids: Optional[torch.LongTensor] = None,
      context_attention_mask: Optional[torch.LongTensor] = None,
      doc_scores: Optional[torch.FloatTensor] = None,
      max_length: Optional[int] = None,
      min_length: Optional[int] = None,
      early_stopping: Optional[bool] = None,
      use_cache: Optional[bool] = None,
      num_beams: Optional[int] = None,
      num_beam_groups: Optional[int] = None,
      diversity_penalty: Optional[float] = None,
      bos_token_id: Optional[int] = None,
      pad_token_id: Optional[int] = None,
      eos_token_id: Optional[int] = None,
      length_penalty: Optional[float] = None,
      no_repeat_ngram_size: Optional[int] = None,
      encoder_no_repeat_ngram_size: Optional[int] = None,
      repetition_penalty: Optional[float] = None,
      bad_words_ids: Optional[List[List[int]]] = None,
      num_return_sequences: Optional[int] = None,
      decoder_start_token_id: Optional[int] = None,
      n_docs: Optional[int] = None,
      prefix_allowed_tokens_fn: Callable[
          [int, torch.Tensor], List[int]
      ] = None,
      logits_processor: Optional[
          LogitsProcessorList
      ] = LogitsProcessorList(),
      renormalize_logits: Optional[bool] = None,
      stopping_criteria: Optional[
          StoppingCriteriaList
      ] = StoppingCriteriaList(),
      forced_bos_token_id: Optional[int] = None,
      forced_eos_token_id: Optional[int] = None,
      remove_invalid_values: Optional[bool] = None,
      exponential_decay_length_penalty: Optional[
          Tuple[Union[int, float]]
      ] = None,
      **model_kwargs,
  ) -> torch.LongTensor:
      """
      Implements RAG token decoding.
  
      Args:
          input_ids (`torch.LongTensor` 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 (`torch.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 (`torch.LongTensor` 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 (`torch.LongTensor` 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 (`torch.FloatTensor` 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__`].
          max_length (`int`, *optional*, defaults to 20):
              The maximum length of the sequence to be generated.
          min_length (`int`, *optional*, defaults to 10):
              The minimum length of the sequence to be generated.
          early_stopping (`bool`, *optional*, defaults to `False`):
              Whether or not to stop the beam search when at least `num_beams` sentences are finished per batch or
              not.
          use_cache: (`bool`, *optional*, defaults to `True`):
              Whether or not the model should use the past last key/values attentions (if applicable to the model) to
              speed up decoding.
          pad_token_id (`int`, *optional*):
              The id of the *padding* token.
          bos_token_id (`int`, *optional*):
              The id of the *beginning-of-sequence* token.
          eos_token_id (`int`, *optional*):
              The id of the *end-of-sequence* token.
          length_penalty (`float`, *optional*, defaults to 1.0):
              Exponential penalty to the length that is used with beam-based generation. It is applied as an exponent
              to the sequence length, which in turn is used to divide the score of the sequence. Since the score is
              the log likelihood of the sequence (i.e. negative), `length_penalty` > 0.0 promotes longer sequences,
              while `length_penalty` < 0.0 encourages shorter sequences.
          no_repeat_ngram_size (`int`, *optional*, defaults to 0):
              If set to int > 0, all ngrams of that size can only occur once.
          encoder_no_repeat_ngram_size (`int`, *optional*, defaults to 0):
              If set to int > 0, all ngrams of that size that occur in the `encoder_input_ids` cannot occur in the
              `decoder_input_ids`.
          bad_words_ids(`List[int]`, *optional*):
              List of token ids that are not allowed to be generated. In order to get the tokens of the words that
              should not appear in the generated text, use `tokenizer.encode(bad_word, add_prefix_space=True)`.
          num_beams (`int`, *optional*, defaults to 1):
              Number of beams for beam search. 1 means no beam search.
          num_beam_groups (`int`, *optional*, defaults to 1):
              Number of groups to divide `num_beams` into in order to ensure diversity among different groups of
              beams. [this paper](https://arxiv.org/pdf/1610.02424.pdf) for more details.
          diversity_penalty (`float`, *optional*, defaults to 0.0):
              This value is subtracted from a beam's score if it generates a token same as any beam from other group
              at a particular time. Note that `diversity_penalty` is only effective if `group beam search` is
              enabled.
          num_return_sequences(`int`, *optional*, defaults to 1):
              The number of independently computed returned sequences for each element in the batch. Note that this
              is not the value we pass to the `generator`'s `[`~generation_utils.GenerationMixin.generate`] function,
              where we set `num_return_sequences` to `num_beams`. decoder_start_token_id (`int`, *optional*): If an
              encoder-decoder model starts decoding with a different token than *bos*, the id of that token.
          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.
          prefix_allowed_tokens_fn: (`Callable[[int, torch.Tensor], List[int]]`, *optional*):
              If provided, this function constraints the beam search to allowed tokens only at each step. If not
              provided no constraint is applied. This function takes 2 arguments `inputs_ids` and the batch ID
              `batch_id`. It has to return a list with the allowed tokens for the next generation step conditioned on
              the previously generated tokens `inputs_ids` and the batch ID `batch_id`. This argument is useful for
              constrained generation conditioned on the prefix, as described in [Autoregressive Entity
              Retrieval](https://arxiv.org/abs/2010.00904).
          logits_processor (`LogitsProcessorList`, *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.
          stopping_criteria (`StoppingCriteriaList`, *optional*):
               Custom stopping criteria that complement the default stopping criteria built from arguments and a
               model's config. If a stopping criteria is passed that is already created with the arguments or a
               model's config an error is thrown.
          forced_bos_token_id (`int`, *optional*):
              The id of the token to force as the first generated token after the `decoder_start_token_id`. Useful
              for multilingual models like [mBART](../model_doc/mbart) where the first generated token needs to be
              the target language token.
          forced_eos_token_id (`int`, *optional*):
              The id of the token to force as the last generated token when `max_length` is reached.
          remove_invalid_values (`bool`, *optional*):
              Whether to remove possible *nan* and *inf* outputs of the model to prevent the generation method to
              crash. Note that using `remove_invalid_values` can slow down generation.
  
      Return:
          `torch.LongTensor` 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`.
      """
      # set default parameters
      n_docs = n_docs if n_docs is not None else self.config.n_docs
      num_beams = (
          num_beams if num_beams is not None else self.config.num_beams
      )
      num_beam_groups = (
          num_beam_groups
          if num_beam_groups is not None
          else self.config.num_beam_groups
      )
      max_length = (
          max_length
          if max_length is not None
          else self.config.max_length
      )
      num_return_sequences = (
          num_return_sequences
          if num_return_sequences is not None
          else self.config.num_return_sequences
      )
      bos_token_id = (
          bos_token_id
          if bos_token_id is not None
          else self.config.generator.bos_token_id
      )
      eos_token_id = (
          eos_token_id
          if eos_token_id is not None
          else self.config.generator.eos_token_id
      )
      pad_token_id = (
          pad_token_id
          if pad_token_id is not None
          else self.config.generator.pad_token_id
      )
      use_cache = (
          use_cache if use_cache is not None else self.config.use_cache
      )
      decoder_start_token_id = (
          decoder_start_token_id
          if decoder_start_token_id is not None
          else self.config.generator.decoder_start_token_id
      )
      remove_invalid_values = (
          remove_invalid_values
          if remove_invalid_values is not None
          else self.config.remove_invalid_values
      )
      exponential_decay_length_penalty = (
          exponential_decay_length_penalty
          if exponential_decay_length_penalty is not None
          else self.config.exponential_decay_length_penalty
      )
  
      # 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.cpu()
              .detach()
              .to(torch.float32)
              .numpy(),
              prefix=self.generator.config.prefix,
              n_docs=n_docs,
              return_tensors="pt",
          )
          (
              context_input_ids,
              context_attention_mask,
              retrieved_doc_embeds,
          ) = (
              out["context_input_ids"],
              out["context_attention_mask"],
              out["retrieved_doc_embeds"],
          )
  
          # set to correct device
          retrieved_doc_embeds = retrieved_doc_embeds.to(
              question_hidden_states
          )
          context_input_ids = context_input_ids.to(input_ids)
          context_attention_mask = context_attention_mask.to(input_ids)
  
          # compute doc_scores
          doc_scores = torch.bmm(
              question_hidden_states.unsqueeze(1),
              retrieved_doc_embeds.transpose(1, 2),
          ).squeeze(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
      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,
          return_dict=True,
      )
  
      input_ids = torch.full(
          (batch_size * num_beams, 1),
          decoder_start_token_id,
          dtype=torch.long,
          device=next(self.parameters()).device,
      )
      input_ids_seq_length = input_ids.shape[-1]
      last_hidden_state = encoder_outputs["last_hidden_state"]
  
      def extend_enc_output(tensor, num_beams=None):
          # split into `batch_size`, `num_beams`, `num_docs`
          tensor = tensor[None, None, :].reshape(
              (batch_size, 1, n_docs) + tensor.shape[1:]
          )
          # repeat same last hidden states over `num_beams` dimension
          tensor = tensor.expand(
              (batch_size, num_beams, n_docs) + tensor.shape[3:]
          )
          # merge `batch_size`, `num_beams`, `num_docs` dims again
          return tensor.reshape(
              (batch_size * num_beams * n_docs,) + tensor.shape[3:]
          )
  
      # correctly extend last_hidden_state and attention mask
      context_attention_mask = extend_enc_output(
          context_attention_mask, num_beams=num_beams
      )
      encoder_outputs["last_hidden_state"] = extend_enc_output(
          last_hidden_state, num_beams=num_beams
      )
  
      doc_scores = doc_scores.repeat_interleave(num_beams, dim=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(
          repetition_penalty=repetition_penalty,
          no_repeat_ngram_size=no_repeat_ngram_size,
          encoder_no_repeat_ngram_size=encoder_no_repeat_ngram_size,
          input_ids_seq_length=input_ids_seq_length,
          encoder_input_ids=context_input_ids,
          bad_words_ids=bad_words_ids,
          min_length=min_length,
          max_length=max_length,
          eos_token_id=eos_token_id,
          forced_bos_token_id=forced_bos_token_id,
          forced_eos_token_id=forced_eos_token_id,
          prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
          num_beams=num_beams,
          num_beam_groups=num_beam_groups,
          diversity_penalty=diversity_penalty,
          remove_invalid_values=remove_invalid_values,
          exponential_decay_length_penalty=exponential_decay_length_penalty,
          logits_processor=logits_processor,
          renormalize_logits=renormalize_logits,
      )
  
      if num_beams == 1:
          if num_return_sequences > 1:
              raise ValueError(
                  f"num_return_sequences has to be 1, but is {num_return_sequences} when doing greedy search."
              )
          return self.greedy_search(
              input_ids,
              logits_processor=pre_processor,
              max_length=max_length,
              pad_token_id=pad_token_id,
              eos_token_id=eos_token_id,
              **model_kwargs,
          )
      elif num_beams > 1:
          length_penalty = (
              length_penalty
              if length_penalty is not None
              else self.config.length_penalty
          )
          early_stopping = (
              early_stopping
              if early_stopping is not None
              else self.config.early_stopping
          )
          if num_return_sequences > num_beams:
              raise ValueError(
                  "`num_return_sequences` has to be smaller or equal to `num_beams`."
              )
          beam_scorer = BeamSearchScorer(
              batch_size=batch_size,
              num_beams=num_beams,
              device=self.device,
              length_penalty=length_penalty,
              do_early_stopping=early_stopping,
              num_beam_hyps_to_keep=num_return_sequences,
          )
          return self.beam_search(
              input_ids,
              beam_scorer,
              logits_processor=pre_processor,
              max_length=max_length,
              pad_token_id=pad_token_id,
              eos_token_id=eos_token_id,
              **model_kwargs,
          )
      else:
          raise ValueError(
              f"`num_beams` has to be an integer strictly superior to 0 (≥ 1), but is {num_beams}"
          )
  
imports: from dataclasses import dataclass
  from typing import Callable, List, Optional, Tuple, Union
  import torch
  from torch import nn
  from ...configuration_utils import PretrainedConfig
  from ...generation_beam_search import BeamSearchScorer
  from ...generation_logits_process import LogitsProcessorList
  from ...generation_stopping_criteria import StoppingCriteriaList
  from ...modeling_outputs import ModelOutput
  from ...modeling_utils import PreTrainedModel
  from ...utils import 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 11 for calling the RagTokenForGeneration obj's underlying member method code with example usage: obj.generate(input_ids, attention_mask, context_input_ids, context_attention_mask, doc_scores, max_length, min_length, early_stopping, use_cache, num_beams, num_beam_groups, diversity_penalty, bos_token_id, pad_token_id, eos_token_id, length_penalty, no_repeat_ngram_size, encoder_no_repeat_ngram_size, repetition_penalty, bad_words_ids, num_return_sequences, decoder_start_token_id, n_docs, prefix_allowed_tokens_fn, logits_processor, renormalize_logits, stopping_criteria, forced_bos_token_id, forced_eos_token_id, remove_invalid_values, exponential_decay_length_penalty) and returns: tensor, self, self