UhjddlZddlmZddlmZmZmZmZmZddl Z ddl m cm Z ddlZ ddl m Z ddlmZmZmZddlmZddlmZdd lmZdd lmZmZmZmZmZdd lm Z dd l!m"Z"m#Z#m$Z$dd l%m&Z&ddl'm(Z(m)Z)e#rddl*m+Z+ddl,m-Z-ddl.m/Z/ne0Z-e$jbe2Z3GddeZ4 dPde jjde jjdee jjdee jjdee jje jje jje6ee jjee jjff dZ7de jjde jjde6de6de jjf dZ8Gdd e jrjtZ; dPd!ee jjd"ee6fd#Z<Gd$d%e-Z=Gd&d'e j|Z?Gd(d)e j|Z@Gd*d+e(ZA dQd,d-d.e jjde jjd/e jjdee jd0ee6e6fd1e6d2e6d3eeCdeee jje jjfee jjffd4ZDe jfd,d-d.e jjd5e=d!e jjd"e6d0ee6e6fd1e6d2e6d6e jdee jjfd7ZGd,d-d.e jjde jjd/e jjdee jd0ee6e6fd1e6d2e6dee jjfd8ZHeGeDeHd9ZIGd:d-e j|ZJGd;de ZLe"Gd?d@eLZMGdAdBe j|ZNe"dCDGdEdFeLZOe"dGDGdHdIeLZPe"dJDGdKdLeLZQe"GdMdNeLZRgdOZSy)RN) nullcontext)DictLiteralOptionalTupleUnion)nn)BCEWithLogitsLossCrossEntropyLossMSELoss)ACT2FN)PretrainedConfig)_prepare_4d_attention_mask)BaseModelOutputMaskedLMOutputQuestionAnsweringModelOutputSequenceClassifierOutputTokenClassifierOutput)PreTrainedModel)auto_docstringis_flash_attn_2_availablelogging)is_triton_available)GemmaRotaryEmbeddingapply_rotary_pos_emb) flash_attn_varlen_qkvpacked_func)RotaryEmbedding) apply_rotaryceZdZdZdZdgZ ddedffd ZfdZxZ S) ModernBertConfiga This is the configuration class to store the configuration of a [`ModernBertModel`]. It is used to instantiate an ModernBert model according to the specified arguments, defining the model architecture. Instantiating a configuration with the defaults will yield a similar configuration to that of the ModernBERT-base. e.g. [answerdotai/ModernBERT-base](https://huggingface.co/answerdotai/ModernBERT-base) Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the documentation from [`PretrainedConfig`] for more information. Args: vocab_size (`int`, *optional*, defaults to 50368): Vocabulary size of the ModernBert model. Defines the number of different tokens that can be represented by the `inputs_ids` passed when calling [`ModernBertModel`] hidden_size (`int`, *optional*, defaults to 768): Dimension of the hidden representations. intermediate_size (`int`, *optional*, defaults to 1152): Dimension of the MLP representations. num_hidden_layers (`int`, *optional*, defaults to 22): Number of hidden layers in the Transformer decoder. num_attention_heads (`int`, *optional*, defaults to 12): Number of attention heads for each attention layer in the Transformer decoder. hidden_activation (`str` or `function`, *optional*, defaults to `"gelu"`): The non-linear activation function (function or string) in the decoder. Will default to `"gelu"` if not specified. max_position_embeddings (`int`, *optional*, defaults to 8192): The maximum sequence length that this model might ever be used with. initializer_range (`float`, *optional*, defaults to 0.02): The standard deviation of the truncated_normal_initializer for initializing all weight matrices. initializer_cutoff_factor (`float`, *optional*, defaults to 2.0): The cutoff factor for the truncated_normal_initializer for initializing all weight matrices. norm_eps (`float`, *optional*, defaults to 1e-05): The epsilon used by the rms normalization layers. norm_bias (`bool`, *optional*, defaults to `False`): Whether to use bias in the normalization layers. pad_token_id (`int`, *optional*, defaults to 50283): Padding token id. eos_token_id (`int`, *optional*, defaults to 50282): End of stream token id. bos_token_id (`int`, *optional*, defaults to 50281): Beginning of stream token id. cls_token_id (`int`, *optional*, defaults to 50281): Classification token id. sep_token_id (`int`, *optional*, defaults to 50282): Separation token id. global_rope_theta (`float`, *optional*, defaults to 160000.0): The base period of the global RoPE embeddings. attention_bias (`bool`, *optional*, defaults to `False`): Whether to use a bias in the query, key, value and output projection layers during self-attention. attention_dropout (`float`, *optional*, defaults to 0.0): The dropout ratio for the attention probabilities. global_attn_every_n_layers (`int`, *optional*, defaults to 3): The number of layers between global attention layers. local_attention (`int`, *optional*, defaults to 128): The window size for local attention. local_rope_theta (`float`, *optional*, defaults to 10000.0): The base period of the local RoPE embeddings. embedding_dropout (`float`, *optional*, defaults to 0.0): The dropout ratio for the embeddings. mlp_bias (`bool`, *optional*, defaults to `False`): Whether to use bias in the MLP layers. mlp_dropout (`float`, *optional*, defaults to 0.0): The dropout ratio for the MLP layers. decoder_bias (`bool`, *optional*, defaults to `True`): Whether to use bias in the decoder layers. classifier_pooling (`str`, *optional*, defaults to `"cls"`): The pooling method for the classifier. Should be either `"cls"` or `"mean"`. In local attention layers, the CLS token doesn't attend to all tokens on long sequences. classifier_dropout (`float`, *optional*, defaults to 0.0): The dropout ratio for the classifier. classifier_bias (`bool`, *optional*, defaults to `False`): Whether to use bias in the classifier. classifier_activation (`str`, *optional*, defaults to `"gelu"`): The activation function for the classifier. deterministic_flash_attn (`bool`, *optional*, defaults to `False`): Whether to use deterministic flash attention. If `False`, inference will be faster but not deterministic. sparse_prediction (`bool`, *optional*, defaults to `False`): Whether to use sparse prediction for the masked language model instead of returning the full dense logits. sparse_pred_ignore_index (`int`, *optional*, defaults to -100): The index to ignore for the sparse prediction. reference_compile (`bool`, *optional*): Whether to compile the layers of the model which were compiled during pretraining. If `None`, then parts of the model will be compiled if 1) `triton` is installed, 2) the model is not on MPS, 3) the model is not shared between devices, and 4) the model is not resized after initialization. If `True`, then the model may be faster in some scenarios. repad_logits_with_grad (`bool`, *optional*, defaults to `False`): When True, ModernBertForMaskedLM keeps track of the logits' gradient when repadding for output. This only applies when using Flash Attention 2 with passed labels. Otherwise output logits always have a gradient. Examples: ```python >>> from transformers import ModernBertModel, ModernBertConfig >>> # Initializing a ModernBert style configuration >>> configuration = ModernBertConfig() >>> # Initializing a model from the modernbert-base style configuration >>> model = ModernBertModel(configuration) >>> # Accessing the model configuration >>> configuration = model.config ``` modernbertpast_key_valuesclassifier_poolingclsmeanc$ "t%|d| || ||d|$||_||_||_||_||_||_||_| |_ | |_ | |_ ||_ ||_ ||_||_||_||_||_||_||_||_||_||_||_||_||_||_| |_|!|_|"|_|#|_|j.dvrtAd|j.dy)N) pad_token_id bos_token_id eos_token_id cls_token_id sep_token_idr&zQInvalid value for `classifier_pooling`, should be either "cls" or "mean", but is .)!super__init__ vocab_sizemax_position_embeddings hidden_sizeintermediate_sizenum_hidden_layersnum_attention_headsinitializer_rangeinitializer_cutoff_factornorm_eps norm_biasglobal_rope_thetaattention_biasattention_dropouthidden_activationglobal_attn_every_n_layerslocal_attentionlocal_rope_thetaembedding_dropoutmlp_bias mlp_dropout decoder_biasr%classifier_dropoutclassifier_biasclassifier_activationdeterministic_flash_attnsparse_predictionsparse_pred_ignore_indexreference_compilerepad_logits_with_grad ValueError)&selfr3r5r6r7r8r@r4r9r:r;r<r*r,r+r-r.r=r>r?rArBrCrDrErFrGr%rHrIrJrKrLrMrNrOkwargs __class__s& /var/www/catia.catastroantioquia-mas.com/valormas/lib/python3.12/site-packages/transformers/models/modernbert/modular_modernbert.pyr2zModernBertConfig.__init__sTN  %%%%%    %'>$&!2!2#6 !2)B&  "!2,!2!2*D'. 0!2  &("4"4.%:"(@%!2(@%!2&<#  " "/ 9cdhd{d{c||}~  :cHt|}|jdd|S)NrN)r1to_dictpop)rQoutputrSs rTrWzModernBertConfig.to_dicts#" &- rU)#iii gelui g{Gz?@gh㈵>Fikjir_r^gAFr @r`Fr`Tr'r`Fr\FFiNF) __name__ __module__ __qualname____doc__ model_typekeys_to_ignore_at_inferencerr2rW __classcell__rSs@rTr"r"5se NJ#4"5 $"%"#$ 5:$!&!%$IQ8$M29QfrUr"inputsattention_mask position_idslabelsreturnc|jdtj}tj|j dj }t |j j}tjjjtj|dtjd}|jdk(r|j |}n*|j^} } } | | z} |j| g| |}||j |nd} ||j |nd}||||| |fS) a Remove padding from input sequences. Args: inputs: (batch, seqlen, ...) or (batch, seqlen) attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. position_ids: (batch, seqlen), int, position ids labels: (batch, seqlen), int, labels Returns: unpadded_inputs: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask. indices: (total_nnz) cu_seqlens: (batch + 1), the cumulative sequence lengths max_seqlen_in_batch: int unpadded_position_ids: (total_nnz) or None unpadded_labels: (total_nnz) or None dimdtypeF)as_tupler)rrN)sumtorchint32nonzeroflattenintmaxitemr functionalpadcumsumrsshapeview)rkrlrmrnseqlens_in_batchindicesmax_seqlen_in_batch cu_seqlensunpadded_inputsbatchseqlenrestrunpadded_position_idsunpadded_labelss rT_unpad_modernbert_inputrs,.&))b )DmmN224uEMMOG.22499;<$$((6FAUZU`U`)acijJ zz|q ..*73%||v%&++e3d3G<?K?WL0027;]a393Efnn&w/4O GZ1DF[]l llrUrrrcl|jdk(rHtj||z|j|j}|||<|j ||}|S|j ^}}tj||zg||j|jd}|||<|j ||g|}|S)aQ Add padding to sequences. Args: inputs: (total_nnz, ...) or (total_nnz,), where total_nnz = number of tokens selected in attention_mask. indices: (total_nnz) batch: int, batch size seqlen: int, max sequence length Returns: padded_inputs: (batch, seqlen, ...) or (batch, seqlen) rv)rtdevice)rsrxzerosrtrrr)rkrrrrY padded_inputs_rs rT_pad_modernbert_outputr"s$zz|qUV^6<< V w E62   <<DUV^]d]&,,v}}] w# E69D9 rUc\eZdZe ddeej deefdZedZ y)ApplyRotaryEmbUnpadNr max_seqlenc |j}|j\}}}} |ddddfj|d| } t| ||d||dd|j |||||_|S)NrrqrFT)seqlen_offsetsrr interleavedinplace) contiguousrrr save_for_backwardr) ctxqkvcossinrr total_nnz_three_nheadsheaddimqks rTforwardzApplyRotaryEmbUnpad.forwardBsnn.1ii+ 67GBQBZ__YG 4   !!  c3 3# rUc |j\}}}|j}|j\}}}}|ddddfj|d|} t | ||d||j ddd |ddddddfS)NrrqrFT)rrrrr conjugate) saved_tensorsrrrr r) rdorrrrrrrdqks rTbackwardzApplyRotaryEmbUnpad.backwardas"00S* ]]_.0hh+ 67GBQBinnYG4   !~~ 4tT455rUNN) rcrdre staticmethodrrxTensorr|rrr0rUrTrrAsQ .2$(  U\\*  SM <66rUrrrc4tj|||||S)a Arguments: qkv: (total_nnz, 3, nheads, headdim) - input tensor for packed QKV. cos, sin: (seqlen_rotary, rotary_dim / 2) interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of 1st half and 2nd half (GPT-NeoX style). inplace: if True, apply rotary embedding in-place. seqlen_offsets: (batch_size,) or int. Each sequence in x is shifted by this amount. Most commonly used in inference when we have KV cache. cu_seqlens: (batch + 1,) or None max_seqlen: int Return: out: (total_nnz, dim) rotary_dim must be <= headdim Apply rotary embedding to the first rotary_dim of x. )rapply)rrrrrs rTapply_rotary_unpaddedrxs.  $ $S#sJ KKrUc "eZdZdZ d dededeedeejdeejf fd Z ddejd ejdeed e eje ejejfffd Zd efd ZxZS)!ModernBertUnpaddedRotaryEmbeddingzP The rotary position embeddings applied directly to unpadded sequences. rsbaserrrtcvt|||d|d||_||||j|||yyyy)a max_seqlen: if max_seqlen, device, and dtype are provided, we precompute the cos_sin_cache up to max_seqlen. If the max_seqlen, device, or dtype during training/inference differ, the cos_sin_cache will be recomputed during the forward pass. TF)rsrpos_idx_in_fp32rrNrrt)r1r2r_update_cos_sin_cache)rQrsrrrrtrSs rTr2z*ModernBertUnpaddedRotaryEmbedding.__init__sV StT&^cd$  !f&8U=N  & &z& & N>O&8 !rUrrroc|(|j||j|jt||j|j ||}|S)z Apply rotary embedding *inplace* to qkv. qkv: (total_nnz, 3, nheads, headdim) cu_seqlens: (batch + 1,) cumulative sequence lengths max_seqlen: int max seq length in the batch rrr)rrrtr _cos_cached _sin_cached)rQrrrs rTrz)ModernBertUnpaddedRotaryEmbedding.forwardsS  !  & &z#**CII & V#       !!   rUcTd|jd|jd|jS)Nzdim=z, base=z , scale_base=)rsr scale_baserQs rT extra_reprz,ModernBertUnpaddedRotaryEmbedding.extra_reprs(dhhZwtyykt>OPPrU)rbNNNN)rcrdrerfr|floatrrxrrtr2rrrrstrrrirjs@rTrrs$()-'+ O OOSM O & O  $ O.%)  \\LLSM  u||U5<<#=>> ? 2QCQrUrceZdZdZdeffd ZejddejdejfdZ d de ejd e ejdejfd Z xZ S) ModernBertEmbeddingszV Same as BertEmbeddings with a tiny tweak for positional embeddings indexing. configcdt|||_tj|j |j |j|_tj|j |j|j|_ tj|j|_y)N) padding_idxepsbias)r1r2rr Embeddingr3r5r*tok_embeddings LayerNormr;r<normDropoutrDdroprQrrSs rTr2zModernBertEmbeddings.__init__sw   ll6+<+P>P^d^q^qrLL!3!3vO_O_` JJv778 rUTdynamic input_idsroc`|j|j|j|Sr)rrr)rQrs rTcompiled_embeddingsz(ModernBertEmbeddings.compiled_embeddingss%yy4#6#6y#ABCCrU inputs_embedsc|"|j|j|}|S|jjr|j |n.|j|j|j |}|Sr)rrrrNrr)rQrr hidden_statess rTrzModernBertEmbeddings.forwardsu  $ IIdii &>?M;;00((3YYtyy)<)rUr target_dtypec  ||||}|jtjtjfv} | rb|j} |j |}t ||||j r |jnd|j|} | j | } n3t ||||j r |jnd|j|} | j||fS)Nrr`)rr dropout_p deterministic window_size) rtrxfloat16bfloat16rrrr?rKr) rrrrrrBrrsrr convert_dtype orig_dtypeattns rTflash_attention_forwardr$s SZJ GCIIemmU^^%DDMYY ff\"/ !!28//f..s 99'  wwz"/ !!28//f..s 99'   IIb#   rUc v|j||\} } |jddjd\} } } t| | | | \} } |dk7r|}t j | | | |j r |jnd|jddj}|j|d |}|fS) Nrr rvrrrr`)r attn_maskrq) rrrrFscaled_dot_product_attentionrr?rr)rrrlrrmrBrrsrrrrrr r s rTsdpa_attention_forwardrOs  < @HC a+22q29E3%eS#s;JE3(", &&   28//f..s$   1a ""2r3/K >rU)flash_attention_2eagersdpaczeZdZdZd dedeeffd Z d dejdee dejfdZ xZ S) raPerforms multi-headed self attention on a batch of unpadded sequences. If Flash Attention 2 is installed, this module uses Flash Attention to improve throughput. If Flash Attention 2 is not installed, the implementation will use PyTorch's SDPA kernel, which requires padding and unpadding inputs, adding some overhead. See `forward` method for additional details. rlayer_idct|||_||_|j|j zdk7r&t d|jd|j d|j|_|j|_|j |_ |j|j z|_ |j|jz|_ tj|jd|jz|j|_||j zdk7r$|j"dz|j"dzf|_nd|_|j$}|j&}|j"dk7r$|j( |j(}|j"}|j*d k(rt-|j|| |_nt1||j| |_tj|j|j|j|_|jd kDrtj4|jntj6|_t;|_y) NrzThe hidden size (z6) is not a multiple of the number of attention heads ()r rrrr)rsrr)rrsrr`)r1r2rr"r5r8rPr?rK num_headsr all_head_sizer rr>WqkvrArBr=r4rC_attn_implementationrrrrrIdentityout_dropset pruned_heads)rQrr" rope_thetar4rSs rTr2zModernBertAttention.__init__s        : : :a ?#F$6$6#77mntoIoInJJKL "(!9!9(.(G(G%33**f.H.HH !]]T^^;IIf00!d6H6H2HvOdOde f77 71 <$*$:$:a$?AWAW[\A\#]D #+D -- "("@"@   8 +&&2#44 &,&<&< #  & &*= =?MM.EJDO8v4==_ijDO))F..0B0BI^I^_@F@X@X[^@^ 6#;#;3  3rUcfeZdZddedeeffd ZejddejdejfdZ ddejd eejd eejd eejd eejd eedee dejfdZ xZS)ModernBertEncoderLayerrr"ct|||_|dk(rtj|_n;tj |j|j|j|_t|||_ tj |j|j|j|_ t||_y)Nrr)rr")r1r2rr r) attn_normrr5r;r<rrmlp_normrmlprQrr"rSs rTr2zModernBertEncoderLayer.__init__s  q=[[]DN\\&*<*<&//X^XhXhiDN'vI  V%7%7V__SYScScd  (rUTrrrocB|j|j|Sr)r8r7rQrs rT compiled_mlpz#ModernBertEncoderLayer.compiled_mlpsxx m455rUrlrrmrrrc  |j|j|||||||}||dz}|jjr|j |n|j |j |} || z}|f|ddzS)Nrlrrmrrrrrv)rr6rrNr<r8r7) rQrrlrrmrrrr0 mlp_outputs rTrzModernBertEncoderLayer.forwardsyy NN= )) 3%!!/! & Q7 {{,,   m ,$-- 67  & 2 ,qr"222rUr)NNNNNF)rcrdrer"rr|r2rxrrr<rr2rrirjs@rTr4r4s )/ )8C= )U]]4 6%,,65<<6!6 266:37-1$(,13||3!.3&ell3 3 u//0 3 U\\* 3SM3$D>3 3rUr4c eZdZeZdZdZddgZdZdZ dZ de jfdZ e dded eej$d eeeeeeffd effd Zd ZfdZxZS)ModernBertPreTrainedModelmodelTrr4Frc|jjddtjdtffd }|jj |jj t jd|jjzz |jj |jjdzd}t|tr||j|dyt|tr-||j|d ||j|d yt|t r-||j"|d ||j|d yt|t$r||j&|d yt|t(r||j*|d yt|t,t.t0fr||j2|d yt|tj4rW|j6j8j;d |j<%|j<j8j?yyy) Nr rstdctjj|jd| |z|zt |tj r7|j *tjj|j yyy)Nr`)r(rDab)r init trunc_normal_weight isinstancerrzeros_)rrD cutoff_factors rT init_weightz.init_weightsq GG ! !  .3&#% " &")),;;*GGNN6;;/+-rUr]r)inout embedding final_outrQrOrPrRg?) rr:r Modulerr9mathsqrtr7r5rKrrrrrrr'ModernBertPredictionHeaddenseModernBertForMaskedLMdecoder#ModernBertForSequenceClassification ModernBertForTokenClassificationModernBertForQuestionAnswering classifierrrJdatafill_rzero_)rQrrNstdsrMs @rT _init_weightsz'ModernBertPreTrainedModel._init_weightss ==  M 0 0 0++//;;00499S4;;C`C`=`3aa6600$6   f2 3 --tK/@ A  .  4: .  4; /  3 4  T$Z 0  4; /  8 9  d5k 2  5 6 U 4   02RTr s  ))4 +< =  - MM   $ $S ){{&   &&('.rUuse_flash_attention_2 torch_dtype device_mapcheck_device_mapc|j,d|_ |j|tj|d|St |||tj||S#tt f$r d|_Y:wxYw)NrF)rdrehard_check_onlyrf)rcrdrerf)_attn_implementation_internal_check_and_enable_flash_attn_2rxrrP ImportErrorr1_autoset_attn_implementation)r'rrcrdrerfrSs rTrlz6ModernBertPreTrainedModel._autoset_attn_implementation0s  / / 73FF 0 <99 % )$)%5 :w3 "7 !- 4   , <7;4 NTzcResizing token embeddings with `torch.compile` is not supported. Falling back to non-compiled mode.F)r1resize_token_embeddingsrrNrsrt)rQargsrR model_embedsrSs rTrxz1ModernBertPreTrainedModel.resize_token_embeddingsps[w6GG ;; ( (L 8{{,,##y-2DKK )rU)FNNT)rcrdrer" config_classbase_model_prefixsupports_gradient_checkpointing_no_split_modules_supports_flash_attn_2_supports_sdpa_supports_flex_attnr rSrb classmethodr2rrxrtrrrr|rlrvrxrirjs@rTrArAs#L&*#/1IJ!N-)BII-)^',-1;?!%   $ ekk*  U3S#X#678     @B>  rUrAc!eZdZdeffd ZdZdZe ddee jdee jdee jdee jd ee jd ee jd ee jd ee d ee dee dee dee dee deee jdfeffdZde jde de jfdZxZS)ModernBertModelrc t||||_t||_t j t|jDcgc]}t||c}|_ t j|j|j|j|_d|_|j#ycc}w)NrF)r1r2rr embeddingsr ModuleListranger7r4layersrr5r;r< final_normgradient_checkpointing post_initr9s rTr2zModernBertModel.__init__s   .v6mmFKFLdLdFe f( #FH 5 f  ,,v'9'9vU[UeUef&+#  gsCc.|jjSrrrrs rTget_input_embeddingsz$ModernBertModel.get_input_embeddingss---rUc&||j_yrr)rQr s rTset_input_embeddingsz$ModernBertModel.set_input_embeddingss).&rUrrlrrmrrrr batch_sizeseq_lenroutput_hidden_states return_dictro.c j | | n|jj} | | n|jj} | | n|jj} |du|duz r t d| rdnd}| rdnd}|j ||j || ) '||jdd\ n|jdd\ | |jn |j}|(tj f|tj}d}|jjdk(rM||d}|0tj5t|| ^}}}}dddnQt|| ^}}}}n>|&tj | j!d }|j#|| \}}|j%|| }|j&D]t}| r||fz}|j(r/|j*r#|j-|j.||||||| }n|||||||| }|d }| s]t1|dkDsl||dfz}v| r||fz}|j3|}|r't5| }|t7  fd|D}| st7d|||fDSt9|||S#1swYxYw) sliding_window_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing attention on padding or far-away tokens. In ModernBert, only every few layers perform global attention, while the rest perform local attention. This mask is used to avoid attending to far-away tokens in the local attention layers when not using Flash Attention. indices (`torch.Tensor` of shape `(total_unpadded_tokens,)`, *optional*): Indices of the non-padding tokens in the input sequence. Used for unpadding the output. cu_seqlens (`torch.Tensor` of shape `(batch + 1,)`, *optional*): Cumulative sequence lengths of the input sequences. Used to index the unpadded tensors. max_seqlen (`int`, *optional*): Maximum sequence length in the batch excluding padding tokens. Used to unpad input_ids and pad output tensors. batch_size (`int`, *optional*): Batch size of the input sequences. Used to pad the output tensors. seq_len (`int`, *optional*): Sequence length of the input sequences including padding tokens. Used to pad the output tensors. Nz:You must specify exactly one of input_ids or inputs_embedsr0rrFrT)rkrl)rr)r)rrr>rvrkrrrc3<K|]}t|yw)rN)r).0hsrrrs rT z*ModernBertModel.forward.. s(*+"gZ`ghh*sc3&K|] }|| ywrr0)rvs rTrz*ModernBertModel.forward..smq_`_lms)last_hidden_stater attentions)rrruse_return_dictrPrv%warn_if_padding_and_no_attention_maskrrrxonesr2r(no_gradrarange unsqueeze_update_attention_maskrrrr_gradient_checkpointing_func__call__rrrrtupler)rQrrlrrmrrrrrrrrrall_hidden_statesall_self_attentionsrrepadrr encoder_layer layer_outputss ` `` rTrzModernBertModel.forwardsHB2C1N-TXT_T_TqTq$8$D $++JjJj &1%J((J2c |r|jjdk(r'tjdd|j_nF|jjdk7r-tjd|jjdt ||j }t j|jdjd}t j||jz }||jjdzkjdjdj|j}|j|j!t j"|j j$}||fS)Nr zOutputting attentions is only supported with the 'eager' attention implementation, not with "sdpa". Falling back to `attn_implementation="eager"`.rzZOutputting attentions is only supported with the eager attention implementation, not with zT. Consider setting `attn_implementation="eager"`. Setting `output_attentions=False`.rr)rr(rsrtrrtrxrrrabsTrBrr masked_fill logical_notfinfomin)rQrlrglobal_attention_maskrowsdistance window_maskrs rTrz&ModernBertModel._update_attention_masksS {{//69##V4; 011W<## $ @ @AB:: !;>4:: V||177:;EEaH99TDFF]+449 9 D DQ G Q QRS T W WXfXmXm n 4?? @W@W@Y[`[f[fgkgqgq[r[v[vw$&999rU NNNNNNNNNNNNN)rcrdrer"r2rrrrrxrrr|r2rrrrrrirjs@rTrr}s / ./15156:3704*.-1$($(!%,0/3&*D E,,-D !.D &ell3 D u//0 D   - D %,,'D U\\*D SMD SMD #D $D>D 'tnD d^D  uU\\3&'8 9D D L:U\\:VZ:_d_k_k:rUrc\eZdZdeffd Zdej dej fdZxZS)rVrcJt|||_tj|j |j |j |_t|j|_ tj|j |j|j|_y)Nr)r1r2rr rr5rIrWrrJrrr;r<rrs rTr2z!ModernBertPredictionHead.__init__8sq  YYv1163E3EvG]G]^ &667LL!3!3vO_O_` rUrroc`|j|j|j|Sr)rrrWr;s rTrz ModernBertPredictionHead.forward?s#yy$**]";<==rU) rcrdrer"r2rxrrrirjs@rTrVrV7s-a/a>U\\>ell>rUrVzd The ModernBert Model with a decoder head on top that is used for masked language modeling. ) custom_introc"8eZdZdgZdeffd ZdZdejfdZ e jdd e jd e jfd Z e dd ee j d ee jdee jdee jdee jdee jdee jdee jdeedeedeedeedeedeed eee jeffdZxZS)rXzdecoder.weightrctt||||_t||_t ||_tj|j|j|j|_ |jj|_ |jj|_|jy)Nr)r1r2rrrBrVheadr rr5r3rGrYrLrMrrs rTr2zModernBertForMaskedLM.__init__Ks   $V, ,V4 yy!3!3V5F5FVM`M`a !%!>!>(, (L(L% rUc|jSrrYrs rTget_output_embeddingsz+ModernBertForMaskedLM.get_output_embeddingsXs ||rUnew_embeddingsc||_yrr)rQrs rTset_output_embeddingsz+ModernBertForMaskedLM.set_output_embeddings[s % rUTrrYrocB|j|j|Sr)rYr)rQrYs rT compiled_headz#ModernBertForMaskedLM.compiled_head^s||DIIf-..rUrrlrrmrrnrrrrrrrrc H||n|jj}|j|jjdk(r||| | )| '||jdd\} } n|jdd\} } | |j n |j }|(t j| | f|t j}|4t j5t||||\}}}} }}dddnt||||\}}}} }}|j|||||||| | | | | | }|d}|jrK|I|jd}|j|jdd}||jk7}||}||}|jjr|j!|n|j#|j%|}d}|(|j'|||jj( }|jjdk(rN|jj*s| t-nt j5t/||| | }ddd|s |f}||f|zS|St1|||j2|j4 S#1swYxYw#1swYHxYw) rNrrr)rkrlrmrn rrlrrmrrrrrrrrrrrq)r3rlosslogitsrr)rrrvr(rrrxrr2rrrBrLrrMrNrrYr loss_functionr3rOrrrrr)rQrrlrrmrrnrrrrrrrrrRroutputsr mask_tokensrrrYs rTrzModernBertForMaskedLM.forwardbsF&1% K(F{{,,   0 1dii(9:;    %%ffAWAW%XD ;; + +/B B"&++"D"D\a\i\i\k r/vwV`ipq rYF)-)9TGf$ Ev E!//))   m^ r rsJ J JJ!NNNNNNNNNNNNNN)rcrdre_tied_weights_keysr"r2rr rrrxrrrrrrr|r2rrrrrirjs@rTrXrXCs ++ / &BII&U]]4 /ELL/U\\/!/15156:/304)-*.-1$($(!%,0/3&*m E,,-m !.m &ell3 m u||, m   - m &m %,,'m U\\*m SMm SMm #m $D>m 'tnm d^m " uU\\"N2 3#m m rUrXz` The ModernBert Model with a sequence classification head on top that performs pooling. c"eZdZdeffd Ze ddeejdeejdeejdeejdeejdeejd eejd eejd ee d ee d ee dee dee dee de e ejeffdZxZS)rZrcnt|||j|_||_t ||_t ||_tjj|j|_ tj|j|j|_|j!yr)r1r2 num_labelsrrrBrVrrxr rrHrrr5r]rrs rTr2z,ModernBertForSequenceClassification.__init__s   ++ $V, ,V4 HH$$V%>%>? ))F$6$68I8IJ rUrrlrrmrrnrrrrrrrrroc f||n|jj}|j|j|||||||| | | | | | }|d}|jjdk(r |dddf}nQ|jjdk(r8||j dzj d|j dd z }|j|}|j|}|j|}d}||jj|jdk(rd |j_ nl|jdkDrL|jtjk(s|jtjk(rd |j_ nd |j_ |jjd k(rIt!}|jdk(r&||j#|j#}n|||}n|jjd k(r=t%}||j'd|j|j'd}n,|jjd k(rt)}|||}|s |f}||f|zS|St+|||j,|j.S)aB sliding_window_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing attention on padding or far-away tokens. In ModernBert, only every few layers perform global attention, while the rest perform local attention. This mask is used to avoid attending to far-away tokens in the local attention layers when not using Flash Attention. 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). indices (`torch.Tensor` of shape `(total_unpadded_tokens,)`, *optional*): Indices of the non-padding tokens in the input sequence. Used for unpadding the output. cu_seqlens (`torch.Tensor` of shape `(batch + 1,)`, *optional*): Cumulative sequence lengths of the input sequences. Used to index the unpadded tensors. max_seqlen (`int`, *optional*): Maximum sequence length in the batch excluding padding tokens. Used to unpad input_ids and pad output tensors. batch_size (`int`, *optional*): Batch size of the input sequences. Used to pad the output tensors. seq_len (`int`, *optional*): Sequence length of the input sequences including padding tokens. Used to pad the output tensors. Nrrr'r(rqrvrT)rskeepdim regressionsingle_label_classificationmulti_label_classificationr)rrrvrBr%rrwrrr] problem_typerrtrxlongr|r squeezer rr rrr)rQrrlrrmrrnrrrrrrrrrRrr pooled_outputrrloss_fctrYs rTrz+ModernBertForSequenceClassification.forwardsN&1%e 'tne d^e " uU\\"$<< =#e e rUrZzv The ModernBert Model with a token classification head on top, e.g. for Named Entity Recognition (NER) tasks. c"eZdZdeffd Ze ddeejdeejdeejdeejdeejdeejd eejd eejd ee d ee d ee dee dee dee de e ejeffdZxZS)r[rc`t|||j|_t||_t ||_tjj|j|_ tj|j|j|_|jyrr1r2rrrBrVrrxr rrHrrr5r]rrs rTr2z)ModernBertForTokenClassification.__init__Us{   ++$V, ,V4 HH$$V%>%>? ))F$6$68I8IJ rUrrlrrmrrnrrrrrrrrroc||n|jj}|j|j|||||||| | | | | | }|d}|j |}|j |}|j |}d}|I 'tnI d^I uU\\"$99 :!I I rUr[c"eZdZdeffd Ze ddeejdeejdeejdeejdeejdeejd eejd eejd ee d ee d ee dee dee dee de e eje ffdZxZS)r\rc`t|||j|_t||_t ||_tjj|j|_ tj|j|j|_|jyrrrs rTr2z'ModernBertForQuestionAnswering.__init__sy   ++$V, ,V4 HH$$V%>%>? ))F$6$68I8IJ rUrrlrrmstart_positions end_positionsrrrrrrrrroc T||n|jj}|j|j||||||| | | | | | }|d}|j |}|j |}|j |}|jdd\}}|jdj}|jdj}d}|||j||||fi|}|s||f|ddz}||f|zS|St||||j|jS)rN) rlrrmrrrrrrrrrrvrqr)r start_logits end_logitsrr)rrrvrBrrr]splitrrrrrr)rQrrlrrmrrrrrrrrrrrRrrrrrrrYs rTrz&ModernBertForQuestionAnswering.forwardshF&1%K 'tnK d^K " uU\\"$@@ A#K K rUr\)r"rrArXrZr[r\rr1)TrT contextlibrtypingrrrrrrxtorch.nn.functionalr rrtorch.utils.checkpointtorch.nnr r r activationsrconfiguration_utilsrmodeling_attn_mask_utilsrmodeling_outputsrrrrrmodeling_utilsrutilsrrrutils.import_utilsrgemma.modeling_gemmarrflash_attn.flash_attn_interfacerflash_attn.layers.rotaryrflash_attn.ops.triton.rotaryr object get_loggerrcrsr"rr|rrautogradFunctionrrrrSrrrrr2r rrtrrr/rr4rArrVrXrZr[r\__all__r0rUrTrs  "88 AA!3B.GG5MP89O   H %A'AN,0%) &m LL&mLL&m5<<(&m U\\ " &m  5<<u||S(5<<:PRZ[`[g[gRh hi &mR LL \\     \\ >46%..1146v*. $ L& L  L42Q2Qj299<:BII:(_ 4_)." !" "LL" " 5++, " 38_ " " " ~" 5u||+ ,eELL.A AB"\!&(! !(! (!2(! (!  (! 38_ (! (! (!++(! 5<<(!V  !   LL   5++,  38_      5<< H1 $ "!M3"))M3`+3RYY+3\BBBJv:/v:v:r >ryy > H 5H  H V t *Ct  t n W '@W  W tX %>X X v rU