o hI@sdZddlZddlZddlmZddlmZmZmZddl Z ddl m Z ddl m Z ddl mZdd lmZmZmZmZdd lmZdd lmZmZmZdd lmZmZmZmZd dlm Z e!e"Z#dZ$dZ%dZ&dZ'gdZ(ddZ)Gddde j*Z+Gddde j*Z,Gddde j*Z-Gddde j*Z.Gddde 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+eZ5eGd,d-d-eZ6eGd.d/d/eZ7Gd0d1d1e j*Z8Gd2d3d3e j*Z9Gd4d5d5e j*Z:Gd6d7d7e j*Z;Gd8d9d9e j*ZGdd?d?e?Z@ed@e=GdAdBdBe?ZAedCe=GdDdEdEe?ZBedFe=GdGdHdHe?ZCedIe=GdJdKdKe?ZDdLZEedMe=GdNdOdOe?ZFdS)Pz 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) RealmConfig(google/realm-cc-news-pretrained-embedder'google/realm-cc-news-pretrained-encoder&google/realm-cc-news-pretrained-scorerr)rrrz&google/realm-cc-news-pretrained-openqazgoogle/realm-orqa-nq-openqazgoogle/realm-orqa-nq-readerzgoogle/realm-orqa-wq-openqazgoogle/realm-orqa-wq-readerc Csfzddl}ddl}ddl}Wn tytdwtj|}t d||j |}g}g} |D] \} } t d| d| |j || } | | | | q6t|| D]\} } t|tryd| vryt d| d |jjd q\| d s| d rt|tr| d d} | dd} | d s| d rt|tr| d d} | drt|trdnd} | d| d} | d| d} | d| d} | d| d} | d| d} | dr*t|trdnd}| d|d} | d|d } | d!|d"} | d#|d$} | d%|d} | d&|d$} n"| d'rLt|tr8dnd}| d(|d } | d)|d"} | d*} td+d,| Drgt dd*| q\|}| D]m}|d-|r{|d.|}n|g}|dd/ks|dd0krt|d1}n4|dd2ks|dd3krt|d4}n z t||d}Wntyt dd*| Yqkwt|d5krt|d6}||}qk|d7dd8krt|d1}n |d/kr| | } z|j!| j!ksJd9|j!d:| j!d;Wnt"y!}z |j#|j!| j!f7_#d}~wwt d<| t$%| |_&q\|S)=z'Load tf checkpoints in a pytorch model.rNzLoading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see https://www.tensorflow.org/install/ for installation instructions.z&Converting TensorFlow checkpoint from zLoading TF weight z with shape readerz Skipping z as it is not z 's parameterbertclszbert/z reader/realm/zcls/z reader/cls/zrealm/zreader/zreader/module/bert/zreader/module/cls/z reader/dense/zqa_outputs/dense_intermediate/zreader/dense_1/zqa_outputs/dense_output/zreader/layer_normalizationzqa_outputs/layer_normalizationzmodule/module/module/z embedder/z!module/module/module/module/bert/zmodule/module/module/LayerNorm/zcls/LayerNorm/zmodule/module/module/dense/z cls/dense/z,module/module/module/module/cls/predictions/zcls/predictions/zmodule/module/module/bert/z%module/module/module/cls/predictions/zmodule/module/zmodule/module/LayerNorm/zmodule/module/dense//css|]}|dvVqdS))adam_vadam_mAdamWeightDecayOptimizerAdamWeightDecayOptimizer_1 global_stepN).0nr%r%^/var/www/html/ai/venv/lib/python3.10/site-packages/transformers/models/realm/modeling_realm.py |s  z+load_tf_weights_in_realm..z [A-Za-z]+_\d+z_(\d+)kernelgammaweight output_biasbetabiasri _embeddingszPointer shape z and array shape z mismatchedzInitialize PyTorch weight )'renumpy tensorflow ImportErrorloggererrorospathabspathinfotrainlist_variables load_variableappendzip isinstance RealmReader __class____name__ startswithRealmForOpenQAreplaceRealmKnowledgeAugEncoder RealmEmbeddersplitanyjoin fullmatchgetattrAttributeErrorlenint transposeshapeAssertionErrorargstorch from_numpydata)modelconfigtf_checkpoint_pathr2nptftf_path init_varsnamesarraysnamerSarray reader_prefixembedder_prefixpointerm_name scope_namesnumer%r%r(load_tf_weights_in_realm:s                     rkc sheZdZdZfddZ     ddeejdeejdeejd eejd e d ej f d d Z 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_idxepsposition_embedding_typeabsolute position_ids)rF) persistenttoken_type_idsdtype)super__init__r Embedding vocab_size hidden_size pad_token_idword_embeddingsmax_position_embeddingsposition_embeddingstype_vocab_sizetoken_type_embeddings LayerNormlayer_norm_epsDropouthidden_dropout_probdropoutrNrpregister_bufferrVarangeexpandzerosrrsizelongselfrZrCr%r(rys  zRealmEmbeddings.__init__Nr input_idsrurr inputs_embedspast_key_values_lengthreturnc Cs|dur |}n|dd}|d}|dur&|jdd|||f}|durPt|drE|jddd|f}||d|} | }n tj|tj|jjd}|durY| |}| |} || } |j dkrp| |} | | 7} | | } || } | S)Nrsrrurrwdevicerq)rrrhasattrrurrVrrrr~rrprrr) rrrurrrr input_shape seq_lengthbuffered_token_type_ids buffered_token_type_ids_expandedr embeddingsrr%r%r(forwards,        zRealmEmbeddings.forward)NNNNr) rD __module__ __qualname____doc__ryrrV LongTensor FloatTensorrQTensorr __classcell__r%r%rr(rls* rlcseZdZdfdd ZdejdejfddZ      dd ejd eejd eejd eejd eejdee e ejdee de ejfddZ Z S)RealmSelfAttentionNcst|j|jdkrt|dstd|jd|jd|j|_t|j|j|_|j|j|_t |j|j|_ t |j|j|_ t |j|j|_ t |j|_|p\t|dd|_|jdksh|jd kry|j|_t d |jd |j|_|j|_dS) Nrembedding_sizezThe hidden size (z6) is not a multiple of the number of attention heads ()rprq relative_keyrelative_key_queryr0r)rxryr|num_attention_headsr ValueErrorrQattention_head_size all_head_sizerLinearquerykeyvaluerattention_probs_dropout_probrrNrprrzdistance_embedding is_decoderrrZrprr%r(rys*   zRealmSelfAttention.__init__xrcCs6|dd|j|jf}||}|ddddS)Nrsrr0rr)rrrviewpermute)rr new_x_shaper%r%r(transpose_for_scoress z'RealmSelfAttention.transpose_for_scoresF hidden_statesattention_mask head_maskencoder_hidden_statesencoder_attention_maskpast_key_valueoutput_attentionscCs||}|du} | r|dur|d} |d} |}nP| r/|||} |||} |}n;|durZ|||} |||} tj|d| gdd} tj|d| gdd} n|||} |||} ||} |du} |jrz| | f}t| | dd}|j dks|j dkr | j d| j d}}| rtj |dtj |j d dd}ntj|tj |j d dd}tj|tj |j d dd}||}|||jd}|j| jd }|j dkrtd | |}||}n|j dkr td | |}td | |}|||}|t|j}|dur||}tjj|dd}||}|dur0||}t|| }|dddd }|dd|jf}||}|rX||fn|f}|jrd||f}|S)Nrrr0dimrsrrrrvzbhld,lrd->bhlrzbhrd,lrd->bhlrr) rrrrrVcatrmatmulrRrprStensorrrrrrrtorweinsummathsqrtrr functionalsoftmaxrr contiguousrr)rrrrrrrrmixed_query_layeris_cross_attention key_layer value_layer query_layer use_cacheattention_scores query_length key_lengthposition_ids_lposition_ids_rdistancepositional_embeddingrelative_position_scoresrelative_position_scores_queryrelative_position_scores_keyattention_probs context_layernew_context_layer_shapeoutputsr%r%r(rsn             zRealmSelfAttention.forwardNNNNNNF) rDrrryrVrrrrrboolrrr%r%rr(rs4 rc8eZdZfddZdejdejdejfddZZS)RealmSelfOutputcsBtt|j|j|_tj|j|jd|_t|j |_ dSNrn) rxryrrr|denserrrrrrrr%r(ryo zRealmSelfOutput.__init__r input_tensorrcC&||}||}|||}|Srrrrrrrr%r%r(ru  zRealmSelfOutput.forwardrDrrryrVrrrr%r%rr(rn $rcseZdZdfdd ZddZ      ddejdeejd eejd eejd eejd ee e ejd ee de ejfddZ Z S)RealmAttentionNcs.tt||d|_t||_t|_dS)Nrp)rxryrrroutputset pruned_headsrrr%r(ry~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)rPrrrrrrrrrrrrunion)rheadsindexr%r%r( prune_headss zRealmAttention.prune_headsFrrrrrrrrc Cs<||||||||}||d|} | f|dd} | S)Nrr)rr) rrrrrrrr self_outputsattention_outputrr%r%r(rs  zRealmAttention.forwardrr) rDrrryrrVrrrrrrrr%r%rr(r}s4 rc2eZdZfddZdejdejfddZZS)RealmIntermediatecsDtt|j|j|_t|jt rt |j|_ dS|j|_ dSr) rxryrrr|intermediate_sizerrA hidden_actstrr intermediate_act_fnrrr%r(rys   zRealmIntermediate.__init__rrcC||}||}|Sr)rrrrr%r%r(r  zRealmIntermediate.forwardrr%r%rr(rs rcr) RealmOutputcsBtt|j|j|_tj|j|jd|_t |j |_ dSr) rxryrrrr|rrrrrrrrr%r(ryrzRealmOutput.__init__rrrcCrrrrr%r%r(rrzRealmOutput.forwardrr%r%rr(rrrcseZdZfddZ      ddejdeejdeejdeejd eejd eeeejd ee d eejfd dZ ddZ Z S) RealmLayercsrt|j|_d|_t||_|j|_|j|_|jr-|js&t|dt|dd|_ t ||_ t ||_ dS)Nrz> should be used as a decoder model if cross attention is addedrqr)rxrychunk_size_feed_forward seq_len_dimr attentionradd_cross_attentionrcrossattentionr intermediaterrrrr%r(rys   zRealmLayer.__init__NFrrrrrrrrc Cs|dur |ddnd}|j|||||d} | d} |jr(| dd} | d} n| dd} d} |jro|durot|dsDtd|d|durN|d dnd}|| ||||||}|d} | |dd} |d} | | } t|j|j|j| }|f| } |jr| | f} | S) Nr0)rrrrrsr z'If `encoder_hidden_states` are passed, z` has to be instantiated with cross-attention layers by setting `config.add_cross_attention=True`r) r rrrr rfeed_forward_chunkrr )rrrrrrrrself_attn_past_key_valueself_attention_outputsrrpresent_key_valuecross_attn_present_key_valuecross_attn_past_key_valuecross_attention_outputs layer_outputr%r%r(rsP         zRealmLayer.forwardcCs||}|||}|Sr)r r)rrintermediate_outputrr%r%r(rs  zRealmLayer.feed_forward_chunkr) rDrrryrVrrrrrrrrr%r%rr(rs4  ArcseZdZfddZ         ddejdeejdeejd eejd eejd eeeejd ee d ee dee dee de eeje ffddZ Z S) RealmEncodercs:t|_tfddtjD|_d|_dS)Ncsg|]}tqSr%)r)r&_rZr%r( )sz)RealmEncoder.__init__..F) rxryrZr ModuleListrangenum_hidden_layerslayergradient_checkpointingrrrr(ry&s   zRealmEncoder.__init__NFTrrrrrpast_key_valuesrroutput_hidden_states return_dictrc  Cs^| rdnd} |r dnd} |r|jjrdnd} |jr%|jr%|r%tdd}|r)dnd}t|jD]^\}}| r;| |f} |durC||nd}|durM||nd}|jrc|jrc||j |||||||}n ||||||||}|d}|rz||df7}|r| |df} |jjr| |df} q0| r| |f} | st dd ||| | | fDSt ||| | | d S) Nr%zZ`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...Frrsrr0css|] }|dur|VqdSrr%)r&vr%r%r(r)nsz'RealmEncoder.forward..)last_hidden_stater r attentionscross_attentions) rZr rtrainingr6 warning_once enumerater_gradient_checkpointing_func__call__tupler )rrrrrrr rrr!r"all_hidden_statesall_self_attentionsall_cross_attentionsnext_decoder_cachei layer_modulelayer_head_maskr layer_outputsr%r%r(r,sz         zRealmEncoder.forward) NNNNNNFFT)rDrrryrVrrrrrrr rrr%r%rr(r%sD      rcr) RealmPoolercs*tt|j|j|_t|_dSr)rxryrrr|rTanh activationrrr%r(rys zRealmPooler.__init__rrcCs(|dddf}||}||}|S)Nr)rr7)rrfirst_token_tensor pooled_outputr%r%r(rs  zRealmPooler.forwardrr%r%rr(r5s r5c@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%) rDrrrr;rVr__annotations__rrrr%r%r%r%r(r:s r:c@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) rDrrrr>rVrr<r?r@r%r%r%r(r=s  r=c@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%)rDrrrrBrVrr<rCrDrE BoolTensorrFrGrrHrIint32rJrrrr%r%r%r%r(rAs #rAc@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) rDrrrrNdictr<rOrVrr%r%r%r(rMs  rMc$eZdZfddZddZZS)RealmPredictionHeadTransformcsVtt|j|j|_t|jtrt |j|_ n|j|_ tj |j|j d|_ dSr) rxryrrr|rrArrr transform_act_fnrrrrr%r(rys  z%RealmPredictionHeadTransform.__init__cCs"||}||}||}|Sr)rrSrrr%r%r(r s   z$RealmPredictionHeadTransform.forwardrDrrryrrr%r%rr(rRs  rRcrQ)RealmLMPredictionHeadcsLtt||_tj|j|jdd|_t t |j|_ |j |j_ dS)NF)r/) rxryrR transformrrr|r{decoder ParameterrVrr/rrr%r(rys  zRealmLMPredictionHead.__init__cCrr)rVrWrr%r%r(r"rzRealmLMPredictionHead.forwardrTr%r%rr(rUs  rUcrQ)RealmOnlyMLMHeadcstt||_dSr)rxryrU predictionsrrr%r(ry)s zRealmOnlyMLMHead.__init__cCs||}|Sr)rZ)rsequence_outputprediction_scoresr%r%r(r-s zRealmOnlyMLMHead.forwardrTr%r%rr(rY(s rYcrQ)RealmScorerProjectioncs>tt||_t|j|j|_tj |j|j d|_ dSr) rxryrUrZrrr|retriever_proj_sizerrrrrr%r(ry3s  zRealmScorerProjection.__init__cCrr)rrrr%r%r(r9rzRealmScorerProjection.forwardrTr%r%rr(r]2s r]crQ)RealmReaderProjectioncsXt||_t|j|jd|_t|jd|_tj |j|j d|_ t |_ dS)Nr0rrn)rxryrZrrr|span_hidden_sizedense_intermediate dense_outputrreader_layer_norm_epslayer_normalizationReLUrelurrr%r(ry@s 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)Nrr)rVrr)widthcurrent_starts current_ends)masksmax_sequence_lenr%r(_spans_given_widthVszRRealmReaderProjection.forward..span_candidates.._spans_given_widthc3s|] }|dVqdS)rNr%)r&w)rmr%r(r)[szIRealmReaderProjection.forward..span_candidates..rrsrr)rSr@rrZmax_span_widthrVr index_select)rkrstartsends start_masks end_masks span_masksr)rmrkrlr(span_candidatesIs "   z6RealmReaderProjection.forward..span_candidatescSd||t|jSN?typerVfinfominmaskrwr%r%r( mask_to_scorehz4RealmReaderProjection.forward..mask_to_scorer0rsrrrorv) rVfloat32rachunkrqrfrdrbsqueezerw)rr block_maskrxrstart_projectionend_projectioncandidate_startscandidate_endscandidate_maskcandidate_start_projectionscandidate_end_projectionscandidate_hidden reader_logitsr%rwr(rHs     zRealmReaderProjection.forwardrTr%r%rr(r_?s r_aH 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|jdur|jj dSdSt|tj rC|jjjd|jjd|j durA|jj|j  dSdSt|tj rX|jj |jj ddSdS)zInitialize the weightsg)meanstdNr{)rArrr,rXnormal_rZinitializer_ranger/zero_rzrmrfill_)rmoduler%r%r( _init_weightss      z"RealmPreTrainedModel._init_weightscGsTg}|D]#}|dur|dq|j}t|dkr"|d|df}||q|S)z.Flatten inputs' shape to (-1, input_shape[-1])Nr0rs)r?rSrPr)rinputsflattened_inputsrrr%r%r(_flatten_inputss   z$RealmPreTrainedModel._flatten_inputsN) rDrrrr config_classrkload_tf_weightsbase_model_prefixrrr%r%r%r(rs rcs^eZdZdZdfdd ZddZddZd d Z dd d ZZ S)RealmBertModelz? Same as the original BertModel but remove docstrings. TcsDt|||_t||_t||_|rt|nd|_| dSr) rxryrZrlrrencoderr5pooler post_init)rrZadd_pooling_layerrr%r(rys    zRealmBertModel.__init__cCs|jjSrrr~rwr%r%r(get_input_embeddingssz#RealmBertModel.get_input_embeddingscCs ||j_dSrrrrr%r%r(set_input_embeddingss z#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)itemsrrr r)rheads_to_prunerrr%r%r( _prune_headsszRealmBertModel._prune_headsNc CsP| dur| n|jj} | dur| n|jj} | dur| n|jj} |jjr-| dur(| n|jj} nd} |dur;|dur;td|durJ||||}n|durW|dd}ntd|\}}|durf|j n|j }| durv| ddj dnd}|durt j |||f|d}|durt |jdr|jjddd|f}|||}|}n t j|t j|d }|||}|jjr|dur|\}}}||f}|durt j ||d}||}nd}|||jj}|j|||||d }|j|||||| | | | | d }|d}|jdur ||nd}| s||f|d dSt|||j|j|j|jd S)NFzDYou cannot specify both input_ids and inputs_embeds at the same timersz5You have to specify either input_ids or inputs_embedsrr0rgrur)rrrrurr) rrrrr rrr!r"r)r$ pooler_outputr rr%r&)rZrr!use_return_dictrrr%warn_if_padding_and_no_attention_maskrrrSrVonesrrrurrrget_extended_attention_maskinvert_attention_mask get_head_maskrrrr r rr%r&)rrrrurrrrrrr rrr!r"r batch_sizerrrrrextended_attention_maskencoder_batch_sizeencoder_sequence_lengthrencoder_hidden_shapeencoder_extended_attention_maskembedding_outputencoder_outputsr[r9r%r%r(rs       zRealmBertModel.forward)T NNNNNNNNNNNNN) rDrrrryrrrrrr%r%rr(rs(  rz`The embedder of REALM outputting projected score that will be used to calculate relevance score.cseZdZdgZfddZddZddZee de e e d  dd e ejd e ejd e ejde ejde ejde ejde ede ede edeee ffddZZS)rIzcls.predictions.decoder.biasc0t|t|j|_t|j|_|dSr)rxryrrZrr]rrrrr%r(ry}s    zRealmEmbedder.__init__cC |jjjSrrrr~rwr%r%r(r z"RealmEmbedder.get_input_embeddingscC||jj_dSrrrr%r%r(rz"RealmEmbedder.set_input_embeddingsbatch_size, sequence_length output_typerNrrrurrrrrr!r"rc Csj| dur| n|jj} |j||||||||| d } | d} || } | s,| f| ddSt| | j| jdS)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 ``` Nrrurrrrrr!r"rr0)r;rr%)rZrrrr:rr%) rrrrurrrrrr!r" realm_outputsrr;r%r%r(rs*! zRealmEmbedder.forward) NNNNNNNNN)rDrr_tied_weights_keysryrrrREALM_INPUTS_DOCSTRINGformatrr:_CONFIG_FOR_DOCrrVrrrrrrrr%r%rr(rIvsJ      rIzoThe scorer of REALM outputting relevance scores representing the score of document candidates (before softmax).c!seZdZdZdfdd Zeedee e d             dde e j de e jd e e j d e e j d e e j d e e jd e e j de e jde e jde e jde ede ede edeee ffddZZS) RealmScorerz Args: query_embedder ([`RealmEmbedder`]): Embedder for input sequences. If not specified, it will use the same embedder as candidate sequences. Ncs8t|t|j|_|dur|n|j|_|dSr)rxryrIrZembedderquery_embedderr)rrZrrr%r(rys   zRealmScorer.__init__rrrrrurrcandidate_input_idscandidate_attention_maskcandidate_token_type_idscandidate_inputs_embedsrrrr!r"rc Cs| dur| n|jj} |dur| durtd|dur"|dur"td|j||||| | | | | d }||||\}}}|j||||| || | | d }|d}|d}|d|jj|jj}t d||}| si|||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.rrrsz bd,bnd->bn)r>r?r@) rZrrrrrrnum_candidatesr^rVrr=)rrrrurrrrrrrrrr!r" query_outputsflattened_input_idsflattened_attention_maskflattened_token_type_idscandidate_outputsr?r@r>r%r%r(rsNI    zRealmScorer.forwardrr)rDrrrryrrrrr=rrrVrrrrrrrr%r%rr(rs^       rzrThe knowledge-augmented encoder of REALM outputting masked language model logits and marginal log-likelihood loss.cseZdZdgZfddZddZddZdd Zd d Ze e d e e ed             ddeejdeejdeejdeejdeejdeejdeejdeejdeejdeedeedeedeee ffddZZS)rHzcls.predictions.decodercrr)rxryrrZrrYrrrrr%r(rycs    z!RealmKnowledgeAugEncoder.__init__cCrrrrwr%r%r(rirz-RealmKnowledgeAugEncoder.get_input_embeddingscCrrrrr%r%r(rlrz-RealmKnowledgeAugEncoder.set_input_embeddingscCrrrrZrWrwr%r%r(get_output_embeddingsorz.RealmKnowledgeAugEncoder.get_output_embeddingscCrrr)rnew_embeddingsr%r%r(set_output_embeddingsrrz.RealmKnowledgeAugEncoder.set_output_embeddingsz+batch_size, num_candidates, sequence_lengthrNrrrurrrrr>labelsmlm_maskrr!r"rc  Cst| dur| n|jj} ||||\} }}|j| |||||| | | d }|d}||}|}d}|dur|dur;td|\}}| durNtj|tj d} n| tj } t dd}| d|jj }|d |jj d}||| ||jj| }|dd}||}|d }tt|| t|  }| s|f|d d }|dur|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 ```NrrzZYou have to specify `relevance_score` when `labels` is specified in order to compute loss.rvnone) reductionrsrr0r)rBlogitsrr%)rZrrrrrrrV ones_likerr}rrr{tiler log_softmax unsqueeze logsumexpnansumsumr rr%)rrrrurrrrr>rrrr!r"rrr joint_outputs joint_outputr\r@masked_lm_lossrrloss_fct mlm_logits mlm_targetsmasked_lm_log_probcandidate_log_probjoint_gold_log_probmarginal_gold_log_probsrr%r%r(rus^9         z RealmKnowledgeAugEncoder.forward) NNNNNNNNNNNN)rDrrrryrrrrrrrrr rrrVrrrrrrrr%r%rr(rH[sd        rHzThe reader of REALM.c#seZdZfddZeedeee d              dde e j de e j de e j d e e j d e e j d e e j d e e j d e e jde e j de e j de e jde ede ede edeeeffddZZS)rBcs>t||j|_t||_t||_t||_| dSr) rxry num_labelsrrrYrr_ qa_outputsrrrr%r(rys     zRealmReader.__init__z!reader_beam_size, sequence_lengthrNrrrurrrrr>rstart_positions end_positions has_answersrr!r"rc$ CsF|dur|n|jj}|durtd|durtd|d|jjkr'td|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}| dur| dur| durd d }d d}|d} | d| } | d| } | }t|}!|||| d|jj| d|jjd}t|}"|||}||d|d}||!tj9}||"tj9}||}|s||||f|dd}#|dur|||||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.rrrsrrocSs\ttt|ddt|d}ttt|ddt|d}tt||dS)zCompute correct span.rrsr)rVeqrrK logical_and)rr gold_starts gold_ends is_gold_start is_gold_endr%r%r(compute_correct_candidatesVsz7RealmReader.forward..compute_correct_candidatescSs@tjfdd}tj||||jddd}tj|dd}||S)z3Loss based on the negative marginal log-likelihood.cSryrzr|rr%r%r(rfrzERealmReader.forward..marginal_log_loss..mask_to_scorervrsr)rVrrrw)r is_correctr log_numeratorlog_denominatorr%r%r(marginal_log_losscsz.RealmReader.forward..marginal_log_loss)rrrrr0) rBrCrDrErFrGrHrIrJrr%)rZrrrrprrreader_beam_sizerVrargmaxmaxvaluesrqclamprKrr}rrrArr%)$rrrrurrrrr>rrrrrr!r"rr[rrrretriever_logitspredicted_block_indexpredicted_candidatepredicted_start predicted_end total_lossrCrDrErFrr ignored_indexany_retriever_correctany_reader_correctrr%r%r(rs&         zRealmReader.forward)NNNNNNNNNNNNNN)rDrrryrrrrrArrrVrrrKrrrrrr%r%rr(rBsb        rBay 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.cseZdZdfdd ZeddZddZee de e e d     dd e ejd e ejd e ejd e ejde edeee ff ddZZS)rFNc s`t|t||_t||_|dtdj |j |j ftj t dd||_|dS)N block_embr%cpu)rrwr)rxryrIrrBrrrVr new_emptynum_block_recordsr^rr retrieverr)rrZrrr%r(rys      zRealmForOpenQA.__init__cCs|jr|jjS|jjSr)r'rZsearcher_beam_sizerrwr%r%r(rsz!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)rr)rrr%r%r(block_embedding_tosz!RealmForOpenQA.block_embedding_toz1, sequence_lengthrrrru answer_idsr"rc Cs|dur|n|jj}|dur|jddkrtd|j|||dd}|d}td|j||jj }tj ||j dd \} } | } tj |jd| d } |j| |||jjd \} } }}||jj }|jtjj|jj d }||jtj| durtj| 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)rrurr"z BD,QD->QBrs)krro) max_lengthrgrzD,BD->B) rrrur>rrrrr")rNrO)!rZrrSrrrVrrrrtopkrrrqrrreader_seq_lenrspecial_tokens_maskr}r logical_not_ logical_and_rurrrrrrGrIrJrM)rrrrurr"question_outputsquestion_projection batch_scoresrretrieved_block_idsretrieved_block_embrrIrJ concat_inputsrretrieved_logitsrNpredicted_blockrOr%r%r(rsV% zRealmForOpenQA.forwardr)NNNN)rDrrrypropertyrrrREALM_FOR_OPEN_QA_DOCSTRINGrrrMrrrVrrrrrrrr%r%rr(rFs0   rF)Grrr8 dataclassesrtypingrrrrVrtorch.nnr activationsr modeling_outputsr r r r modeling_utilsr pytorch_utilsrrrutilsrrrrconfiguration_realmr get_loggerrDr6_EMBEDDER_CHECKPOINT_FOR_DOC_ENCODER_CHECKPOINT_FOR_DOC_SCORER_CHECKPOINT_FOR_DOCr#REALM_PRETRAINED_MODEL_ARCHIVE_LISTrkModulerlrrrrrrrr5r:r=rArMrRrUrYr]r_REALM_START_DOCSTRINGrrrrIrrHrBr*rFr%r%r%r(s        lA2W^1  C 2(N(!