o ۷io@sdZddlZddlmZmZmZddlZddlmZddlm Z m Z m Z ddl m Z ddlmZdd lmZmZmZmZmZmZdd lmZmZdd lmZmZmZdd lmZm Z m!Z!d dl"m#Z#e!$e%Z&Gdddej'Z(dDddZ)Gdddej'Z*Gdddej'Z+Gdddej'Z,Gdddej'Z-Gdddej'Z.Gdddej'Z/Gdd d ej'Z0Gd!d"d"ej'Z1 # dEd$ej'd%ej2d&ej2d'ej2d(eej2d)e3d*e3d+eej2fd,d-Z4Gd.d/d/ej'Z5Gd0d1d1ej'Z6Gd2d3d3eZ7Gd4d5d5ej'Z8eGd6d7d7eZ9eGd8d9d9e9Z:eGd:d;d;e9Z;edd?d?e9ZdS)FzPyTorch MarkupLM model.N)CallableOptionalUnion)nn)BCEWithLogitsLossCrossEntropyLossMSELoss)ACT2FN)GradientCheckpointingLayer)BaseModelOutputBaseModelOutputWithPoolingMaskedLMOutputQuestionAnsweringModelOutputSequenceClassifierOutputTokenClassifierOutput)ALL_ATTENTION_FUNCTIONSPreTrainedModel)apply_chunking_to_forward find_pruneable_heads_and_indicesprune_linear_layer)auto_docstringcan_return_tuplelogging)MarkupLMConfigcs*eZdZdZfddZdddZZS)XPathEmbeddingszConstruct the embeddings from xpath tags and subscripts. We drop tree-id in this version, as its info can be covered by xpath. cstj|_tj|jj|_tj |_ t |_ tj|jdj|_ tdjj|_tfddt|jD|_tfddt|jD|_dS)Ncg|] }tjjqS)r Embeddingmax_xpath_tag_unit_embeddingsxpath_unit_hidden_size.0_configrd/home/ubuntu/vllm_env/lib/python3.10/site-packages/transformers/models/markuplm/modeling_markuplm.py >z,XPathEmbeddings.__init__..crr)rr max_xpath_subs_unit_embeddingsr"r#r&rr(r)Er*)super__init__ max_depthrLinearr" hidden_sizexpath_unitseq2_embeddingsDropouthidden_dropout_probdropoutReLU activationxpath_unitseq2_inner inner2emb ModuleListrangexpath_tag_sub_embeddingsxpath_subs_sub_embeddingsselfr' __class__r&r(r-1s"     zXPathEmbeddings.__init__Nc Csg}g}t|jD](}||j||dddd|f||j||dddd|fq tj|dd}tj|dd}||}||| | |}|S)Ndim) r:r.appendr;r<torchcatr8r4r6r7)r>xpath_tags_seqxpath_subs_seqxpath_tags_embeddingsxpath_subs_embeddingsixpath_embeddingsrrr(forwardKs&(zXPathEmbeddings.forward)NN)__name__ __module__ __qualname____doc__r-rM __classcell__rrr?r(r+s rcCs6||}tj|dd|||}||S)a Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols are ignored. This is modified from fairseq's `utils.make_positions`. Args: x: torch.Tensor x: Returns: torch.Tensor rrB)neintrEcumsumtype_aslong) input_ids padding_idxpast_key_values_lengthmaskincremental_indicesrrr("create_position_ids_from_input_ids^s  r]cs@eZdZdZfddZddZ       d dd ZZS) MarkupLMEmbeddingszGConstruct the embeddings from word, position and token_type embeddings.cst||_tj|j|j|jd|_t|j |j|_ |j |_ t ||_ t|j|j|_tj|j|jd|_t|j|_|jdt|j ddd|j|_tj|j |j|jd|_ dS)N)rYeps position_ids)rrAF) persistent)r,r-r'rr vocab_sizer0 pad_token_idword_embeddingsmax_position_embeddingsposition_embeddingsr.rrLtype_vocab_sizetoken_type_embeddings LayerNormlayer_norm_epsr2r3r4register_bufferrEarangeexpandrYr=r?rr(r-qs    zMarkupLMEmbeddings.__init__cCsN|dd}|d}tj|jd||jdtj|jd}|d|S)z We are provided embeddings directly. We cannot infer which are padded so just generate sequential position ids. Args: inputs_embeds: torch.Tensor Returns: torch.Tensor NrArdtypedevicer)sizerErmrYrWrq unsqueezern)r> inputs_embeds input_shapesequence_lengthrarrr(&create_position_ids_from_inputs_embedss  z9MarkupLMEmbeddings.create_position_ids_from_inputs_embedsNrcCs<|dur |}n|dd}|dur|jn|j} |dur0|dur+t||j|}n||}|dur=tj|tj| d}|durF||}|dur_|j j tj t t ||jgtj| d}|durx|j jtj t t ||jgtj| d}|} ||} ||} |||} | | | | }||}||}|S)NrAro)rrrqr]rYrwrEzerosrWrer' tag_pad_idonestuplelistr. subs_pad_idrgrirLrjr4)r>rXrGrHtoken_type_idsrartrZrurqwords_embeddingsrgrirL embeddingsrrr(rMs8          zMarkupLMEmbeddings.forward)NNNNNNr)rNrOrPrQr-rwrMrRrrr?r(r^ns r^c8eZdZfddZdejdejdejfddZZS)MarkupLMSelfOutputcsBtt|j|j|_tj|j|jd|_t|j |_ dSNr_) r,r-rr/r0denserjrkr2r3r4r=r?rr(r- zMarkupLMSelfOutput.__init__ hidden_states input_tensorreturncC&||}||}|||}|SNrr4rjr>rrrrr(rM  zMarkupLMSelfOutput.forwardrNrOrPr-rETensorrMrRrrr?r(r $rc2eZdZfddZdejdejfddZZS)MarkupLMIntermediatecsDtt|j|j|_t|jt rt |j|_ dS|j|_ dSr) r,r-rr/r0intermediate_sizer isinstance hidden_actstrr intermediate_act_fnr=r?rr(r-s   zMarkupLMIntermediate.__init__rrcC||}||}|Sr)rrr>rrrr(rM  zMarkupLMIntermediate.forwardrrrr?r(rs rcr)MarkupLMOutputcsBtt|j|j|_tj|j|jd|_t |j |_ dSr) r,r-rr/rr0rrjrkr2r3r4r=r?rr(r-rzMarkupLMOutput.__init__rrrcCrrrrrrr(rMrzMarkupLMOutput.forwardrrrr?r(rrrcr)MarkupLMPoolercs*tt|j|j|_t|_dSr)r,r-rr/r0rTanhr6r=r?rr(r-s zMarkupLMPooler.__init__rrcCs(|dddf}||}||}|S)Nr)rr6)r>rfirst_token_tensor pooled_outputrrr(rMs  zMarkupLMPooler.forwardrrrr?r(rs rcr)MarkupLMPredictionHeadTransformcsVtt|j|j|_t|jtrt |j|_ n|j|_ tj |j|j d|_ dSr) r,r-rr/r0rrrrr transform_act_fnrjrkr=r?rr(r-s  z(MarkupLMPredictionHeadTransform.__init__rrcCs"||}||}||}|Sr)rrrjrrrr(rMs   z'MarkupLMPredictionHeadTransform.forwardrrrr?r(rs  rcs,eZdZfddZddZddZZS)MarkupLMLMPredictionHeadcsLtt||_tj|j|jdd|_t t |j|_ |j |j_ dS)NF)bias) r,r-r transformrr/r0rcdecoder ParameterrErxrr=r?rr(r-"s  z!MarkupLMLMPredictionHead.__init__cCs|j|j_dSr)rrr>rrr( _tie_weights/sz%MarkupLMLMPredictionHead._tie_weightscCrr)rrrrrr(rM2rz MarkupLMLMPredictionHead.forward)rNrOrPr-rrMrRrrr?r(r!s  rcr)MarkupLMOnlyMLMHeadcstt||_dSr)r,r-r predictionsr=r?rr(r-:s zMarkupLMOnlyMLMHead.__init__sequence_outputrcCs||}|Sr)r)r>rprediction_scoresrrr(rM>s zMarkupLMOnlyMLMHead.forwardrrrr?r(r9s rmodulequerykeyvalueattention_maskscalingr4 head_maskc Kst||dd|} |dur'|ddddddd|jdf} | | } tjj| dtjd|j } tjj | ||j d} |durM| | dddd} t| |} | dd } | | fS)Nr rA)rCrp)ptrainingr)rEmatmul transposeshaper functionalsoftmaxfloat32torpr4rview contiguous) rrrrrrr4rkwargs attn_weights causal_mask attn_outputrrr(eager_attention_forwardDs & rc sZeZdZfddZ   d dejdeejdeejdeed e ejf d d Z Z S) MarkupLMSelfAttentioncst|j|jdkrt|dstd|jd|jd||_|j|_t|j|j|_|j|j|_ t |j|j |_ t |j|j |_ t |j|j |_t |j|_|j|_|jd|_dS)Nrembedding_sizezThe hidden size (z6) is not a multiple of the number of attention heads ()g)r,r-r0num_attention_headshasattr ValueErrorr'rTattention_head_size all_head_sizerr/rrrr2attention_probs_dropout_probr4attention_dropoutrr=r?rr(r-as"  zMarkupLMSelfAttention.__init__NFrrroutput_attentionsrcKs|jdd}g|d|jR}|||dd}|||dd} |||dd} t} |jj dkrCt |jj } | ||| | |f|j sOdn|j |j |d|\} } | jg|dR} |rp| | f}|S| f}|S)NrArreagerr)r4rr)rrrrrrrrr'_attn_implementationrrrrreshaper)r>rrrrrru hidden_shape query_states key_states value_statesattention_interfacerroutputsrrr(rMvs4    zMarkupLMSelfAttention.forwardNNF) rNrOrPr-rErr FloatTensorboolr{rMrRrrr?r(r`s  rc sbeZdZfddZddZ   ddejdeejd eejd ee d e ejf d d Z Z S)MarkupLMAttentioncs*tt||_t||_t|_dSr)r,r-rr>routputset pruned_headsr=r?rr(r-s    zMarkupLMAttention.__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)NrrrB)lenrr>rrrrrrrrrrunion)r>headsindexrrr( prune_headss zMarkupLMAttention.prune_headsNFrrrrrc Ks@|j|f|||d|}||d|}|f|dd}|SN)rrrrr)r>r) r>rrrrr self_outputsattention_outputrrrr(rMszMarkupLMAttention.forwardr) rNrOrPr-rrErrrrr{rMrRrrr?r(rs" rc sbeZdZfddZ   ddejdeejdeejdeed e ejf d d Z d d Z Z S) MarkupLMLayercs:t|j|_d|_t||_t||_t||_ dS)Nr) r,r-chunk_size_feed_forward seq_len_dimr attentionr intermediaterrr=r?rr(r-s   zMarkupLMLayer.__init__NFrrrrrc KsP|j|f|||d|}|d}|dd}t|j|j|j|} | f|}|Sr)rrfeed_forward_chunkrr) r>rrrrrself_attention_outputsrr layer_outputrrr(rMs   zMarkupLMLayer.forwardcCs||}|||}|Sr)rr)r>rintermediate_outputrrrr(rs  z MarkupLMLayer.feed_forward_chunkr) rNrOrPr-rErrrrr{rMrrRrrr?r(rs"   rcszeZdZfddZe     ddejdeejdeejd ee d ee d ee d e e eje ffd dZ ZS)MarkupLMEncodercs:t|_tfddtjD|_d|_dS)Ncsg|]}tqSr)r)r$rKr&rr(r)sz,MarkupLMEncoder.__init__..F) r,r-r'rr9r:num_hidden_layerslayergradient_checkpointingr=r?r&r(r-s   zMarkupLMEncoder.__init__NFTrrrroutput_hidden_states return_dictrcKs|rdnd}|r dnd} t|jD].\} } |r||f}|dur$|| nd} | d||| |d|} | d}|r?| | df} q|rG||f}t||| dS)Nr)rrrrrr)last_hidden_stater attentions) enumeraterr )r>rrrrrrrall_hidden_statesall_self_attentionsrK layer_modulelayer_head_mask layer_outputsrrr(rMs2   zMarkupLMEncoder.forward)NNFFT)rNrOrPr-rrErrrrrr{r rMrRrrr?r(rs.  rcsJeZdZUeed<dZddZedee e e j fffdd Z ZS)MarkupLMPreTrainedModelr'markuplmcCst|tjr |jjjd|jjd|jdur|jj dSdSt|tj rC|jjjd|jjd|j durA|jj|j  dSdSt|tj rX|jj |jj ddSt|tre|jj dSdS)zInitialize the weightsr)meanstdN?)rrr/weightdatanormal_r'initializer_rangerzero_r rYrjfill_r)r>rrrr( _init_weights)s       z%MarkupLMPreTrainedModel._init_weightspretrained_model_name_or_pathcstj|g|Ri|Sr)r,from_pretrained)clsr  model_argsrr?rr(r;sz'MarkupLMPreTrainedModel.from_pretrained)rNrOrPr__annotations__base_model_prefixr  classmethodrrrosPathLikerrRrrr?r(r#s *rcseZdZdfdd ZddZddZdd Zee dd e e j d e e j d e e j de e j de e j de e j de e j de e j de e de e de e deeeffddZZS) MarkupLMModelTcsDt|||_t||_t||_|rt|nd|_| dS)zv add_pooling_layer (bool, *optional*, defaults to `True`): Whether to add a pooling layer N) r,r-r'r^rrencoderrpooler post_init)r>r'add_pooling_layerr?rr(r-Cs    zMarkupLMModel.__init__cCs|jjSrrrerrrr(get_input_embeddingsSsz"MarkupLMModel.get_input_embeddingscCs ||j_dSrr)r>rrrr(set_input_embeddingsVs z"MarkupLMModel.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)itemsrrrr)r>heads_to_prunerrrrr( _prune_headsYszMarkupLMModel._prune_headsNrXrGrHrr~rarrtrrrrc Cs| dur| n|jj} | dur| n|jj} | dur| n|jj} |dur*|dur*td|dur9||||} n|durF|dd} ntd|durQ|jn|j} |dur_tj | | d}|durltj | tj | d}| d d}|j |jd }d |d }|dur|dkr| d  d  d d}||jjdddd}n|dkr| d d d}|j t|jd }ndg|jj}|j||||||d }|j|||| | dd}|d }|jdur||nd}t|||j|jdS)a xpath_tags_seq (`torch.LongTensor` of shape `(batch_size, sequence_length, config.max_depth)`, *optional*): Tag IDs for each token in the input sequence, padded up to config.max_depth. xpath_subs_seq (`torch.LongTensor` of shape `(batch_size, sequence_length, config.max_depth)`, *optional*): Subscript IDs for each token in the input sequence, padded up to config.max_depth. Examples: ```python >>> from transformers import AutoProcessor, MarkupLMModel >>> processor = AutoProcessor.from_pretrained("microsoft/markuplm-base") >>> model = MarkupLMModel.from_pretrained("microsoft/markuplm-base") >>> html_string = " Page Title " >>> encoding = processor(html_string, return_tensors="pt") >>> outputs = model(**encoding) >>> last_hidden_states = outputs.last_hidden_state >>> list(last_hidden_states.shape) [1, 4, 768] ```NzDYou cannot specify both input_ids and inputs_embeds at the same timerAz5You have to specify either input_ids or inputs_embeds)rqrorr)rprgr)rXrGrHrar~rtT)rrrr)r pooler_outputrr)r'rruse_return_dictr%warn_if_padding_and_no_attention_maskrrrqrErzrxrWrsrrprCrnrnext parametersrrrr rr)r>rXrGrHrr~rarrtrrrrurqextended_attention_maskembedding_outputencoder_outputsrrrrr(rMash'     zMarkupLMModel.forward)T) NNNNNNNNNNN)rNrOrPr-rrr rrrrE LongTensorrrrr{r rMrRrrr?r(r@sV     rc!seZdZfddZee             ddeejdeejdeejdeejdeejd eejd eejd eejd eejd eejdee dee dee de e eje ffddZ ZS)MarkupLMForQuestionAnsweringcs@t||j|_t|dd|_t|j|j|_| dSNF)r) r,r- num_labelsrrrr/r0 qa_outputsrr=r?rr(r-s  z%MarkupLMForQuestionAnswering.__init__NrXrGrHrr~rarrtstart_positions end_positionsrrrrc Cs| dur| n|jj} |j||||||||| | dd }|d}||}|jddd\}}|d}|d}d}| dur| durt| dkrQ| d} t| dkr^| d} |d}| d|| d|t |d}||| }||| }||d }t ||||j |j d S) ae xpath_tags_seq (`torch.LongTensor` of shape `(batch_size, sequence_length, config.max_depth)`, *optional*): Tag IDs for each token in the input sequence, padded up to config.max_depth. xpath_subs_seq (`torch.LongTensor` of shape `(batch_size, sequence_length, config.max_depth)`, *optional*): Subscript IDs for each token in the input sequence, padded up to config.max_depth. Examples: ```python >>> from transformers import AutoProcessor, MarkupLMForQuestionAnswering >>> import torch >>> processor = AutoProcessor.from_pretrained("microsoft/markuplm-base-finetuned-websrc") >>> model = MarkupLMForQuestionAnswering.from_pretrained("microsoft/markuplm-base-finetuned-websrc") >>> html_string = " My name is Niels " >>> question = "What's his name?" >>> encoding = processor(html_string, questions=question, return_tensors="pt") >>> with torch.no_grad(): ... outputs = model(**encoding) >>> answer_start_index = outputs.start_logits.argmax() >>> answer_end_index = outputs.end_logits.argmax() >>> predict_answer_tokens = encoding.input_ids[0, answer_start_index : answer_end_index + 1] >>> processor.decode(predict_answer_tokens).strip() 'Niels' ```NT rGrHrr~rarrtrrrrrrArB) ignore_indexr)loss start_logits end_logitsrr)r'r"rr-splitsqueezerrrrclamp_rrrr)r>rXrGrHrr~rarrtr.r/rrrrrlogitsr3r4 total_loss ignored_indexloss_fct start_lossend_lossrrr(rMsN0          z$MarkupLMForQuestionAnswering.forward) NNNNNNNNNNNNN)rNrOrPr-rrrrErrrr{rrMrRrrr?r(r*s\       r*zC MarkupLM Model with a `token_classification` head on top. ) custom_introceZdZfddZee            ddeejdeejdeejdeejdeejd eejd eejd eejd eejd ee dee dee de e eje ffddZ ZS)MarkupLMForTokenClassificationcsbt||j|_t|dd|_|jdur|jn|j}t||_ t |j |j|_ | dSr+)r,r-r,rrclassifier_dropoutr3rr2r4r/r0 classifierrr>r'rAr?rr(r->s   z'MarkupLMForTokenClassification.__init__NrXrGrHrr~rarrtlabelsrrrrc  Cs| dur| n|jj} |j||||||||| | dd } | d}||}d}| dur:t}||d|jj| d}t||| j| j dS)a xpath_tags_seq (`torch.LongTensor` of shape `(batch_size, sequence_length, config.max_depth)`, *optional*): Tag IDs for each token in the input sequence, padded up to config.max_depth. xpath_subs_seq (`torch.LongTensor` of shape `(batch_size, sequence_length, config.max_depth)`, *optional*): Subscript IDs for each token in the input sequence, padded up to config.max_depth. labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`. Examples: ```python >>> from transformers import AutoProcessor, AutoModelForTokenClassification >>> import torch >>> processor = AutoProcessor.from_pretrained("microsoft/markuplm-base") >>> processor.parse_html = False >>> model = AutoModelForTokenClassification.from_pretrained("microsoft/markuplm-base", num_labels=7) >>> nodes = ["hello", "world"] >>> xpaths = ["/html/body/div/li[1]/div/span", "/html/body/div/li[1]/div/span"] >>> node_labels = [1, 2] >>> encoding = processor(nodes=nodes, xpaths=xpaths, node_labels=node_labels, return_tensors="pt") >>> with torch.no_grad(): ... outputs = model(**encoding) >>> loss = outputs.loss >>> logits = outputs.logits ```NTr0rrAr2r8rr) r'r"rrBrrr,rrr)r>rXrGrHrr~rarrtrDrrrrrrr2r;rrr(rMLs:. z&MarkupLMForTokenClassification.forward NNNNNNNNNNNN)rNrOrPr-rrrrErrrr{rrMrRrrr?r(r@7sV      r@z MarkupLM Model transformer with a sequence classification/regression head on top (a linear layer on top of the pooled output) e.g. for GLUE tasks. cr?)!MarkupLMForSequenceClassificationcsdt||j|_||_t||_|jdur|jn|j}t ||_ t |j |j|_ |dSr)r,r-r,r'rrrAr3rr2r4r/r0rBrrCr?rr(r-s    z*MarkupLMForSequenceClassification.__init__NrXrGrHrr~rarrtrDrrrrc  CsJ| dur| n|jj} |j||||||||| | dd } | d}||}||}d}| dur|jjdurX|jdkr>d|j_n|jdkrT| jtj ksO| jtj krTd|j_nd|j_|jjdkrvt }|jdkrp|| | }n+||| }n%|jjdkrt }||d|j| d}n|jjdkrt}||| }t||| j| jd S) a xpath_tags_seq (`torch.LongTensor` of shape `(batch_size, sequence_length, config.max_depth)`, *optional*): Tag IDs for each token in the input sequence, padded up to config.max_depth. xpath_subs_seq (`torch.LongTensor` of shape `(batch_size, sequence_length, config.max_depth)`, *optional*): Subscript IDs for each token in the input sequence, padded up to config.max_depth. labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If `config.num_labels > 1` a classification loss is computed (Cross-Entropy). Examples: ```python >>> from transformers import AutoProcessor, AutoModelForSequenceClassification >>> import torch >>> processor = AutoProcessor.from_pretrained("microsoft/markuplm-base") >>> model = AutoModelForSequenceClassification.from_pretrained("microsoft/markuplm-base", num_labels=7) >>> html_string = " Page Title " >>> encoding = processor(html_string, return_tensors="pt") >>> with torch.no_grad(): ... outputs = model(**encoding) >>> loss = outputs.loss >>> logits = outputs.logits ```NTr0r regressionsingle_label_classificationmulti_label_classificationrArE)r'r"rr4rB problem_typer,rprErWrTrr6rrrrrr)r>rXrGrHrr~rarrtrDrrrrrr8r2r;rrr(rMsT-     "       z)MarkupLMForSequenceClassification.forwardrF)rNrOrPr-rrrrErrrr{rrMrRrrr?r(rGsV      rG)r*rGr@rr)r)rN)?rQrtypingrrrrErtorch.nnrrr activationsr modeling_layersr modeling_outputsr r rrrrmodeling_utilsrr pytorch_utilsrrrutilsrrrconfiguration_markuplmr get_loggerrNloggerModulerr]r^rrrrrrrrfloatrrrrrrrr*r@rG__all__rrrr(s~       3c ;.)1 ma r