o iD@sDdZddlZddlmZddlmZmZddlZddlmZddl m Z ddl m Z dd l mZdd lmZmZmZdd lmZdd lmZmZmZdd lmZmZmZmZddlmZe e!Z"eeddGdddeZ#Gdddej$Z%Gdddej$Z&Gdddej$Z'Gdddej$Z(Gdddej$Z)Gdddej$Z*Gd d!d!ej$Z+Gd"d#d#ej$Z,Gd$d%d%ej$Z-Gd&d'd'eZ.Gd(d)d)ej$Z/Gd*d+d+ej$Z0Gd,d-d-ej$Z1eGd.d/d/eZ2eGd0d1d1e2Z3eGd2d3d3e2Z4ed4dGd5d6d6e2Z5ed7dGd8d9d9e2Z6gd:Z7dS);zPyTorch Bros model.N) dataclass)OptionalUnion)nn)CrossEntropyLoss)ACT2FN)GradientCheckpointingLayer)"BaseModelOutputWithCrossAttentions,BaseModelOutputWithPoolingAndCrossAttentionsTokenClassifierOutput)PreTrainedModel)apply_chunking_to_forward find_pruneable_heads_and_indicesprune_linear_layer) ModelOutputauto_docstringcan_return_tuplelogging) BrosConfigz@ Base class for outputs of token classification models. ) custom_introc@steZdZUdZdZeejed<dZ eejed<dZ eejed<dZ ee ejed<dZ ee ejed<dS)BrosSpadeOutputa loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): Classification loss. initial_token_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`): Classification scores for entity initial tokens (before SoftMax). subsequent_token_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, sequence_length+1)`): Classification scores for entity sequence tokens (before SoftMax). Nlossinitial_token_logitssubsequent_token_logits hidden_states attentions)__name__ __module__ __qualname____doc__rrtorch FloatTensor__annotations__rrrtuplerr&r&Z/home/ubuntu/.local/lib/python3.10/site-packages/transformers/models/bros/modeling_bros.pyr)s  rc2eZdZfddZdejdejfddZZS)BrosPositionalEmbedding1Dcs@t|j|_ddtd|jd|j}|d|dS)Nri'g@inv_freq)super__init__dim_bbox_sinusoid_emb_1dr"arangeregister_buffer)selfconfigr+ __class__r&r'r-Cs z"BrosPositionalEmbedding1D.__init__pos_seqreturncCsX|}|\}}}||||d|jddd|jd}tj||gdd}|S)Nrdim)sizeviewr+r.r"catsincos)r1r5seq_sizeb1b2b3 sinusoid_inppos_embr&r&r'forwardMs  (z!BrosPositionalEmbedding1D.forwardrrr r-r"TensorrF __classcell__r&r&r3r'r)@s  r)cr()BrosPositionalEmbedding2Dcs*t|j|_t||_t||_dSN)r,r-dim_bboxr) x_pos_emb y_pos_embr1r2r3r&r'r-Vs  z"BrosPositionalEmbedding2D.__init__bboxr6cCsdg}t|jD]!}|ddkr|||d|fq|||d|fqtj|dd}|S)Nr7r.r8r9)rangerLappendrMrNr"r=)r1rPstacki bbox_pos_embr&r&r'rF]s z!BrosPositionalEmbedding2D.forwardrGr&r&r3r'rJUs rJcs,eZdZfddZdejfddZZS)BrosBboxEmbeddingscs.tt||_tj|j|jdd|_dS)NF)bias) r,r-rJbbox_sinusoid_embrLineardim_bbox_sinusoid_emb_2ddim_bbox_projectionbbox_projectionrOr3r&r'r-is  zBrosBboxEmbeddings.__init__rPcCs\|dd}|dddddddf|dddddddf}||}||}|S)Nrr) transposerXr\)r1rPbbox_tbbox_posrUr&r&r'rFns 8  zBrosBboxEmbeddings.forwardrGr&r&r3r'rVhs rVc sbeZdZdZfddZ    d deejdeejdeejdeejd ejf d d ZZ S) BrosTextEmbeddingszGConstruct the embeddings from word, position and token_type embeddings.csttj|j|j|jd|_t|j|j|_ t|j |j|_ tj |j|j d|_ t|j|_t|dd|_|dt|jd|jdtj|jtj|jjdd d dS) N) padding_idxepsposition_embedding_typeabsolute position_ids)rr8token_type_idsdtypedeviceF) persistent)r,r-r Embedding vocab_size hidden_size pad_token_idword_embeddingsmax_position_embeddingsposition_embeddingstype_vocab_sizetoken_type_embeddings LayerNormlayer_norm_epsDropouthidden_dropout_probdropoutgetattrrdr0r"r/expandzerosrfr;longrjrOr3r&r'r-zs"  zBrosTextEmbeddings.__init__N input_idsrgrf inputs_embedsr6c Cs|dur |}n|dd}|d}|dur$|jddd|f}|durNt|drC|jddd|f}||d|}|}n tj|tj|jjd}|durW| |}| |} || } |j dkrn| |} | | 7} | | } || } | S)Nr8rrgrrhre)r;rfhasattrrgr{r"r|r}rjrprtrdrrrury) r1r~rgrfr input_shape seq_lengthbuffered_token_type_ids buffered_token_type_ids_expandedrt embeddingsrrr&r&r'rFs,        zBrosTextEmbeddings.forward)NNNN) rrr r!r-rr"rHrFrIr&r&r3r'r`ws$ r`cszeZdZfddZ     ddejdejdeejdeejd eejd eejd eejd eejfd dZZ S)BrosSelfAttentioncst|j|jdkrt|dstd|jd|jd|j|_t|j|j|_|j|j|_t |j|j|_ t |j|j|_ t |j|j|_ t |j|_t|dd|_|jdksf|jd krw|j|_t d |jd |j|_|j|_dS) Nrembedding_sizezThe hidden size (z6) is not a multiple of the number of attention heads ()rdre relative_keyrelative_key_queryr7r)r,r-rnnum_attention_headsr ValueErrorintattention_head_size all_head_sizerrYquerykeyvaluerwattention_probs_dropout_probryrzrdrqrldistance_embedding is_decoderrOr3r&r'r-s&   zBrosSelfAttention.__init__NFrrUattention_mask head_maskencoder_hidden_statesencoder_attention_maskoutput_attentionsr6cCs|jdd|j|jf}|||dd} |du} | r8|||dd} |||dd} |}n|||dd} |||dd} t | | dd} |j dksd|j dkr| d}tj |tj |jddd}tj |tj |jddd}||}|||jd}|j| jd }|j dkrtd | |}| |} n|j dkrtd | |}td | |}| ||} | j\}}}}|||||}|gd }td | |f}| |} | t|j} |dur| |} tjdd| }||}|dur ||}t || }|dddd}| dd|jf}|j|}|r5||fn|f}|jr@|d}|S)Nrr8rr7rrrh)rizbhld,lrd->bhlrzbhrd,lrd->bhlr)r7rrrzbnid,bijd->bnijr9rrK)shaperrrr<r]rrr"matmulrdr;r/r}rjrrqtorieinsumpermutemathsqrtrSoftmaxry contiguousrr)r1rrUrrrrr hidden_shape query_layeris_cross_attention key_layer value_layerattention_scoresrposition_ids_lposition_ids_rdistancepositional_embeddingrelative_position_scoresrelative_position_scores_queryrelative_position_scores_key batch_sizen_headd_headbbox_pos_scoresattention_probs context_layernew_context_layer_shapeoutputsr&r&r'rFsX          zBrosSelfAttention.forwardNNNNF) rrr r-r"rHrr%rFrIr&r&r3r'rs0  rc8eZdZfddZdejdejdejfddZZS)BrosSelfOutputcsBtt|j|j|_tj|j|jd|_t|j |_ dSNrb) r,r-rrYrndenserurvrwrxryrOr3r&r'r-& zBrosSelfOutput.__init__r input_tensorr6cC&||}||}|||}|SrKrryrur1rrr&r&r'rF,  zBrosSelfOutput.forwardrGr&r&r3r'r% $rcseZdZfddZddZ     ddejdejd eejd eejd eejd eejd eede ejfddZ Z S) BrosAttentioncs*tt||_t||_t|_dSrK)r,r-rr1routputset pruned_headsrOr3r&r'r-4s    zBrosAttention.__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)Nrrr9)lenrr1rrrrrrrrrrunion)r1headsindexr&r&r' prune_heads:s zBrosAttention.prune_headsNFrrUrrrrrr6c Cs>|j|||||||d}||d|} | f|dd} | S)NrrUrrrrrrr)r1r) r1rrUrrrrr self_outputsattention_outputrr&r&r'rFOs  zBrosAttention.forwardr) rrr r-rr"rHrboolr%rFrIr&r&r3r'r3s2  rcr()BrosIntermediatecsDtt|j|j|_t|jt rt |j|_ dS|j|_ dSrK) r,r-rrYrnintermediate_sizer isinstance hidden_actstrrintermediate_act_fnrOr3r&r'r-is   zBrosIntermediate.__init__rr6cCs||}||}|SrK)rr)r1rr&r&r'rFqs  zBrosIntermediate.forwardrGr&r&r3r'rhs rcr) BrosOutputcsBtt|j|j|_tj|j|jd|_t |j |_ dSr) r,r-rrYrrnrrurvrwrxryrOr3r&r'r-xrzBrosOutput.__init__rrr6cCrrKrrr&r&r'rF~rzBrosOutput.forwardrGr&r&r3r'rwrrcseZdZfddZ     ddejdejdeejdeejd eejd eejd eed e ejfd dZ ddZ Z S) BrosLayercsnt|j|_d|_t||_|j|_|j|_|jr+|js&t|dt||_ t ||_ t ||_ dS)Nrz> should be used as a decoder model if cross attention is added)r,r-chunk_size_feed_forward seq_len_dimr attentionradd_cross_attention Exceptioncrossattentionr intermediaterrrOr3r&r'r-s    zBrosLayer.__init__NFrrUrrrrrr6c Cs|j|||||d}|d} |jr|dd} n|dd} |jrI|durIt|dr2td|d|j| |||||d} | d} | | dd} t|j|j|j| } | f| } |jr_| d } | S) N)rUrrrrrr8rz'If `encoder_hidden_states` are passed, z` has to be instantiated with cross-attention layers by setting `config.add_cross_attention=True`)rrrrrrK) rrrrrrfeed_forward_chunkrr) r1rrUrrrrrself_attention_outputsrrcross_attention_outputs layer_outputr&r&r'rFsH     zBrosLayer.forwardcCs||}|||}|SrK)rr)r1rintermediate_outputrr&r&r'rs  zBrosLayer.feed_forward_chunkr) rrr r-r"rHrr#rr%rFrrIr&r&r3r'rs2  8rcseZdZfddZe       ddejdejdeejd eejd eejd eejd ee d ee dee de e eje ffddZ ZS) BrosEncodercs4t|_tfddtjD|_dS)Ncsg|]}tqSr&)r).0_r2r&r' sz(BrosEncoder.__init__..)r,r-r2r ModuleListrQnum_hidden_layerslayerrOr3rr'r-s $zBrosEncoder.__init__NFTrrUrrrrroutput_hidden_states return_dictr6c  Cs|rdnd} |r dnd} |r|jjrdnd} t|jD]8\} }|r&| |f} |dur.|| nd}||||||||d}|d}|rS| |df} |jjrS| |df} q|r[| |f} t|| | | dS)Nr&rrrr7)last_hidden_staterrcross_attentions)r2r enumeraterr )r1rrUrrrrrrrall_hidden_statesall_self_attentionsall_cross_attentionsrT layer_modulelayer_head_mask layer_outputsr&r&r'rFs<    zBrosEncoder.forward)NNNNFFT)rrr r-rr"rHrr#rrr%r rFrIr&r&r3r'rs>    rcr() BrosPoolercs*tt|j|j|_t|_dSrK)r,r-rrYrnrTanh activationrOr3r&r'r- s zBrosPooler.__init__rr6cCs(|dddf}||}||}|S)Nr)rr)r1rfirst_token_tensor pooled_outputr&r&r'rFs  zBrosPooler.forwardrGr&r&r3r'r s rcr()BrosRelationExtractorcst|j|_|j|_|j|_|j|_t|j|_ t |j|j|j|_ t |j|j|j|_ t td|j|_dS)Nr)r,r- n_relationsrnbackbone_hidden_sizehead_hidden_sizeclassifier_dropout_probrrwdroprYrr Parameterr"r| dummy_noderOr3r&r'r-s zBrosRelationExtractor.__init__rrc Cs|||}|jdd|dd}tj||gdd}|||}| |d|d|j |j }| |d|d|j |j }t | dddd| dddd}|S)Nrraxisr7r)rrr unsqueezerepeatr;r"r=rr<rrrr)r1rr dummy_vecrelation_scorer&r&r'rF)s zBrosRelationExtractor.forwardrGr&r&r3r'rs rc@s*eZdZUeed<dZdejfddZdS)BrosPreTrainedModelr2brosmodulecCs|jj}t|tjr"|jjjd|d|jdur |jj dSdSt|tj rC|jjjd|d|j durA|jj|j  dSdSt|tj rX|jj |jj ddSt|trhtjj|j|ddSdS)zInitialize the weightsr*)meanstdNg?)r)r2initializer_rangerrrYweightdatanormal_rWzero_rlrarufill_rinitr)r1rrr&r&r' _init_weightsAs"       z!BrosPreTrainedModel._init_weightsN) rrr rr$base_model_prefixrModulerr&r&r&r'r <s r cseZdZdfdd 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 j de e de e de e de ee j effddZZS) BrosModelTcsNt|||_t||_t||_t||_|rt |nd|_ | dS)zv add_pooling_layer (bool, *optional*, defaults to `True`): Whether to add a pooling layer N) r,r-r2r`rrVbbox_embeddingsrencoderrpooler init_weights)r1r2add_pooling_layerr3r&r'r-Ws     zBrosModel.__init__cCs|jjSrKrrp)r1r&r&r'get_input_embeddingsgszBrosModel.get_input_embeddingscCs ||j_dSrKr")r1rr&r&r'set_input_embeddingsjs zBrosModel.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)r1heads_to_prunerrr&r&r' _prune_headsmszBrosModel._prune_headsNr~rPrrgrfrrrrrrrr6c  Cs*| dur| n|jj} | dur| n|jj} | dur| n|jj} |dur*|dur*td|dur3|} n|dur@|dd} ntd|durLtd| \}}|durW|jn|j}|duretj| |d}|durt |j dr|j j ddd|f}| ||}|}n tj | tj|d}||| |}|jjr|dur|\}}}||f}| durtj||d} || }nd}|||jj}|j ||||d }|jdd kr|ddddgd f}||jj}||}|j||||||| | d d }|d}|jdur||nd}t|||j|j|jdS)a bbox ('torch.FloatTensor' of shape '(batch_size, num_boxes, 4)'): Bounding box coordinates for each token in the input sequence. Each bounding box is a list of four values (x1, y1, x2, y2), where (x1, y1) is the top left corner, and (x2, y2) is the bottom right corner of the bounding box. Examples: ```python >>> import torch >>> from transformers import BrosProcessor, BrosModel >>> processor = BrosProcessor.from_pretrained("jinho8345/bros-base-uncased") >>> model = BrosModel.from_pretrained("jinho8345/bros-base-uncased") >>> encoding = processor("Hello, my dog is cute", add_special_tokens=False, return_tensors="pt") >>> bbox = torch.tensor([[[0, 0, 1, 1]]]).repeat(1, encoding["input_ids"].shape[-1], 1) >>> encoding["bbox"] = bbox >>> outputs = model(**encoding) >>> last_hidden_states = outputs.last_hidden_state ```NzDYou cannot specify both input_ids and inputs_embeds at the same timer8z5You have to specify either input_ids or inputs_embedszYou have to specify bbox)rjrgrh)r~rfrgr)rrr7rr7rrrT)rUrrrrrrrr)r pooler_outputrrr)r2rruse_return_dictrr;rjr"onesrrrgr{r|r}get_extended_attention_maskrinvert_attention_mask get_head_maskrr bbox_scalerrrr rrr)r1r~rPrrgrfrrrrrrrrrrrjrrextended_attention_maskencoder_batch_sizeencoder_sequence_lengthrencoder_hidden_shapeencoder_extended_attention_maskembedding_output scaled_bboxbbox_position_embeddingsencoder_outputssequence_outputrr&r&r'rFus|(       zBrosModel.forward)T NNNNNNNNNNNN)rrr r-r#r$r'rrrr"rHrrr%r rFrIr&r&r3r'rUs\     rceZdZdgZfddZee            ddeej deej deej deej d eej d eej d eej d eej d eej dee dee dee de e ej e ffddZZS)BrosForTokenClassificationrcs^t||j|_t||_t|dr|jn|j}t ||_ t |j |j|_ |dSNclassifier_dropout)r,r- num_labelsrrrr>rxrrwryrYrn classifierr r1r2r>r3r&r'r-s    z#BrosForTokenClassification.__init__Nr~rPrbbox_first_token_maskrgrfrrlabelsrrrr6c  Cs| dur| n|jj} |j|||||||| | dd } | d}||}||}d}| durXt}|durK|d}||d|j|| d|}n ||d|j| d}t||| j | j dS)a bbox ('torch.FloatTensor' of shape '(batch_size, num_boxes, 4)'): Bounding box coordinates for each token in the input sequence. Each bounding box is a list of four values (x1, y1, x2, y2), where (x1, y1) is the top left corner, and (x2, y2) is the bottom right corner of the bounding box. bbox_first_token_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to indicate the first token of each bounding box. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. Examples: ```python >>> import torch >>> from transformers import BrosProcessor, BrosForTokenClassification >>> processor = BrosProcessor.from_pretrained("jinho8345/bros-base-uncased") >>> model = BrosForTokenClassification.from_pretrained("jinho8345/bros-base-uncased") >>> encoding = processor("Hello, my dog is cute", add_special_tokens=False, return_tensors="pt") >>> bbox = torch.tensor([[[0, 0, 1, 1]]]).repeat(1, encoding["input_ids"].shape[-1], 1) >>> encoding["bbox"] = bbox >>> outputs = model(**encoding) ```NT) rPrrgrfrrrrrrr8rlogitsrr) r2r*rryr@rr<r?r rr)r1r~rPrrBrgrfrrrCrrrrr9rErloss_fctr&r&r'rFs>-   z"BrosForTokenClassification.forwardr:rrr "_keys_to_ignore_on_load_unexpectedr-rrrr"rHrrr%r rFrIr&r&r3r'r<sX       r<a Bros Model with a token classification head on top (initial_token_layers and subsequent_token_layer on top of the hidden-states output) e.g. for Named-Entity-Recognition (NER) tasks. The initial_token_classifier is used to predict the first token of each entity, and the subsequent_token_classifier is used to predict the subsequent tokens within an entity. Compared to BrosForTokenClassification, this model is more robust to serialization errors since it predicts next token from one token. c!seZdZdgZfddZee             ddeej deej deej deej d eej d eej d eej d eej d eej deej dee dee dee de e ej e ffddZZS)!BrosSpadeEEForTokenClassificationrc st|||_|j|_|j|_|j|_t||_t |dr"|j n|j }t t |t |j|jt |t |j|j|_t||_|dSr=)r,r-r2r?rrnrrrrr>rxr SequentialrwrYinitial_token_classifierrsubsequent_token_classifierr rAr3r&r'r-hs    z*BrosSpadeEEForTokenClassification.__init__Nr~rPrrBrgrfrrinitial_token_labelssubsequent_token_labelsrrrr6c Cs| dur| n|jj} |j|||||||| | dd }|d}|dd}||dd}|||d}d|}|j\}}|j }t j |t |dg |gdd}||dddddft |jj}t ||dj |t jd}||dddddft |jj}|d}d}| dur| durt}| d} |dur|d}||d|j|| |}n ||d|j| }| d} ||d|d|| |}||}t||||j|jd S) a> bbox ('torch.FloatTensor' of shape '(batch_size, num_boxes, 4)'): Bounding box coordinates for each token in the input sequence. Each bounding box is a list of four values (x1, y1, x2, y2), where (x1, y1) is the top left corner, and (x2, y2) is the bottom right corner of the bounding box. bbox_first_token_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to indicate the first token of each bounding box. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. initial_token_labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for the initial token classification. subsequent_token_labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for the subsequent token classification. Examples: ```python >>> import torch >>> from transformers import BrosProcessor, BrosSpadeEEForTokenClassification >>> processor = BrosProcessor.from_pretrained("jinho8345/bros-base-uncased") >>> model = BrosSpadeEEForTokenClassification.from_pretrained("jinho8345/bros-base-uncased") >>> encoding = processor("Hello, my dog is cute", add_special_tokens=False, return_tensors="pt") >>> bbox = torch.tensor([[[0, 0, 1, 1]]]).repeat(1, encoding["input_ids"].shape[-1], 1) >>> encoding["bbox"] = bbox >>> outputs = model(**encoding) ```NT r~rPrrgrfrrrrrrrrrjrir8)rrrrr)r2r*rr]rrKrLsqueezerrjr"r=r|rr masked_fillfinforimineyer<rr?rrr)r1r~rPrrBrgrfrrrMrNrrrrlast_hidden_statesrrinv_attention_maskrmax_seq_lengthrjinvalid_token_maskself_token_masksubsequent_token_maskrrFinitial_token_losssubsequent_token_lossr&r&r'rFsj2  &     z)BrosSpadeEEForTokenClassification.forward) NNNNNNNNNNNNN)rrr rHr-rrrr"rHrrr%rrFrIr&r&r3r'rI\s^      rIz Bros Model with a token classification head on top (a entity_linker layer on top of the hidden-states output) e.g. for Entity-Linking. The entity_linker is used to predict intra-entity links (one entity to another entity). cr;)!BrosSpadeELForTokenClassificationrcsxt|||_|j|_|j|_|j|_t||_t |dr"|j n|j t ||_ |dSt ||_ |dSr=)r,r-r2r?rrnrrrrr>rxr entity_linkerr rOr3r&r'r-s      z*BrosSpadeELForTokenClassification.__init__Nr~rPrrBrgrfrrrCrrrr6c  CsR| dur| n|jj} |j|||||||| | dd } | d}|dd}|||d}d}| durt}|j\}}|j }t ||dj |t j d}|d}t j|t j|dgt j |dgdd }||dddddft |jj}||dddddft |jj}||d|d|| d|}t||| j| jd S) a bbox ('torch.FloatTensor' of shape '(batch_size, num_boxes, 4)'): Bounding box coordinates for each token in the input sequence. Each bounding box is a list of four values (x1, y1, x2, y2), where (x1, y1) is the top left corner, and (x2, y2) is the bottom right corner of the bounding box. bbox_first_token_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to indicate the first token of each bounding box. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. Examples: ```python >>> import torch >>> from transformers import BrosProcessor, BrosSpadeELForTokenClassification >>> processor = BrosProcessor.from_pretrained("jinho8345/bros-base-uncased") >>> model = BrosSpadeELForTokenClassification.from_pretrained("jinho8345/bros-base-uncased") >>> encoding = processor("Hello, my dog is cute", add_special_tokens=False, return_tensors="pt") >>> bbox = torch.tensor([[[0, 0, 1, 1]]]).repeat(1, encoding["input_ids"].shape[-1], 1) >>> encoding["bbox"] = bbox >>> outputs = model(**encoding) ```NTrOrrrPr8rhrrD)r2r*rr]rr_rQrrrjr"rUrrr<r=r|rRrSrirTr rr)r1r~rPrrBrgrfrrrCrrrrrVrErrFrrXrjrZmaskr&r&r'rF sL,   (($z)BrosSpadeELForTokenClassification.forwardr:rGr&r&r3r'r^sX      r^)r rr<rIr^)8r!r dataclassesrtypingrrr"rtorch.nnr activationsrmodeling_layersr modeling_outputsr r r modeling_utilsr pytorch_utilsrrrutilsrrrrconfiguration_brosr get_loggerrloggerrrr)rJrVr`rrrrrrrrrr rr<rIr^__all__r&r&r&r'sd        Am5M9""d  o