o Zh][@sNdZddlZddlmZddlmZmZmZmZm Z ddl Z ddl m Z ddl m Z ddlmZmZmZmZdd lmZdd lmZdd lmZmZd d lmZd dlmZeeZ eGdddeZ!eGdddeZ"eddGdddeZ#eGddde#Z$eddGddde#Z%eddGddde#eZ&gdZ'dS) zRAG model implementation.N) dataclass)CallableListOptionalTupleUnion)nn)PretrainedConfig)GenerationConfigGenerationMixinLogitsProcessorListStoppingCriteriaList) ModelOutput)PreTrainedModel)auto_docstringlogging) RagConfig) RagRetrieverc@seZdZUdZdZeejed<dZ eejed<dZ eejed<dZ ee ejed<dZ eejed<dZeejed<dZeejed <dZeejed <dZeejed <dZeeejd fed <dZeeejd fed<dZeejed<dZeeejd fed<dZeeejd fed<dZeeejd fed<dZeeejd fed<dZeeejd fed<dS)RetrievAugLMMarginOutputa Base class for retriever augmented marginalized models outputs. Args: loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): Language modeling loss. logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): Prediction scores of the language modeling head. The score is possibly marginalized over all documents for each vocabulary token. 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`. past_key_values (`List[torch.FloatTensor]`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): List of `torch.FloatTensor` of length `config.n_layers`, with each tensor of shape `(2, batch_size, num_heads, sequence_length, embed_size_per_head)`). Contains precomputed hidden-states (key and values in the attention blocks) of the decoder that can be used (see `past_key_values` input) to speed up sequential decoding. retrieved_doc_embeds (`torch.FloatTensor` of shape `(batch_size, config.n_docs, hidden_size)`, *optional*, returned when *output_retrieved=True*): Embedded documents retrieved by the retriever. Is used with `question_encoder_last_hidden_state` to compute the `doc_scores`. retrieved_doc_ids (`torch.LongTensor` of shape `(batch_size, config.n_docs)`, *optional*, returned when *output_retrieved=True*): The indexes of the embedded documents retrieved by the retriever. 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. question_encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): Sequence of hidden states at the output of the last layer of the question encoder pooled output of the model. question_enc_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings and one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden states of the question encoder at the output of each layer plus the initial embedding outputs. question_enc_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights of the question encoder, after the attention softmax, used to compute the weighted average in the self-attention heads. generator_enc_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): Sequence of hidden-states at the output of the last layer of the generator encoder of the model. generator_enc_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings and one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden states of the generator encoder at the output of each layer plus the initial embedding outputs. generator_enc_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights of the generator encoder, after the attention softmax, used to compute the weighted average in the self-attention heads. generator_dec_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings and one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden states of the generator decoder at the output of each layer plus the initial embedding outputs. generator_dec_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights of the generator decoder, after the attention softmax, used to compute the weighted average in the self-attention heads. generator_cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Cross-attentions weights of the generator decoder, after the attention softmax, used to compute the weighted average in the cross-attention heads. Nlosslogits doc_scorespast_key_valuesretrieved_doc_embedsretrieved_doc_idscontext_input_idscontext_attention_mask"question_encoder_last_hidden_state.question_enc_hidden_statesquestion_enc_attentionsgenerator_enc_last_hidden_stategenerator_enc_hidden_statesgenerator_enc_attentionsgenerator_dec_hidden_statesgenerator_dec_attentionsgenerator_cross_attentions)__name__ __module__ __qualname____doc__rrtorch FloatTensor__annotations__rrrrrr LongTensorrrrr rr!r"r#r$r%r&r'r0r0S/var/www/auris/lib/python3.10/site-packages/transformers/models/rag/modeling_rag.pyr$s& Jrc@sneZdZUdZdZeejed<dZ eejed<dZ ee ejed<dZ eejed<dZ eejed<dZeejed<dZeejed <dZeejed <dZeeejd fed <dZeeejd fed <dZeejed<dZeeejd fed<dZeeejd fed<dZeeejd fed<dZeeejd fed<dZeeejd fed<dS)RetrievAugLMOutputa; Args: logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): Prediction scores of the language modeling head. The score is possibly marginalized over all documents for each vocabulary token. 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`. past_key_values (`List[torch.FloatTensor]`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): List of `torch.FloatTensor` of length `config.n_layers`, with each tensor of shape `(2, batch_size, num_heads, sequence_length, embed_size_per_head)`). Contains precomputed hidden-states (key and values in the attention blocks) of the decoder that can be used (see `past_key_values` input) to speed up sequential decoding. retrieved_doc_embeds (`torch.FloatTensor` of shape `(batch_size, config.n_docs, hidden_size)`, *optional*, returned when *output_retrieved=True*): Embedded documents retrieved by the retriever. Is used with `question_encoder_last_hidden_state` to compute the `doc_scores`. retrieved_doc_ids (`torch.LongTensor` of shape `(batch_size, config.n_docs)`, *optional*, returned when *output_retrieved=True*): The indexes of the embedded documents retrieved by the retriever. 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. question_encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): Sequence of hidden states at the output of the last layer of the question encoder pooled output of the model. question_enc_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings and one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden states of the question encoder at the output of each layer plus the initial embedding outputs. question_enc_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights of the question encoder, after the attention softmax, used to compute the weighted average in the self-attention heads. generator_enc_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): Sequence of hidden-states at the output of the last layer of the generator encoder of the model. generator_enc_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings and one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden states of the generator encoder at the output of each layer plus the initial embedding outputs. generator_enc_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights of the generator encoder, after the attention softmax, used to compute the weighted average in the self-attention heads. generator_dec_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings and one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden states of the generator decoder at the output of each layer plus the initial embedding outputs. generator_dec_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights of the generator decoder, after the attention softmax, used to compute the weighted average in the self-attention heads. generator_cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Cross-attentions weights of the generator decoder, after the attention softmax, used to compute the weighted average in the cross-attention heads. Nrrrrrrrr.r r!r"r#r$r%r&r')r(r)r*r+rrr,r-r.rrrrrr/rrrr rr!r"r#r$r%r&r'r0r0r0r1r2s$ Fr2a RAG models were released with the paper [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://arxiv.org/abs/2005.11401) by Patrick Lewis, Ethan Perez, Aleksandra Piktus et al. RAG is a retriever augmented model and encapsulate three components: a question encoder, a dataset retriever and a generator, the encoder and generator are trainable while the retriever is just an indexed dataset. )Z custom_introc @sJeZdZeZdZdZdZe   d de e de e de de fdd Z dS) RagPreTrainedModelragTN.question_encoder_pretrained_model_name_or_path'generator_pretrained_model_name_or_path retrieverreturncKsdd|D}dd|D}|D]}|d|=q|D]}|d|=q"|dd}|durh|dus>> from transformers import RagModel >>> # initialize a RAG from two pretrained models. >>> model = RagModel.from_pretrained_question_encoder_generator( ... "facebook/dpr-question_encoder-single-nq-base", "google-t5/t5-small" ... ) >>> # saving model after fine-tuning >>> model.save_pretrained("./rag") >>> # load fine-tuned model >>> model = RagModel.from_pretrained("./rag") ```cS,i|]\}}|dr|tdd|qS)question_encoder_N startswithlen.0argumentvaluer0r0r1 0 zQRagPreTrainedModel.from_pretrained_question_encoder_generator..cSr9) generator_Nr;r>r0r0r1rB6rCr:rDmodelNznIf `model` is not defined as an argument, a `question_encoder_pretrained_model_name_or_path` has to be defined AutoModelconfig) AutoConfigZreturn_unused_kwargsTzqIf `generator_model` is not defined as an argument, a `generator_pretrained_model_name_or_path` has to be definedAutoModelForSeq2SeqLM)question_encoder generatorrIr7) itemskeyspopauto.modeling_autorHZauto.configuration_autorJZfrom_pretrainedrLgetr'from_question_encoder_generator_configsrI)clsr5r6r7kwargsZkwargs_question_encoderZkwargs_generatorkeyrMrHrJZquestion_encoder_configrNrLZgenerator_configrIr0r0r1*from_pretrained_question_encoder_generatorsxD               z=RagPreTrainedModel.from_pretrained_question_encoder_generator)NNN)r(r)r*r config_classZbase_model_prefixZ_supports_flash_attn_2Z_supports_sdpa classmethodrstrrrrXr0r0r0r1r3s$ r3c"seZdZ    ddeedeedeedeeffdd Ze              ddee j d ee j d ee e e j d ee j d ee jd ee e e j dee j dee j dee j deedeedeedeedeedee e j effddZZS)RagModelNrIrMrNr7c s|dus|dur |dusJd|dur!tj|j|jfi|}nt||js2Jd|d|jt||durHddlm}| |j }|durXddlm }| |j }||_ |j durst|tspJdt|j d ||_ ||_ ||_ d|_d |_dS)  question_encoder (`PreTrainedModel`, *optional*): The model responsible for encoding the question into hidden states for retrieval. generator (`PreTrainedModel`, *optional*): The model responsible for generating text based on retrieved documents. retriever (`RagRetriever`, *optional*): The component responsible for retrieving documents from a knowledge base given the encoded question. NzQEither a configuration or an question_encoder and a generator has to be provided.zconfig: z has to be of type rFrGrKz`self.retriever` is of type z&, but should be of type `RagRetriever`F)rrTrI isinstancerYsuper__init__rRrHZ from_configrMrLrNr7rtype ctx_encodercontext_encoder_training)selfrIrMrNr7rVrHrL __class__r0r1r`|s6"        zRagModel.__init__ input_idsattention_maskencoder_outputsdecoder_input_idsdecoder_attention_maskrrrr use_cacheoutput_attentionsoutput_hidden_statesoutput_retrievedn_docsr8c Csr|dur|n|jj}| dur| n|jj} | dur| n|jj} | dur$| n|jj} | dur.| n|jj} |jduoF|dupB| dupB|duoF|du}|dur|r|j||dd}|d}|j||j dt j d |j jj|dd}|jr|d |d |d |d |d |df\}} }}}}| |}| |} | |}| |}|j||ddj}|d||jd}t |d|ddd}nM|d |d |d |df\}} }}| |}| |}| |} t |d|ddd}n|dusJd| dus Jd|dusJd|dusJd|jd|dks4Jd|d|jdd|dur@|j|dd}|durL|j|dd}|j || ||||| | dd }|shd}d}d}d}d}n|j}|j}|rt| s|d}d} d}d}td)id|jd|d|jd |d | d |d|d |d!|d"|d#|jd$|jd%|j d&|j!d'|j"d(|j#S)*ay input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. [`RagConfig`], used to initialize the model, specifies which generator to use, it also specifies a compatible generator tokenizer. Use that tokenizer class to obtain the indices. [What are input IDs?](../glossary#input-ids) encoder_outputs (`tuple(tuple(torch.FloatTensor)`, *optional*) Tuple consists of (`generator_enc_last_hidden_state`, *optional*: `generator_enc_hidden_states`, *optional*: `generator_enc_attentions`). `generator_enc_last_hidden_state` of shape `(batch_size, n_docs * sequence_length, hidden_size)` is a sequence of hidden-states at the output of the last layer of the generator's encoder. Used by the ([`RagModel`]) model during decoding. decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): Provide for generation tasks. `None` by default, construct as per instructions for the generator model you're using with your RAG instance. decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*): Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also be used by default. 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` `doc_scores` has to be provided to the forward pass. `doc_scores` can be computed via `question_encoder_last_hidden_state` and `retrieved_doc_embeds`, see examples for more information. 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 was 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_attention_mask` has to be provided to the forward pass. `context_attention_mask` are returned by [`~RagRetriever.__call__`]. output_retrieved (`bool`, *optional*): Whether or not to return the `retrieved_doc_embeds`, `retrieved_doc_ids`, `context_input_ids` and `context_attention_mask`. See returned tensors for more detail. n_docs (`int`, *optional*): The number of documents to retrieve. Example: ```python >>> from transformers import AutoTokenizer, RagRetriever, RagModel >>> import torch >>> tokenizer = AutoTokenizer.from_pretrained("facebook/rag-token-base") >>> retriever = RagRetriever.from_pretrained( ... "facebook/rag-token-base", index_name="exact", use_dummy_dataset=True ... ) >>> # initialize with RagRetriever to do everything in one forward call >>> model = RagModel.from_pretrained("facebook/rag-token-base", retriever=retriever) >>> inputs = tokenizer("How many people live in Paris?", return_tensors="pt") >>> outputs = model(input_ids=inputs["input_ids"]) ```NT)rh return_dictrcpudevicedtypeptprefixrpZreturn_tensorsrrrZtokenized_doc_idsZtokenized_doc_attention_maskZdoc_idsrrFzMake sure that `context_input_ids` are passed, if no `retriever` is set. Alternatively, you can set a retriever using the `set_retriever(...)` function.zMake sure that `context_attention_mask` are passed, if no `retriever` is set. Alternatively, you can set a retriever using the `set_retriever(...)` function.zMake sure that `doc_scores` are passed, if no `retriever` is set. Alternatively, you can set a retriever using the `set_retriever(...)` function.z^Make sure that `doc_scores` are passed when passing `encoder_outputs` to the forward function.M The first dimension of `context_input_ids` should be a multiple of `n_docs`= , but is .dim) rgrhrirjrkrrlrmrqNrrrrrr r!r"r#r$r%r&r'r0)$rIrprlrmrnror7rMdetachtor,float32numpyrNrxrcrbZ pooler_outputviewshapebmm unsqueeze transposesqueezerepeat_interleave hidden_statesZ attentionsr2rrZencoder_last_hidden_stateZencoder_hidden_statesZencoder_attentionsZdecoder_hidden_statesZdecoder_attentionsZcross_attentions)rdrgrhrirjrkrrrrrlrmrnrorpZhas_to_retrieveZquestion_enc_outputsrZretriever_outputsrZretrieved_doc_input_idsZretrieved_doc_attention_maskrZ gen_outputsr r!r0r0r1forwards&I                       zRagModel.forwardNNNN)NNNNNNNNNNNNNN)r(r)r*rr rrr`rr,r/Tensorrr- BoolTensorboolintrr2r __classcell__r0r0rer1r\zsx2     r\zu A RAG-sequence model implementation. It performs RAG-sequence specific marginalization in the forward pass. c&seZdZ    d2deedeedeedeeffdd Zdefdd Zd efd d Z e                 d3d ee j dee j deeee j dee j dee jdeeee j dee j dee j dee jdeedeedeedeedeedeedee j deedef$dd Zed!d"Zed#d$Zed%d&Ze          d4d ee j dee j dee j dee j dee jd'eed(eed)eedeede j fd*d+Z d5d.d/Zed0d1ZZS)6RagSequenceForGenerationNrIrMrNr7c b|dus|dur |dusJd|dur tj|j|jfi|}t|t||||d|_dSr]NzHEither a configuration or an encoder and a generator has to be provided.)rIrMrNr7rrTrIr_r`r\r4rdrIrMrNr7rVrer0r1r`s z!RagSequenceForGeneration.__init__cC ||j_dSrr4r7rdr7r0r0r1 set_retriever z&RagSequenceForGeneration.set_retrieverrbcCd|j_||j_dSNTr4rcrbrdrbr0r0r1 set_context_encoder_for_training z9RagSequenceForGeneration.set_context_encoder_for_trainingrgrhrirjrkrrrrrlrmrnroexclude_bos_score reduce_losslabelsrpr8cKs6|dur|n|jj}|dur|n|jj}|dur|n|jj}|dur*|dur(|}d} |j|||||||| || | | | |d}d}|durS|j|j|j|||jj||d}t did|d|jd|jd|j d |j d |j d |j d |jd |jd|jd|jd|jd|jd|jd|jd|jd|jS)a3 input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. [`RagConfig`], used to initialize the model, specifies which generator to use, it also specifies a compatible generator tokenizer. Use that tokenizer class to obtain the indices. [What are input IDs?](../glossary#input-ids) encoder_outputs (`tuple(tuple(torch.FloatTensor)`, *optional*) Tuple consists of (`generator_enc_last_hidden_state`, *optional*: `generator_enc_hidden_states`, *optional*: `generator_enc_attentions`). `generator_enc_last_hidden_state` of shape `(batch_size, n_docs * sequence_length, hidden_size)` is a sequence of hidden-states at the output of the last layer of the generator's encoder. Used by the ([`RagModel`]) model during decoding. decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): Provide for generation tasks. `None` by default, construct as per instructions for the generator model you're using with your RAG instance. decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*): Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also be used by default. 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 was 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_attention_mask` has to be provided to the forward pass. `context_attention_mask` 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` `doc_scores` has to be provided to the forward pass. `doc_scores` can be computed via `question_encoder_last_hidden_state` and `retrieved_doc_embeds`, see examples for more information. output_retrieved (`bool`, *optional*): Whether or not to return the `retrieved_doc_embeds`, `retrieved_doc_ids`, `context_input_ids` and `context_attention_mask`. See returned tensors for more detail. exclude_bos_score (`bool`, *optional*): Only relevant if `labels` is passed. If `True`, the score of the BOS token is disregarded when computing the loss. reduce_loss (`bool`, *optional*): Only relevant if `labels` is passed. If `True`, the NLL loss is reduced using the `torch.Tensor.sum` operation. n_docs (`int`, *optional*): The number of documents to retrieve. Example: ```python >>> from transformers import AutoTokenizer, RagRetriever, RagSequenceForGeneration >>> import torch >>> tokenizer = AutoTokenizer.from_pretrained("facebook/rag-sequence-nq") >>> retriever = RagRetriever.from_pretrained( ... "facebook/rag-sequence-nq", index_name="exact", use_dummy_dataset=True ... ) >>> # initialize with RagRetriever to do everything in one forward call >>> model = RagSequenceForGeneration.from_pretrained("facebook/rag-token-nq", retriever=retriever) >>> inputs = tokenizer("How many people live in Paris?", return_tensors="pt") >>> targets = tokenizer(text_target="In Paris, there are 10 million people.", return_tensors="pt") >>> input_ids = inputs["input_ids"] >>> labels = targets["input_ids"] >>> outputs = model(input_ids=input_ids, labels=labels) >>> # or use retriever separately >>> model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-nq", use_dummy_dataset=True) >>> # 1. Encode >>> question_hidden_states = model.question_encoder(input_ids)[0] >>> # 2. Retrieve >>> docs_dict = retriever(input_ids.numpy(), question_hidden_states.detach().numpy(), return_tensors="pt") >>> doc_scores = torch.bmm( ... question_hidden_states.unsqueeze(1), docs_dict["retrieved_doc_embeds"].float().transpose(1, 2) ... ).squeeze(1) >>> # 3. Forward to generator >>> outputs = model( ... context_input_ids=docs_dict["context_input_ids"], ... context_attention_mask=docs_dict["context_attention_mask"], ... doc_scores=doc_scores, ... decoder_input_ids=labels, ... ) ```NFrgrhrirjrkrrrrrlrmrnrorp)repsilonrrprrrrrrrrrr r!r"r#r$r%r&r'r0)rIrprrr4get_nllrrlabel_smoothingrrrrrrrr r!r"r#r$r%r&r')rdrgrhrirjrkrrrrrlrmrnrorrrrprVoutputsrr0r0r1rsg      z RagSequenceForGeneration.forwardcC|jjSrrrdr0r0r1r7az"RagSequenceForGeneration.retrievercCrrr4rNrr0r0r1rNerz"RagSequenceForGeneration.generatorcCrrr4rMrr0r0r1rMirz)RagSequenceForGeneration.question_encoderdo_deduplicationnum_return_sequences num_beamsc KsT| dur| n|jj} |dur|n|jj}|dur|n|jj} |dur$|n|jj}|dus4|dus4Jd|jdurd|durd|j||dd} |j|| jdt j d |j jj | ddd }||}g} || d <|| d <d| d <|dur{|jdn|jd| }t|D]}||| |d | }|j j|fi| }|rt tdd|D}|jd}|dur|||d |d }|||dd}nC|dusJd|dusJd||d }||| |d | }||d }|||d ddf}||d }|||||dd}|d | d }| ||q|j| |jj jdS)a Implements RAG sequence "thorough" decoding. Read the [`~generation.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.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 (`Dict[str, Any]`, *optional*): Additional kwargs will be passed to [`~generation.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`. Nz= At least one of input_ids or context_input_ids must be givenrhrrrrsrvrwrrrrhrcSsi|] }t||qSr0)r[tolist)r?kr0r0r1rBsz5RagSequenceForGeneration.generate..T)rrzMake sure that `context_attention_mask` are passed, if no `input_ids` is set. Alternatively, you can set a retriever using the `set_retriever(...)` function.zMake sure that `doc_scores` are passed, if no `input_ids` is set. Alternatively, you can set a retriever using the `set_retriever(...)` function.)rrrrrr) pad_token_id)rIrprrrr7rMrrr,rrrNrxrrangegeneratestacklistvaluesrepeatZtopkappend _cat_and_padr)rdrgrhrrrrrrrp model_kwargsZnum_doc_return_sequencesquestion_hidden_statesZhypos batch_sizeindexZgenerator_input_idsZoutput_sequencesZnum_candidatesZ new_input_idsrZindividual_input_idsZindividual_attention_maskZindividual_doc_scoresZ top_cand_indsr0r0r1rms~A       z!RagSequenceForGeneration.generateFcsBtddddfjddjjjgd|dur#|njj}jj p/jjj }|duo@dddf | } fdd} t j j|dd|jd||d|d} t j j|dddd} | ddddddddf} | ddddddddf}| ddddddddf}tj| || |gdd}ddd|dd|ksJ|jdd}|jdd d }| ||\}}|r| r|ddddddfdn|d}|d}|d}|d}| }| }|r|}|}||d}d ||||}|S) NrrcDjjj}|r||d||d|d|dfSNrryeqrIrNranyZ masked_fill_rll smooth_objZpad_maskrdtargetr0r1 _mask_pads   z4RagSequenceForGeneration.get_nll.._mask_padsryr}rFr~rTr~Zkeepdim?)r,catnewrfill_rIrNrrp bos_token_idrallr functional log_softmaxrsizerrr~gathersum logsumexp)rd seq_logitsrrrrrrprZuse_bosr seq_logprobs doc_logprobsZfirst_token_scoresZsecond_token_scores remainder rag_logprobsrrnll_loss smooth_losseps_irr0rr1rs@2"   2   z RagSequenceForGeneration.get_nllcCsv|dtdd|Dtdd|D|}d}|D]}|||||jdd|jdf<||jd7}q|S)NrcSg|]}|jdqS)rrr?tr0r0r1 Bz9RagSequenceForGeneration._cat_and_pad..cSr)rrrr0r0r1rBrr)rrmaxrr)Ztensorsroutputindrr0r0r1r?s0$z%RagSequenceForGeneration._cat_and_padrNNNNNNNNNNNNNNNNN) NNNNNNNNN)FrFN)r(r)r*rr rrr`rrrr,r/rrrr-rrrrpropertyr7rNrMno_gradrr staticmethodrrr0r0rer1rs      !       ;rzo A RAG-token model implementation. It performs RAG-token specific marginalization in the forward pass. c&s`eZdZ    d?deedeedeedeeffdd Zdefdd Zd efd d Z      d@d dZ e ddZ e ddZ e ddZeddZdAddZe                 dBdeejdeejdeeeejdeejdeejdeeeejdeejd eejd!eejd"eed#eed$eed%eed&eed'eed(eejd)eed*ef$d+d,Zeddddddddeef deejdeejdeejd eejd!eejd)eed-ee d.ee!eejge"efd/eed0eed*ejfd1d2Z#d3d4Z$d5d6Z%d7d8Z&dAd9d:Z'dCd=d>Z(Z)S)DRagTokenForGenerationNrIrMrNr7c rrrrrer0r1r`Qs zRagTokenForGeneration.__init__cCrrrrr0r0r1rorz#RagTokenForGeneration.set_retrieverrbcCrrrrr0r0r1rrrz6RagTokenForGeneration.set_context_encoder_for_trainingc Ks4|dur|ddddf}d||||||d|d S)NryT) rgrirrrjrrldo_marginalizerpr0) rdrjrrhrlrirrprVr0r0r1prepare_inputs_for_generationvs z3RagTokenForGeneration.prepare_inputs_for_generationcCrrrrr0r0r1r7rzRagTokenForGeneration.retrievercCrrrrr0r0r1rNrzRagTokenForGeneration.generatorcCrrrrr0r0r1rMrz&RagTokenForGeneration.question_encodercs8ddd}|D]}|tfdd|Df7}q|S)zeReorders cache for generation. BART-inspired but we need to take care of the extra dimension for docscSs^|jd|jd}|jd|g|jddR}|d|}|jdg|jddR}|S)NrryrrF)rrZ index_select)rZ new_orderrpresultr0r0r1_reorder_stackeds  z>RagTokenForGeneration._reorder_cache.._reorder_stackedr0c3s"|] }||jVqdSr)rrt)r?Z past_staterbeam_idxr0r1 s z7RagTokenForGeneration._reorder_cache..)tuple)rrZreordered_pastZ layer_pastr0rr1_reorder_cachesz$RagTokenForGeneration._reorder_cachecCsp|dur|n|jj}tjj|dd|jd||d|d}tj|dd}|| d d}tj |ddS)Nryr}rr) rIrprrrrrrr,rr)rdrrrprrZ log_prob_sumr0r0r1 marginalizesz!RagTokenForGeneration.marginalizergrhrirjrkrrrrrlrmrnrorrrrpr8cKsX|dur|n|jj}|dur|n|jj}|dur|n|jj}|dur*|dur(|}d} |j|||||||| || | | | |d}d}|j}|dur[|dusLJ|j|j|j|||jj|d}|re| ||j|}t did|d|d|jd|j d |j d |j d |jd |jd |jd|jd|jd|jd|jd|jd|jd|jd|jS)a input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. [`RagConfig`], used to initialize the model, specifies which generator to use, it also specifies a compatible generator tokenizer. Use that tokenizer class to obtain the indices. [What are input IDs?](../glossary#input-ids) encoder_outputs (`tuple(tuple(torch.FloatTensor)`, *optional*) Tuple consists of (`generator_enc_last_hidden_state`, *optional*: `generator_enc_hidden_states`, *optional*: `generator_enc_attentions`). `generator_enc_last_hidden_state` of shape `(batch_size, n_docs * sequence_length, hidden_size)` is a sequence of hidden-states at the output of the last layer of the generator's encoder. Used by the ([`RagModel`]) model during decoding. decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): Provide for generation tasks. `None` by default, construct as per instructions for the generator model you're using with your RAG instance. decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*): Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also be used by default. 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 was 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_attention_mask` has to be provided to the forward pass. `context_attention_mask` 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` `doc_scores` has to be provided to the forward pass. `doc_scores` can be computed via `question_encoder_last_hidden_state` and `retrieved_doc_embeds`, see examples for more information. output_retrieved (`bool`, *optional*): Whether or not to return the `retrieved_doc_embeds`, `retrieved_doc_ids`, `context_input_ids` and `context_attention_mask`. See returned tensors for more detail. do_marginalize (`bool`, *optional*): If `True`, the logits are marginalized over all documents by making use of `torch.nn.functional.log_softmax`. reduce_loss (`bool`, *optional*): Only relevant if `labels` is passed. If `True`, the NLL loss is reduced using the `torch.Tensor.sum` operation. n_docs (`int`, *optional*): The number of documents to retrieve. Example: ```python >>> from transformers import AutoTokenizer, RagRetriever, RagTokenForGeneration >>> import torch >>> tokenizer = AutoTokenizer.from_pretrained("facebook/rag-token-nq") >>> retriever = RagRetriever.from_pretrained( ... "facebook/rag-token-nq", index_name="exact", use_dummy_dataset=True ... ) >>> # initialize with RagRetriever to do everything in one forward call >>> model = RagTokenForGeneration.from_pretrained("facebook/rag-token-nq", retriever=retriever) >>> inputs = tokenizer("How many people live in Paris?", return_tensors="pt") >>> targets = tokenizer(text_target="In Paris, there are 10 million people.", return_tensors="pt") >>> input_ids = inputs["input_ids"] >>> labels = targets["input_ids"] >>> outputs = model(input_ids=input_ids, labels=labels) >>> # or use retriever separately >>> model = RagTokenForGeneration.from_pretrained("facebook/rag-token-nq", use_dummy_dataset=True) >>> # 1. Encode >>> question_hidden_states = model.question_encoder(input_ids)[0] >>> # 2. Retrieve >>> docs_dict = retriever(input_ids.numpy(), question_hidden_states.detach().numpy(), return_tensors="pt") >>> doc_scores = torch.bmm( ... question_hidden_states.unsqueeze(1), docs_dict["retrieved_doc_embeds"].float().transpose(1, 2) ... ).squeeze(1) >>> # 3. Forward to generator >>> outputs = model( ... context_input_ids=docs_dict["context_input_ids"], ... context_attention_mask=docs_dict["context_attention_mask"], ... doc_scores=doc_scores, ... decoder_input_ids=labels, ... ) >>> # or directly generate >>> generated = model.generate( ... context_input_ids=docs_dict["context_input_ids"], ... context_attention_mask=docs_dict["context_attention_mask"], ... doc_scores=doc_scores, ... ) >>> generated_string = tokenizer.batch_decode(generated, skip_special_tokens=True) ```NFr)rrrprrrrrrrrrr r!r"r#r$r%r&r'r0)rIrprrr4rrrrrrrrrrrrr r!r"r#r$r%r&r')rdrgrhrirjrkrrrrrlrmrnrorrrrprVrrrr0r0r1rso       zRagTokenForGeneration.forwardgeneration_configprefix_allowed_tokens_fnlogits_processorstopping_criteriac  s|dur|j}t|}|jd&i| } | dddu} ||| dur(n|jj|jdur|dur|j ||dd}|j|| j dt j d|jjjdd}|d |d |d }}}| |}| |}| |}t |d |d d d }|jddksJdd|jdd|jd|jj}|||dd}t j|jd f|jt jt|jd}|jd}|d}d'fdd }|||jd}|||jd|d<|j|jdd}|| d<|| d<|| d<| d<|j ||||| |jd}|j!|| d}|jd kr2|j"d kr#t#d|j"d |j$|f|||d!dd"| S|jd krQ|j"|jkrCt#d#|j%|f|||d!d$| St#d%|j)(a 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`. NrhrrrrrsrvrwrrrrrFrzr{r|T)rgrhrq)rurtrylast_hidden_statecsl|ddddfdf|jdd}||f|jdd}||f|jddS)Nrr )Zreshaperexpand)Ztensorrrrpr0r1extend_enc_outputs,"z9RagTokenForGeneration.generate..extend_enc_output)rr}rrirp)rinput_ids_seq_lengthZencoder_input_idsrrrt)rrz)num_return_sequences has to be 1, but is z when doing greedy search.F)rrr synced_gpusstreamerzA`num_return_sequences` has to be smaller or equal to `num_beams`.)rrrruH`num_beams` has to be an integer strictly superior to 0 (≥ 1), but is r0r)&rcopydeepcopyupdaterSZ_prepare_special_tokensrIrpr7rMrrr,rrrNrxrrrrrr4Z get_encoderfullrdecoder_start_token_idlongnext parametersrtrZ_get_logits_processorZ_get_stopping_criteriar ValueErrorZ_sampleZ _beam_search)rdrgrhrrrrprrrrrVrZkwargs_has_attention_maskroutrencoderrirrrZ pre_processorZprepared_stopping_criteriar0rr1rksQ                 zRagTokenForGeneration.generatecC |jjSr)r4rNget_input_embeddingsrr0r0r1r7rz*RagTokenForGeneration.get_input_embeddingscCrr)r4rNget_output_embeddingsrr0r0r1r:rz+RagTokenForGeneration.get_output_embeddingscCs|jj|Sr)r4rNset_output_embeddings)rdZnew_embeddingsr0r0r1r=sz+RagTokenForGeneration.set_output_embeddingscCsX|dur|jj}||j}|ddddf|ddddf<||dddf<|S)zCShift input ids one token to the right, and pad with start_token_idNryrr)rIr Z new_zerosrclone)rdrgZstart_token_idZshifted_input_idsr0r0r1shift_tokens_right@s  (z(RagTokenForGeneration.shift_tokens_rightFrcs |dur|njj}tddddfjddjjjgdfdd} |||} d | ksDJ|j dd} |j ddd} || | \} } | d} | d} | } | } |rs| } | } ||d} d || | | }|S) Nrrcrrrrrr0r1rPrz1RagTokenForGeneration.get_nll.._mask_padsryrTrr)rIrpr,rrrrrNrrrr~rrr)rdrrrrrrprrrrrrrrr0rr1rIs*2   zRagTokenForGeneration.get_nllr)NNNNNNrr)FrN)*r(r)r*rr rrr`rrrrr7rNrMrrrrr,r/r-rrrrrrrrr rr rrrrrrrrrr0r0rer1rKs             -    L  r)r\r3rr)(r+r  dataclassesrtypingrrrrrr,rZconfiguration_utilsr Z generationr r r rZmodeling_outputsrZmodeling_utilsrutilsrrZconfiguration_ragrZ retrieval_ragrZ get_loggerr(loggerrr2r3r\rr__all__r0r0r0r1sV        ^Y 3 "