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__`].
      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 has 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.
      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.
      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:
      `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,
      n_docs: Optional[int] = None,
      generation_config: Optional[GenerationConfig] = None,
      prefix_allowed_tokens_fn: Callable[
          [int, torch.Tensor], List[int]
      ] = None,
      logits_processor: Optional[
          LogitsProcessorList
      ] = LogitsProcessorList(),
      stopping_criteria: Optional[
          StoppingCriteriaList
      ] = StoppingCriteriaList(),
      **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__`].
          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 has 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.
          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.
          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:
          `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`.
      """
      # 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.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 * generation_config.num_beams, 1),
          generation_config.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=generation_config.num_beams
      )
      encoder_outputs["last_hidden_state"] = extend_enc_output(
          last_hidden_state, num_beams=generation_config.num_beams
      )
  
      doc_scores = doc_scores.repeat_interleave(
          generation_config.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(
          generation_config=generation_config,
          input_ids_seq_length=input_ids_seq_length,
          encoder_input_ids=context_input_ids,
          prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
          logits_processor=logits_processor,
      )
  
      if generation_config.num_beams == 1:
          if generation_config.num_return_sequences > 1:
              raise ValueError(
                  f"num_return_sequences has to be 1, but is {generation_config.num_return_sequences} when doing"
                  " greedy search."
              )
          return self.greedy_search(
              input_ids,
              logits_processor=pre_processor,
              max_length=generation_config.max_length,
              pad_token_id=generation_config.pad_token_id,
              eos_token_id=generation_config.eos_token_id,
              **model_kwargs,
          )
      elif generation_config.num_beams > 1:
          if (
              generation_config.num_return_sequences
              > generation_config.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=generation_config.num_beams,
              device=self.device,
              length_penalty=generation_config.length_penalty,
              do_early_stopping=generation_config.early_stopping,
              num_beam_hyps_to_keep=generation_config.num_return_sequences,
              max_length=generation_config.max_length,
          )
          return self.beam_search(
              input_ids,
              beam_scorer,
              logits_processor=pre_processor,
              max_length=generation_config.max_length,
              pad_token_id=generation_config.pad_token_id,
              eos_token_id=generation_config.eos_token_id,
              **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: import copy
  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 import BeamSearchScorer, GenerationConfig, LogitsProcessorList, 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, n_docs, generation_config, prefix_allowed_tokens_fn, logits_processor, stopping_criteria) and returns: tensor, self, self