U ,-eܖ@sdZddlZddlmZddlmZmZmZddlZddl Zddlm Z ddl m Z m Z mZddlmZmZdd lmZmZdd lmZdd lmZdd lmZmZmZmZmZmZd dl m!Z!e"e#Z$dZ%dZ&ddgZ'eGdddeZ(eGdddeZ)eGdddeZ*eGdddeZ+eGdddeZ,eGdddeZ-eGdddeZ.eGd d!d!eZ/eGd"d#d#eZ0eGd$d%d%eZ1Gd&d'd'e j2Z3Gd(d)d)e j2Z4Gd*d+d+e j2Z5Gd,d-d-e j2Z6Gd.d/d/e j2Z7Gd0d1d1e j2Z8Gd2d3d3e j2Z9Gd4d5d5e j2Z:Gd6d7d7e j2Z;Gd8d9d9e j2ZGd>d?d?eZ?d@Z@dAZAedBe@GdCdDdDe?ZBdEdFZCGdGdHdHe j2ZDedIe@GdJdKdKe?ZEedLe@GdMdNdNe?ZFedOe@GdPdQdQe?ZGedRe@GdSdTdTe?ZHedUe@GdVdWdWe?ZIedXe@GdYdZdZe?ZJed[e@Gd\d]d]e?ZKed^e@Gd_d`d`e?ZLdS)azPyTorch LUKE model.N) dataclass)OptionalTupleUnion)nn)BCEWithLogitsLossCrossEntropyLossMSELoss)ACT2FNgelu)BaseModelOutputBaseModelOutputWithPooling)PreTrainedModel)apply_chunking_to_forward) ModelOutputadd_code_sample_docstringsadd_start_docstrings%add_start_docstrings_to_model_forwardloggingreplace_return_docstrings) LukeConfigrzstudio-ousia/luke-basezstudio-ousia/luke-largec@s6eZdZUdZdZejed<dZe e ejed<dS)BaseLukeModelOutputWithPoolinga Base class for outputs of the LUKE model. Args: last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): Sequence of hidden-states at the output of the last layer of the model. entity_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, entity_length, hidden_size)`): Sequence of entity hidden-states at the output of the last layer of the model. pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size)`): Last layer hidden-state of the first token of the sequence (classification token) further processed by a Linear layer and a Tanh activation function. 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. entity_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, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each layer plus the initial entity 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 + entity_length, sequence_length + entity_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. Nentity_last_hidden_stateentity_hidden_states __name__ __module__ __qualname____doc__rtorch FloatTensor__annotations__rrrr$r$g/var/www/html/Darija-Ai-Train/env/lib/python3.8/site-packages/transformers/models/luke/modeling_luke.pyr5s rc@s6eZdZUdZdZejed<dZe e ejed<dS)BaseLukeModelOutputa# Base class for model's outputs, with potential hidden states and attentions. Args: last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): Sequence of hidden-states at the output of the last layer of the model. entity_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, entity_length, hidden_size)`): Sequence of entity hidden-states at the output of the last layer of the model. 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. entity_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, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each layer plus the initial entity 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. Nrrrr$r$r$r%r&Ts r&c@seZdZUdZdZeejed<dZ eejed<dZ eejed<dZ ejed<dZ ejed<dZ eeejed<dZeeejed <dZeeejed <dS) LukeMaskedLMOutputa> Base class for model's outputs, with potential hidden states and attentions. Args: loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): The sum of masked language modeling (MLM) loss and entity prediction loss. mlm_loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): Masked language modeling (MLM) loss. mep_loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): Masked entity prediction (MEP) loss. logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). entity_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): Prediction scores of the entity prediction head (scores for each entity vocabulary token before SoftMax). 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. entity_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, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each layer plus the initial entity 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. Nlossmlm_lossmep_losslogits entity_logits hidden_statesr attentions)rrrr r(rr!r"r#r)r*r+r,r-rrr.r$r$r$r%r'ss r'c@steZdZUdZdZeejed<dZ ejed<dZ ee ejed<dZ ee ejed<dZ ee ejed<dS)EntityClassificationOutputay Outputs of entity classification models. Args: loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): Classification loss. logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`): Classification scores (before SoftMax). 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. entity_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, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each layer plus the initial entity 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. Nr(r+r-rr.rrrr r(rr!r"r#r+r-rrr.r$r$r$r%r/s r/c@steZdZUdZdZeejed<dZ ejed<dZ ee ejed<dZ ee ejed<dZ ee ejed<dS)EntityPairClassificationOutputa~ Outputs of entity pair classification models. Args: loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): Classification loss. logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`): Classification scores (before SoftMax). 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. entity_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, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each layer plus the initial entity 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. Nr(r+r-rr.r0r$r$r$r%r1s r1c@steZdZUdZdZeejed<dZ ejed<dZ ee ejed<dZ ee ejed<dZ ee ejed<dS)EntitySpanClassificationOutputa Outputs of entity span classification models. Args: loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): Classification loss. logits (`torch.FloatTensor` of shape `(batch_size, entity_length, config.num_labels)`): Classification scores (before SoftMax). 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. entity_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, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each layer plus the initial entity 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. Nr(r+r-rr.r0r$r$r$r%r2s r2c@steZdZUdZdZeejed<dZ ejed<dZ ee ejed<dZ ee ejed<dZ ee ejed<dS)LukeSequenceClassifierOutputa Outputs of sentence classification models. Args: loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): Classification (or regression if config.num_labels==1) loss. logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`): Classification (or regression if config.num_labels==1) scores (before SoftMax). 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, if the model has an embedding layer, + 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 optional initial embedding outputs. entity_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, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each layer plus the initial entity 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. Nr(r+r-rr.r0r$r$r$r%r3s r3c@steZdZUdZdZeejed<dZ ejed<dZ ee ejed<dZ ee ejed<dZ ee ejed<dS)LukeTokenClassifierOutputa Base class for outputs of token classification models. Args: loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided) : Classification loss. logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`): Classification scores (before SoftMax). 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, if the model has an embedding layer, + 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 optional initial embedding outputs. entity_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, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each layer plus the initial entity 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. Nr(r+r-rr.r0r$r$r$r%r4s r4c@seZdZUdZdZeejed<dZ ejed<dZ ejed<dZ ee ejed<dZ ee ejed<dZee ejed<dS) LukeQuestionAnsweringModelOutputay Outputs of question answering models. Args: loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): Total span extraction loss is the sum of a Cross-Entropy for the start and end positions. start_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length)`): Span-start scores (before SoftMax). end_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length)`): Span-end scores (before SoftMax). 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, if the model has an embedding layer, + 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 optional initial embedding outputs. entity_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, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each layer plus the initial entity 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. Nr( start_logits end_logitsr-rr.)rrrr r(rr!r"r#r6r7r-rrr.r$r$r$r%r5?s r5c@steZdZUdZdZeejed<dZ ejed<dZ ee ejed<dZ ee ejed<dZ ee ejed<dS)LukeMultipleChoiceModelOutputa  Outputs of multiple choice models. Args: loss (`torch.FloatTensor` of shape *(1,)*, *optional*, returned when `labels` is provided): Classification loss. logits (`torch.FloatTensor` of shape `(batch_size, num_choices)`): *num_choices* is the second dimension of the input tensors. (see *input_ids* above). Classification scores (before SoftMax). 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, if the model has an embedding layer, + 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 optional initial embedding outputs. entity_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, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each layer plus the initial entity 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. Nr(r+r-rr.r0r$r$r$r%r8ds r8cs2eZdZdZfddZd ddZddZZS) LukeEmbeddingszV Same as BertEmbeddings with a tiny tweak for positional embeddings indexing. csttj|j|j|jd|_t|j|j|_ t|j |j|_ tj |j|j d|_ t|j|_|j|_tj|j|j|jd|_ dS)N padding_idxZeps)super__init__r Embedding vocab_size hidden_sizeZ pad_token_idword_embeddingsmax_position_embeddingsposition_embeddingstype_vocab_sizetoken_type_embeddings LayerNormlayer_norm_epsDropouthidden_dropout_probdropoutr;selfconfig __class__r$r%r>s zLukeEmbeddings.__init__Nc Cs|dkr0|dk r&t||j|j}n ||}|dk rB|}n|dd}|dkrptj|tj|j jd}|dkr| |}| |}| |}|||}| |}||}|S)Ndtypedevice)"create_position_ids_from_input_idsr;torT&create_position_ids_from_inputs_embedssizer!zeroslong position_idsrBrDrFrGrK) rM input_idstoken_type_idsr[ inputs_embeds input_shaperDrF embeddingsr$r$r%forwards"        zLukeEmbeddings.forwardcCsN|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 NrQrrRr)rXr!Zaranger;rZrT unsqueezeexpand)rMr^r_Zsequence_lengthr[r$r$r%rWs  z5LukeEmbeddings.create_position_ids_from_inputs_embeds)NNNN)rrrr r>rarW __classcell__r$r$rOr%r9s  !r9cs<eZdZedfdd ZdejejejdddZZS) LukeEntityEmbeddingsrNcst||_tj|j|jdd|_|j|jkrHtj |j|jdd|_ t|j |j|_ t|j |j|_tj|j|jd|_t|j|_dS)Nrr:Fbiasr<)r=r>rNrr?entity_vocab_sizeentity_emb_sizeentity_embeddingsrALinearentity_embedding_denserCrDrErFrGrHrIrJrKrLrOr$r%r>s  zLukeEntityEmbeddings.__init__N) entity_idsr[r]c Cs|dkrt|}||}|jj|jjkr6||}||jdd}|dk | d}||}tj |dd}||j ddjdd}| |}|||}| |}||}|S)Nr)minrQdimgHz>)r!Z zeros_likerkrNrjrArmrDclamptype_asrbsumrFrGrK) rMrnr[r]rkrDZposition_embedding_maskrFr`r$r$r%ras       zLukeEntityEmbeddings.forward)N) rrrrr>r! LongTensorrardr$r$rOr%resrecs.eZdZfddZddZd ddZZS) LukeSelfAttentioncst|j|jdkr@t|ds@td|jfd|jd|j|_t|j|j|_|j|j|_|j |_ t |j|j|_ t |j|j|_ t |j|j|_|j rt |j|j|_t |j|j|_t |j|j|_t |j|_dS)NrZembedding_sizezThe hidden size z4 is not a multiple of the number of attention heads .)r=r>rAnum_attention_headshasattr ValueErrorintattention_head_size all_head_sizeuse_entity_aware_attentionrrlquerykeyvalue w2e_query e2w_query e2e_queryrIZattention_probs_dropout_probrKrLrOr$r%r>s" zLukeSelfAttention.__init__cCs6|dd|j|jf}|j|}|ddddS)NrQrrr )rXryr}viewpermute)rMxZ new_x_shaper$r$r%transpose_for_scoress z&LukeSelfAttention.transpose_for_scoresNFc Cs|d}|dkr|}ntj||gdd}|||}|||} |jr|dk r|||} |||} || |} || |} |ddddd|ddf}|ddddd|ddf}|dddd|dddf}|dddd|dddf}t | | dd}t | | dd}t | | dd}t | | dd}tj||gdd}tj||gdd}tj||gdd}n$|||}t || dd}|t |j}|dk r||}tjj|dd}||}|dk r||}t || }|dddd}|dd|jf}|j|}|ddd|ddf}|dkrd}n|dd|dddf}|r|||f}n||f}|S)NrrqrQrpr rr)rXr!catrrrrrrrrmatmulZ transposemathsqrtr}r functionalZsoftmaxrKr contiguousr~r) rMword_hidden_statesrattention_mask head_maskoutput_attentions word_sizeconcat_hidden_statesZ key_layerZ value_layerZw2w_query_layerZw2e_query_layerZe2w_query_layerZe2e_query_layerZ w2w_key_layerZ e2w_key_layerZ w2e_key_layerZ e2e_key_layerZw2w_attention_scoresZw2e_attention_scoresZe2w_attention_scoresZe2e_attention_scoresZword_attention_scoresZentity_attention_scoresZattention_scoresZ query_layerZattention_probsZ context_layerZnew_context_layer_shapeZoutput_word_hidden_statesZoutput_entity_hidden_statesoutputsr$r$r%rasV            zLukeSelfAttention.forward)NNF)rrrr>rrardr$r$rOr%rws   rwcs4eZdZfddZejejejdddZZS)LukeSelfOutputcsBtt|j|j|_tj|j|jd|_t|j |_ dSNr<) r=r>rrlrAdenserGrHrIrJrKrLrOr$r%r>ls zLukeSelfOutput.__init__r- input_tensorreturncCs&||}||}|||}|SNrrKrGrMr-rr$r$r%rars  zLukeSelfOutput.forwardrrrr>r!ZTensorrardr$r$rOr%rks rcs.eZdZfddZddZd ddZZS) LukeAttentioncs*tt||_t||_t|_dSr)r=r>rwrMroutputsetZ pruned_headsrLrOr$r%r>zs   zLukeAttention.__init__cCs tddSNz4LUKE does not support the pruning of attention headsNotImplementedError)rMZheadsr$r$r% prune_headsszLukeAttention.prune_headsNFcCs|d}||||||}|dkr2|d}|} n(tj|dddd}tj||gdd} ||| } | ddd|ddf} |dkrd} n| dd|dddf} | | f|dd} | S)Nrrrrq)rXrMr!rr)rMrrrrrrZ self_outputsZconcat_self_outputsrattention_outputZword_attention_outputZentity_attention_outputrr$r$r%ras(  zLukeAttention.forward)NNF)rrrr>rrardr$r$rOr%rys  rcs0eZdZfddZejejdddZZS)LukeIntermediatecsBtt|j|j|_t|jt r6t |j|_ n|j|_ dSr) r=r>rrlrAintermediate_sizer isinstance hidden_actstrr intermediate_act_fnrLrOr$r%r>s   zLukeIntermediate.__init__r-rcCs||}||}|Sr)rrrMr-r$r$r%ras  zLukeIntermediate.forwardrr$r$rOr%rs rcs4eZdZfddZejejejdddZZS) LukeOutputcsBtt|j|j|_tj|j|jd|_t |j |_ dSr) r=r>rrlrrArrGrHrIrJrKrLrOr$r%r>s zLukeOutput.__init__rcCs&||}||}|||}|Srrrr$r$r%ras  zLukeOutput.forwardrr$r$rOr%rs rcs.eZdZfddZd ddZddZZS) LukeLayercs:t|j|_d|_t||_t||_t||_ dSNr) r=r>chunk_size_feed_forward seq_len_dimr attentionr intermediaterrrLrOr$r%r>s    zLukeLayer.__init__NFc Cs|d}|j|||||d}|dkr0|d}ntj|dddd}|dd} t|j|j|j|} | ddd|ddf} |dkrd} n| dd|dddf} | | f| } | S)Nrrrrrq)rXrr!rrfeed_forward_chunkrr) rMrrrrrrZself_attention_outputsZconcat_attention_outputr layer_outputZword_layer_outputZentity_layer_outputr$r$r%ras0    zLukeLayer.forwardcCs||}|||}|Sr)rr)rMrZintermediate_outputrr$r$r%rs  zLukeLayer.feed_forward_chunk)NNF)rrrr>rarrdr$r$rOr%rs   %rcs&eZdZfddZdddZZS) LukeEncodercs:t|_tfddtjD|_d|_dS)Ncsg|] }tqSr$)r).0_rfr$r% sz(LukeEncoder.__init__..F) r=r>rNrZ ModuleListrangenum_hidden_layerslayergradient_checkpointingrLrOrfr%r>s  zLukeEncoder.__init__NFTcs|rdnd}|rdnd} r dnd} t|jD]\} } |rN||f}| |f} |dk r^|| nd} |jr|jrfdd}tjj|| |||| }n| |||| }|d}|dk r|d}r.| |df} q.|r||f}| |f} |s tdd||| || fDSt||| || d S) Nr$csfdd}|S)Ncs|fSrr$)inputs)modulerr$r%custom_forwardszJLukeEncoder.forward..create_custom_forward..custom_forwardr$)rrrrr%create_custom_forwardsz2LukeEncoder.forward..create_custom_forwardrrrcss|]}|dk r|VqdSrr$rvr$r$r% ;sz&LukeEncoder.forward..)last_hidden_stater-r.rr) enumeraterrZtrainingr!utils checkpointtupler&)rMrrrrroutput_hidden_states return_dictZall_word_hidden_statesZall_entity_hidden_statesZall_self_attentionsiZ layer_moduleZlayer_head_maskrZ layer_outputsr$rr%ras`          zLukeEncoder.forward)NNFFTrrrr>rardr$r$rOr%rs rcs0eZdZfddZejejdddZZS) LukePoolercs*tt|j|j|_t|_dSr)r=r>rrlrArZTanh activationrLrOr$r%r>Qs zLukePooler.__init__rcCs(|dddf}||}||}|S)Nr)rr)rMr-Zfirst_token_tensor pooled_outputr$r$r%raVs  zLukePooler.forwardrr$r$rOr%rPs rcs$eZdZfddZddZZS)EntityPredictionHeadTransformcsVtt|j|j|_t|jt r6t |j|_ n|j|_ tj |j|j d|_ dSr)r=r>rrlrArjrrrrr transform_act_fnrGrHrLrOr$r%r>`s   z&EntityPredictionHeadTransform.__init__cCs"||}||}||}|Sr)rrrGrr$r$r%rais   z%EntityPredictionHeadTransform.forwardrr$r$rOr%r_s rcs$eZdZfddZddZZS)EntityPredictionHeadcsHt||_t||_tj|j|jdd|_ t t |j|_ dS)NFrg)r=r>rNr transformrrlrjridecoder Parameterr!rYrhrLrOr$r%r>qs   zEntityPredictionHead.__init__cCs||}|||j}|Sr)rrrhrr$r$r%raxs zEntityPredictionHead.forwardrr$r$rOr%rps rc@s>eZdZdZeZdZdZddgZe j dddZ d d d Z d S)LukePreTrainedModelz An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models. lukeTrrercCst|tjr:|jjjd|jjd|jdk r|jj n~t|tj r|j dkr^|jj n|jjjd|jjd|j dk r|jj|j  n&t|tj r|jj |jjddS)zInitialize the weightsg)ZmeanZstdNr?)rrrlweightdataZnormal_rNZinitializer_rangerhZzero_r?Z embedding_dimr;rGZfill_)rMrr$r$r% _init_weightss       z!LukePreTrainedModel._init_weightsFcCst|tr||_dSr)rrr)rMrrr$r$r%_set_gradient_checkpointings z/LukePreTrainedModel._set_gradient_checkpointingN)F) rrrr r config_classZbase_model_prefixZsupports_gradient_checkpointingZ_no_split_modulesrModulerrr$r$r$r%rsra> This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads etc.) This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and behavior. Parameters: config ([`LukeConfig`]): 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. a> 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) entity_ids (`torch.LongTensor` of shape `(batch_size, entity_length)`): Indices of entity tokens in the entity vocabulary. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. entity_attention_mask (`torch.FloatTensor` of shape `(batch_size, entity_length)`, *optional*): Mask to avoid performing attention on padding entity token indices. Mask values selected in `[0, 1]`: - 1 for entity tokens that are **not masked**, - 0 for entity tokens that are **masked**. entity_token_type_ids (`torch.LongTensor` of shape `(batch_size, entity_length)`, *optional*): Segment token indices to indicate first and second portions of the entity token inputs. Indices are selected in `[0, 1]`: - 0 corresponds to a *portion A* entity token, - 1 corresponds to a *portion B* entity token. entity_position_ids (`torch.LongTensor` of shape `(batch_size, entity_length, max_mention_length)`, *optional*): Indices of positions of each input entity in the position embeddings. Selected in the range `[0, config.max_position_embeddings - 1]`. 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. 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**. 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. zThe bare LUKE model transformer outputting raw hidden-states for both word tokens and entities without any specific head on top.cseZdZdeedfdd ZddZddZd d Zd d Z d dZ e e deeeddeejeejeejeejeejeejeejeejeejeejeeeeeeeeefdddZejeejdddZZS) LukeModelT)rNadd_pooling_layercsNt|||_t||_t||_t||_|rrNr9r`rerkrencoderrpooler post_init)rMrNrrOr$r%r>s    zLukeModel.__init__cCs|jjSrr`rBrMr$r$r%get_input_embeddingsszLukeModel.get_input_embeddingscCs ||j_dSrrrMrr$r$r%set_input_embeddingsszLukeModel.set_input_embeddingscCs|jjSrrkrr$r$r%get_entity_embeddingsszLukeModel.get_entity_embeddingscCs ||j_dSrrrr$r$r%set_entity_embeddingsszLukeModel.set_entity_embeddingscCs tddSrr)rMZheads_to_pruner$r$r% _prune_headsszLukeModel._prune_headsbatch_size, sequence_length output_typerN)r\rr]r[rnentity_attention_maskentity_token_type_idsentity_position_idsrr^rrrrc Cs| dk r | n|jj} | dk r | n|jj} | dk r4| n|jj} |dk rV| dk rVtdn@|dk rt||||}n"| dk r| dd}ntd|\}}|dk r|jn| j}|dkrtj ||f|d}|dkrtj |tj |d}|dk r6|d}|dkrtj ||f|d}|dkr6tj ||ftj |d}| | |jj } |j|||| d}|||}|dkrtd}n||||}|j|||| | | | d }|d }|jdk r||nd}| s||f|ddSt|||j|j|j|jd S) u Returns: Examples: ```python >>> from transformers import AutoTokenizer, LukeModel >>> tokenizer = AutoTokenizer.from_pretrained("studio-ousia/luke-base") >>> model = LukeModel.from_pretrained("studio-ousia/luke-base") # Compute the contextualized entity representation corresponding to the entity mention "Beyoncé" >>> text = "Beyoncé lives in Los Angeles." >>> entity_spans = [(0, 7)] # character-based entity span corresponding to "Beyoncé" >>> encoding = tokenizer(text, entity_spans=entity_spans, add_prefix_space=True, return_tensors="pt") >>> outputs = model(**encoding) >>> word_last_hidden_state = outputs.last_hidden_state >>> entity_last_hidden_state = outputs.entity_last_hidden_state # Input Wikipedia entities to obtain enriched contextualized representations of word tokens >>> text = "Beyoncé lives in Los Angeles." >>> entities = [ ... "Beyoncé", ... "Los Angeles", ... ] # Wikipedia entity titles corresponding to the entity mentions "Beyoncé" and "Los Angeles" >>> entity_spans = [ ... (0, 7), ... (17, 28), ... ] # character-based entity spans corresponding to "Beyoncé" and "Los Angeles" >>> encoding = tokenizer( ... text, entities=entities, entity_spans=entity_spans, add_prefix_space=True, return_tensors="pt" ... ) >>> outputs = model(**encoding) >>> word_last_hidden_state = outputs.last_hidden_state >>> entity_last_hidden_state = outputs.entity_last_hidden_state ```NzDYou cannot specify both input_ids and inputs_embeds at the same timerQz5You have to specify either input_ids or inputs_embeds)rTrRr)r\r[r]r^)rrrrrr)r pooler_outputr-r.rr)rNrruse_return_dictr{Z%warn_if_padding_and_no_attention_maskrXrTr!ZonesrYrZZ get_head_maskrr`get_extended_attention_maskrkrrrr-r.rr)rMr\rr]r[rnrrrrr^rrrr_Z batch_sizeZ seq_lengthrTZentity_seq_lengthZword_embedding_outputextended_attention_maskZentity_embedding_outputZencoder_outputssequence_outputrr$r$r%rasp9          zLukeModel.forward)word_attention_maskrcCs|}|dk rtj||gdd}|dkrH|dddddddf}n8|dkrn|ddddddf}ntd|jd|j|jd}d |t|jj}|S) ac Makes broadcastable attention and causal masks so that future and masked tokens are ignored. Arguments: word_attention_mask (`torch.LongTensor`): Attention mask for word tokens with ones indicating tokens to attend to, zeros for tokens to ignore. entity_attention_mask (`torch.LongTensor`, *optional*): Attention mask for entity tokens with ones indicating tokens to attend to, zeros for tokens to ignore. Returns: `torch.Tensor` The extended attention mask, with a the same dtype as `attention_mask.dtype`. NrQrqr rz&Wrong shape for attention_mask (shape ))rSr) r!rrrr{shaperVrSZfinforo)rMrrrrr$r$r%rs  z%LukeModel.get_extended_attention_mask)T) NNNNNNNNNNNNN)rrrrboolr>rrrrrrLUKE_INPUTS_DOCSTRINGformatrr_CONFIG_FOR_DOCrr!rvr"rrrarrdr$r$rOr%rsR     rcCs2||}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 rrq)ner|r!ZcumsumrtrZ)r\r;maskZincremental_indicesr$r$r%rUs rUcs0eZdZdZfddZddZddZZS) LukeLMHeadz*Roberta Head for masked language modeling.csdtt|j|j|_tj|j|jd|_t|j|j |_ t t |j |_|j|j _dSr)r=r>rrlrArrGrH layer_normr@rrr!rYrhrLrOr$r%r>s  zLukeLMHead.__init__cKs*||}t|}||}||}|Sr)rr r r)rMfeatureskwargsrr$r$r%ras    zLukeLMHead.forwardcCs*|jjjjdkr|j|j_n |jj|_dS)Nmeta)rrhrTtyperr$r$r% _tie_weightss zLukeLMHead._tie_weights)rrrr r>rarrdr$r$rOr%r s  r z The LUKE model with a language modeling head and entity prediction head on top for masked language modeling and masked entity prediction. cseZdZdddgZfddZfddZdd Zd d Zee d e e e d deejeejeejeejeejeejeejeejeejeejeejeejeeeeeeeee fdddZZS)LukeForMaskedLMzlm_head.decoder.weightzlm_head.decoder.biasz!entity_predictions.decoder.weightcs@t|t||_t||_t||_t |_ | dSr) r=r>rrr lm_headrentity_predictionsrrloss_fnrrLrOr$r%r>s      zLukeForMaskedLM.__init__cs$t||jj|jjjdSr)r= tie_weightsZ_tie_or_clone_weightsrrrrkrrOr$r%rs zLukeForMaskedLM.tie_weightscCs|jjSrrrrr$r$r%get_output_embeddingssz%LukeForMaskedLM.get_output_embeddingscCs ||j_dSrr)rMZnew_embeddingsr$r$r%set_output_embeddingssz%LukeForMaskedLM.set_output_embeddingsrrN)r\rr]r[rnrrrlabels entity_labelsrr^rrrrcCs0|dk r |n|jj}|j||||||||| | | |dd }d}d}||j}| dk r| |j} ||d|jj | d}|dkr|}d}d}|j dk r| |j }| dk r||d|jj | d}|dkr|}n||}|st dd||||||j|j|jfDSt||||||j|j|jdS)aS 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]` entity_labels (`torch.LongTensor` of shape `(batch_size, entity_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]` Returns: NT r\rr]r[rnrrrrr^rrrrQcss|]}|dk r|VqdSrr$rr$r$r%rds z*LukeForMaskedLM.forward..)r(r)r*r+r,r-rr.)rNrrrrrVrTrrr@rrrirr-rr.r')rMr\rr]r[rnrrrrrrr^rrrrr(r)r+r*r,r$r$r%rasn"    zLukeForMaskedLM.forward)NNNNNNNNNNNNNNN)rrrZ_tied_weights_keysr>rrrrrrrr'rrr!rvr"rrrrardr$r$rOr%rsP     rz The LUKE model with a classification head on top (a linear layer on top of the hidden state of the first entity token) for entity classification tasks, such as Open Entity. 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 j e e j e e j e e j e ee ee eeeefdddZZS) LukeForEntityClassificationcsJt|t||_|j|_t|j|_t |j |j|_ | dSr r=r>rr num_labelsrrIrJrKrlrA classifierrrLrOr$r%r>s   z$LukeForEntityClassification.__init__rrNr\rr]r[rnrrrrr^rrrrrcCs|dk r |n|jj}|j||||||||| | | | dd }|jdddddf}||}||}d}| dk r| |j} | jdkrt j || }n t j | d| d|}|stdd|||j|j|jfDSt|||j|j|jd S) u labels (`torch.LongTensor` of shape `(batch_size,)` or `(batch_size, num_labels)`, *optional*): Labels for computing the classification loss. If the shape is `(batch_size,)`, the cross entropy loss is used for the single-label classification. In this case, labels should contain the indices that should be in `[0, ..., config.num_labels - 1]`. If the shape is `(batch_size, num_labels)`, the binary cross entropy loss is used for the multi-label classification. In this case, labels should only contain `[0, 1]`, where 0 and 1 indicate false and true, respectively. Returns: Examples: ```python >>> from transformers import AutoTokenizer, LukeForEntityClassification >>> tokenizer = AutoTokenizer.from_pretrained("studio-ousia/luke-large-finetuned-open-entity") >>> model = LukeForEntityClassification.from_pretrained("studio-ousia/luke-large-finetuned-open-entity") >>> text = "Beyoncé lives in Los Angeles." >>> entity_spans = [(0, 7)] # character-based entity span corresponding to "Beyoncé" >>> inputs = tokenizer(text, entity_spans=entity_spans, return_tensors="pt") >>> outputs = model(**inputs) >>> logits = outputs.logits >>> predicted_class_idx = logits.argmax(-1).item() >>> print("Predicted class:", model.config.id2label[predicted_class_idx]) Predicted class: person ```NTrrrrQcss|]}|dk r|VqdSrr$rr$r$r%rsz6LukeForEntityClassification.forward..r(r+r-rr.)rNrrrrKrrVrTndimrr cross_entropy binary_cross_entropy_with_logitsrrtrr-rr.r/rMr\rr]r[rnrrrrr^rrrrrfeature_vectorr+r(r$r$r%rasH.     z#LukeForEntityClassification.forward)NNNNNNNNNNNNNN)rrrr>rrrrr/rrr!rvr"rrrrardr$r$rOr%rsD   rz The LUKE model with a classification head on top (a linear layer on top of the hidden states of the two entity tokens) for entity pair classification tasks, such as TACRED. 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 j e e j e e j e e j e ee ee eeeefdddZZS) LukeForEntityPairClassificationcsPt|t||_|j|_t|j|_t |j d|jd|_ | dS)NrFrrLrOr$r%r>s   z(LukeForEntityPairClassification.__init__rrNrcCs|dk r |n|jj}|j||||||||| | | | dd }tj|jdddddf|jdddddfgdd}||}||}d}| dk r| |j } | j dkrt j || }n t j |d| d|}|stdd |||j|j|jfDSt|||j|j|jd S) u" labels (`torch.LongTensor` of shape `(batch_size,)` or `(batch_size, num_labels)`, *optional*): Labels for computing the classification loss. If the shape is `(batch_size,)`, the cross entropy loss is used for the single-label classification. In this case, labels should contain the indices that should be in `[0, ..., config.num_labels - 1]`. If the shape is `(batch_size, num_labels)`, the binary cross entropy loss is used for the multi-label classification. In this case, labels should only contain `[0, 1]`, where 0 and 1 indicate false and true, respectively. Returns: Examples: ```python >>> from transformers import AutoTokenizer, LukeForEntityPairClassification >>> tokenizer = AutoTokenizer.from_pretrained("studio-ousia/luke-large-finetuned-tacred") >>> model = LukeForEntityPairClassification.from_pretrained("studio-ousia/luke-large-finetuned-tacred") >>> text = "Beyoncé lives in Los Angeles." >>> entity_spans = [ ... (0, 7), ... (17, 28), ... ] # character-based entity spans corresponding to "Beyoncé" and "Los Angeles" >>> inputs = tokenizer(text, entity_spans=entity_spans, return_tensors="pt") >>> outputs = model(**inputs) >>> logits = outputs.logits >>> predicted_class_idx = logits.argmax(-1).item() >>> print("Predicted class:", model.config.id2label[predicted_class_idx]) Predicted class: per:cities_of_residence ```NTrrrrqrQcss|]}|dk r|VqdSrr$rr$r$r%r[sz:LukeForEntityPairClassification.forward..r )rNrrr!rrrKrrVrTr!rrr"r#rrtrr-rr.r1r$r$r$r%rasN1.     z'LukeForEntityPairClassification.forward)NNNNNNNNNNNNNN)rrrr>rrrrr1rrr!rvr"rrrrardr$r$rOr%r&sD   r&z The LUKE model with a span classification head on top (a linear layer on top of the hidden states output) for tasks such as named entity recognition. 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 j e e j e e j e e j e e j e e j e ee ee eeeefdddZZS) LukeForEntitySpanClassificationcsNt|t||_|j|_t|j|_t |j d|j|_ | dS)Nr rrLrOr$r%r>rs   z(LukeForEntitySpanClassification.__init__rrN)r\rr]r[rnrrrentity_start_positionsentity_end_positionsrr^rrrrrcCs|dk r |n|jj}|j||||||||| | ||dd }|jd}| ddd|} | j|jjkrt| |jj} t |jd| }| ddd|} | j|jjkr| |jj} t |jd| }t j |||j gdd}| |}||}d}| dk rT| |j} | jdkr4tj|d|j| d}n tj|d| d|}|s|tdd |||j|j|jfDSt|||j|j|jd S) u' entity_start_positions (`torch.LongTensor`): The start positions of entities in the word token sequence. entity_end_positions (`torch.LongTensor`): The end positions of entities in the word token sequence. labels (`torch.LongTensor` of shape `(batch_size, entity_length)` or `(batch_size, entity_length, num_labels)`, *optional*): Labels for computing the classification loss. If the shape is `(batch_size, entity_length)`, the cross entropy loss is used for the single-label classification. In this case, labels should contain the indices that should be in `[0, ..., config.num_labels - 1]`. If the shape is `(batch_size, entity_length, num_labels)`, the binary cross entropy loss is used for the multi-label classification. In this case, labels should only contain `[0, 1]`, where 0 and 1 indicate false and true, respectively. Returns: Examples: ```python >>> from transformers import AutoTokenizer, LukeForEntitySpanClassification >>> tokenizer = AutoTokenizer.from_pretrained("studio-ousia/luke-large-finetuned-conll-2003") >>> model = LukeForEntitySpanClassification.from_pretrained("studio-ousia/luke-large-finetuned-conll-2003") >>> text = "Beyoncé lives in Los Angeles" # List all possible entity spans in the text >>> word_start_positions = [0, 8, 14, 17, 21] # character-based start positions of word tokens >>> word_end_positions = [7, 13, 16, 20, 28] # character-based end positions of word tokens >>> entity_spans = [] >>> for i, start_pos in enumerate(word_start_positions): ... for end_pos in word_end_positions[i:]: ... entity_spans.append((start_pos, end_pos)) >>> inputs = tokenizer(text, entity_spans=entity_spans, return_tensors="pt") >>> outputs = model(**inputs) >>> logits = outputs.logits >>> predicted_class_indices = logits.argmax(-1).squeeze().tolist() >>> for span, predicted_class_idx in zip(entity_spans, predicted_class_indices): ... if predicted_class_idx != 0: ... print(text[span[0] : span[1]], model.config.id2label[predicted_class_idx]) Beyoncé PER Los Angeles LOC ```NTrrQrprrqcss|]}|dk r|VqdSrr$rr$r$r%rsz:LukeForEntitySpanClassification.forward..r )rNrrrrXrbrcrTrVr!gatherrrrKrr!rrr"rrr#rtrr-rr.r2)rMr\rr]r[rnrrrr(r)rr^rrrrrrAZ start_statesZ end_statesr%r+r(r$r$r%ra~sZA        z'LukeForEntitySpanClassification.forward)NNNNNNNNNNNNNNNN)rrrr>rrrrr2rrr!rvr"rrrrardr$r$rOr%r'jsL   r'z The LUKE Model transformer with a sequence classification/regression head on top (a linear layer on top of the pooled output) e.g. for GLUE tasks. cseZdZfddZeedeee e dd e e j e e je e j e e j e e j e e je e j e e j e e je e je e je ee ee eeee fdddZZS) LukeForSequenceClassificationcsZt||j|_t||_t|jdk r2|jn|j|_ t |j |j|_ | dSrr=r>rrrrrIclassifier_dropoutrJrKrlrArrrLrOr$r%r>s  z&LukeForSequenceClassification.__init__rrrrNrcCs|dk r |n|jj}|j||||||||| | | | dd }|j}||}||}d}| dk rJ| |j} |jjdkr|j dkrd|j_n4|j dkr| j t j ks| j t j krd|j_nd|j_|jjdkrt}|j dkr||| }n ||| }nN|jjdkr,t}||d|j | d}n|jjdkrJt}||| }|srtd d |||j|j|jfDSt|||j|j|jd S) a 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). NTrrZ regressionZsingle_label_classificationZmulti_label_classificationrQcss|]}|dk r|VqdSrr$rr$r$r%r^sz8LukeForSequenceClassification.forward..r )rNrrrrKrrVrTZ problem_typerrSr!rZr|r squeezerrrrr-rr.r3)rMr\rr]r[rnrrrrr^rrrrrrr+r(loss_fctr$r$r%rasd       "     z%LukeForSequenceClassification.forward)NNNNNNNNNNNNNN)rrrr>rrrr_CHECKPOINT_FOR_DOCr3rrr!rvr"rrrrardr$r$rOr%r+sL  r+z The LUKE Model with a token classification head on top (a linear layer on top of the hidden-states output). To solve Named-Entity Recognition (NER) task using LUKE, `LukeForEntitySpanClassification` is more suitable than this class. cseZdZfddZeedeee e dd e e j e e je e j e e j e e j e e je e j e e j e e je e je e je ee ee eeee fdddZZS) LukeForTokenClassificationcs^t||j|_t|dd|_t|jdk r6|jn|j|_ t |j |j|_ | dSNF)rr,rLrOr$r%r>vs z#LukeForTokenClassification.__init__rr.NrcCs|dk r |n|jj}|j||||||||| | | | dd }|j}||}||}d}| dk r| |j} t}|| d|j | d}|st dd|||j |j |jfDSt|||j |j |jdS)J labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): Labels for computing the multiple choice classification loss. Indices should be in `[0, ..., num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See `input_ids` above) NTrrQcss|]}|dk r|VqdSrr$rr$r$r%rsz5LukeForTokenClassification.forward..r )rNrrrrKrrVrTrrrrr-rr.r4)rMr\rr]r[rnrrrrr^rrrrrrr+r(r0r$r$r%rasF   z"LukeForTokenClassification.forward)NNNNNNNNNNNNNN)rrrr>rrrrr1r4rrr!rvr"rrrrardr$r$rOr%r2msL  r2z The LUKE Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear layers on top of the hidden-states output to compute `span start logits` and `span end logits`). cseZdZfddZeedeee e dd e e j e e je e j e e je e j e e je e j e e j e e je e je e j e e j e ee ee eeee fdddZZS) LukeForQuestionAnsweringcs@t||j|_t|dd|_t|j|j|_| dSr3) r=r>rrrrrlrA qa_outputsrrLrOr$r%r>s  z!LukeForQuestionAnswering.__init__rr.N)r\rr]r[rnrrrrr^start_positions end_positionsrrrrcCsJ|dk r |n|jj}|j||||||||| | | |dd }|j}||}|jddd\}}|d}|d}d}| dk r| dk rt| dkr| d} t| dkr| d} |d}| d|| d|t |d}||| }||| }||d }|s0t d d ||||j |j |jfDSt||||j |j |jd S) a start_positions (`torch.LongTensor` of shape `(batch_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 `(batch_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. NTrrrQrqr)Z ignore_indexrcss|]}|dk r|VqdSrr$rr$r$r%r.sz3LukeForQuestionAnswering.forward..)r(r6r7r-rr.)rNrrrr6splitr/lenrXZclamp_rrr-rr.r5)rMr\rr]r[rnrrrrr^r7r8rrrrrr+r6r7Z total_lossZ ignored_indexr0Z start_lossZend_lossr$r$r%rash"             z LukeForQuestionAnswering.forward)NNNNNNNNNNNNNNN)rrrr>rrrrr1r5rrr!rvr"rrrrardr$r$rOr%r5sP  r5z The LUKE Model with a multiple choice classification head on top (a linear layer on top of the pooled output and a softmax) e.g. for RocStories/SWAG tasks. cseZdZfddZeedeee e dd e e j e e je e j e e j e e j e e je e j e e j e e je e je e je ee ee eeee fdddZZS) LukeForMultipleChoicecsPt|t||_t|jdk r*|jn|j|_t |j d|_ | dSr) r=r>rrrrIr-rJrKrlrArrrLrOr$r%r>Ms  zLukeForMultipleChoice.__init__z(batch_size, num_choices, sequence_lengthr.NrcCs|dk r |n|jj}|dk r&|jdn| jd}|dk rJ|d|dnd}|dk rh|d|dnd}|dk r|d|dnd}|dk r|d|dnd}| dk r| d| d| dnd} |dk r|d|dnd}|dk r|d|dnd}|dk r(|d|dnd}|dk rP|d|d|dnd}|j||||||||| | | | dd }|j}||}||}|d|}d}| dk r| |j } t }||| }|st dd|||j |j|jfDSt|||j |j|jd S) r4NrrQrpTrcss|]}|dk r|VqdSrr$rr$r$r%rsz0LukeForMultipleChoice.forward..r )rNrrrrXrrrKrrVrTrrr-rr.r8)rMr\rr]r[rnrrrrr^rrrrZ num_choicesrrr+Zreshaped_logitsr(r0r$r$r%raYs~       zLukeForMultipleChoice.forward)NNNNNNNNNNNNNN)rrrr>rrrrr1r8rrr!rvr"rrrrardr$r$rOr%r;EsL  r;)Mr r dataclassesrtypingrrrr!Ztorch.utils.checkpointrZtorch.nnrrr Z activationsr r Zmodeling_outputsr rZmodeling_utilsrZ pytorch_utilsrrrrrrrrZconfiguration_lukerZ get_loggerrloggerrr1Z"LUKE_PRETRAINED_MODEL_ARCHIVE_LISTrr&r'r/r1r2r3r4r5r8rr9rerwrrrrrrrrrrZLUKE_START_DOCSTRINGrrrUr rrr&r'r+r2r5r;r$r$r$r%s       *!!$#I(r04U!KIlq jYp