a h @svdZddlZddlmZmZmZddlZddlZddlmZddl m Z m Z m Z ddl mZddlmZdd lmZmZmZmZmZmZdd lmZmZdd lmZmZmZdd lm Z m!Z!m"Z"d dl#m$Z$e"%e&Z'Gdddej(Z)d=ddZ*Gdddej(Z+Gdddej(Z,Gdddej(Z-Gdddej(Z.Gdddej(Z/Gdddej(Z0Gdd d ej(Z1Gd!d"d"ej(Z2d>ej(ej3ej3ej3eej3e4e4eej3d$d%d&Z5Gd'd(d(ej(Z6Gd)d*d*ej(Z7Gd+d,d,eZ8Gd-d.d.ej(Z9e Gd/d0d0eZ:e Gd1d2d2e:Z;e Gd3d4d4e:Zgd<Z?dS)?zPyTorch 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)Ncsg|]}tjjqS)r EmbeddingZmax_xpath_tag_unit_embeddingsxpath_unit_hidden_size.0_configrj/var/www/html/assistant/venv/lib/python3.9/site-packages/transformers/models/markuplm/modeling_markuplm.py ?sz,XPathEmbeddings.__init__..csg|]}tjjqSr)rrZmax_xpath_subs_unit_embeddingsr r!r$rr&r'Fs)super__init__ max_depthrLinearr hidden_sizeZxpath_unitseq2_embeddingsDropouthidden_dropout_probdropoutZReLU activationxpath_unitseq2_inner inner2emb ModuleListrangexpath_tag_sub_embeddingsxpath_subs_sub_embeddingsselfr% __class__r$r&r)2s"    zXPathEmbeddings.__init__Nc Csg}g}t|jD]P}||j||dddd|f||j||dddd|fqtj|dd}tj|dd}||}||| | |}|S)Ndim) r4r*appendr5r6torchcatr2r/r0r1)r8xpath_tags_seqxpath_subs_seqZxpath_tags_embeddingsZxpath_subs_embeddingsixpath_embeddingsrrr&forwardLs&(zXPathEmbeddings.forward)NN)__name__ __module__ __qualname____doc__r)rE __classcell__rrr9r&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 rr<)neintr?ZcumsumZtype_aslong) input_ids padding_idxpast_key_values_lengthmaskZincremental_indicesrrr&"create_position_ids_from_input_ids_s rRcs2eZdZdZfddZddZd 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)rOeps position_ids)rr;F) persistent)r(r)r%rr vocab_sizer,Z pad_token_idword_embeddingsZmax_position_embeddingsposition_embeddingsr*rrDZtype_vocab_sizetoken_type_embeddings LayerNormlayer_norm_epsr-r.r/Zregister_bufferr?arangeexpandrOr7r9rr&r)rs    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 Nr;rdtypedevicer)sizer?r^rOrMrb unsqueezer_)r8 inputs_embeds input_shapeZsequence_lengthrVrrr&&create_position_ids_from_inputs_embedss z9MarkupLMEmbeddings.create_position_ids_from_inputs_embedsNrcCs<|dur|}n|dd}|dur0|jn|j} |dur`|durVt||j|}n ||}|durztj|tj| d}|dur||}|dur|j j tj t t ||jgtj| d}|dur|j jtj t t ||jgtj| d}|} ||} ||} |||} | | | | }||}||}|S)Nr;r`)rcrbrRrOrgr?zerosrMrYr%Z tag_pad_idonestuplelistr*Z subs_pad_idrZr[rDr\r/)r8rNrArBtoken_type_idsrVrerPrfrbZwords_embeddingsrZr[rD embeddingsrrr&rEs8           zMarkupLMEmbeddings.forward)NNNNNNr)rFrGrHrIr)rgrErJrrr9r&rSos rScs4eZdZfddZejejejdddZZS)MarkupLMSelfOutputcsBtt|j|j|_tj|j|jd|_t|j |_ dSNrT) r(r)rr+r,denser\r]r-r.r/r7r9rr&r)s zMarkupLMSelfOutput.__init__ hidden_states input_tensorreturncCs&||}||}|||}|SNrpr/r\r8rrrsrrr&rEs  zMarkupLMSelfOutput.forwardrFrGrHr)r?TensorrErJrrr9r&rns rncs0eZdZfddZejejdddZZS)MarkupLMIntermediatecsBtt|j|j|_t|jt r6t |j|_ n|j|_ dSru) r(r)rr+r,intermediate_sizerp isinstance hidden_actstrr intermediate_act_fnr7r9rr&r)s   zMarkupLMIntermediate.__init__rrrtcCs||}||}|Sru)rprr8rrrrr&rEs  zMarkupLMIntermediate.forwardrxrrr9r&rzs rzcs4eZdZfddZejejejdddZZS)MarkupLMOutputcsBtt|j|j|_tj|j|jd|_t |j |_ dSro) r(r)rr+r{r,rpr\r]r-r.r/r7r9rr&r)s zMarkupLMOutput.__init__rqcCs&||}||}|||}|Srurvrwrrr&rEs  zMarkupLMOutput.forwardrxrrr9r&rs rcs0eZdZfddZejejdddZZS)MarkupLMPoolercs*tt|j|j|_t|_dSru)r(r)rr+r,rpZTanhr0r7r9rr&r)s zMarkupLMPooler.__init__rcCs(|dddf}||}||}|S)Nr)rpr0)r8rrZfirst_token_tensor pooled_outputrrr&rEs  zMarkupLMPooler.forwardrxrrr9r&rs rcs0eZdZfddZejejdddZZS)MarkupLMPredictionHeadTransformcsVtt|j|j|_t|jtr6t |j|_ n|j|_ tj |j|j d|_ dSro) r(r)rr+r,rpr|r}r~r transform_act_fnr\r]r7r9rr&r)s   z(MarkupLMPredictionHeadTransform.__init__rcCs"||}||}||}|Sru)rprr\rrrr&rEs   z'MarkupLMPredictionHeadTransform.forwardrxrrr9r&rs rcs,eZdZfddZddZddZZS)MarkupLMLMPredictionHeadcsLtt||_tj|j|jdd|_t t |j|_ |j |j_ dS)NF)bias) r(r)r transformrr+r,rXdecoder Parameterr?rhrr7r9rr&r)#s   z!MarkupLMLMPredictionHead.__init__cCs|j|j_dSru)rrr8rrr& _tie_weights0sz%MarkupLMLMPredictionHead._tie_weightscCs||}||}|Sru)rrrrrr&rE3s  z MarkupLMLMPredictionHead.forward)rFrGrHr)rrErJrrr9r&r"s rcs0eZdZfddZejejdddZZS)MarkupLMOnlyMLMHeadcstt||_dSru)r(r)r predictionsr7r9rr&r);s zMarkupLMOnlyMLMHead.__init__)sequence_outputrtcCs||}|Sru)r)r8rprediction_scoresrrr&rE?s zMarkupLMOnlyMLMHead.forwardrxrrr9r&r:s r)modulequerykeyvalueattention_maskscalingr/ head_maskc Kst||dd|} |durN|ddddddd|jdf} | | } tjj| dtjd|j } tjj | ||j d} |dur| | dddd} t| |} | dd } | | fS)Nr r;)r=ra)ptrainingr)r?matmul transposeshaperZ functionalZsoftmaxZfloat32torar/rview contiguous) rrrrrrr/rkwargs attn_weightsZ causal_mask attn_outputrrr&eager_attention_forwardEs & rcsLeZdZfddZdejeejeejeee ejdddZ Z S) MarkupLMSelfAttentioncst|j|jdkr>t|ds>td|jd|jd||_|j|_t|j|j|_|j|j|_ t |j|j |_ t |j|j |_ t |j|j |_t |j|_|j|_|jd|_dS)NrZembedding_sizezThe hidden size (z6) is not a multiple of the number of attention heads ()g)r(r)r,num_attention_headshasattr ValueErrorr%rLattention_head_size all_head_sizerr+rrrr-Zattention_probs_dropout_probr/attention_dropoutrr7r9rr&r)bs"  zMarkupLMSelfAttention.__init__NFrrrroutput_attentionsrtcKs|jdd}g|d|jR}|||dd}|||dd} |||dd} t} |jj dkrt |jj } | ||| | |f|j sdn|j |j |d|\} } | jg|dR} |r| | fn| f}|S)Nr;rreagerr)r/rr)rrrrrrrrr%Z_attn_implementationrrrrZreshaper)r8rrrrrrrfZ hidden_shapeZ query_statesZ key_statesZ value_statesZattention_interfacerroutputsrrr&rEws0    zMarkupLMSelfAttention.forward)NNF) rFrGrHr)r?ryr FloatTensorboolrjrErJrrr9r&ras rcsTeZdZfddZddZd ejeejeejee e ejddd Z Z S) MarkupLMAttentioncs*tt||_t||_t|_dSru)r(r)rr8rnoutputset pruned_headsr7r9rr&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)Nrrr<)lenrr8rrrrrrrrrprunion)r8headsindexrrr& prune_headss zMarkupLMAttention.prune_headsNFrc Ks@|j|f|||d|}||d|}|f|dd}|SN)rrrrr)r8r) r8rrrrrrZ self_outputsattention_outputrrrr&rEszMarkupLMAttention.forward)NNF) rFrGrHr)rr?ryrrrrjrErJrrr9r&rs rcsTeZdZfddZd ejeejeejeee ejdddZ dd Z Z S) MarkupLMLayercs:t|j|_d|_t||_t||_t||_ dS)Nr) r(r)chunk_size_feed_forward seq_len_dimr attentionrz intermediaterrr7r9rr&r)s    zMarkupLMLayer.__init__NFrc KsP|j|f|||d|}|d}|dd}t|j|j|j|} | f|}|Sr)rrfeed_forward_chunkrr) r8rrrrrrZself_attention_outputsrr layer_outputrrr&rEs   zMarkupLMLayer.forwardcCs||}|||}|Sru)rr)r8rZintermediate_outputrrrr&rs  z MarkupLMLayer.feed_forward_chunk)NNF) rFrGrHr)r?ryrrrrjrErrJrrr9r&rs  rc sdeZdZfddZed ejeejeejee ee ee e e eje fdddZ ZS) MarkupLMEncodercs:t|_tfddtjD|_d|_dS)Ncsg|] }tqSr)r)r"rCr$rr&r'z,MarkupLMEncoder.__init__..F) r(r)r%rr3r4num_hidden_layerslayerZgradient_checkpointingr7r9r$r&r)s  zMarkupLMEncoder.__init__NFT)rrrrroutput_hidden_states return_dictrtcKs|rdnd}|rdnd} t|jD]\\} } |r8||f}|durH|| nd} | f||| |d|} | d}|r"| | df} q"|r||f}t||| dS)Nr)rrrrrrr)last_hidden_staterr attentions) enumeraterr )r8rrrrrrrrZall_hidden_statesZall_self_attentionsrCZ layer_moduleZlayer_head_maskZ layer_outputsrrr&rEs0    zMarkupLMEncoder.forward)NNFFT)rFrGrHr)rr?ryrrrrrjr rErJrrr9r&rs  rcsJeZdZUeed<dZddZeee e e j fdfdd Z ZS)MarkupLMPreTrainedModelr%markuplmcCst|tjr:|jjjd|jjd|jdur|jj n~t|tj rz|jjjd|jjd|j dur|jj|j  n>t|tj r|jj |jj dnt|tr|jj dS)zInitialize the weightsr)meanZstdN?)r|rr+weightdataZnormal_r%Zinitializer_rangerZzero_rrOr\Zfill_r)r8rrrr& _init_weights*s       z%MarkupLMPreTrainedModel._init_weights)pretrained_model_name_or_pathcstj|g|Ri|Sru)r(from_pretrained)clsrZ model_argsrr9rr&r<sz'MarkupLMPreTrainedModel.from_pretrained)rFrGrHr__annotations__Zbase_model_prefixr classmethodrrr~osPathLikerrJrrr9r&r$s rcseZdZdfdd ZddZddZdd Zeede e j e e j e e j e e j e e j e e j e e j e e j e e e e e e eeefd d d ZZS) MarkupLMModelTcsDt|||_t||_t||_|r2t|nd|_| dS)zv add_pooling_layer (bool, *optional*, defaults to `True`): Whether to add a pooling layer N) r(r)r%rSrmrencoderrpooler post_init)r8r%add_pooling_layerr9rr&r)Ds    zMarkupLMModel.__init__cCs|jjSrurmrYrrrr&get_input_embeddingsTsz"MarkupLMModel.get_input_embeddingscCs ||j_dSrur)r8rrrr&set_input_embeddingsWsz"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)r8Zheads_to_prunerrrrr& _prune_headsZszMarkupLMModel._prune_headsN) rNrArBrrlrVrrerrrrtc Cs| dur | n|jj} | dur | n|jj} | dur4| n|jj} |durV|durVtdn@|durt||||} n"|dur|dd} ntd|dur|jn|j} |durtj | | d}|durtj | tj | d}| d d}|j |jd }d |d }|dur|dkrP| d  d  d d}||jjdddd}n$|dkrt| 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 timer;z5You have to specify either input_ids or inputs_embeds)rbr`rr)rargr)rNrArBrVrlreT)rrrr)rZ pooler_outputrrr)r%rruse_return_dictrZ%warn_if_padding_and_no_attention_maskrcrbr?rirhrMrdrrar=r_rnext parametersrmrrr rrr)r8rNrArBrrlrVrrerrrrfrbZextended_attention_maskZembedding_outputZencoder_outputsrrrrr&rEbsh'     zMarkupLMModel.forward)T) NNNNNNNNNNN)rFrGrHr)rrrrrrr?Z LongTensorrrrrjr rErJrrr9r&rAs> rcseZdZfddZeedeejeejeejeejeejeejeejeejeejeejee ee ee e e eje fdddZ ZS)MarkupLMForQuestionAnsweringcs@t||j|_t|dd|_t|j|j|_| dSNF)r) r(r) num_labelsrrrr+r, qa_outputsrr7r9rr&r)s  z%MarkupLMForQuestionAnswering.__init__N)rNrArBrrlrVrrestart_positions end_positionsrrrrtc Cs"| dur | n|jj} |j||||||||| | dd }|d}||}|jddd\}}|d}|d}d}| dur | dur t| dkr| 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 rArBrrlrVrrerrrrrr;r<)Z ignore_indexr)loss start_logits end_logitsrrr)r%rrrsplitsqueezerrrcZclamp_rrrrr)r8rNrArBrrlrVrrerrrrrrrlogitsrrZ total_lossZ ignored_indexloss_fctZ start_lossZend_lossrrr&rEsN0          z$MarkupLMForQuestionAnswering.forward) NNNNNNNNNNNNN)rFrGrHr)rrrr?ryrrrjrrErJrrr9r&rs@ rzC MarkupLM Model with a `token_classification` head on top. )Z custom_introcseZdZfddZeedeejeejeejeejeejeejeejeejeejee ee ee e e eje fd ddZ ZS)MarkupLMForTokenClassificationcsbt||j|_t|dd|_|jdur2|jn|j}t||_ t |j |j|_ | dSr)r(r)rrrclassifier_dropoutr.rr-r/r+r, classifierrr8r%rr9rr&r)?s  z'MarkupLMForTokenClassification.__init__N rNrArBrrlrVrrelabelsrrrrtc  Cs| dur | n|jj} |j||||||||| | dd } | d}||}d}| durtt}||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 ```NTrrr;rrrrr) r%rrrrrrrrrr)r8rNrArBrrlrVrrerrrrrrrrrrrr&rEMs:. z&MarkupLMForTokenClassification.forward) NNNNNNNNNNNN)rFrGrHr)rrrr?ryrrrjrrErJrrr9r&r8s< rz 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. cseZdZfddZeedeejeejeejeejeejeejeejeejeejee ee ee e e eje fd ddZ ZS)!MarkupLMForSequenceClassificationcsdt||j|_||_t||_|jdur4|jn|j}t ||_ t |j |j|_ |dSru)r(r)rr%rrrr.rr-r/r+r,rrrr9rr&r)s   z*MarkupLMForSequenceClassification.__init__Nrc  CsP| dur | n|jj} |j||||||||| | dd } | d}||}||}d}| dur<|jjdur|jdkr~d|j_n4|jdkr| jtj ks| jtj krd|j_nd|j_|jjdkrt }|jdkr|| | }n ||| }nN|jjdkrt }||d|j| d}n|jjdkr 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 ```NTrrZ regressionZsingle_label_classificationZmulti_label_classificationr;r)r%rrr/rZ problem_typerrar?rMrLrrrrrrrrr)r8rNrArBrrlrVrrerrrrrrrrrrrr&rEsT-      "     z)MarkupLMForSequenceClassification.forward) NNNNNNNNNNNN)rFrGrHr)rrrr?ryrrrjrrErJrrr9r&rs< r)rrrrr)r)rN)@rIrtypingrrrr?Ztorch.utils.checkpointrZtorch.nnrrrZ activationsr Zmodeling_layersr Zmodeling_outputsr r rrrrZmodeling_utilsrrZ pytorch_utilsrrrutilsrrrZconfiguration_markuplmrZ get_loggerrFloggerModulerrRrSrnrzrrrrrryfloatrrrrrrrrrr__all__rrrr&sp      3 c ;.)1 mar