a h\@shdZddlZddlmZddlmZmZmZddlZddlm Z ddl m Z m Z ddl mZdd lmZmZmZmZdd lmZdd lmZdd lmZmZd dlmZd dlmZee Z!eeddGdddeZ"eeGdddeZ#eddeGdddeZ$eGddde$Z%eddGddde$Z&eddGdd d e$eZ'gd!Z(dS)"zRAG model implementation.N) dataclass)CallableOptionalUnion)nn)CacheEncoderDecoderCache)PretrainedConfig)GenerationConfigGenerationMixinLogitsProcessorListStoppingCriteriaList) ModelOutput)PreTrainedModel)auto_docstringlogging) RagConfig) RagRetrieverzI Base class for retriever augmented marginalized models outputs. )Z custom_introc@szeZdZUdZdZeejed<dZ eejed<dZ eejed<dZ ee ed<dZ eejed<dZeejed<dZeejed <dZeejed <dZeejed <dZeeejd fed <dZeeejd fed<dZeejed<dZeeejd fed<dZeeejd fed<dZeeejd fed<dZeeejd fed<dZeeejd fed<dS)RetrievAugLMMarginOutputa 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 (`Cache`, *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 tupler!r"r#r$r%r&r'r1r1`/var/www/html/assistant/venv/lib/python3.9/site-packages/transformers/models/rag/modeling_rag.pyr%s$ Grc@sheZdZUdZdZeejed<dZ eejed<dZ ee ed<dZ eejed<dZ eejed<dZeejed<dZeejed <dZeejed <dZeeejd fed <dZeeejd fed <dZeejed<dZeeejd fed<dZeeejd fed<dZeeejd fed<dZeeejd fed<dZeeejd fed<dS)RetrievAugLMOutputa7 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 (`Cache`, *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 r0r!r"r#r$r%r&r'r1r1r1r2r3s" Er3a RAG models were released with the paper [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://huggingface.co/papers/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. c@sDeZdZUeed<dZdZdZede e e e e e dddZ dS) RagPreTrainedModelconfigragTN).question_encoder_pretrained_model_name_or_path'generator_pretrained_model_name_or_path retrieverreturncKsdd|D}dd|D}|D]}|d|=q(|D]}|d|=q<|dd}|dur|duspJdd d lm} d |vrd d lm} | j|fi|d di\} }| |d <| j|fi|}|dd} | durH|dusJdd dlm} d |vr6d d lm} | j|fi|d di\}}||d <| j|fi|} |d }|durtt j |j | j fi|}||| ||dS)a Instantiates an question encoder and a generator from one or two base classes of the library from pretrained model checkpoints. The model is set in evaluation mode by default using `model.eval()` (Dropout modules are deactivated). To train the model, you need to first set it back in training mode with `model.train()`. Params: question_encoder_pretrained_model_name_or_path (`str`, *optional*, defaults to `None`): Information necessary to initiate the question encoder. Can be either: - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co. - A path to a *directory* containing model weights saved using [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`. - A path or url to a *tensorflow index checkpoint file* (e.g, `./tf_model/model.ckpt.index`). In this case, `from_tf` should be set to `True` and a configuration object should be provided as `config` argument. This loading path is slower than converting the TensorFlow checkpoint in a PyTorch model using the provided conversion scripts and loading the PyTorch model afterwards. generator_pretrained_model_name_or_path (`str`, *optional*, defaults to `None`): Information necessary to initiate the generator. Can be either: - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co. - A path to a *directory* containing model weights saved using [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`. - A path or url to a *tensorflow index checkpoint file* (e.g, `./tf_model/model.ckpt.index`). In this case, `from_tf` should be set to `True` and a configuration object should be provided as `config` argument. This loading path is slower than converting the TensorFlow checkpoint in a PyTorch model using the provided conversion scripts and loading the PyTorch model afterwards. model_args (remaining positional arguments, *optional*): All remaining positional arguments will be passed to the underlying model's `__init__` method. retriever ([`RagRetriever`], *optional*): The retriever to use. kwwargs (remaining dictionary of keyword arguments, *optional*): Can be used to update the configuration object (after it being loaded) and initiate the model (e.g., `output_attentions=True`). - To update the question_encoder configuration, use the prefix *question_encoder_* for each configuration parameter. - To update the generator configuration, use the prefix *generator_* for each configuration parameter. - To update the parent model configuration, do not use a prefix for each configuration parameter. Behaves differently depending on whether a `config` is provided or automatically loaded. Example: ```python >>> 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") ```cSs,i|]$\}}|dr|tdd|qS)question_encoder_N startswithlen.0argumentvaluer1r1r2 4s zQRagPreTrainedModel.from_pretrained_question_encoder_generator..cSs,i|]$\}}|dr|tdd|qS) generator_Nr<r?r1r1r2rC:s r;rDmodelNznIf `model` is not defined as an argument, a `question_encoder_pretrained_model_name_or_path` has to be defined AutoModelr5) AutoConfigZreturn_unused_kwargsTzqIf `generator_model` is not defined as an argument, a `generator_pretrained_model_name_or_path` has to be definedAutoModelForSeq2SeqLM)question_encoder generatorr5r9) itemspopauto.modeling_autorHZauto.configuration_autorIZfrom_pretrainedrKgetr'from_question_encoder_generator_configsr5)clsr7r8r9kwargsZkwargs_question_encoderZkwargs_generatorkeyrLrHrIZquestion_encoder_configrMrKZgenerator_configr5r1r1r2*from_pretrained_question_encoder_generatorsxD                z=RagPreTrainedModel.from_pretrained_question_encoder_generator)NNN)r(r)r*rr.Zbase_model_prefixZ_supports_flash_attnZ_supports_sdpa classmethodrstrrrrVr1r1r1r2r4s r4cseZdZdeeeeeeeedfdd Zed ee j ee j ee e e j ee j ee jeeee j ee j ee j eeeeeeeeeeee e j efdddZZS) RagModelNr5rLrMr9c s|dus |dur|dus Jd|durBtj|j|jfi|}n"t||jsdJd|d|jt||durddlm}| |j }|durddlm }| |j }||_ |j durt|tsJdt|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 rFrGrJz`self.retriever` is of type z&, but should be of type `RagRetriever`F)rrRr5 isinstanceZ config_classsuper__init__rPrH from_configrLrKrMr9rtype ctx_encodercontext_encoder_training)selfr5rLrMr9rTrHrK __class__r1r2r^s6"       zRagModel.__init__) input_idsattention_maskencoder_outputsdecoder_input_idsdecoder_attention_maskrrrr use_cacheoutput_attentionsoutput_hidden_statesoutput_retrievedn_docsr:cCs4|dur |n|jj}| dur | n|jj} | dur4| n|jj} | durH| n|jj} | dur\| n|jj} |jduo|dus| dus|duo|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}nb|d |d |d |df\}} }}| |}| |}| |} t |d|ddd}n6|dusJd| dusJd|dus*Jd|dus>> 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)rg return_dictrcpudevicedtypeptprefixroZreturn_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) rfrgrhrirjrrkrlrpN)rrrrrrrrr r!r"r#r$r%r&r')$r5rorkrlrmrnr9rLdetachtor,float32numpyrMrwrbraZ pooler_outputviewshapebmm unsqueeze transposesqueezerepeat_interleave hidden_statesZ attentionsr3rrZencoder_last_hidden_stateZencoder_hidden_statesZencoder_attentionsZdecoder_hidden_statesZdecoder_attentionsZcross_attentions)rcrfrgrhrirjrrrrrkrlrmrnroZhas_to_retrieveZquestion_enc_outputsrZretriever_outputsrZretrieved_doc_input_idsZretrieved_doc_attention_maskrZ gen_outputsr r!r1r1r2forwardsI                    zRagModel.forward)NNNN)NNNNNNNNNNNNNN)r(r)r*rr rrr^rr,r/Tensorr0r- BoolTensorrboolintrr3r __classcell__r1r1rdr2rY~sT2rYzu A RAG-sequence model implementation. It performs RAG-sequence specific marginalization in the forward pass. cs~eZdZdeeeeeeeedfdd ZedddZedd d Z e dee j ee j eeee j ee j ee jeeee j ee j ee jeeeeeeeeeeeeee j eeed d d ZeddZeddZeddZe dee j ee j ee j ee j ee jeeeeeeeee j d ddZd ddZeddZZS)!RagSequenceForGenerationNrZc sb|dus |dur|dus Jd|dur@tj|j|jfi|}t|t||||d|_dSr[NzHEither a configuration or an encoder and a generator has to be provided.rZrrRr5r]r^rYr6rcr5rLrMr9rTrdr1r2r^s z!RagSequenceForGeneration.__init__r9cCs ||j_dSr~r6r9rcr9r1r1r2 set_retrieversz&RagSequenceForGeneration.set_retrieverracCsd|j_||j_dSNTr6rbrarcrar1r1r2 set_context_encoder_for_trainingsz9RagSequenceForGeneration.set_context_encoder_for_training)rfrgrhrirjrrrrrkrlrmrnexclude_bos_score reduce_losslabelsror:cKs|dur |n|jj}|dur |n|jj}|dur4|n|jj}|durT|durP|}d} |j|||||||| || | | | |d}d}|dur|j|j|j|||jj||d}t ||j|j|j |j |j |j |j|j|j|j|j|j|j|j|j|jdS)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, ... ) ```NFrfrgrhrirjrrrrrkrlrmrnro)repsilonrrorrrrrrrrrr r!r"r#r$r%r&r')r5rorrr6get_nllrrlabel_smoothingrrrrrrrr r!r"r#r$r%r&r')rcrfrgrhrirjrrrrrkrlrmrnrrrrorToutputsrr1r1r2rsjg z RagSequenceForGeneration.forwardcCs|jjSr~rrcr1r1r2r9esz"RagSequenceForGeneration.retrievercCs|jjSr~r6rMrr1r1r2rMisz"RagSequenceForGeneration.generatorcCs|jjSr~r6rLrr1r1r2rLmsz)RagSequenceForGeneration.question_encoder) rfrgrrrdo_deduplicationnum_return_sequences num_beamsror:c Ks`| dur | n|jj} |dur |n|jj}|dur4|n|jj} |durH|n|jj}|dush|dushJd|jdur|dur|j||dd} |j|| jdt j d |j jj | ddd }||}g} || d <|| d <d| d <|dur|jdn |jd| }t|D]<}||| |d | }|j j|fi| }|r`t tdd|D}|jd}|dur|||d |d }|||dd}n|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 givenrgrrqrrrurvrrrrgrcSsi|]}t||qSr1)rXtolist)r@kr1r1r2rCz5RagSequenceForGeneration.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)r5rorrrr9rLrrr,rrrMrwrrangegeneratestacklistvaluesrepeatZtopkappend _cat_and_padr)rcrfrgrrrrrrro model_kwargsZnum_doc_return_sequencesquestion_hidden_statesZhypos batch_sizeindexZgenerator_input_idsZoutput_sequencesZnum_candidatesZ new_input_idsrZindividual_input_idsZindividual_attention_maskZindividual_doc_scoresZ top_cand_indsr1r1r2rqs~A       z!RagSequenceForGeneration.generateFcsHtddddfjddjjjgd|durF|njj}jj p^jjj }|duodddf | } 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) NrrcsDjjj}|r0||d||d|d|dfSNrrxeqr5rMranyZ masked_fill_rll smooth_objZpad_maskrctargetr1r2 _mask_padss   z4RagSequenceForGeneration.get_nll.._mask_padsrxr|rFr}rTr}Zkeepdim?)r,catnewrfill_r5rMrro bos_token_idrallr functional log_softmaxrsizerrr}gathersum logsumexp)rc seq_logitsrrrrrrorZuse_bosr seq_logprobs doc_logprobsZfirst_token_scoresZsecond_token_scores remainder rag_logprobsrrnll_loss smooth_losseps_irr1rr2rs@2"   6   z RagSequenceForGeneration.get_nllcCsv|dtdd|Dtdd|D|}d}|D]6}|||||jdd|jdf<||jd7}q:|S)NrcSsg|]}|jdqS)rrr@tr1r1r2 Frz9RagSequenceForGeneration._cat_and_pad..cSsg|]}|jdqS)rrrr1r1r2rFrr)rrmaxrr)Ztensorsroutputindrr1r1r2rCs0$z%RagSequenceForGeneration._cat_and_pad)NNNN)NNNNNNNNNNNNNNNNN) NNNNNNNNN)FrFN) r(r)r*rr rrr^rrrr,r/rr0rrr-rrrrpropertyr9rMrLno_gradrr staticmethodrrr1r1rdr2rs!    ;rzo A RAG-token model implementation. It performs RAG-token specific marginalization in the forward pass. cseZdZd+eeeeeeeedfdd ZedddZedd d Z d,d d Z e d dZ e ddZ e ddZeddZd-ddZed.eejeejeeeejeejeejeeeejeejeejeeeeeeeeeeeeeejeeedddZeddddddddee f eejeejeejeejeejeeee!ee"eejge#efeeee ejd ddZ$ddZ%dd Z&d!d"Z'd#d$Z(d/d%d&Z)d0d)d*Z*Z+S)1RagTokenForGenerationNrZc sb|dus |dur|dus Jd|dur@tj|j|jfi|}t|t||||d|_dSrrrrdr1r2r^Us zRagTokenForGeneration.__init__rcCs ||j_dSr~rrr1r1r2rssz#RagTokenForGeneration.set_retrieverrcCsd|j_||j_dSrrrr1r1r2rvsz6RagTokenForGeneration.set_context_encoder_for_trainingc Ks4|dur|ddddf}d||||||d|d S)NrxT) rfrhrrrirrkdo_marginalizeror1) rcrirrgrkrhrrorTr1r1r2prepare_inputs_for_generationzs z3RagTokenForGeneration.prepare_inputs_for_generationcCs|jjSr~rrr1r1r2r9szRagTokenForGeneration.retrievercCs|jjSr~rrr1r1r2rMszRagTokenForGeneration.generatorcCs|jjSr~rrr1r1r2rLsz&RagTokenForGeneration.question_encodercsLddd}|D]"}|tfdd|Df7}qt|trHt|}|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)NrrxrrF)rrZ index_select)rZ new_orderroresultr1r1r2_reorder_stackeds  z>RagTokenForGeneration._reorder_cache.._reorder_stackedr1c3s |]}||jVqdSr~)rrs)r@Z past_staterbeam_idxr1r2 rz7RagTokenForGeneration._reorder_cache..)r0r\r Zfrom_legacy_cache)rrZreordered_pastZ layer_pastr1rr2_reorder_caches  z$RagTokenForGeneration._reorder_cachecCsp|dur |n|jj}tjj|dd|jd||d|d}tj|dd}|| d d}tj |ddS)Nrxr|rr) r5rorrrrrrr,rr)rcrrrorrZ log_prob_sumr1r1r2 marginalizesz!RagTokenForGeneration.marginalize)rfrgrhrirjrrrrrkrlrmrnrrrror:cKs|dur |n|jj}|dur |n|jj}|dur4|n|jj}|durT|durP|}d} |j|||||||| || | | | |d}d}|j}|dur|dusJ|j|j|j|||jj|d}|r| ||j|}t |||j|j |j |j |j|j|j|j|j|j|j|j|j|j|jdS)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)rrror)r5rorrr6rrrrrrrrrrrrr r!r"r#r$r%r&r')rcrfrgrhrirjrrrrrkrlrmrnrrrrorTrrrr1r1r2rspo  zRagTokenForGeneration.forward) rfrgrrrrogeneration_configprefix_allowed_tokens_fnlogits_processorstopping_criteriar:c  s|dur|j}t|}|jfi| } | dddu} ||| durPn|jj|jdur|dur|j ||dd}|j|| j dt j d|jjjdd}|d |d |d }}}| |}| |}| |}t |d |d d d }|jddks8Jdd|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}|j"|| d|jd|j#d d|jd kr|j$d krlt%d |j$d!|j&|f|||d"dd#| S|jd kr|j$|jkrt%d$|j'|f|||d"d%| St%d&|jdS)(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://huggingface.co/papers/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`. NrgrrrqrrrurvrrrrrFryrzr{T)rfrgrp)rtrsrxlast_hidden_statecsl|ddddfdf|jdd}||f|jdd}||f|jddS)Nrr)Zreshaperexpand)Ztensorrrror1r2extend_enc_outputs,z9RagTokenForGeneration.generate..extend_enc_output)rr|rrhro)rinput_ids_seq_lengthZencoder_input_idsrrrs)rr)Zassistant_modelrZmax_cache_lengthz)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 )N)(rcopydeepcopyupdaterQZ_prepare_special_tokensr5ror9rLrrr,rrrMrwrrrrrr6Z get_encoderfullrdecoder_start_token_idlongnext parametersrsrZ_get_logits_processorZ_get_stopping_criteriaZ_prepare_cache_for_generation max_lengthr ValueErrorZ_sampleZ _beam_search)rcrfrgrrrrorrrrrTrZkwargs_has_attention_maskroutrencoderrhrrrZ pre_processorZprepared_stopping_criteriar1rr2rqsQ                   zRagTokenForGeneration.generatecCs|||}|Sr~)r)rcrrr1r1r2_temporary_reorder_cacheGs z.RagTokenForGeneration._temporary_reorder_cachecCs |jjSr~)r6rMget_input_embeddingsrr1r1r2r Nsz*RagTokenForGeneration.get_input_embeddingscCs |jjSr~)r6rMget_output_embeddingsrr1r1r2r Qsz+RagTokenForGeneration.get_output_embeddingscCs|jj|Sr~)r6rMset_output_embeddings)rcZnew_embeddingsr1r1r2r Tsz+RagTokenForGeneration.set_output_embeddingscCsX|dur|jj}||j}|ddddf|ddddf<||dddf<|S)zCShift input ids one token to the right, and pad with start_token_idNrxrr)r5rZ new_zerosrclone)rcrfZstart_token_idZshifted_input_idsr1r1r2shift_tokens_rightWs  (z(RagTokenForGeneration.shift_tokens_rightFrcs |dur |njj}tddddfjddjjjgdfdd} |||} d | ksJ|j dd} |j ddd} || | \} } | d} | d} | } | } |r| } | } ||d} d || | | }|S) NrrcsDjjj}|r0||d||d|d|dfSrrrrr1r2rgs   z1RagTokenForGeneration.get_nll.._mask_padsrxrTrr)r5ror,rrrrrMrrrr}rrr)rcrrrrrrorrrrrrrrr1rr2r`s*2   zRagTokenForGeneration.get_nll)NNNN)NNNNNN)N)NNNNNNNNNNNNNNNNN)N)FrN),r(r)r*rr rrr^rrrrr9rMrLrrrrr,r/r-r0rrrrrrrrr rr rrrr r r r rrrr1r1rdr2rOs      -V r)rYr4rr))r+r dataclassesrtypingrrrr,rZ cache_utilsrr Zconfiguration_utilsr Z generationr r r rZmodeling_outputsrZmodeling_utilsrutilsrrZconfiguration_ragrZ retrieval_ragrZ get_loggerr(loggerrr3r4rYrr__all__r1r1r1r2s`        [X 35