docstr: Implements RAG sequence "thorough" decoding. Read the [`~generation_utils.GenerationMixin.generate`]`
  documentation for more information on how to set other generate input parameters.
  
  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.
      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 is not initialized with a `retriever` or `input_ids` is not given, `context_input_ids` and
          `context_attention_mask` have to be provided to the forward pass. They 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 is not initialized with a `retriever` or `input_ids` is not given, `doc_scores` has to be
          provided to the forward pass. `doc_scores` are returned by [`~RagRetriever.__call__`].
      do_deduplication (`bool`, *optional*):
          Whether or not to deduplicate the generations from different context documents for a given input. Has
          to be set to `False` if used while training with distributed backend.
      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`.
      num_beams (`int`, *optional*, defaults to 1):
          Number of beams for beam search. 1 means no beam search.
      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.
      kwargs:
          Additional kwargs will be passed to [`~generation_utils.GenerationMixin.generate`].
  
  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,
      do_deduplication: Optional[bool] = None,  # defaults to True
      num_return_sequences: Optional[int] = None,  # defaults to 1
      num_beams: Optional[int] = None,  # defaults to 1
      n_docs: Optional[int] = None,
      **model_kwargs,
  ) -> torch.LongTensor:
      """
      Implements RAG sequence "thorough" decoding. Read the [`~generation_utils.GenerationMixin.generate`]`
      documentation for more information on how to set other generate input parameters.
  
      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.
          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 is not initialized with a `retriever` or `input_ids` is not given, `context_input_ids` and
              `context_attention_mask` have to be provided to the forward pass. They 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 is not initialized with a `retriever` or `input_ids` is not given, `doc_scores` has to be
              provided to the forward pass. `doc_scores` are returned by [`~RagRetriever.__call__`].
          do_deduplication (`bool`, *optional*):
              Whether or not to deduplicate the generations from different context documents for a given input. Has
              to be set to `False` if used while training with distributed backend.
          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`.
          num_beams (`int`, *optional*, defaults to 1):
              Number of beams for beam search. 1 means no beam search.
          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.
          kwargs:
              Additional kwargs will be passed to [`~generation_utils.GenerationMixin.generate`].
  
      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`.
      """
  
      n_docs = n_docs if n_docs is not None else self.config.n_docs
      do_deduplication = (
          do_deduplication
          if do_deduplication is not None
          else self.config.do_deduplication
      )
      num_doc_return_sequences = (
          num_return_sequences
          if num_return_sequences is not None
          else self.config.num_return_sequences
      )
      num_beams = (
          num_beams if num_beams is not None else self.config.num_beams
      )
  
      assert (
          input_ids is not None or context_input_ids is not None
      ), " At least one of input_ids or context_input_ids must be given"
  
      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]
          context_input_ids = 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"]
  
          # set to correct device
          context_input_ids = context_input_ids.to(input_ids)
  
      hypos = []
      model_kwargs["num_beams"] = num_beams
      model_kwargs["num_return_sequences"] = num_beams
      model_kwargs["attention_mask"] = None
  
      batch_size = (
          input_ids.shape[0]
          if input_ids is not None
          else context_input_ids.shape[0] // n_docs
      )
  
      for index in range(batch_size):
          # first, generate beams from documents:
          generator_input_ids = context_input_ids[
              index * n_docs : (index + 1) * n_docs
          ]  # (n_docs, max_len)
  
          output_sequences = self.generator.generate(
              generator_input_ids,
              **model_kwargs,
          )  # n_docs * n_beam, tgt_len
          if do_deduplication:
              # do_deduplication, max_output_len
              output_sequences = torch.stack(
                  list(
                      {
                          str(k.tolist()): k for k in output_sequences
                      }.values()
                  )
              )
  
          num_candidates = output_sequences.shape[
              0
          ]  # after deduplication, this number can be less than n_docs*n_beam
  
          # then, run model forwards to get nll scores:
          if input_ids is not None:
              new_input_ids = input_ids[index : index + 1].repeat(
                  num_candidates, 1
              )
              outputs = self(
                  new_input_ids,
                  labels=output_sequences,
                  exclude_bos_score=True,
              )
          else:  # input_ids is None, need context_input_ids/mask and doc_scores
              assert context_attention_mask is not None, (
                  "Make sure that `context_attention_mask` are passed, if no `input_ids` 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 `input_ids` is set. Alternatively, you can set a"
                  " retriever using the `set_retriever(...)` function."
              )
  
              individual_input_ids = generator_input_ids.repeat(
                  num_candidates, 1
              )  # (num_candidates*n_docs, max_len)
  
              individual_attention_mask = context_attention_mask[
                  index * n_docs : (index + 1) * n_docs
              ]
              individual_attention_mask = (
                  individual_attention_mask.repeat(num_candidates, 1)
              )
  
              individual_doc_scores = doc_scores[
                  index : (index + 1), :
              ]  # doc_scores.shape = [batch, n_docs]
              individual_doc_scores = individual_doc_scores.repeat(
                  num_candidates, 1
              )  # [num_candidates, n_docs]
  
              outputs = self(
                  context_input_ids=individual_input_ids,
                  context_attention_mask=individual_attention_mask,
                  doc_scores=individual_doc_scores,
                  labels=output_sequences,
                  exclude_bos_score=True,
              )
  
          top_cand_inds = (-outputs["loss"]).topk(
              num_doc_return_sequences
          )[1]
  
          # add hypothesis
          hypos.append(output_sequences[top_cand_inds])
  
      return self._cat_and_pad(
          hypos, pad_token_id=self.config.generator.pad_token_id
      )
  
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 8 for calling the RagSequenceForGeneration obj's underlying member method code with example usage: obj.generate(input_ids, attention_mask, context_input_ids, context_attention_mask, doc_scores, do_deduplication, num_return_sequences, num_beams, n_docs) and returns: self