rh\dZddlZddlmZddlmZmZmZddlZddlm Z ddl m Z m Z ddl mZdd lmZmZmZmZdd lmZdd lmZdd lmZmZd dlmZd dlmZej>e Z!eedGddeZ"eeGddeZ#edeGddeZ$eGdde$Z%edGdde$Z&edGdd e$eZ'gd!Z(y)"zRAG model implementation.N) dataclass)CallableOptionalUnion)nn)CacheEncoderDecoderCache)PretrainedConfig)GenerationConfigGenerationMixinLogitsProcessorListStoppingCriteriaList) ModelOutput)PreTrainedModel)auto_docstringlogging) RagConfig) RagRetrieverzI Base class for retriever augmented marginalized models outputs. ) custom_introceZdZUdZdZeejed<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<y)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__rrrr rr LongTensorr r!r"r#tupler$r%r&r'r(r)r*w/var/www/html/ai-insurance-compliance-backend/venv/lib/python3.12/site-packages/transformers/models/rag/modeling_rag.pyrr%sEN)-D(5$$ %,*.FHU&& '..2J**+2'+OXe_+8<(5#4#45<48x 0 01848x 0 0189=HU%5%56=FJ&1B1B(CJJNu/@/@#/E)F GNGKXeE,=,=s,B&CDKCG#Xe.?.?%@GKO%0A0A30F*G!HOHLhuU->->-C'DELKO%0A0A30F*G!HOHLhuU->->-C'DELJNu/@/@#/E)F GNr5rceZdZUdZdZeejed<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<y)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. Nrrrrrr r!r".r#r$r%r&r'r(r)r*)r+r,r-r.rrr/r0r1rrr rrr2r r!r"r#r3r$r%r&r'r(r)r*r4r5r6r8r8sCJ+/FHU&& '..2J**+2'+OXe_+8<(5#4#45<48x 0 01848x 0 0189=HU%5%56=FJ&1B1B(CJJNu/@/@#/E)F GNGKXeE,=,=s,B&CDKCG#Xe.?.?%@GKO%0A0A30F*G!HOHLhuU->->-C'DELKO%0A0A30F*G!HOHLhuU->->-C'DELJNu/@/@#/E)F GNr5r8a 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 ZeZdZUeed<dZdZdZe d de e de e de de fd Z y) RagPreTrainedModelconfigragTN.question_encoder_pretrained_model_name_or_path'generator_pretrained_model_name_or_path retrieverreturnc ^|jDcic]%\}}|jdr|tdd|'}}}|jDcic]%\}}|jdr|tdd|'}}}|D]} |d| z= |D]} |d| z= |jdd} | K|Jdddlm} d|vr%dd lm} | j|fi|d d i\} }| |d<| j|fi|} |jdd}|K|Jd dd lm }d|vr%dd lm} | j|fi|d d i\}}||d<|j|fi|}|jd}|+tj| j|jfi|}|| |||Scc}}wcc}}w)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") ```question_encoder_N generator_modelznIf `model` is not defined as an argument, a `question_encoder_pretrained_model_name_or_path` has to be defined AutoModelr;) AutoConfigreturn_unused_kwargsTzqIf `generator_model` is not defined as an argument, a `generator_pretrained_model_name_or_path` has to be definedAutoModelForSeq2SeqLM)question_encoder generatorr;r?)items startswithlenpopauto.modeling_autorGauto.configuration_autorHfrom_pretrainedrKgetr'from_question_encoder_generator_configsr;)clsr=r>r?kwargsargumentvaluekwargs_question_encoderkwargs_generatorkeyrLrGrHquestion_encoder_configrMrKgenerator_configr;s r6*from_pretrained_question_encoder_generatorz=RagPreTrainedModel.from_pretrained_question_encoder_generatorsoL$*<<># %""#67 S,-/ 0% 7# # $*<<> %""<0 S&( )5 0  + 2C*S01 2# +C|c)* + 366wE  #AM  M 766@C]:C]C]BD-D*.D@')@ 5L'18y88> BY  %(($7  :F ! F C//@5OZ5O5O;6?O6fj62 "2.> *=-==7;KI H% >FF '')9)9=CF$4 RXdmnnO#  s *F#*F))NNN)r+r,r-rr1base_model_prefix_supports_flash_attn_supports_sdpa classmethodrstrrrr`r4r5r6r:r:st NIMAE"& Jo8@ Jo2:#Jo Jo  JoJor5r:c"eZdZ ddeedeedeedeeffd Ze ddee jdee jdee e e jd ee jd ee jd eed ee jd ee jdee jdeedeedeedeedeedee e jeffdZxZS)RagModelr;rLrMr?c H| ||Jd|,tj|j|jfi|}n-t||jsJd|d|jt |||!ddlm}|j|j}|!ddlm }|j|j}||_ |j5t|tsJdt|jd ||_ ||_ ||_ d|_d |_y)  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 rErFrJz`self.retriever` is of type z&, but should be of type `RagRetriever`F)rrVr; isinstance config_classsuper__init__rRrG from_configrLrKrMr?rtype ctx_encodercontext_encoder_training) selfr;rLrMr?rXrGrK __class__s r6rmzRagModel.__init__s< !&6&ByG\ _ ] >FF '')9)9=CFfd&7&78 sHVHL_`d`q`q_r:s s8   # 6(44V5L5LM    B-99&:J:JKI" >> %i6 .tDNN/C.DDjk 6'DN 0"(-%r5 input_idsattention_maskencoder_outputsdecoder_input_idsdecoder_attention_maskrrr r! use_cacheoutput_attentionsoutput_hidden_statesoutput_retrievedn_docsr@c ||n|jj}| | n|jj} | | n|jj} | | n|jj} | | n|jj } |j duxr|duxs | duxs|duxr|du}|)|r |j||d}|d}|j ||jjdtjj|jjj|d}|jr|d |d |d |d |d |df\}} }}}}|j|}| j|} |j|}|j|}|j!||dj"}|j%d||j&d}tj(|j+d|j-ddj/d}n|d |d |d |df\}} }}|j|}|j|}| j|} tj(|j+d|j-ddj/d}n|Jd| Jd|Jd|Jd|j&d|zdk(sJd|d|j&dd||j1|d}||j1|d}|j|| ||||| | d }|s d}d}d}d}d}nj2}|j4}|r| sd}d} d}d}t7d)id|j8d|d|j:d |d | d dd d!|d"|d#|j<d$|j>d%|j@d&|jBd'|jDd(|jFS)*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)ru return_dictrcpudevicedtypeptprefixr}return_tensorsr r!rtokenized_doc_idstokenized_doc_attention_maskdoc_idsrrEzMake 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) rtrurvrwrxrryrzrNrrrrr"r#r$r%r&r'r(r)r*r4)$r;r}ryrzr{r|r?rLdetachtor/float32numpyrMrrqrp pooler_outputviewshapebmm unsqueeze transposesqueezerepeat_interleave hidden_states attentionsr8rrencoder_last_hidden_stateencoder_hidden_statesencoder_attentionsdecoder_hidden_statesdecoder_attentionscross_attentions)rrrtrurvrwrxrrr r!ryrzr{r|r}has_to_retrievequestion_enc_outputsr"retriever_outputsrretrieved_doc_input_idsretrieved_doc_attention_maskr gen_outputsr#r$s r6forwardzRagModel.forwardsR"-4;;3E3E!*!6IDKK>0077!#' %3%!00**=>)*BC)*@A)*=>)*HI))4 ).,/4)):(<()*BC)*@A))4 jf%'=?SUf,@+B+BCe+f((9(<(d $D>d 'tnd #4.d  d uU\\"$66 7!d d r5rgzu A RAG-sequence model implementation. It performs RAG-sequence specific marginalization in the forward pass. c&reZdZ d%deedeedeedeeffd ZdefdZdefdZ e d&d ee jd ee jd eeee jd ee jd ee jdeedee jdee jdee j"deedeedeedeedeedeedee jdeedef$dZedZedZedZe j4 d'd ee jd ee jdee jdee jdee j"deed eed!eedeede jfd"Z d(d#Zed$ZxZS))RagSequenceForGenerationr;rLrMr?c | ||Jd|+tj|j|jfi|}t||t |||||_yriNzHEither a configuration or an encoder and a generator has to be provided.)r;rLrMr?rrVr;rlrmrgr<rrr;rLrMr?rXrss r6rmz!RagSequenceForGeneration.__init__s| !&6&ByG\ V ] >FF '')9)9=CF  6>> 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, ... ) ```NFrtrurvrwrxr r!rrryrzr{r|r})repsilonrr}rrrrr r!rrr"r#r$r%r&r'r(r)r*r4)r;r}rrr<get_nllrrlabel_smoothingrrr r!rrr"r#r$r%r&r'r(r)r*)rrrtrurvrwrxrr r!rryrzr{r|rrrr}rXoutputsrs r6rz RagSequenceForGeneration.forwardsN"-4;;3E3E1B1N-TXT_T_TqTq%0%> )) $33  &77  $+#A#A  ")!=!= &77 07/Y/Y (/'I'I %,$C$C -4,S,S )0(K(K &-%E%E )0(K(K &-%E%E! "(/'I'I#  r5c.|jjSrrrrs r6r?z"RagSequenceForGeneration.retrieverexx!!!r5c.|jjSrr<rMrs r6rMz"RagSequenceForGeneration.generatorirr5c.|jjSrr<rLrs r6rLz)RagSequenceForGeneration.question_encodermxx(((r5do_deduplicationnum_return_sequences num_beamsc | | n|jj} ||n|jj}||n|jj} ||n|jj}| |Jd|j ||j ||d} |j || jjdtjj|jjj| dd }|j|}g} || d <|| d <d| d <||jdn|jd| z}t|D]T}||| z|d z| z}|jj |fi| }|rRtj"t%|Dcic]}t'|j)|c}j+}|jd}|$|||d zj-|d }|||d}nq|Jd|Jd|j-|d }||| z|d z| z}|j-|d }|||d zddf}|j-|d }|||||d}|d j/| d }| j1||W|j3| |jjj4Scc}w)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 givenrurrrrrr rrrurT)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.)r r!rrrr) pad_token_id)r;r}rrrr?rLrrr/rrrMrrrangegeneratestacklistretolistvaluesrepeattopkappend _cat_and_padr)rrrtrur r!rrrrr} model_kwargsnum_doc_return_sequencesquestion_hidden_stateshypos batch_sizeindexgenerator_input_idsoutput_sequencesknum_candidates new_input_idsrindividual_input_idsindividual_attention_maskindividual_doc_scores top_cand_indss r6rz!RagSequenceForGeneration.generateqsxB"-4;;3E3E/?/K+QUQ\Q\QmQm$8$D $++JjJj !"+!6IDKK> %*;*C%)%:%:9Uc%:%def%g " $&--/22%u}}2U[[]~~,,33# !/! " !# !2 4 4Y ? $- [!/8 +,)- %&+4+@Y__Q'FWF]F]^_F`djFj :&3 :E"3EFNeaiSYEY"Z 6t~~66#    #(;;tQa4bAS_a5G4b4i4i4k/l#m -33N $ )%%!) < C CNTU V }5EY]^-9T9"-J- (;'A'A"A($-C56>UZ]^U^bhTh,i),E,L,L^]^,_)(25EAI3F3I(J%(=(D(D^UV(W%&:+D4+&* &fo-334LMaPM LL)-8 9g3 :j  T[[5J5J5W5W XXW5cs#!K#ctjddddfjjddj j j jgd||nj j}j jxs j j j}|duxr&dddfj|j} fd} tjj|dj|jd|z|d|j!d} tjj|dj#dj#d} | ddddddddf} | ddddddddf}| ddddddddf}tj| || z|gd}j#dj#dj%d|ddj'|j'k(sJ|j)d}|j+dd }| ||\}}|r| r|ddddddfj+dn|j+d}|j+d}|j-d}|j-d}| }| }|r |j+}|j+}||j!dz }d |z |z||zz}|S) Nrrcjjjj}|j r$|j |d|j |d|j d|j dfSNreqr;rMrany masked_fill_rll smooth_objpad_maskrrtargets r6 _mask_padsz4RagSequenceForGeneration.get_nll.._mask_padsfyy!6!6!C!CDH||~#.''#6::b>:#5#5b#99 9r5rrrErrTrkeepdim?)r/catnewrfill_r;rMrr} bos_token_idrallr functional log_softmaxrsizerrrgathersum logsumexp)rr seq_logitsrrrrrr}r use_bosr seq_logprobs doc_logprobsfirst_token_scoressecond_token_scores remainder rag_logprobsrrnll_loss smooth_losseps_irs` ` r6rz RagSequenceForGeneration.get_nlls AqrE]FJJv||A:@@AVAVAcAcd egh "-4;;3E3E{{//U4;;3H3H3U3U d*Rvad||/L/P/P/R :}}000DII   Q 6 )62zr7J }}000CMMbQ[[\^_ *!QA+6*1a1a<8 Aqr1- yy"46IL6XZc!djkl !!!$..r299!VQJzz||//1111  Rv 6!%%"d%; #B 3J%6'R1ab\  a rvvay^^A& \\!_))!, 3!k ||~H%//+K,++B//g )EK,?? r5c x|djt|Dcgc]}|jdc}t|Dcgc]}|jdc}j |}d}|D]<}|||||jdzd|jdf<||jdz }>|Scc}wcc}w)Nrr)r rrmaxr )tensorsrtoutputinds r6rz%RagSequenceForGeneration._cat_and_padCs AJNN3G^ $D>^ 'tn^ #4.^ $D>^ d^!^ "))*#^ $ %^ ( ")^ ^ @""""))U]]_15598<=A26+/.2#' $TYE,,-TY!!1!12TY$E$4$45 TY !))9)9 : TY U../ TY#4.TY'smTYC=TY TY   TYTYnos9vr5rzo A RAG-token model implementation. It performs RAG-token specific marginalization in the forward pass. c&eZdZ d.deedeedeedeeffd ZdefdZdefd Z d/d Z e d Z e d Z e d ZedZd0dZe d1deej(deej*deeeej.deej(deej0deedeej(deej(deej*deedeedeedeedeedeedeej(d eed!ef$d"Zej<ddddddddee f deej(deej(deej(deej(deej*d eed#ee!d$ee"eej.ge#efd%eed&ee d!ej(fd'Z$d(Z%d)Z&d*Z'd+Z(d0d,Z)d2d-Z*xZ+S)3RagTokenForGenerationNr;rLrMr?c | ||Jd|+tj|j|jfi|}t||t |||||_yrrrs r6rmzRagTokenForGeneration.__init__Us| !&6&ByG\ V ] >FF '')9)9=CF  6RagTokenForGeneration._reorder_cache.._reorder_stackeds"((+yq/AAF.M..r6TMz7RagTokenForGeneration._reorder_cache..s)pWa&z8;;z?P?P3QRps,/)r3rjr from_legacy_cache)rr<reordered_past layer_pastr8s ` @r6_reorder_cachez$RagTokenForGeneration._reorder_caches_ ) J peopp N  o': ;0BB>RNr5c|||n|jj}tjj |dj |j d|z|d|jd}tj|d}||jdjdz}tj|dS)Nrrrr) r;r}rr rrrrr/rr)rrrrr}rr log_prob_sums r6 marginalizez!RagTokenForGeneration.marginalizes!-4;;3E3E}}000DII   Q 6 )62zr7J ((; #l&<&>> 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)rrr}rrrrr r!rrr"r#r$r%r&r'r(r)r*r4)r;r}r.rr<rrrrrDrrr r!rrr"r#r$r%r&r'r(r)r*)rrrtrurvrwrxrr r!rryrzr{r|r.rrr}rXrrrs r6rzRagTokenForGeneration.forwards"^"-4;;3E3E+9+E4;;KeKe%0%.extend_enc_outputsD$M*22J63JV\\Z[Z\M]3]^F]]J 6#BV\\RSRTEU#UVF>>: #9F#B"Dv||TUTVGW"WX Xr5)rrrrvr})rFinput_ids_seq_lengthencoder_input_idsrGrHr)rFrI)assistant_modelrmax_cache_lengthrz)num_return_sequences has to be 1, but is z when doing greedy search.F)rHrIrF synced_gpusstreamerzA`num_return_sequences` has to be smaller or equal to `num_beams`.)rHrIrFrUuH`num_beams` has to be an integer strictly superior to 0 (≥ 1), but is r4r)(rFcopydeepcopyupdaterU_prepare_special_tokensr;r}r?rLrrr/rrrMrrrrrrr< get_encoderfullrdecoder_start_token_idlongnext parametersrr_get_logits_processor_get_stopping_criteria_prepare_cache_for_generation max_lengthr ValueError_sample _beam_search)rrrtrur r!rr}rFrGrHrIrXrkwargs_has_attention_maskroutrencoderrvrQrKrP pre_processorprepared_stopping_criteriars ` @r6rzRagTokenForGeneration.generateqsb  $ $ 6 6  MM*;</(//9&9 $0$4$45Et$LTX$X! $$%68QR"-4;;3E3E >> %*;*C%)%:%:9Uc%:%def%g "..&--/22%u}}2U[[]~~,,33# !C'(,-*+8L5 $8#:#:;Q#R 1 4 4Y ? %;%>%>y%I "#9#C#CA#FH\HfHfghjkHlmuuJ"''*V39 [\b[cd!''*+1 . 9 ',,Q/69 (($$002!,=NdrvwJJ +55 5q 9  4 4**)*11  )r2+,?@ Y"33IUfUpUp!q/@ ):)D)D0 +, 112C2M2MST1U &0 \"*9 &')? %&!' X22/!5/%=-## 3 &*%@%@/CT&A& " **    q).99A=## +   & &! + 559 ?@Q@f@f?gh&& 4<<!."<"3!   ( (1 , 558I8S8SS !dee$4$$!."<"3!    Z[l[v[vZwx r5c*|j||}|Sr)rA)rrrr<s r6_temporary_reorder_cachez.RagTokenForGeneration._temporary_reorder_cacheHs--oxHr5cJ|jjjSr)r<rMget_input_embeddingsrs r6rpz*RagTokenForGeneration.get_input_embeddingsOsxx!!6688r5cJ|jjjSr)r<rMget_output_embeddingsrs r6rrz+RagTokenForGeneration.get_output_embeddingsRsxx!!7799r5cL|jjj|Sr)r<rMset_output_embeddings)rrnew_embeddingss r6rtz+RagTokenForGeneration.set_output_embeddingsUsxx!!77GGr5c||jj}|j|j}|ddddfj |ddddf<||dddf<|S)zCShift input ids one token to the right, and pad with start_token_idNrrr)r;r] new_zerosrclone)rrrtstart_token_idshifted_input_idss r6shift_tokens_rightz(RagTokenForGeneration.shift_tokens_rightXsh  !![[??N%// @#,QV#4#:#:#<!QR% "0!Q$  r5c||njj}tjddddfj j ddj jjjgdfd}j|||}jdj|jk(sJ|jd} |jdd} || | \} } | jd} | jd} | } | } |r | j} | j} ||jdz } d|z | z| | zz}|S) Nrrcjjjj}|j r$|j |d|j |d|j d|j dfSrrrs r6rz1RagTokenForGeneration.get_nll.._mask_padshrr5rrTrr)r;r}r/rr rr rMrrDrrrrr)rrrrrrrr}rrrrrrrrs` ` r6rzRagTokenForGeneration.get_nllasa!-4;;3E3E AqrE]FJJv||A:@@AVAVAcAcd egh  :'' JG !!"%zz||//1111  Rv 6!%%"d%; #B 3J VVAY^^A& 3!k ||~H%//+K,++B//g )EK,?? r5r)NNNNNNrr$)FrN),r+r,r-rr rrrmrrr/r%r?rMrLr'rArDrr/r2r0r3rrr rrrrr&rrr rrrrnrprrrtr{rrrs@r6r)r)Os.26:/3,0 x)*x#?3xO, x L) x<'|'+O+ :""""))* 4156:@D8<=A+/8<=A26$(,0/3+/)-&*-1 $%j E,,-j !!2!23j "%ell(;"<= j $E$4$45 j !))9)9 : j "%j $E$4$45j !))9)9 :j U../j D>j $D>j 'tnj #4.j !j d^!j "))*#j $ %j ( ")j j XU]]_15598<=A26 $8rs] !,, 53ff+-,('   H %  XO{XO  XOv UOUO UOpQoQoQohX !X X v m1m  m`  o.o  od br5