U 9%e^J@sdZddlZddlZddlmZddlmZmZmZddl Z ddl m Z ddl m Z ddl mZdd lmZmZmZmZdd lmZdd lmZmZmZdd lmZmZmZmZd dlm Z e!e"Z#dZ$dZ%dZ&dZ'ddddddddgZ(ddZ)Gddde j*Z+Gddde j*Z,Gddde j*Z-Gd d!d!e j*Z.Gd"d#d#e j*Z/Gd$d%d%e j*Z0Gd&d'd'e j*Z1Gd(d)d)e j*Z2Gd*d+d+e j*Z3eGd,d-d-eZ4eGd.d/d/eZ5eGd0d1d1eZ6eGd2d3d3eZ7Gd4d5d5e j*Z8Gd6d7d7e j*Z9Gd8d9d9e j*Z:Gd:d;d;e j*Z;GdZ=d?Z>Gd@dAdAeZ?GdBdCdCe?Z@edDe=GdEdFdFe?ZAedGe=GdHdIdIe?ZBedJe=GdKdLdLe?ZCedMe=GdNdOdOe?ZDdPZEedQe=GdRdSdSe?ZFdS)Tz PyTorch REALM model.N) dataclass)OptionalTupleUnion)nn)CrossEntropyLoss)ACT2FN))BaseModelOutputWithPastAndCrossAttentions,BaseModelOutputWithPoolingAndCrossAttentionsMaskedLMOutput ModelOutput)PreTrainedModel)apply_chunking_to_forward find_pruneable_heads_and_indicesprune_linear_layer)add_start_docstrings%add_start_docstrings_to_model_forwardloggingreplace_return_docstrings) RealmConfigz(google/realm-cc-news-pretrained-embedderz'google/realm-cc-news-pretrained-encoderz&google/realm-cc-news-pretrained-scorerrz&google/realm-cc-news-pretrained-openqazgoogle/realm-orqa-nq-openqazgoogle/realm-orqa-nq-readerzgoogle/realm-orqa-wq-openqazgoogle/realm-orqa-wq-readerc Cszddl}ddl}ddl}Wn tk r<tdYnXtj|}t d||j |}g}g} |D]@\} } t d| d| |j || } | | | | qrt|| D]\} } t|trd| krt d| d |jjd q| d s| d r4t|tr4| d d} | dd} | d sL| d rdt|trd| d d} | drt|trdnd} | d| d} | d| d} | d| d} | d| d} | d| d} | drlt|trdnd}| d|d} | d|d } | d!|d"} | d#|d$} | d%|d} | d&|d$} nD| d'rt|trdnd}| d(|d } | d)|d"} | d*} td+d,| Drt dd*| q|}| D]}|d-|r|d.|}n|g}|dd/ks0|dd0kr|sz+load_tf_weights_in_realm..z [A-Za-z]+_\d+z_(\d+)ZkernelgammaweightZ output_biasbetabiasriZ _embeddingszPointer shape z and array shape z mismatchedzInitialize PyTorch weight )'renumpyZ tensorflow ImportErrorloggererrorospathabspathinfotrainZlist_variablesZ load_variableappendzip isinstance RealmReader __class____name__ startswithRealmForOpenQAreplaceRealmKnowledgeAugEncoder RealmEmbeddersplitanyjoin fullmatchgetattrAttributeErrorlenint transposeshapeAssertionErrorargstorchZ from_numpydata)modelconfigZtf_checkpoint_pathr&nptfZtf_pathZ init_varsnamesZarraysnamerDarrayZ reader_prefixZembedder_prefixZpointerZm_nameZ scope_namesnumerrrload_tf_weights_in_realm:s      $  $               rRcsTeZdZdZfddZd eejeejeejeeje ej dddZ Z S) RealmEmbeddingszGConstruct the embeddings from word, position and token_type embeddings.csttj|j|j|jd|_t|j|j|_ t|j |j|_ tj |j|j d|_ t|j|_t|dd|_|jdt|jddd|jd tj|jtjd dddS) N) padding_idxZepsposition_embedding_typeabsolute position_ids)rF) persistenttoken_type_idsdtype)super__init__r Embedding vocab_size hidden_sizeZ pad_token_idword_embeddingsmax_position_embeddingsposition_embeddingsZtype_vocab_sizetoken_type_embeddings LayerNormlayer_norm_epsDropouthidden_dropout_probdropoutr?rVregister_bufferrGarangeexpandzerosrXsizelongselfrJr4rrr_s" zRealmEmbeddings.__init__Nr) input_idsr[rX inputs_embedspast_key_values_lengthreturnc Cs|dk r|}n|dd}|d}|dkrL|jdd|||f}|dkrt|dr|jddd|f}||d|} | }ntj|tj|jjd}|dkr| |}| |} || } |j dkr| |} | | 7} | | } || } | S)NrYrr[rr]devicerW)rprXhasattrr[rnrGrorqrzrcrfrVrergrk) rsrur[rXrvrw input_shape seq_lengthbuffered_token_type_ids buffered_token_type_ids_expandedrf embeddingsrerrrforwards,        zRealmEmbeddings.forward)NNNNr) r5 __module__ __qualname____doc__r_rrG LongTensor FloatTensorrBTensorr __classcell__rrrtrrSs rSc seZdZd fdd ZejejdddZd ejeejeejeejeejee e ejee e ejdd d Z Z S) RealmSelfAttentionNcst|j|jdkr>t|ds>td|jd|jd|j|_t|j|j|_|j|j|_t |j|j|_ t |j|j|_ t |j|j|_ t |j|_|pt|dd|_|jdks|jd kr|j|_t d |jd |j|_|j|_dS) NrZembedding_sizezThe hidden size (z6) is not a multiple of the number of attention heads ()rVrW relative_keyrelative_key_queryr%r)r^r_rbnum_attention_headsr{ ValueErrorrBattention_head_size all_head_sizerLinearquerykeyvalueriZattention_probs_dropout_probrkr?rVrdr`distance_embedding is_decoderrsrJrVrtrrr_s* zRealmSelfAttention.__init__)xrxcCs6|dd|j|jf}||}|ddddS)NrYrr%rr)rprrviewpermute)rsrZ new_x_shaperrrtranspose_for_scoress z'RealmSelfAttention.transpose_for_scoresF hidden_statesattention_mask head_maskencoder_hidden_statesencoder_attention_maskpast_key_valueoutput_attentionsrxcCs||}|dk } | r4|dk r4|d} |d} |}n| r^|||} |||} |}nv|dk r|||} |||} tj|d| gdd} tj|d| gdd} n |||} |||} ||} |dk } |jr| | f}t| | dd}|j dks |j dkr| j d| j d}}| r^tj |dtj |j d dd}ntj|tj |j d dd}tj|tj |j d dd}||}|||jd}|j| jd }|j dkrtd | |}||}n4|j dkrtd | |}td | |}|||}|t|j}|dk r:||}tjj|dd}||}|dk rf||}t|| }|dddd }|dd|jf}||}|r||fn|f}|jr||f}|S)Nrrr%dimrYrrryr\zbhld,lrd->bhlrzbhrd,lrd->bhlrr) rrrrrGcatrmatmulrCrVrDtensorrqrzrrmrrdtor]einsummathsqrtrrZ functionalZsoftmaxrkr contiguousrpr)rsrrrrrrrZmixed_query_layerZis_cross_attentionZ key_layerZ value_layerZ query_layer use_cacheZattention_scoresZ query_lengthZ key_lengthZposition_ids_lZposition_ids_rZdistanceZpositional_embeddingZrelative_position_scoresZrelative_position_scores_queryZrelative_position_scores_keyZattention_probsZ context_layerZnew_context_layer_shapeoutputsrrrrsp             zRealmSelfAttention.forward)N)NNNNNF) r5rrr_rGrrrrrboolrrrrrtrrs$rcs4eZdZfddZejejejdddZZS)RealmSelfOutputcsBtt|j|j|_tj|j|jd|_t|j |_ dSNrU) r^r_rrrbdensergrhrirjrkrrrtrrr_os zRealmSelfOutput.__init__r input_tensorrxcCs&||}||}|||}|SNrrkrgrsrrrrrrus  zRealmSelfOutput.forwardr5rrr_rGrrrrrrtrrns rc sveZdZd fdd ZddZd ejeejeejeejeejee e ejee e ejddd Z Z S) RealmAttentionNcs.tt||d|_t||_t|_dS)NrV)r^r_rrsroutputset pruned_headsrrtrrr_~s  zRealmAttention.__init__cCst|dkrdSt||jj|jj|j\}}t|jj||j_t|jj||j_t|jj ||j_ t|j j |dd|j _ |jjt||j_|jj|jj|j_ |j ||_dS)Nrrr)rArrsrrrrrrrrrrunion)rsheadsindexrrr prune_headss zRealmAttention.prune_headsFrc Cs<||||||||}||d|} | f|dd} | S)Nrr)rsr) rsrrrrrrrZ self_outputsattention_outputrrrrrs  zRealmAttention.forward)N)NNNNNF) r5rrr_rrGrrrrrrrrrrtrr}s$rcs0eZdZfddZejejdddZZS)RealmIntermediatecsBtt|j|j|_t|jt r6t |j|_ n|j|_ dSr) r^r_rrrbintermediate_sizerr2 hidden_actstrr intermediate_act_fnrrrtrrr_s   zRealmIntermediate.__init__rrxcCs||}||}|Sr)rrrsrrrrrs  zRealmIntermediate.forwardrrrrtrrs rcs4eZdZfddZejejejdddZZS) RealmOutputcsBtt|j|j|_tj|j|jd|_t |j |_ dSr) r^r_rrrrbrrgrhrirjrkrrrtrrr_s zRealmOutput.__init__rcCs&||}||}|||}|Srrrrrrrs  zRealmOutput.forwardrrrrtrrs rc steZdZfddZd ejeejeejeejeejeeeejee eejdddZ dd Z Z S) RealmLayercsrt|j|_d|_t||_|j|_|j|_|jrZ|jsLt|dt|dd|_ t ||_ t ||_ dS)Nrz> should be used as a decoder model if cross attention is addedrWr)r^r_chunk_size_feed_forward seq_len_dimr attentionradd_cross_attentionrcrossattentionr intermediaterrrrrtrrr_s   zRealmLayer.__init__NFrc Cs|dk r|ddnd}|j|||||d} | d} |jrP| dd} | d} n | dd} d} |jr|dk rt|dstd|d|dk r|d dnd}|| ||||||}|d} | |dd} |d} | | } t|j|j|j| }|f| } |jr| | f} | S) Nr%rrrrrYrz'If `encoder_hidden_states` are passed, z` has to be instantiated with cross-attention layers by setting `config.add_cross_attention=True`r) rrr{rrrfeed_forward_chunkrr)rsrrrrrrrZself_attn_past_key_valueZself_attention_outputsrrZpresent_key_valueZcross_attn_present_key_valueZcross_attn_past_key_valueZcross_attention_outputs layer_outputrrrrsV         zRealmLayer.forwardcCs||}|||}|Sr)rr)rsrZintermediate_outputrrrrrs  zRealmLayer.feed_forward_chunk)NNNNNF) r5rrr_rGrrrrrrrrrrrtrrs$  ArcseZdZfddZd ejeejeejeejeejeeeejee ee ee ee e eeje fd ddZ Z S) RealmEncodercs:t|_tfddtjD|_d|_dS)Ncsg|] }tqSr)r)r_rJrr )sz)RealmEncoder.__init__..F) r^r_rJrZ ModuleListrangenum_hidden_layerslayergradient_checkpointingrrrtrrr_&s  zRealmEncoder.__init__NFT) rrrrrpast_key_valuesrroutput_hidden_states return_dictrxc  st| rdnd} rdnd} r(|jjr(dnd} |jrJ|jrJ|rJtdd}|rRdnd}t|jD]\}}| rv| |f} |dk r||nd}|dk r||nd|jr|jrևfdd}tj j |||||||}n||||||}|d}|r||df7}r`| |df} |jjr`| |d f} q`| r@| |f} | sbt d d ||| | | fDSt ||| | | d S) NrzZ`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...Fcsfdd}|S)Ncs|fSrr)inputs)modulerrrrcustom_forwardOszKRealmEncoder.forward..create_custom_forward..custom_forwardr)rrr)rrcreate_custom_forwardNsz3RealmEncoder.forward..create_custom_forwardrrYrr%css|]}|dk r|VqdSrr)rvrrrr ssz'RealmEncoder.forward..)last_hidden_staterr attentionscross_attentions) rJrrtrainingr)Z warning_once enumeraterrGutils checkpointtupler )rsrrrrrrrrrrZall_hidden_statesZall_self_attentionsZall_cross_attentionsZnext_decoder_cacheiZ layer_moduleZlayer_head_maskrZ layer_outputsrrrr,sv          zRealmEncoder.forward) NNNNNNFFT)r5rrr_rGrrrrrrr rrrrrtrr%s. rcs0eZdZfddZejejdddZZS) RealmPoolercs*tt|j|j|_t|_dSr)r^r_rrrbrZTanh activationrrrtrrr_s zRealmPooler.__init__rcCs(|dddf}||}||}|S)Nr)rr)rsrZfirst_token_tensor pooled_outputrrrrs  zRealmPooler.forwardrrrrtrrs rc@sLeZdZUdZdZejed<dZe e ejed<dZ e e ejed<dS)RealmEmbedderOutputa* Outputs of [`RealmEmbedder`] models. Args: projected_score (`torch.FloatTensor` of shape `(batch_size, config.retriever_proj_size)`): Projected score. 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 + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus the initial embedding outputs. 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 after the attention softmax, used to compute the weighted average in the self-attention heads. Nprojected_scorerr) r5rrrrrGr__annotations__rrrrrrrrrs rc@s<eZdZUdZdZejed<dZejed<dZ ejed<dS)RealmScorerOutputa' Outputs of [`RealmScorer`] models. Args: relevance_score (`torch.FloatTensor` of shape `(batch_size, config.num_candidates)`): The relevance score of document candidates (before softmax). query_score (`torch.FloatTensor` of shape `(batch_size, config.retriever_proj_size)`): Query score derived from the query embedder. candidate_score (`torch.FloatTensor` of shape `(batch_size, config.num_candidates, config.retriever_proj_size)`): Candidate score derived from the embedder. Nrelevance_score query_scorecandidate_score) r5rrrrrGrrrrrrrrrs  rc@seZdZUdZdZejed<dZejed<dZ ejed<dZ ej ed<dZ ej ed<dZ ejed<dZejed <dZejed <dZejed <dZeeejed <dZeeejed <dS)RealmReaderOutputa+ Outputs of [`RealmReader`] models. Args: loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `start_positions`, `end_positions`, `has_answers` are provided): Total loss. retriever_loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `start_positions`, `end_positions`, `has_answers` are provided): Retriever loss. reader_loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `start_positions`, `end_positions`, `has_answers` are provided): Reader loss. retriever_correct (`torch.BoolTensor` of shape `(config.searcher_beam_size,)`, *optional*): Whether or not an evidence block contains answer. reader_correct (`torch.BoolTensor` of shape `(config.reader_beam_size, num_candidates)`, *optional*): Whether or not a span candidate contains answer. block_idx (`torch.LongTensor` of shape `()`): The index of the retrieved evidence block in which the predicted answer is most likely. candidate (`torch.LongTensor` of shape `()`): The index of the retrieved span candidates in which the predicted answer is most likely. start_pos (`torch.IntTensor` of shape `()`): Predicted answer starting position in *RealmReader*'s inputs. end_pos (`torch.IntTensor` of shape `()`): Predicted answer ending position in *RealmReader*'s inputs. 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 + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus the initial embedding outputs. 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 after the attention softmax, used to compute the weighted average in the self-attention heads. Nlossretriever_loss reader_lossretriever_correctreader_correct block_idx candidate start_posend_posrr)r5rrrrrGrrrrr BoolTensorrrrrrZint32rrrrrrrrrrs #rc@s,eZdZUdZdZeed<dZej ed<dS)RealmForOpenQAOutputz Outputs of [`RealmForOpenQA`] models. Args: reader_output (`dict`): Reader output. predicted_answer_ids (`torch.LongTensor` of shape `(answer_sequence_length)`): Predicted answer ids. N reader_outputpredicted_answer_ids) r5rrrrdictrrrGrrrrrrs  rcs$eZdZfddZddZZS)RealmPredictionHeadTransformcsVtt|j|j|_t|jtr6t |j|_ n|j|_ tj |j|j d|_ dSr) r^r_rrrbrr2rrr transform_act_fnrgrhrrrtrrr_ s   z%RealmPredictionHeadTransform.__init__cCs"||}||}||}|Sr)rr rgrrrrrs   z$RealmPredictionHeadTransform.forwardr5rrr_rrrrrtrr s r cs$eZdZfddZddZZS)RealmLMPredictionHeadcsLtt||_tj|j|jdd|_t t |j|_ |j |j_ dS)NF)r$) r^r_r  transformrrrbradecoder ParameterrGror$rrrtrrr_s   zRealmLMPredictionHead.__init__cCs||}||}|Sr)r rrrrrr's  zRealmLMPredictionHead.forwardr rrrtrr s r cs$eZdZfddZddZZS)RealmOnlyMLMHeadcstt||_dSr)r^r_r  predictionsrrrtrrr_.s zRealmOnlyMLMHead.__init__cCs||}|Sr)r)rssequence_outputprediction_scoresrrrr2s zRealmOnlyMLMHead.forwardr rrrtrr-s rcs$eZdZfddZddZZS)RealmScorerProjectioncs>tt||_t|j|j|_tj |j|j d|_ dSr) r^r_r rrrrbretriever_proj_sizerrgrhrrrtrrr_8s  zRealmScorerProjection.__init__cCs||}||}|Sr)rrgrrrrr>s  zRealmScorerProjection.forwardr rrrtrr7s rcs$eZdZfddZddZZS)RealmReaderProjectioncsXt||_t|j|jd|_t|jd|_tj |j|j d|_ t |_ dS)Nr%rrU)r^r_rJrrrbZspan_hidden_sizedense_intermediate dense_outputrgZreader_layer_norm_epslayer_normalizationZReLUrelurrrtrrr_Es  zRealmReaderProjection.__init__csfdd}tjfdd}|}|jddd\}}||\}}} tj|d|d } tj|d|d } | | } | } | } | d} | || | j d 7} | ||fS) Ncsj\}fddtfddtjjD\}}t|d}t|d}tjd|d}tjd|d}||}|||fS)aK Generate span candidates. Args: masks: [num_retrievals, max_sequence_len] Returns: starts: [num_spans] ends: [num_spans] span_masks: [num_retrievals, num_spans] whether spans locate in evidence block. cs6tj|djd}tj|djd}||fS)Nrrz)rGrmrz)widthZcurrent_startsZ current_ends)masksmax_sequence_lenrr_spans_given_width[szRRealmReaderProjection.forward..span_candidates.._spans_given_widthc3s|]}|dVqdS)rNr)rw)rrrr `szIRealmReaderProjection.forward..span_candidates..rrYrr)rDr1rrJmax_span_widthrGr index_select)rrZstartsZendsZ start_masksZ end_masksZ span_masksrs)rrrrspan_candidatesNs "  z6RealmReaderProjection.forward..span_candidatescSsd||t|jSN?typerGZfinfominmaskr]rrr mask_to_scoremsz4RealmReaderProjection.forward..mask_to_scorer%rYrrr!r\) rGfloat32rchunkr#rrrsqueezer])rsr block_maskr%r-Zstart_projectionZend_projectioncandidate_startscandidate_endsZcandidate_maskZcandidate_start_projectionsZcandidate_end_projectionsZcandidate_hidden reader_logitsrr$rrMs    zRealmReaderProjection.forwardr rrrtrrDs raH This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and behavior. Parameters: config ([`RealmConfig`]): Model configuration class with all the parameters of the model. Initializing with a config file does not load the weights associated with the model, only the configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights. a5 Args: input_ids (`torch.LongTensor` of shape `({0})`): Indices of input sequence tokens in the vocabulary. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) attention_mask (`torch.FloatTensor` of shape `({0})`, *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) token_type_ids (`torch.LongTensor` of shape `({0})`, *optional*): Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0, 1]`: - 0 corresponds to a *sentence A* token, - 1 corresponds to a *sentence B* token. [What are token type IDs?](../glossary#token-type-ids) position_ids (`torch.LongTensor` of shape `({0})`, *optional*): Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, config.max_position_embeddings - 1]`. [What are position IDs?](../glossary#position-ids) head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_size)`, *optional*): Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This is useful if you want more control over how to convert *input_ids* indices into associated vectors than the model's internal embedding lookup matrix. output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for more detail. return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. c@s,eZdZdZeZeZdZddZ ddZ dS)RealmPreTrainedModelz An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models. realmcCst|tjr:|jjjd|jjd|jdk r|jj nft|tj rz|jjjd|jjd|j dk r|jj|j  n&t|tj r|jj |jj ddS)zInitialize the weightsg)meanZstdNr')r2rrr"rHZnormal_rJZinitializer_ranger$Zzero_r`rTrgZfill_)rsrrrr _init_weightss      z"RealmPreTrainedModel._init_weightscGsTg}|D]F}|dkr |dq|j}t|dkrD|d|df}||q|S)z.Flatten inputs' shape to (-1, input_shape[-1])Nr%rY)r0rDrAr)rsrZflattened_inputsrr|rrr_flatten_inputss   z$RealmPreTrainedModel._flatten_inputsN) r5rrrr config_classrRZload_tf_weightsZbase_model_prefixr8r9rrrrr5s r5c sDeZdZdZdfdd ZddZddZd d Zdd d ZZ S)RealmBertModelz? Same as the original BertModel but remove docstrings. TcsDt|||_t||_t||_|r2t|nd|_| dSr) r^r_rJrSrrencoderrpooler post_init)rsrJZadd_pooling_layerrtrrr_s    zRealmBertModel.__init__cCs|jjSrrrcr$rrrget_input_embeddingssz#RealmBertModel.get_input_embeddingscCs ||j_dSrr?rsrrrrset_input_embeddingssz#RealmBertModel.set_input_embeddingscCs*|D]\}}|jj|j|qdS)z Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base class PreTrainedModel N)itemsr<rrr)rsZheads_to_prunerrrrr _prune_headsszRealmBertModel._prune_headsNc Cs^| dk r | n|jj} | dk r | n|jj} | dk r4| n|jj} |jjrZ| dk rP| n|jj} nd} |dk rx|dk rxtdn@|dk r||||}n"|dk r|dd}ntd|\}}|dk r|j n|j }| dk r| ddj dnd}|dkrt j |||f|d}|dkrft |jdrT|jjddd|f}|||}|}nt j|t j|d }|||}|jjr|dk r|\}}}||f}|dkrt j ||d}||}nd}|||jj}|j|||||d }|j|||||| | | | | d }|d}|jdk r$||nd}| sB||f|d dSt|||j|j|j|jd S)NFzDYou cannot specify both input_ids and inputs_embeds at the same timerYz5You have to specify either input_ids or inputs_embedsrr%rr[ry)rurXr[rvrw) rrrrrrrrrr)r pooler_outputrrrr)rJrruse_return_dictrrrZ%warn_if_padding_and_no_attention_maskrprzrDrGZonesr{rr[rnrorqZget_extended_attention_maskZinvert_attention_maskZ get_head_maskrr<r=r rrrr)rsrurr[rXrrvrrrrrrrr| batch_sizer}rzrwr~rZextended_attention_maskZencoder_batch_sizeZencoder_sequence_lengthrZencoder_hidden_shapeZencoder_extended_attention_maskZembedding_outputZencoder_outputsrrrrrr s          zRealmBertModel.forward)T) NNNNNNNNNNNNN) r5rrrr_r@rBrDrrrrrtrr;s&  r;z`The embedder of REALM outputting projected score that will be used to calculate relevance score.cseZdZdgZfddZddZddZee de e e d de eje eje eje eje eje eje ee ee eeee fd d d ZZS)r:zcls.predictions.decoder.biascs0t|t|j|_t|j|_|dSr)r^r_r;rJr6rrr>rrrtrrr_s   zRealmEmbedder.__init__cCs |jjjSrr6rrcr$rrrr@sz"RealmEmbedder.get_input_embeddingscCs||jj_dSrrHrArrrrBsz"RealmEmbedder.set_input_embeddingsbatch_size, sequence_length output_typer:N) rurr[rXrrvrrrrxc Csn| dk r | n|jj} |j||||||||| d } | d} || } | sX| f| ddSt| | j| jdSdS)a Returns: Example: ```python >>> from transformers import AutoTokenizer, RealmEmbedder >>> import torch >>> tokenizer = AutoTokenizer.from_pretrained("google/realm-cc-news-pretrained-embedder") >>> model = RealmEmbedder.from_pretrained("google/realm-cc-news-pretrained-embedder") >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt") >>> outputs = model(**inputs) >>> projected_score = outputs.projected_score ``` Nrr[rXrrvrrrrr%)rrr)rJrFr6rrrr) rsrurr[rXrrvrrrZ realm_outputsrErrrrrs*!  zRealmEmbedder.forward) NNNNNNNNN)r5rr_tied_weights_keysr_r@rBrREALM_INPUTS_DOCSTRINGformatrr_CONFIG_FOR_DOCrrGrrrrrrrrrrtrr:{s6    r:zoThe scorer of REALM outputting relevance scores representing the score of document candidates (before softmax).cseZdZdZd fdd Zeedee e dd e e j e e je e j e e j e e j e e je e j e e je e je e je ee ee eeee fddd ZZS) RealmScorerz Args: query_embedder ([`RealmEmbedder`]): Embedder for input sequences. If not specified, it will use the same embedder as candidate sequences. Ncs8t|t|j|_|dk r$|n|j|_|dSr)r^r_r:rJembedderquery_embedderr>)rsrJrTrtrrr_s  zRealmScorer.__init__rIrJ)rurr[rXcandidate_input_idscandidate_attention_maskcandidate_token_type_idscandidate_inputs_embedsrrvrrrrxc Cs| dk r | n|jj} |dkr,| dkr,td|dkrD|dkrDtd|j||||| | | | | d }||||\}}}|j||||| || | | d }|d}|d}|d|jj|jj}t d||}| s|||fSt |||dS) a candidate_input_ids (`torch.LongTensor` of shape `(batch_size, num_candidates, sequence_length)`): Indices of candidate input sequence tokens in the vocabulary. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) candidate_attention_mask (`torch.FloatTensor` of shape `(batch_size, num_candidates, 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) candidate_token_type_ids (`torch.LongTensor` of shape `(batch_size, num_candidates, sequence_length)`, *optional*): Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0, 1]`: - 0 corresponds to a *sentence A* token, - 1 corresponds to a *sentence B* token. [What are token type IDs?](../glossary#token-type-ids) candidate_inputs_embeds (`torch.FloatTensor` of shape `(batch_size * num_candidates, sequence_length, hidden_size)`, *optional*): Optionally, instead of passing `candidate_input_ids` you can choose to directly pass an embedded representation. This is useful if you want more control over how to convert *candidate_input_ids* indices into associated vectors than the model's internal embedding lookup matrix. Returns: Example: ```python >>> import torch >>> from transformers import AutoTokenizer, RealmScorer >>> tokenizer = AutoTokenizer.from_pretrained("google/realm-cc-news-pretrained-scorer") >>> model = RealmScorer.from_pretrained("google/realm-cc-news-pretrained-scorer", num_candidates=2) >>> # batch_size = 2, num_candidates = 2 >>> input_texts = ["How are you?", "What is the item in the picture?"] >>> candidates_texts = [["Hello world!", "Nice to meet you!"], ["A cute cat.", "An adorable dog."]] >>> inputs = tokenizer(input_texts, return_tensors="pt") >>> candidates_inputs = tokenizer.batch_encode_candidates(candidates_texts, max_length=10, return_tensors="pt") >>> outputs = model( ... **inputs, ... candidate_input_ids=candidates_inputs.input_ids, ... candidate_attention_mask=candidates_inputs.attention_mask, ... candidate_token_type_ids=candidates_inputs.token_type_ids, ... ) >>> relevance_score = outputs.relevance_score ```Nz5You have to specify either input_ids or input_embeds.zJYou have to specify either candidate_input_ids or candidate_inputs_embeds.rLrrYz bd,bnd->bn)rrr) rJrFrrTr9rSrnum_candidatesrrGrr)rsrurr[rXrUrVrWrXrrvrrrZ query_outputsflattened_input_idsflattened_attention_maskflattened_token_type_idsZcandidate_outputsrrrrrrrsVI    zRealmScorer.forward)N) NNNNNNNNNNNNN)r5rrrr_rrOrPrrrQrrGrrrrrrrrrrtrrRsB   rRzrThe knowledge-augmented encoder of REALM outputting masked language model logits and marginal log-likelihood loss.cseZdZdgZfddZddZddZdd Zd d Ze e d e e ed deejeejeejeejeejeejeejeejeejeeeeeeeee fd ddZZS)r9zcls.predictions.decodercs0t|t|j|_t|j|_|dSr)r^r_r;rJr6rrr>rrrtrrr_hs   z!RealmKnowledgeAugEncoder.__init__cCs |jjjSrrHr$rrrr@nsz-RealmKnowledgeAugEncoder.get_input_embeddingscCs||jj_dSrrHrArrrrBqsz-RealmKnowledgeAugEncoder.set_input_embeddingscCs |jjjSrrrrr$rrrget_output_embeddingstsz.RealmKnowledgeAugEncoder.get_output_embeddingscCs||jj_dSrr])rsZnew_embeddingsrrrset_output_embeddingswsz.RealmKnowledgeAugEncoder.set_output_embeddingsz+batch_size, num_candidates, sequence_lengthrJN) rurr[rXrrvrlabelsmlm_maskrrrrxc  Csz| dk r | n|jj} ||||\} }}|j| |||||| | | d }|d}||}|}d}|dk r6|dkrxtd|\}}| dkrtj|tj d} n | tj } t dd}| d|jj }|d |jj d}||| ||jj| }|dd}||}|d }tt|| t|  }| sf|f|d d }|dk rb|f|S|St|||j|jd S) a relevance_score (`torch.FloatTensor` of shape `(batch_size, num_candidates)`, *optional*): Relevance score derived from RealmScorer, must be specified if you want to compute the masked language modeling loss. labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]` mlm_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid calculating joint loss on certain positions. If not specified, the loss will not be masked. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. Returns: Example: ```python >>> import torch >>> from transformers import AutoTokenizer, RealmKnowledgeAugEncoder >>> tokenizer = AutoTokenizer.from_pretrained("google/realm-cc-news-pretrained-encoder") >>> model = RealmKnowledgeAugEncoder.from_pretrained( ... "google/realm-cc-news-pretrained-encoder", num_candidates=2 ... ) >>> # batch_size = 2, num_candidates = 2 >>> text = [["Hello world!", "Nice to meet you!"], ["The cute cat.", "The adorable dog."]] >>> inputs = tokenizer.batch_encode_candidates(text, max_length=10, return_tensors="pt") >>> outputs = model(**inputs) >>> logits = outputs.logits ```NrLrzZYou have to specify `relevance_score` when `labels` is specified in order to compute loss.r\none)Z reductionrYrr%rM)rlogitsrr)rJrFr9r6rrrprGZ ones_liker.r)rrraZtilerYZ log_softmax unsqueeze logsumexpZnansumsumr rr)rsrurr[rXrrvrr`rarrrrZr[r\Z joint_outputsZ joint_outputrrZmasked_lm_lossrGr}Zloss_fctZ mlm_logitsZ mlm_targetsZmasked_lm_log_probZcandidate_log_probZjoint_gold_log_probZmarginal_gold_log_probsrrrrrzsf9          z RealmKnowledgeAugEncoder.forward) NNNNNNNNNNNN)r5rrrNr_r@rBr^r_rrOrPrr rQrrGrrrrrrrrrrtrr9`sJ   r9zThe reader of REALM.cseZdZfddZeedeee dd e e j e e j e e j e e j e e j e e j e e j e e je e j e e j e e je ee ee eeeefdddZZS) r3cs>t||j|_t||_t||_t||_| dSr) r^r_Z num_labelsr;r6rrr qa_outputsr>rrrtrrr_s     zRealmReader.__init__z!reader_beam_size, sequence_lengthrJN)rurr[rXrrvrr1start_positions end_positions has_answersrrrrxc$ CsL|dk r |n|jj}|dkr$td|dkr4td|d|jjkrNtd|j||||||| | |d }|d}|||d|jj\}}}t |d|jjd}||7}t tj |dd j }t tj |dd j }tj |d|d }tj |d|d }d}d}d}d}d}| dk r| dk r| dk rd d }d d}|d} | d| } | d| } | }t|}!|||| d|jj| d|jjd}t|}"|||}||d|d}||!tj9}||"tj9}||}|s*||||f|dd}#|dk r&|||||f|#S|#St||||||||||j|jd S)ar relevance_score (`torch.FloatTensor` of shape `(searcher_beam_size,)`, *optional*): Relevance score, which must be specified if you want to compute the logits and marginal log loss. block_mask (`torch.BoolTensor` of shape `(searcher_beam_size, sequence_length)`, *optional*): The mask of the evidence block, which must be specified if you want to compute the logits and marginal log loss. start_positions (`torch.LongTensor` of shape `(searcher_beam_size,)`, *optional*): Labels for position (index) of the start of the labelled span for computing the token classification loss. Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence are not taken into account for computing the loss. end_positions (`torch.LongTensor` of shape `(searcher_beam_size,)`, *optional*): Labels for position (index) of the end of the labelled span for computing the token classification loss. Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence are not taken into account for computing the loss. has_answers (`torch.BoolTensor` of shape `(searcher_beam_size,)`, *optional*): Whether or not the evidence block has answer(s). Returns: NzCYou have to specify `relevance_score` to calculate logits and loss.zOYou have to specify `block_mask` to separate question block and evidence block.rzQThe input sequence length must be greater than or equal to config.max_span_width.rLrrYrr!cSs\ttt|ddt|d}ttt|ddt|d}tt||dS)zCompute correct span.rrYr)rGeqrdr< logical_and)r2r3 gold_starts gold_endsZ is_gold_startZ is_gold_endrrrcompute_correct_candidates[s  z7RealmReader.forward..compute_correct_candidatescSs@tjfdd}tj||||jddd}tj|dd}||S)z3Loss based on the negative marginal log-likelihood.cSsd||t|jSr&r(r+rrrr-kszERealmReader.forward..marginal_log_loss..mask_to_scorer\rYr)rGr.rer])rcZ is_correctr-Z log_numeratorZlog_denominatorrrrmarginal_log_losshsz.RealmReader.forward..marginal_log_loss)r2r3rmrnr%) rrrrrrrrrrr)rJrFrrpr"r6rgreader_beam_sizerGrdZargmaxmaxvaluesr#clampr<rr)r.r7rrr)$rsrurr[rXrrvrr1rhrirjrrrrrr4r2r3Zretriever_logitsZpredicted_block_indexZpredicted_candidateZpredicted_startZ predicted_endZ total_lossrrrrrorpZ ignored_indexZany_retriever_correctZany_reader_correctrrrrrs&          zRealmReader.forward)NNNNNNNNNNNNNN)r5rrr_rrOrPrrrQrrGrrrrrrrrrrrtrr3sD   r3ay Args: input_ids (`torch.LongTensor` of shape `({0})`): Indices of input sequence tokens in the vocabulary. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) attention_mask (`torch.FloatTensor` of shape `({0})`, *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) token_type_ids (`torch.LongTensor` of shape `({0})`, *optional*): Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0, 1]`: - 0 corresponds to a *sentence A* token, - 1 corresponds to a *sentence B* token (should not be used in this model by design). [What are token type IDs?](../glossary#token-type-ids) answer_ids (`list` of shape `(num_answers, answer_length)`, *optional*): Answer ids for computing the marginal log-likelihood loss. Indices should be in `[-1, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-1` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]` return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. z?`RealmForOpenQA` for end-to-end open domain question answering.c seZdZd fdd ZeddZddZee de e e d de eje eje eje eje eeee fd d d ZZS)r7Nc s`t|t||_t||_|dtdj |j |j ftj t dd||_|dS)N block_embrcpu)rpr]rz)r^r_r:rSr3rrlrGroZ new_emptyZnum_block_recordsrr.rz retrieverr>)rsrJrwrtrrr_s     zRealmForOpenQA.__init__cCs|jr|jjS|jjSr)rrJsearcher_beam_sizerqr$rrrrxsz!RealmForOpenQA.searcher_beam_sizecCs|j||_dS)zSend `self.block_emb` to a specific device. Args: device (`str` or `torch.device`): The device to which `self.block_emb` will be sent. N)rur)rsrzrrrblock_embedding_tosz!RealmForOpenQA.block_embedding_toz1, sequence_lengthrJ)rurr[ answer_idsrrxc Cs|dk r |n|jj}|dk r2|jddkr2td|j|||dd}|d}td|j||jj }tj ||j dd \} } | } tj |jd| d } |j| |||jjd \} } }}||jj }|jtjj|jj d }||jtj| dk rDtj| tj|jj d } tj| tj|jj d } tj|tj|jj d }td| | |jj }|j|jd|jj|jd|jj|jd|jj||| | |dd }|j|j}||j|jd}|s||fSt ||dS)a Returns: Example: ```python >>> import torch >>> from transformers import RealmForOpenQA, RealmRetriever, AutoTokenizer >>> retriever = RealmRetriever.from_pretrained("google/realm-orqa-nq-openqa") >>> tokenizer = AutoTokenizer.from_pretrained("google/realm-orqa-nq-openqa") >>> model = RealmForOpenQA.from_pretrained("google/realm-orqa-nq-openqa", retriever=retriever) >>> question = "Who is the pioneer in modern computer science?" >>> question_ids = tokenizer([question], return_tensors="pt") >>> answer_ids = tokenizer( ... ["alan mathison turing"], ... add_special_tokens=False, ... return_token_type_ids=False, ... return_attention_mask=False, ... ).input_ids >>> reader_output, predicted_answer_ids = model(**question_ids, answer_ids=answer_ids, return_dict=False) >>> predicted_answer = tokenizer.decode(predicted_answer_ids) >>> loss = reader_output.loss ```Nrrz'The batch_size of the inputs must be 1.T)rur[rrz BD,QD->QBrY)krr!) max_lengthrryzD,BD->B) rurr[rr1rjrhrir)rr)!rJrFrDrrSrGrrurrzZtopkrxr0r#rwrvZreader_seq_lenrZspecial_tokens_maskr)rZ logical_not_Z logical_and_r[rrqrurqrrrrr)rsrurr[rzrZquestion_outputsZquestion_projectionZ batch_scoresrZretrieved_block_idsZretrieved_block_embrjrrZ concat_inputsr1Zretrieved_logitsrZpredicted_blockrrrrrsf%   zRealmForOpenQA.forward)N)NNNN)r5rrr_propertyrxryrREALM_FOR_OPEN_QA_DOCSTRINGrPrrrQrrGrrrrrrrrrrtrr7s$    r7)Grrr+ dataclassesrtypingrrrrGrZtorch.nnrZ activationsr Zmodeling_outputsr r r r Zmodeling_utilsrZ pytorch_utilsrrrrrrrrZconfiguration_realmrZ get_loggerr5r)Z_EMBEDDER_CHECKPOINT_FOR_DOCZ_ENCODER_CHECKPOINT_FOR_DOCZ_SCORER_CHECKPOINT_FOR_DOCrQZ#REALM_PRETRAINED_MODEL_ARCHIVE_LISTrRModulerSrrrrrrrrrrrrr r rrrZREALM_START_DOCSTRINGrOr5r;r:rRr9r3r~r7rrrrs        lA2Wc1  C 2(N(!