Hi @PereLluis13.
First, we do need to update the documentation, thanks for pointing it out! The current models has two forward modes corresponding to the two shapes you mentioned: without and with marginalization, controlled by the do_marginalize flag, which is set to True in the generate function:

  
      github.com
  
  
    huggingface/transformers/blob/13c185771847370d695b8eee3cbf12f4edc2111c/src/transformers/modeling_rag.py#L1086 5

    decoder_input_ids=None,
    decoder_attention_mask=None,
    past_key_values=None,
    context_input_ids=None,
    context_attention_mask=None,
    doc_scores=None,
    use_cache=None,
    output_attentions=None,
    output_hidden_states=None,
    output_retrieved=None,
    do_marginalize=None,
    reduce_loss=None,
    labels=None,
    **kwargs  # needs kwargs for generation
):
    r"""
    do_marginalize (:obj:`bool`, `optional`):
        If :obj:`True`, the logits are marginalized over all documents
        by making use of ``torch.nn.functional.log_softmax``.
    reduce_loss (:obj:`bool`, `optional`):
        Only relevant if ``labels`` is passed.



  
  
    
    
  
  


You are right that getting the posterior token-level probabilities (used in the demo) currently requires an additional forward pass (on the decoder side only, you can re-use the encoder output and retrieved documents), as you can see here:

  
      github.com
  
  
    yjernite/transformers/blob/291cea09c3951cc3d042661913fff2acd7148812/examples/rag/app_rag_demo.py#L142 5

            answer_scores += [answer_score]
            explained_gen = []
            ans_id_list = answer_ids[0].tolist()
            if tokenizer.generator.eos_token_id in ans_id_list[1:]:
                s_len = ans_id_list[1:].index(tokenizer.generator.eos_token_id)
            else:
                s_len = len(ans_id_list) - 1
            for i in range(1, s_len):
                tid = ans_id_list[i + 1]
                token = tokenizer.generator._convert_id_to_token(tid).replace("Ġ", "_")
                step_probs = eval_probs[:, i, tid].tolist()
                explained_gen += [(token, [(di, p, p * doc_probs[di]) for di, p in enumerate(step_probs)])]
            explanations += [explained_gen]
    return {
        "answers": [(te, -sc) for te, sc in zip(answer_texts, answer_scores)],
        "documents": support_docs,
        "explanations": explanations,
    }


@st.cache(allow_output_mutation=True)



  
  
    
    
  
  


We’ve been going back and forth on returning the scores in the generate function, it will likely be available in a future PR.
The demo uses the retrieval scores for the documents, which corresponds to the priors for generation. To get the posteriors, you can use Bayes rule with the config.n_docs doc-level log-likelihoods obtained with do_marginalize=False:
p(d|q,a) \proportional p(a|d,q) \times p(d|q)