Ë
    ÛÓ^jâ  ã                   ó`  — d Z ddlmZ ddlZddlZddlmZ ddlmZm	Z	m
Z
 ddlmZ ddlmZ dd	lmZmZmZmZmZmZ dd
lmZ ddlmZmZmZ ddlmZ  ej>                  e «      Z!dZ" G d„ dejF                  «      Z$ G d„ dejF                  «      Z%dejL                  de'de'dejL                  fd„Z( G d„ dejF                  «      Z) G d„ dejF                  «      Z* G d„ dejF                  «      Z+ G d„ dejF                  «      Z,	 dFd ejL                  d!e'd"e'd#e-d$e-dejL                  fd%„Z. G d&„ d'ejF                  «      Z/ G d(„ d)ejF                  «      Z0e G d*„ d+e«      «       Z1 G d,„ d-ejF                  «      Z2 ed.¬/«      e G d0„ d1e«      «       «       Z3 ed2¬/«       G d3„ d4e1«      «       Z4e G d5„ d6e1«      «       Z5 ed7¬/«       G d8„ d9e1«      «       Z6e G d:„ d;e1«      «       Z7 ed<¬/«       G d=„ d>e1«      «       Z8e G d?„ d@e1«      «       Z9e G dA„ dBe1«      «       Z:e G dC„ dDe1«      «       Z;g dE¢Z<y)Gz!PyTorch Funnel Transformer model.é    )Ú	dataclassN)Únn)ÚBCEWithLogitsLossÚCrossEntropyLossÚMSELossé   )Úinitialization)ÚACT2FN)ÚBaseModelOutputÚMaskedLMOutputÚMultipleChoiceModelOutputÚQuestionAnsweringModelOutputÚSequenceClassifierOutputÚTokenClassifierOutput)ÚPreTrainedModel)ÚModelOutputÚauto_docstringÚloggingé   )ÚFunnelConfigg    €„.Ac                   óˆ   ‡ — e Zd Zdeddfˆ fd„Z	 ddej                  dz  dej                  dz  dej                  fd„Zˆ xZS )	ÚFunnelEmbeddingsÚconfigÚreturnNc                 ó@  •— t         ‰| �  «        t        j                  |j                  |j
                  |j                  ¬«      | _        t        j                  |j                  |j                  ¬«      | _        t        j                  |j                  «      | _        y )N)Úpadding_idx©Úeps)ÚsuperÚ__init__r   Ú	EmbeddingÚ
vocab_sizeÚhidden_sizeÚpad_token_idÚword_embeddingsÚ	LayerNormÚd_modelÚlayer_norm_epsÚ
layer_normÚDropoutÚhidden_dropoutÚdropout©Úselfr   Ú	__class__s     €úu/var/www/ramen.bs-engineer-server.com/venv/lib/python3.12/site-packages/transformers/models/funnel/modeling_funnel.pyr    zFunnelEmbeddings.__init__-   sh   ø€ Ü‰ÑÔÜ!Ÿ|™|¨F×,=Ñ,=¸v×?QÑ?QÐ_e×_rÑ_rÔsˆÔÜŸ,™, v§~¡~¸6×;PÑ;PÔQˆŒÜ—z‘z &×"7Ñ"7Ó8ˆ�ó    Ú	input_idsÚinputs_embedsc                 óp   — |€| j                  |«      }| j                  |«      }| j                  |«      }|S ©N)r%   r)   r,   )r.   r2   r3   Ú
embeddingss       r0   ÚforwardzFunnelEmbeddings.forward3   s<   € ð Ð Ø ×0Ñ0°Ó;ˆMØ—_‘_ ]Ó3ˆ
Ø—\‘\ *Ó-ˆ
ØÐr1   ©NN©	Ú__name__Ú
__module__Ú__qualname__r   r    ÚtorchÚTensorr7   Ú__classcell__©r/   s   @r0   r   r   ,   sQ   ø„ ð9˜|ð 9°õ 9ð [_ñØŸ™¨Ñ,ðØDIÇLÁLÐSWÑDWðà	�‰÷r1   r   c                   ó˜  ‡ — e Zd ZU dZdZeed<   deddfˆ fd„Z	 	 d de	j                  d	e	j                  dz  d
e	j                  dz  dee	j                     fd„Zd
e	j                  de	j                  fd„Zdede	j                  de	j                  dee	j                     eee	j                        z  fd„Zde	j                  defd„Zd!de	j                  dedede	j                  fd„Zde	j                  ee	j                     z  ee	j                     z  deee   z  ee   z  de	j                  fd„Z	 d"de	j                  ee	j                     z  ee	j                     z  dedede	j                  fd„Zdee	j                     dee	j                  ee	j                     f   fd„Zdee	j                     dee	j                     fd„Zˆ xZS )#ÚFunnelAttentionStructurez>
    Contains helpers for `FunnelRelMultiheadAttention `.
    é   Úcls_token_type_idr   r   Nc                 óÎ   •— t         ‰| �  «        || _        t        j                  |j
                  «      | _        t        j                  |j
                  «      | _        d | _        y r5   )	r   r    r   r   r*   r+   Úsin_dropoutÚcos_dropoutÚpooling_multr-   s     €r0   r    z!FunnelAttentionStructure.__init__D   sM   ø€ Ü‰ÑÔØˆŒÜŸ:™: f×&;Ñ&;Ó<ˆÔÜŸ:™: f×&;Ñ&;Ó<ˆÔð !ˆÕr1   r3   Úattention_maskÚtoken_type_idsc                 ób  — d| _         |j                  d«      x| _        }| j                  ||j                  |j
                  «      }|�| j                  |«      nd}| j                  j                  r7t        j                  j                  |j                  |dz
  |dz
  g«      d«      nd}||||fS )zCReturns the attention inputs associated to the inputs of the model.r   N)r   r   r   r   )rH   ÚsizeÚseq_lenÚget_position_embedsÚdtypeÚdeviceÚtoken_type_ids_to_matr   Úseparate_clsr   Ú
functionalÚpadÚnew_ones)r.   r3   rI   rJ   rM   Úposition_embedsÚtoken_type_matÚcls_masks           r0   Úinit_attention_inputsz.FunnelAttentionStructure.init_attention_inputsM   s³   € ð ˆÔØ!.×!3Ñ!3°AÓ!6Ð6ˆŒ�wØ×2Ñ2°7¸M×<OÑ<OÐQ^×QeÑQeÓfˆØGUÐGa˜×3Ñ3°NÔCÐgkˆð �{‰{×'Ò'ô �M‰M×Ñ˜m×4Ñ4°gÀ±kÀ7ÈQÁ;Ð5OÓPÐR^Ô_àð 	ð
   °ÀÐJÐJr1   c                 ó‚   — |dd…dd…df   |dd…df   k(  }|| j                   k(  }|dd…dd…df   |dd…df   z  }||z  S )z-Convert `token_type_ids` to `token_type_mat`.N)rD   )r.   rJ   rW   Úcls_idsÚcls_mats        r0   rQ   z.FunnelAttentionStructure.token_type_ids_to_mata   sY   € à'ªª1¨d¨
Ñ3°~ÂaÈÀgÑ7NÑNˆà  D×$:Ñ$:Ñ:ˆØš!šQ ˜*Ñ%¨²°4°Ñ(8Ñ8ˆØ˜Ñ'Ð'r1   rM   rO   rP   c                 óz  — | j                   j                  }| j                   j                  dk(  �rEt        j                  d|dt        j
                  |¬«      j                  |«      }t        j                  d|dz  dt        j
                  |¬«      j                  |«      }dd||dz  z  z  z  }|dd…df   |d   z  }t        j                  |«      }	| j                  |	«      }
t        j                  |«      }| j                  |«      }t        j                  |
|
gd	¬
«      }t        j                  ||	gd	¬
«      }t        j                  ||gd	¬
«      }t        j                  |	 |gd	¬
«      }||||fS t        j                  d|dz  dt        j
                  |¬«      j                  |«      }dd||dz  z  z  z  }t        j                  | dz  |dz  dt        j
                  |¬«      j                  |«      }|dz  }|dd…df   |d   z  }| j                  t        j                  |«      «      }	| j                  t        j                  |«      «      }t        j                  |	|gd	¬
«      }t        j                  d|t        j
                  |¬«      j                  |«      }|}g }t        d| j                   j                  «      D ]í  }|dk(  rd}ns| j                  ||«      }d|dz
  z  }| j                  |||d¬«      }|dd…df   |z   }|j!                  |j#                  d«      |«      }t        j$                  |d|«      }|}d|z  }| j                  ||«      }|dd…df   |z   }|j!                  |j#                  d«      |«      }t        j$                  |d|«      }|j'                  ||g«       Œï |S )a  
        Create and cache inputs related to relative position encoding. Those are very different depending on whether we
        are using the factorized or the relative shift attention:

        For the factorized attention, it returns the matrices (phi, pi, psi, omega) used in the paper, appendix A.2.2,
        final formula.

        For the relative shift attention, it returns all possible vectors R used in the paper, appendix A.2.1, final
        formula.

        Paper link: https://huggingface.co/papers/2006.03236
        Ú
factorizedr   ç      ð?©rO   rP   rC   r   i'  Néÿÿÿÿ©Údim)Úshift)r   r'   Úattention_typer=   ÚarangeÚint64ÚtoÚsinrF   ÚcosrG   ÚcatÚrangeÚ
num_blocksÚstride_pool_posÚrelative_posÚexpandrL   ÚgatherÚappend)r.   rM   rO   rP   r'   Úpos_seqÚfreq_seqÚinv_freqÚsinusoidÚ	sin_embedÚsin_embed_dÚ	cos_embedÚcos_embed_dÚphiÚpsiÚpiÚomegaÚ
rel_pos_idÚzero_offsetÚ	pos_embedÚposÚ
pooled_posÚposition_embeds_listÚblock_indexÚposition_embeds_poolingÚstrideÚrel_posÚposition_embeds_no_poolings                               r0   rN   z,FunnelAttentionStructure.get_position_embedsi   s|  € ð —+‘+×%Ñ%ˆØ�;‰;×%Ñ%¨Ó5ô —l‘l 1 g¨s¼%¿+¹+ÈfÔU×XÑXÐY^Ó_ˆGÜ—|‘| A w°!¡|°SÄÇÁÐTZÔ[×^Ñ^Ð_dÓeˆHØ˜E h°'¸Q±,Ñ&?Ñ@ÑAˆHØšq $˜wÑ'¨(°4©.Ñ8ˆHÜŸ	™	 (Ó+ˆIØ×*Ñ*¨9Ó5ˆKÜŸ	™	 (Ó+ˆIØ×*Ñ*¨9Ó5ˆKä—)‘)˜[¨+Ð6¸BÔ?ˆCÜ—)‘)˜Y¨	Ð2¸Ô;ˆCÜ—‘˜K¨Ð5¸2Ô>ˆBÜ—I‘I 	˜z¨9Ð5¸2Ô>ˆEØ˜˜S %Ð(Ð(ô —|‘| A w°!¡|°SÄÇÁÐTZÔ[×^Ñ^Ð_dÓeˆHØ˜E h°'¸Q±,Ñ&?Ñ@ÑAˆHäŸ™ w h°¡l°G¸a±KÀÌEÏKÉKÐ`fÔg×jÑjÐkpÓqˆJØ! A™+ˆKØ!¢! T 'Ñ*¨X°d©^Ñ;ˆHØ×(Ñ(¬¯©°8Ó)<Ó=ˆIØ×(Ñ(¬¯©°8Ó)<Ó=ˆIÜŸ	™	 9¨iÐ"8¸bÔAˆIä—,‘,˜q '´·±ÀVÔL×OÑOÐPUÓVˆCØˆJØ#%Ð Ü$ Q¨¯©×(>Ñ(>Ó?ò c�ð  !Ò#Ø.2Ñ+à!%×!5Ñ!5°c¸;Ó!G�Jð  ;°¡?Ñ3�FØ"×/Ñ/°°V¸ZÈqÐ/ÓQ�GØ%¢a¨ gÑ.°Ñ<�GØ%Ÿn™n¨W¯\©\¸!«_¸gÓF�GÜ.3¯l©l¸9ÀaÈÓ.QÐ+ð !�Ø˜K™�Ø×+Ñ+¨C°Ó8�à!¢! T 'Ñ*¨[Ñ8�Ø!Ÿ.™.¨¯©°a«¸'ÓB�Ü-2¯\©\¸)ÀQÈÓ-PÐ*à$×+Ñ+Ð-GÐI`Ð,aÕbð9cð: (Ð'r1   Úpos_idr…   c                 óì   — | j                   j                  rW|j                  d|z   dz   g«      }| j                   j                  r|dd n|dd }t	        j
                  ||ddd…   gd«      S |ddd…   S )ze
        Pool `pos_id` while keeping the cls token separate (if `config.separate_cls=True`).
        rC   r   ra   Nr   )r   rR   Ú
new_tensorÚtruncate_seqr=   rk   )r.   rŠ   r…   Úcls_posÚpooled_pos_ids        r0   rn   z(FunnelAttentionStructure.stride_pool_pos¹   s€   € ð �;‰;×#Ò#ð
 ×'Ñ'¨1¨k©>Ð):¸QÑ)>Ð(?Ó@ˆGØ,0¯K©K×,DÒ,D˜F 1 R™LÈ&ÐQRÐQSÈ*ˆMÜ—9‘9˜g }±S°q°SÑ'9Ð:¸AÓ>Ð>à™#˜A˜#‘;Ðr1   r‚   r‡   rd   c                 óÖ   — |€|}|d   |d   z
  }||j                   d   z  }|||z  z   }|d   |d   z
  }t        j                  ||dz
  | t        j                  |j                  ¬«      S )zV
        Build the relative positional vector between `pos` and `pooled_pos`.
        r   ra   r   r`   )Úshaper=   rf   ÚlongrP   )	r.   r‚   r‡   rƒ   rd   Ú	ref_pointÚ
num_removeÚmax_distÚmin_dists	            r0   ro   z%FunnelAttentionStructure.relative_posÈ   s   € ð ÐØˆJà˜q‘M C¨¡FÑ*ˆ	Ø˜Z×-Ñ-¨aÑ0Ñ0ˆ
Ø˜z¨FÑ2Ñ2ˆØ˜a‘= 3 r¡7Ñ*ˆä�|‰|˜H h°¡l°V°GÄ5Ç:Á:ÐVY×V`ÑV`ÔaÐar1   ÚtensorÚaxisc                 ól  ‡ ‡— |€yt        ‰t        t        f«      r‰D ]  }‰ j                  ||«      }Œ |S t        |t        t        f«      r t	        |«      ˆˆ fd„|D «       «      S ‰|j
                  z  Š‰ j                  j                  r#‰ j                  j                  rt        ddd«      nt        ddd«      }t        t        d«      g‰z  |gz   «      }‰ j                  j                  rBt        t        d«      g‰z  t        dd«      gz   «      }t        j                  ||   |g‰¬«      }||   S )zT
        Perform pooling by stride slicing the tensor along the given axis.
        Nc              3   óB   •K  — | ]  }‰j                  |‰«      –— Œ y ­wr5   )Ústride_pool)Ú.0Úxr˜   r.   s     €€r0   ú	<genexpr>z7FunnelAttentionStructure.stride_pool.<locals>.<genexpr>é   s   øè ø€ ÒJ¸a × 0Ñ 0°°D× 9ÑJùs   ƒra   rC   r   )r˜   )Ú
isinstanceÚlistÚtupler›   ÚtypeÚndimr   rR   r�   Úslicer=   rk   )r.   r—   r˜   ÚaxÚ
axis_sliceÚ	enc_sliceÚ	cls_slices   ` `    r0   r›   z$FunnelAttentionStructure.stride_poolÖ   s'  ù€ ð ˆ>Øô �dœT¤5˜MÔ*Øò 6�Ø×)Ñ)¨&°"Ó5‘ð6àˆMô �fœu¤d˜mÔ,Ø”4˜“<ÔJÀ6ÔJÓJÐJð 	�—‘Ñˆð #'§+¡+×":Ò":¸t¿{¹{×?WÒ?WŒE�$˜˜AÔÔ]bÐcgÐimÐopÓ]qð 	ô œ5 ›;˜-¨$Ñ.°*°Ñ=Ó>ˆ	Ø�;‰;×#Ò#Üœu T›{˜m¨dÑ2´e¸DÀ!³nÐ5EÑEÓFˆIÜ—Y‘Y  yÑ 1°6Ð:ÀÔFˆFØ�iÑ Ð r1   Úmodec                 óî  ‡ ‡‡‡— ‰€yt        ‰t        t        f«      r t        ‰«      ˆˆ ˆˆfd„‰D «       «      S ‰ j                  j
                  rE‰ j                  j                  r‰dd…dd…f   n‰}t        j                  ‰dd…dd…f   |gd¬«      Š‰j                  }|dk(  r‰dd…ddd…df   Šn|dk(  r‰dd…ddd…dd…f   Š‰dfŠ‰dk(  r$t        j                  j                  ‰‰‰d	¬
«      Šn_‰dk(  r$t        j                  j                  ‰‰‰d	¬
«      Šn6‰dk(  r&t        j                  j                  ‰ ‰‰d	¬
«       Šnt        d«      ‚|dk(  r‰dd…ddd…df   S |dk(  r	‰dd…df   S ‰S )z3Apply 1D pooling to a tensor of size [B x T (x H)].Nc              3   óF   •K  — | ]  }‰j                  ‰‰‰¬ «      –— Œ y­w))r©   r‡   N)Úpool_tensor)rœ   r�   r©   r.   r‡   r—   s     €€€€r0   rž   z7FunnelAttentionStructure.pool_tensor.<locals>.<genexpr>   s$   øè ø€ ÒcÐWX × 0Ñ 0°¸dÈ6Ð 0× RÑcùs   ƒ!ra   r   rb   rC   r   ÚmeanT)r‡   Ú	ceil_modeÚmaxÚminz0The supported modes are 'mean', 'max' and 'min'.r   )rŸ   r¡   r    r¢   r   rR   r�   r=   rk   r£   r   rS   Ú
avg_pool2dÚ
max_pool2dÚNotImplementedError)r.   r—   r©   r‡   Úsuffixr£   s   ````  r0   r¬   z$FunnelAttentionStructure.pool_tensor÷   s{  û€ ð ˆ>Øô �fœu¤d˜mÔ,Ø”4˜“<ÖcÐ\bÔcÓcÐcà�;‰;×#Ò#Ø'+§{¡{×'?Ò'?�VšA˜s ˜s˜F’^ÀVˆFÜ—Y‘Y ¢q¨"¨1¨" u¡¨vÐ6¸AÔ>ˆFà�{‰{ˆØ�1Š9ØšA˜t¢Q¨Ð,Ñ-‰FØ�QŠYØšA˜t¢Qª˜MÑ*ˆFà˜!�ˆà�6Š>Ü—]‘]×-Ñ-¨f°fÀVÐW[Ð-Ó\‰FØ�UŠ]Ü—]‘]×-Ñ-¨f°fÀVÐW[Ð-Ó\‰FØ�UŠ]Ü—m‘m×.Ñ.°¨w¸ÀvÐY]Ð.Ó^Ð^‰Fä%Ð&XÓYÐYà�1Š9Øš!˜Q¢ 1˜*Ñ%Ð%Ø�QŠYØš!˜Q˜$‘<ÐØˆr1   Úattention_inputsc                 ó”  — |\  }}}}| j                   j                  r€| j                   j                  dk(  r| j                  |dd d«      |dd z   }| j                  |d«      }| j                  |d«      }| j	                  || j                   j
                  ¬«      }n¢| xj                  dz  c_        | j                   j                  dk(  r| j                  |d«      }| j                  |ddg«      }| j                  |ddg«      }| j	                  |d¬«      }| j	                  || j                   j
                  ¬«      }||||f}||fS )zTPool `output` and the proper parts of `attention_inputs` before the attention layer.r^   NrC   r   r   ©r©   r°   )r   Úpool_q_onlyre   r›   r¬   Úpooling_typerH   )r.   Úoutputrµ   rV   rW   rI   rX   s          r0   Úpre_attention_poolingz.FunnelAttentionStructure.pre_attention_pooling  sM  € ð EUÑAˆ˜¨¸Ø�;‰;×"Ò"Ø�{‰{×)Ñ)¨\Ò9Ø"&×"2Ñ"2°?À2ÀAÐ3FÈÓ"JÈ_Ð]^Ð]_ÐM`Ñ"`�Ø!×-Ñ-¨n¸aÓ@ˆNØ×'Ñ'¨°!Ó4ˆHØ×%Ñ% f°4·;±;×3KÑ3KÐ%ÓL‰Fà×Ò Ñ"ÕØ�{‰{×)Ñ)¨\Ò9Ø"&×"2Ñ"2°?ÀAÓ"F�Ø!×-Ñ-¨n¸qÀ!¸fÓEˆNØ×'Ñ'¨°1°a°&Ó9ˆHØ!×-Ñ-¨nÀ5Ð-ÓIˆNØ×%Ñ% f°4·;±;×3KÑ3KÐ%ÓLˆFØ+¨^¸^ÈXÐVÐØÐ'Ð'Ð'r1   c                 óL  — |\  }}}}| j                   j                  r€| xj                  dz  c_        | j                   j                  dk(  r|dd | j	                  |dd d«      z   }| j	                  |d«      }| j	                  |d«      }| j                  |d¬«      }||||f}|S )zFPool the proper parts of `attention_inputs` after the attention layer.rC   r^   Nr   r   r°   r·   )r   r¸   rH   re   r›   r¬   )r.   rµ   rV   rW   rI   rX   s         r0   Úpost_attention_poolingz/FunnelAttentionStructure.post_attention_pooling3  s¶   € àDTÑAˆ˜¨¸Ø�;‰;×"Ò"Ø×Ò Ñ"ÕØ�{‰{×)Ñ)¨\Ò9Ø"1°"°1Ð"5¸×8HÑ8HÈÐYZÐY[ÐI\Ð^_Ó8`Ñ"`�Ø!×-Ñ-¨n¸aÓ@ˆNØ×'Ñ'¨°!Ó4ˆHØ!×-Ñ-¨nÀ5Ð-ÓIˆNØ+¨^¸^ÈXÐVÐØÐr1   r8   ©Nr   )r­   rC   )r:   r;   r<   Ú__doc__rD   ÚintÚ__annotations__r   r    r=   r>   r¡   rY   rQ   rO   rP   r    rN   rn   ro   r›   Ústrr¬   r»   r½   r?   r@   s   @r0   rB   rB   =   s8  ø… ñð Ð�sÓð!˜|ð !°õ !ð /3Ø.2ñ	Kà—|‘|ðKð Ÿ™ tÑ+ðKð Ÿ™ tÑ+ð	Kð
 
ˆu�|‰|Ñ	óKð((°E·L±Lð (ÀUÇ\Á\ó (ðN(ØðN(Ø#(§;¡;ðN(Ø8=¿¹ðN(à	ˆu�|‰|Ñ	˜t D¨¯©Ñ$6Ñ7Ñ	7óN(ð` e§l¡lð Àó ñb §¡ð b°cð bÐSVð bÐ_d×_kÑ_kó bð!à—‘˜u U§\¡\Ñ2Ñ2°T¸%¿,¹,Ñ5GÑGð!ð �E˜#‘JÑ  c¡Ñ*ð!ð 
�‰ó	!ðD rsñ$Ø—l‘l U¨5¯<©<Ñ%8Ñ8¸4ÀÇÁÑ;MÑMð$ØUXð$Øknð$à	�‰ó$ðL(Ø(-¨e¯l©lÑ(;ð(à	ˆu�|‰|˜U 5§<¡<Ñ0Ð0Ñ	1ó(ð, °u¸U¿\¹\Ñ7Jð  ÈuÐUZ×UaÑUaÑOb÷  r1   rB   Úpositional_attnÚcontext_lenrd   r   c                 óÊ   — | j                   \  }}}}t        j                  | ||||g«      } | d d …d d …|d …d d …f   } t        j                  | |||||z
  g«      } | dd |…f   } | S )N.)r‘   r=   Úreshape)rÃ   rÄ   rd   Ú
batch_sizeÚn_headrM   Úmax_rel_lens          r0   Ú_relative_shift_gatherrÊ   A  s   € Ø/>×/DÑ/DÑ,€J�˜ ô —m‘m O°jÀ&È+ÐW^Ð5_Ó`€OØ%¢aª¨E©F²A oÑ6€OÜ—m‘m O°jÀ&È'ÐS^ÐafÑSfÐ5gÓh€OØ% c¨<¨K¨<Ð&7Ñ8€OØÐr1   c                   óÔ   ‡ — e Zd Zdededdfˆ fd„Zdd„Zdd„Z	 ddej                  d	ej                  d
ej                  de
ej                     dede
ej                  df   fd„Zˆ xZS )ÚFunnelRelMultiheadAttentionr   r…   r   Nc                 óJ  •— t         ‰| �  «        || _        || _        |j                  |j
                  |j                  }}}t        j                  |j                  «      | _	        t        j                  |j                  «      | _
        t        j                  |||z  d¬«      | _        t        j                  |||z  «      | _        t        j                  |||z  «      | _        t        j                  t!        j"                  ||g«      «      | _        t        j                  t!        j"                  ||g«      «      | _        t        j                  t!        j"                  |||g«      «      | _        t        j                  t!        j"                  ||g«      «      | _        t        j                  t!        j"                  d||g«      «      | _        t        j                  ||z  |«      | _        t        j0                  ||j2                  ¬«      | _        d|dz  z  | _        y )NF)ÚbiasrC   r   r_   g      à?)r   r    r   r…   r'   rÈ   Úd_headr   r*   r+   Úattention_dropoutÚLinearÚq_headÚk_headÚv_headÚ	Parameterr=   ÚzerosÚr_w_biasÚr_r_biasÚr_kernelÚr_s_biasÚ	seg_embedÚ	post_projr&   r(   r)   Úscale)r.   r   r…   r'   rÈ   rÏ   r/   s         €r0   r    z$FunnelRelMultiheadAttention.__init__R  s€  ø€ Ü‰ÑÔØˆŒØ&ˆÔØ"(§.¡.°&·-±-ÀÇÁ˜�ˆä Ÿj™j¨×)>Ñ)>Ó?ˆÔÜ!#§¡¨F×,DÑ,DÓ!EˆÔä—i‘i ¨°&©¸uÔEˆŒÜ—i‘i ¨°&©Ó9ˆŒÜ—i‘i ¨°&©Ó9ˆŒäŸ™¤U§[¡[°&¸&Ð1AÓ%BÓCˆŒÜŸ™¤U§[¡[°&¸&Ð1AÓ%BÓCˆŒÜŸ™¤U§[¡[°'¸6À6Ð1JÓ%KÓLˆŒÜŸ™¤U§[¡[°&¸&Ð1AÓ%BÓCˆŒÜŸ™¤e§k¡k°1°f¸fÐ2EÓ&FÓGˆŒäŸ™ 6¨F¡?°GÓ<ˆŒÜŸ,™, w°F×4IÑ4IÔJˆŒØ˜F C™KÑ(ˆ�
r1   c                 ó~  — | j                   j                  dk(  rŽ|\  }}}}| j                  | j                  z  }	| j                  }
t        j                  d||	z   |
«      }||dd…df   z  }||dd…df   z  }t        j                  d||«      t        j                  d||«      z   }nŽ|j                  d   |k7  rdnd}|| j                     |dz
     }| j                  | j                  z  }| j                  }
t        j                  d||
«      }t        j                  d||z   |«      }t        |||«      }|�||z  }|S )	z5Relative attention score for the positional encodingsr^   zbinh,dnh->bindNzbind,jd->bnijr   rC   ztd,dnh->tnhzbinh,tnh->bnit)
r   re   rØ   rÝ   rÙ   r=   Úeinsumr‘   r…   rÊ   )r.   rV   rÒ   rÄ   rX   r{   r}   r|   r~   ÚuÚw_rÚq_r_attentionÚq_r_attention_1Úq_r_attention_2rÃ   rd   ÚrÚvÚr_heads                      r0   Úrelative_positional_attentionz9FunnelRelMultiheadAttention.relative_positional_attentioni  sK  € ð �;‰;×%Ñ%¨Ò5ð #2ÑˆC��S˜%à—‘ §
¡
Ñ*ˆAà—-‘-ˆCô "ŸL™LÐ)9¸6ÀA¹:ÀsÓKˆMØ+¨c²!°T°'©lÑ:ˆOØ+¨b²°D°©kÑ9ˆOô $Ÿl™l¨?¸OÈSÓQÔTY×T`ÑT`Ø °%óUñ ‰Oð  Ÿ™ a™¨KÒ7‘A¸QˆEð   × 0Ñ 0Ñ1°%¸!±)Ñ<ˆAà—‘ §
¡
Ñ*ˆAà—-‘-ˆCô —\‘\ -°°CÓ8ˆFä#Ÿl™lÐ+;¸VÀa¹ZÈÓPˆOä4°_ÀkÐSXÓYˆOàÐØ˜xÑ'ˆOØÐr1   c                 óÎ  — |€y|j                   \  }}}| j                  | j                  z  }t        j                  d||z   | j
                  «      }|dd…df   j                  ||j                   d   ||g«      }t        j                  |dd¬«      \  }	}
t        j                  ||
j                  |j                   «      |	j                  |j                   «      «      }|�||z  }|S )z/Relative attention score for the token_type_idsNr   zbind,snd->bnisrC   r   ra   rb   )	r‘   rÚ   rÝ   r=   rß   rÛ   rp   ÚsplitÚwhere)r.   rW   rÒ   rX   rÇ   rM   rÄ   rÚ   Útoken_type_biasÚdiff_token_typeÚsame_token_typeÚtoken_type_attns               r0   Úrelative_token_type_attentionz9FunnelRelMultiheadAttention.relative_token_type_attention“  sã   € àÐ!ØØ+9×+?Ñ+?Ñ(ˆ
�G˜[ð —=‘= 4§:¡:Ñ-ˆô  Ÿ,™,Ð'7¸À(Ñ9JÈDÏNÉNÓ[ˆà'ª¨4¨Ñ0×7Ñ7¸ÀVÇ\Á\ÐRSÁ_ÐV]Ð_jÐ8kÓlˆä+0¯;©;°ÈÈrÔ+RÑ(ˆ˜äŸ+™+Ø˜O×2Ñ2°>×3GÑ3GÓHÈ/×J`ÑJ`Ðao×auÑauÓJvó
ˆð ÐØ˜xÑ'ˆOØÐr1   ÚqueryÚkeyÚvaluerµ   Úoutput_attentions.c                 ó  — |\  }}}}	|j                   \  }
}}|j                   d   }| j                  j                  | j                  j                  }}| j	                  |«      j                  |
|||«      }| j                  |«      j                  |
|||«      }| j                  |«      j                  |
|||«      }|| j                  z  }| j                  | j                  z  }t        j                  d||z   |«      }| j                  ||||	«      }| j                  |||	«      }||z   |z   }|j                  }|j                  «       }|�%|t         d|d d …d d f   j                  «       z
  z  z
  }t        j"                  |d|¬«      }| j%                  |«      }t        j                  d||«      }| j'                  |j)                  |
|||z  «      «      }| j+                  |«      }| j-                  ||z   «      }|r||fS |fS )Nr   zbind,bjnd->bnijra   )rc   rO   zbnij,bjnd->bind)r‘   r   rÈ   rÏ   rÒ   ÚviewrÓ   rÔ   rÝ   r×   r=   rß   rè   rð   rO   ÚfloatÚINFÚsoftmaxrÐ   rÜ   rÆ   r+   r)   )r.   rñ   rò   ró   rµ   rô   rV   rW   rI   rX   rÇ   rM   Ú_rÄ   rÈ   rÏ   rÒ   rÓ   rÔ   r×   Úcontent_scorerÃ   rï   Ú
attn_scorerO   Ú	attn_probÚattn_vecÚattn_outrº   s                                r0   r7   z#FunnelRelMultiheadAttention.forward«  sþ  € ð EUÑAˆ˜¨¸à!&§¡Ñˆ
�G˜QØ—i‘i ‘lˆØŸ™×+Ñ+¨T¯[©[×-?Ñ-?�ˆð —‘˜UÓ#×(Ñ(¨°W¸fÀfÓMˆà—‘˜SÓ!×&Ñ& z°;ÀÈÓOˆØ—‘˜UÓ#×(Ñ(¨°[À&È&ÓQˆà˜$Ÿ*™*Ñ$ˆà—=‘= 4§:¡:Ñ-ˆäŸ™Ð%6¸ÀÑ8IÈ6ÓRˆØ×<Ñ<¸_ÈfÐVaÐckÓlˆØ×<Ñ<¸^ÈVÐU]Ó^ˆð # _Ñ4°ÑFˆ
ð × Ñ ˆØ×%Ñ%Ó'ˆ
àÐ%Ø#¤c¨Q°ÂÀ4ÈÀÑ1N×1TÑ1TÓ1VÑ-VÑ&WÑWˆJä—M‘M *°"¸EÔBˆ	Ø×*Ñ*¨9Ó5ˆ	ô —<‘<Ð 1°9¸fÓEˆð —>‘> (×"2Ñ"2°:¸wÈÐQWÉÓ"XÓYˆØ×&Ñ& xÓ0ˆà—‘ ¨Ñ!1Ó2ˆÙ&7�˜	Ð"ÐF¸f¸YÐFr1   r5   ©F)r:   r;   r<   r   rÀ   r    rè   rð   r=   r>   r¡   Úboolr7   r?   r@   s   @r0   rÌ   rÌ   Q  s–   ø„ ð)˜|ð )¸#ð )À$õ )ó.(óTð< #(ñ3Gà�|‰|ð3Gð �\‰\ð3Gð �|‰|ð	3Gð
   §¡Ñ-ð3Gð  ð3Gð 
ˆu�|‰|˜SÐ Ñ	!÷3Gr1   rÌ   c                   ó`   ‡ — e Zd Zdeddfˆ fd„Zdej                  dej                  fd„Zˆ xZS )ÚFunnelPositionwiseFFNr   r   Nc                 óü  •— t         ‰| �  «        t        j                  |j                  |j
                  «      | _        t        |j                     | _	        t        j                  |j                  «      | _        t        j                  |j
                  |j                  «      | _        t        j                  |j                  «      | _        t        j                  |j                  |j                   «      | _        y r5   )r   r    r   rÑ   r'   Úd_innerÚlinear_1r
   Ú
hidden_actÚactivation_functionr*   Úactivation_dropoutÚlinear_2r+   r,   r&   r(   r)   r-   s     €r0   r    zFunnelPositionwiseFFN.__init__â  sž   ø€ Ü‰ÑÔÜŸ	™	 &§.¡.°&·.±.ÓAˆŒÜ#)¨&×*;Ñ*;Ñ#<ˆÔ Ü"$§*¡*¨V×-FÑ-FÓ"GˆÔÜŸ	™	 &§.¡.°&·.±.ÓAˆŒÜ—z‘z &×"7Ñ"7Ó8ˆŒÜŸ,™, v§~¡~°v×7LÑ7LÓMˆ�r1   Úhiddenc                 óÔ   — | j                  |«      }| j                  |«      }| j                  |«      }| j                  |«      }| j	                  |«      }| j                  ||z   «      S r5   )r  r  r	  r
  r,   r)   )r.   r  Úhs      r0   r7   zFunnelPositionwiseFFN.forwardë  s^   € Ø�M‰M˜&Ó!ˆØ×$Ñ$ QÓ'ˆØ×#Ñ# AÓ&ˆØ�M‰M˜!ÓˆØ�L‰L˜‹OˆØ�‰˜v¨™zÓ*Ð*r1   r9   r@   s   @r0   r  r  á  s4   ø„ ðN˜|ð N°õ Nð+˜eŸl™lð +¨u¯|©|÷ +r1   r  c                   óˆ   ‡ — e Zd Zdededdfˆ fd„Z	 ddej                  dej                  dej                  d	ede	f
d
„Z
ˆ xZS )ÚFunnelLayerr   r…   r   Nc                 ód   •— t         ‰| �  «        t        ||«      | _        t	        |«      | _        y r5   )r   r    rÌ   Ú	attentionr  Úffn)r.   r   r…   r/   s      €r0   r    zFunnelLayer.__init__õ  s(   ø€ Ü‰ÑÔÜ4°V¸[ÓIˆŒÜ(¨Ó0ˆ�r1   rñ   rò   ró   rô   c                 ón   — | j                  |||||¬«      }| j                  |d   «      }|r||d   fS |fS )N©rô   r   r   )r  r  )r.   rñ   rò   ró   rµ   rô   Úattnrº   s           r0   r7   zFunnelLayer.forwardú  sH   € ð �~‰~˜e S¨%Ð1AÐUfˆ~ÓgˆØ—‘˜$˜q™'Ó"ˆÙ$5�˜˜Q™Ð ÐD¸F¸9ÐDr1   r   )r:   r;   r<   r   rÀ   r    r=   r>   r  r¡   r7   r?   r@   s   @r0   r  r  ô  si   ø„ ð1˜|ð 1¸#ð 1À$õ 1ð #(ñ
Eà�|‰|ð
Eð �\‰\ð
Eð �|‰|ð	
Eð  ð
Eð 
÷
Er1   r  c                   ó¦   ‡ — e Zd Zdeddfˆ fd„Z	 	 	 	 	 ddej                  dej                  dz  dej                  dz  ded	ed
edee	z  fd„Z
ˆ xZS )ÚFunnelEncoderr   r   Nc                 óT  •— t         ‰| �  «        || _        t        |«      | _        t        j                  t        |j                  «      D ���cg c];  \  }}t        j                  t        |«      D �cg c]  }t        ||«      ‘Œ c}«      ‘Œ= c}}}«      | _        y c c}w c c}}}w r5   )r   r    r   rB   Úattention_structurer   Ú
ModuleListÚ	enumerateÚblock_sizesrl   r  Úblocks)r.   r   r…   Ú
block_sizerú   r/   s        €r0   r    zFunnelEncoder.__init__  s†   ø€ Ü‰ÑÔØˆŒÜ#;¸FÓ#CˆÔ Ü—m‘mô 09¸×9KÑ9KÓ/L÷ð á+�K ô —‘ÌÈzÓIZÖ[ÀAœ{¨6°;Õ?Ò[Õ\ôó
ˆ�ùâ[ùôs   Á$B#Á3BÂ	B#ÂB#r3   rI   rJ   rô   Úoutput_hidden_statesÚreturn_dictc           
      ó  — |j                  |«      }| j                  j                  |||¬«      }|}|r|fnd }	|rdnd }
t        | j                  «      D �]  \  }}|j                  d«      | j                  j                  rdndkD  }|xr |dkD  }|r| j                  j                  ||«      \  }}t        |«      D ]¥  \  }}t        | j                  j                  |   «      D ]{  }|dk(  xr	 |dk(  xr |}|r}| j                  j                  r|n|x}}n|x}x}} ||||||¬«      }|d   }|r| j                  j                  |«      }|r|
|dd  z   }
|sŒv|	|fz   }	Œ} Œ§ �Œ |st        d„ ||	|
fD «       «      S t        ||	|
¬«      S )	N©rI   rJ   © r   rC   r   r  c              3   ó&   K  — | ]	  }|€Œ|–— Œ y ­wr5   r#  ©rœ   ræ   s     r0   rž   z(FunnelEncoder.forward.<locals>.<genexpr>B  ó   è ø€ Òa˜qÐSTÑS`œÑaùó   ‚Š©Úlast_hidden_stateÚhidden_statesÚ
attentions)Útype_asr  rY   r  r  rL   r   rR   r»   rl   Úblock_repeatsr¸   r½   r¡   r   )r.   r3   rI   rJ   rô   r  r   rµ   r  Úall_hidden_statesÚall_attentionsr…   ÚblockÚpooling_flagÚpooled_hiddenÚlayer_indexÚlayerÚrepeat_indexÚ
do_poolingrñ   rò   ró   Úlayer_outputs                          r0   r7   zFunnelEncoder.forward  sà  € ð (×/Ñ/°Ó>ˆØ×3Ñ3×IÑIØØ)Ø)ð Jó 
Ðð
 ˆá0D˜]Ñ,È$ÐÙ0™°dˆä"+¨D¯K©KÓ"8ó 	JÑˆK˜Ø!Ÿ;™; q›>°$·+±+×2JÒ2J©QÐPQÑRˆLØ'Ò;¨K¸!©OˆLÙØ26×2JÑ2J×2`Ñ2`ØÐ,ó3Ñ/�Ð/ô '0°Ó&6ò JÑ"�˜UÜ$)¨$¯+©+×*CÑ*CÀKÑ*PÓ$Qò J�LØ".°!Ñ"3Ò!\¸+ÈÑ:JÒ!\ÐP\�JÙ!Ø -˜Ø04·±×0GÒ0G¡fÈ]ÐZ˜™eà.4Ð4˜Ð4  eÙ#(¨°°UÐ<LÐ`qÔ#r�LØ)¨!™_�FÙ!Ø+/×+CÑ+C×+ZÑ+ZÐ[kÓ+lÐ(á(Ø)7¸,ÀqÀrÐ:JÑ)J˜Ú+Ø,=ÀÀ	Ñ,IÑ)ñJòJð	Jñ2 ÜÑa VÐ->ÀÐ$OÔaÓaÐaÜ°ÐGXÐesÔtÐtr1   ©NNFFT©r:   r;   r<   r   r    r=   r>   r  r¡   r   r7   r?   r@   s   @r0   r  r    s“   ø„ ð	
˜|ð 	
°õ 	
ð /3Ø.2Ø"'Ø%*Ø ñ0uà—|‘|ð0uð Ÿ™ tÑ+ð0uð Ÿ™ tÑ+ð	0uð
  ð0uð #ð0uð ð0uð 
�Ñ	 ÷0ur1   r  r�   r‡   Ú
target_lenrR   r�   c           	      ó6  — |dk(  r| S |r| dd…dd…f   }| dd…dd…f   } t        j                  | |d¬«      }|rT|r)t        j                  j	                  |ddd|dz
  ddf«      }|dd…d|dz
  …f   }t        j
                  |gd¬«      }|S |dd…d|…f   }|S )z{
    Upsample tensor `x` to match `target_len` by repeating the tokens `stride` time on the sequence length dimension.
    r   N)Úrepeatsrc   r   rb   )r=   Úrepeat_interleaver   rS   rT   rk   )r�   r‡   r:  rR   r�   Úclsrº   s          r0   Úupsampler?  F  sÁ   € ð �‚{ØˆÙØ’�2�A�2�‰hˆØŠa�‘ˆe‰HˆÜ×$Ñ$ Q°¸AÔ>€FÙÙÜ—]‘]×&Ñ& v°°1°a¸À!¹ÀQÈÐ/JÓKˆFØšÐ+˜Z¨!™^Ð+Ð+Ñ,ˆÜ—‘˜C ˜=¨aÔ0ˆð €Mð š˜;˜J˜;˜Ñ'ˆØ€Mr1   c                   ó¾   ‡ — e Zd Zdeddfˆ fd„Z	 	 	 	 	 ddej                  dej                  dej                  dz  dej                  dz  d	ed
ededee	z  fd„Z
ˆ xZS )ÚFunnelDecoderr   r   Nc           	      óä   •— t         ‰| �  «        || _        t        |«      | _        t        j                  t        |j                  «      D �cg c]  }t        |d«      ‘Œ c}«      | _
        y c c}w )Nr   )r   r    r   rB   r  r   r  rl   Únum_decoder_layersr  Úlayers)r.   r   rú   r/   s      €r0   r    zFunnelDecoder.__init__]  sR   ø€ Ü‰ÑÔØˆŒÜ#;¸FÓ#CˆÔ Ü—m‘mÄUÈ6×KdÑKdÓEeÖ$fÀ¤[°¸Õ%;Ò$fÓgˆ�ùÒ$fs   ÁA-Úfinal_hiddenÚfirst_block_hiddenrI   rJ   rô   r  r   c                 óè  — t        |dt        | j                  j                  «      dz
  z  |j                  d   | j                  j
                  | j                  j                  ¬«      }||z   }	|r|	fnd }
|rdnd }| j                  j                  |	||¬«      }| j                  D ]'  } ||	|	|	||¬«      }|d   }	|r||dd  z   }|sŒ"|
|	fz   }
Œ) |st        d„ |	|
|fD «       «      S t        |	|
|¬	«      S )
NrC   r   )r‡   r:  rR   r�   r#  r"  r  r   c              3   ó&   K  — | ]	  }|€Œ|–— Œ y ­wr5   r#  r%  s     r0   rž   z(FunnelDecoder.forward.<locals>.<genexpr>‰  r&  r'  r(  )r?  Úlenr   r  r‘   rR   r�   r  rY   rD  r¡   r   )r.   rE  rF  rI   rJ   rô   r  r   Úupsampled_hiddenr  r.  r/  rµ   r4  r7  s                  r0   r7   zFunnelDecoder.forwardc  s%  € ô $ØØœ˜TŸ[™[×4Ñ4Ó5¸Ñ9Ñ:Ø)×/Ñ/°Ñ2ØŸ™×1Ñ1ØŸ™×1Ñ1ô
Ðð "Ð$6Ñ6ˆÙ)=˜V™IÀ4ÐÙ0™°dˆà×3Ñ3×IÑIØØ)Ø)ð Jó 
Ðð —[‘[ò 	BˆEÙ  ¨°Ð9IÐ]nÔoˆLØ! !‘_ˆFá Ø!/°,¸q¸rÐ2BÑ!B�Ú#Ø$5¸¸	Ñ$AÑ!ð	Bñ ÜÑa VÐ->ÀÐ$OÔaÓaÐaÜ°ÐGXÐesÔtÐtr1   r8  r9  r@   s   @r0   rA  rA  \  s¥   ø„ ðh˜|ð h°õ hð /3Ø.2Ø"'Ø%*Ø ñ'uà—l‘lð'uð "ŸL™Lð'uð Ÿ™ tÑ+ð	'uð
 Ÿ™ tÑ+ð'uð  ð'uð #ð'uð ð'uð 
�Ñ	 ÷'ur1   rA  c                   ód   ‡ — e Zd ZdZdeddfˆ fd„Zdej                  dej                  fd„Zˆ xZ	S )ÚFunnelDiscriminatorPredictionszEPrediction module for the discriminator, made up of two dense layers.r   r   Nc                 óØ   •— t         ‰| �  «        || _        t        j                  |j
                  |j
                  «      | _        t        j                  |j
                  d«      | _        y r¾   )r   r    r   r   rÑ   r'   ÚdenseÚdense_predictionr-   s     €r0   r    z'FunnelDiscriminatorPredictions.__init__�  sF   ø€ Ü‰ÑÔØˆŒÜ—Y‘Y˜vŸ~™~¨v¯~©~Ó>ˆŒ
Ü "§	¡	¨&¯.©.¸!Ó <ˆÕr1   Údiscriminator_hidden_statesc                 ó¬   — | j                  |«      }t        | j                  j                     |«      }| j	                  |«      j                  d«      }|S )Nra   )rN  r
   r   r  rO  Úsqueeze)r.   rP  r*  Úlogitss       r0   r7   z&FunnelDiscriminatorPredictions.forward–  sJ   € ØŸ
™
Ð#>Ó?ˆÜ˜tŸ{™{×5Ñ5Ñ6°}ÓEˆØ×&Ñ& }Ó5×=Ñ=¸bÓAˆØˆr1   )
r:   r;   r<   r¿   r   r    r=   r>   r7   r?   r@   s   @r0   rL  rL  �  s4   ø„ ÙOð=˜|ð =°õ =ð°5·<±<ð ÀEÇLÁL÷ r1   rL  c                   óV   ‡ — e Zd ZU eed<   dZ ej                  «       ˆ fd„«       Zˆ xZ	S )ÚFunnelPreTrainedModelr   Úfunnelc                 ó’  •— t         ‰| �  |«       |j                  j                  }|j	                  d«      dk7  rÇt        |dd «      �‹| j                  j                  €>|j                  j                  \  }}t        j                  dt        ||z   «      z  «      }n| j                  j                  }t        j                  |j                  |¬«       t        |dd «      �!t        j                  |j                   d«       y y |dk(  �r
t        j"                  |j$                  | j                  j&                  ¬	«       t        j"                  |j(                  | j                  j&                  ¬	«       t        j"                  |j*                  | j                  j&                  ¬	«       t        j"                  |j,                  | j                  j&                  ¬	«       t        j"                  |j.                  | j                  j&                  ¬	«       y |d
k(  r±| j                  j                  €dn| j                  j                  }t        j                  |j0                  j                  |¬«       |j0                  j2                  �At        j4                  |j0                  j                  |j0                  j2                     «       y y y )NrÑ   ra   Úweightr_   )ÚstdrÎ   g        rÌ   )Úbr   )r   Ú_init_weightsr/   r:   ÚfindÚgetattrr   Úinitializer_stdrX  r‘   ÚnpÚsqrtr÷   ÚinitÚnormal_Ú	constant_rÎ   Úuniform_r×   Úinitializer_rangerØ   rÙ   rÚ   rÛ   r%   r   Úzeros_)r.   ÚmoduleÚ	classnameÚfan_outÚfan_inrY  r/   s         €r0   r[  z#FunnelPreTrainedModel._init_weights¢  sç  ø€ ä‰Ñ˜fÔ%Ø×$Ñ$×-Ñ-ˆ	Ø�>‰>˜(Ó# rÒ)Ü�v˜x¨Ó.Ð:Ø—;‘;×.Ñ.Ð6Ø&,§m¡m×&9Ñ&9‘O�G˜VÜŸ'™' #¬¨f°wÑ.>Ó(?Ñ"?Ó@‘CàŸ+™+×5Ñ5�CÜ—‘˜VŸ]™]°Õ4Ü�v˜v tÓ,Ð8Ü—‘˜vŸ{™{¨CÕ0ð 9àÐ7Ó7Ü�M‰M˜&Ÿ/™/¨T¯[©[×-JÑ-JÕKÜ�M‰M˜&Ÿ/™/¨T¯[©[×-JÑ-JÕKÜ�M‰M˜&Ÿ/™/¨T¯[©[×-JÑ-JÕKÜ�M‰M˜&Ÿ/™/¨T¯[©[×-JÑ-JÕKÜ�M‰M˜&×*Ñ*¨d¯k©k×.KÑ.KÖLØÐ,Ò,ØŸ™×4Ñ4Ð<‘#À$Ç+Á+×B]ÑB]ˆCÜ�L‰L˜×/Ñ/×6Ñ6¸CÕ@Ø×%Ñ%×1Ñ1Ð=Ü—‘˜F×2Ñ2×9Ñ9¸&×:PÑ:P×:\Ñ:\Ñ]Õ^ð >ð -r1   )
r:   r;   r<   r   rÁ   Úbase_model_prefixr=   Úno_gradr[  r?   r@   s   @r0   rU  rU  �  s+   ø… àÓØ Ðà€U‡]�]ƒ_ó_ó ô_r1   rU  c                   ód   ‡ — e Zd Zdededdfˆ fd„Zdej                  dej                  fd„Zˆ xZ	S )ÚFunnelClassificationHeadr   Ún_labelsr   Nc                 ó  •— t         ‰| �  «        t        j                  |j                  |j                  «      | _        t        j                  |j                  «      | _        t        j                  |j                  |«      | _	        y r5   )
r   r    r   rÑ   r'   Úlinear_hiddenr*   r+   r,   Ú
linear_out)r.   r   ro  r/   s      €r0   r    z!FunnelClassificationHead.__init__¾  sU   ø€ Ü‰ÑÔÜŸY™Y v§~¡~°v·~±~ÓFˆÔÜ—z‘z &×"7Ñ"7Ó8ˆŒÜŸ)™) F§N¡N°HÓ=ˆ�r1   r  c                 ó’   — | j                  |«      }t        j                  |«      }| j                  |«      }| j	                  |«      S r5   )rq  r=   Útanhr,   rr  )r.   r  s     r0   r7   z FunnelClassificationHead.forwardÄ  s=   € Ø×#Ñ# FÓ+ˆÜ—‘˜FÓ#ˆØ—‘˜fÓ%ˆØ�‰˜vÓ&Ð&r1   )
r:   r;   r<   r   rÀ   r    r=   r>   r7   r?   r@   s   @r0   rn  rn  ½  s8   ø„ ð>˜|ð >°sð >¸tõ >ð'˜eŸl™lð '¨u¯|©|÷ 'r1   rn  z2
    Output type of [`FunnelForPreTraining`].
    )Úcustom_introc                   ó¾   — e Zd ZU dZdZej                  dz  ed<   dZej                  dz  ed<   dZ	e
ej                     dz  ed<   dZe
ej                     dz  ed<   y)ÚFunnelForPreTrainingOutputa1  
    loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`):
        Total loss of the ELECTRA-style objective.
    logits (`torch.FloatTensor` of shape `(batch_size, sequence_length)`):
        Prediction scores of the head (scores for each token before SoftMax).
    NÚlossrS  r*  r+  )r:   r;   r<   r¿   rx  r=   ÚFloatTensorrÁ   rS  r*  r¡   r+  r#  r1   r0   rw  rw  Ë  sg   … ñð &*€Dˆ%×
Ñ
˜dÑ
"Ó)Ø'+€FˆE×Ñ Ñ$Ó+Ø59€M�5˜×*Ñ*Ñ+¨dÑ2Ó9Ø26€J��e×'Ñ'Ñ(¨4Ñ/Ô6r1   rw  z¨
    The base Funnel Transformer Model transformer outputting raw hidden-states without upsampling head (also called
    decoder) or any task-specific head on top.
    c                   óN  ‡ — e Zd Zdeddfˆ fd„Zdej                  fd„Zdej                  ddfd„Ze		 	 	 	 	 	 	 	 dde
j                  dz  d	e
j                  dz  d
e
j                  dz  de
j                  dz  de
j                  dz  dedz  dedz  dedz  deez  fd„«       Zˆ xZS )ÚFunnelBaseModelr   r   Nc                 ó„   •— t         ‰| �  |«       t        |«      | _        t	        |«      | _        | j                  «        y r5   )r   r    r   r6   r  ÚencoderÚ	post_initr-   s     €r0   r    zFunnelBaseModel.__init__æ  s4   ø€ Ü‰Ñ˜Ô ä*¨6Ó2ˆŒÜ$ VÓ,ˆŒð 	�‰Õr1   c                 ó.   — | j                   j                  S r5   ©r6   r%   ©r.   s    r0   Úget_input_embeddingsz$FunnelBaseModel.get_input_embeddingsï  ó   € Ø�‰×.Ñ.Ð.r1   Únew_embeddingsc                 ó&   — || j                   _        y r5   r€  ©r.   r„  s     r0   Úset_input_embeddingsz$FunnelBaseModel.set_input_embeddingsò  ó   € Ø*8ˆ�‰Õ'r1   r2   rI   rJ   Úposition_idsr3   rô   r  r   c	                 óV  — |�|n| j                   j                  }|�|n| j                   j                  }|�|n| j                   j                  }|�|�t	        d«      ‚|�#| j                  ||«       |j                  «       }
n!|�|j                  «       d d }
nt	        d«      ‚|�|j                  n|j                  }|€t        j                  |
|¬«      }|€&t        j                  |
t        j                  |¬«      }| j                  ||¬«      }| j                  ||||||¬«      }|S )NúDYou cannot specify both input_ids and inputs_embeds at the same timera   ú5You have to specify either input_ids or inputs_embeds©rP   r`   ©r3   ©rI   rJ   rô   r  r   )r   rô   r  r   Ú
ValueErrorÚ%warn_if_padding_and_no_attention_maskrL   rP   r=   ÚonesrÖ   r’   r6   r}  )r.   r2   rI   rJ   r‰  r3   rô   r  r   ÚkwargsÚinput_shaperP   Úencoder_outputss                r0   r7   zFunnelBaseModel.forwardõ  s;  € ð 2CÐ1NÑ-ÐTX×T_ÑT_×TqÑTqÐà$8Ð$DÑ È$Ï+É+×JjÑJjð 	ð &1Ð%<‘kÀ$Ç+Á+×BYÑBYˆàÐ  ]Ð%>ÜÐcÓdÐdØÐ"Ø×6Ñ6°yÀ.ÔQØ#Ÿ.™.Ó*‰KØÐ&Ø'×,Ñ,Ó.¨s°Ð3‰KäÐTÓUÐUà%.Ð%:�×!Ò!À×@TÑ@TˆàÐ!Ü"ŸZ™Z¨¸FÔCˆNØÐ!Ü"Ÿ[™[¨¼E¿J¹JÈvÔVˆNàŸ™¨	À˜ÓOˆàŸ,™,ØØ)Ø)Ø/Ø!5Ø#ð 'ó 
ˆð Ðr1   ©NNNNNNNN©r:   r;   r<   r   r    r   r!   r‚  r‡  r   r=   r>   r  r¡   r   r7   r?   r@   s   @r0   r{  r{  ß  s  ø„ ð˜|ð °õ ð/ b§l¡ló /ð9°2·<±<ð 9ÀDó 9ð ð *.Ø.2Ø.2Ø,0Ø-1Ø)-Ø,0Ø#'ñ.à—<‘< $Ñ&ð.ð Ÿ™ tÑ+ð.ð Ÿ™ tÑ+ð	.ð
 —l‘l TÑ)ð.ð —|‘| dÑ*ð.ð   $™;ð.ð # T™kð.ð ˜D‘[ð.ð 
�Ñ	 ò.ó ô.r1   r{  c                   ó.  ‡ — e Zd Zdeddfˆ fd„Zdej                  fd„Zdej                  ddfd„Ze		 	 	 	 	 	 	 dde
j                  dz  d	e
j                  dz  d
e
j                  dz  de
j                  dz  dedz  dedz  dedz  deez  fd„«       Zˆ xZS )ÚFunnelModelr   r   Nc                 ó²   •— t         ‰| �  |«       || _        t        |«      | _        t        |«      | _        t        |«      | _        | j                  «        y r5   )
r   r    r   r   r6   r  r}  rA  Údecoderr~  r-   s     €r0   r    zFunnelModel.__init__)  sG   ø€ Ü‰Ñ˜Ô ØˆŒÜ*¨6Ó2ˆŒÜ$ VÓ,ˆŒÜ$ VÓ,ˆŒð 	�‰Õr1   c                 ó.   — | j                   j                  S r5   r€  r�  s    r0   r‚  z FunnelModel.get_input_embeddings3  rƒ  r1   r„  c                 ó&   — || j                   _        y r5   r€  r†  s     r0   r‡  z FunnelModel.set_input_embeddings6  rˆ  r1   r2   rI   rJ   r3   rô   r  r   c           	      óÊ  — |�|n| j                   j                  }|�|n| j                   j                  }|�|n| j                   j                  }|�|�t	        d«      ‚|�#| j                  ||«       |j                  «       }	n!|�|j                  «       d d }	nt	        d«      ‚|�|j                  n|j                  }
|€t        j                  |	|
¬«      }|€&t        j                  |	t        j                  |
¬«      }| j                  ||¬«      }| j                  ||||d|¬«      }| j                  |d	   |d
   | j                   j                  d	      |||||¬«      }|s6d	}|d	   f}|r|d
z  }||d
   ||   z   fz   }|r|d
z  }||d   ||   z   fz   }|S t!        |d	   |r|j"                  |j"                  z   nd |r|j$                  |j$                  z   ¬«      S d ¬«      S )Nr‹  ra   rŒ  r�  r`   rŽ  Tr�  r   r   )rE  rF  rI   rJ   rô   r  r   rC   r(  )r   rô   r  r   r�  r‘  rL   rP   r=   r’  rÖ   r’   r6   r}  r›  r  r   r*  r+  )r.   r2   rI   rJ   r3   rô   r  r   r“  r”  rP   r•  Údecoder_outputsÚidxÚoutputss                  r0   r7   zFunnelModel.forward9  s?  € ð 2CÐ1NÑ-ÐTX×T_ÑT_×TqÑTqÐà$8Ð$DÑ È$Ï+É+×JjÑJjð 	ð &1Ð%<‘kÀ$Ç+Á+×BYÑBYˆàÐ  ]Ð%>ÜÐcÓdÐdØÐ"Ø×6Ñ6°yÀ.ÔQØ#Ÿ.™.Ó*‰KØÐ&Ø'×,Ñ,Ó.¨s°Ð3‰KäÐTÓUÐUà%.Ð%:�×!Ò!À×@TÑ@TˆàÐ!Ü"ŸZ™Z¨¸FÔCˆNØÐ!Ü"Ÿ[™[¨¼E¿J¹JÈvÔVˆNàŸ™¨	À˜ÓOˆàŸ,™,ØØ)Ø)Ø/Ø!%Ø#ð 'ó 
ˆð Ÿ,™,Ø(¨Ñ+Ø.¨qÑ1°$·+±+×2IÑ2IÈ!Ñ2LÑMØ)Ø)Ø/Ø!5Ø#ð 'ó 
ˆñ ØˆCØ& qÑ)Ð+ˆGÙ#Ø�q‘�Ø! _°QÑ%7¸/È#Ñ:NÑ%NÐ$PÑP�Ù Ø�q‘�Ø! _°QÑ%7¸/È#Ñ:NÑ%NÐ$PÑP�ØˆNäØ-¨aÑ0á#ð +×8Ñ8¸?×;XÑ;XÒXàÙTe˜×2Ñ2°_×5OÑ5OÑOô
ð 	
ð
 lpô
ð 	
r1   )NNNNNNNr—  r@   s   @r0   r™  r™  '  sõ   ø„ ð˜|ð °õ ð/ b§l¡ló /ð9°2·<±<ð 9ÀDó 9ð ð *.Ø.2Ø.2Ø-1Ø)-Ø,0Ø#'ñH
à—<‘< $Ñ&ðH
ð Ÿ™ tÑ+ðH
ð Ÿ™ tÑ+ð	H
ð
 —|‘| dÑ*ðH
ð   $™;ðH
ð # T™kðH
ð ˜D‘[ðH
ð 
�Ñ	 òH
ó ôH
r1   r™  zŒ
    Funnel Transformer model with a binary classification head on top as used during pretraining for identifying
    generated tokens.
    c                   ó
  ‡ — e Zd Zdeddfˆ fd„Ze	 	 	 	 	 	 	 	 ddej                  dz  dej                  dz  dej                  dz  dej                  dz  d	ej                  dz  d
edz  dedz  dedz  de	e
z  fd„«       Zˆ xZS )ÚFunnelForPreTrainingr   r   Nc                 ó„   •— t         ‰| �  |«       t        |«      | _        t	        |«      | _        | j                  «        y r5   )r   r    r™  rV  rL  Údiscriminator_predictionsr~  r-   s     €r0   r    zFunnelForPreTraining.__init__Œ  s3   ø€ Ü‰Ñ˜Ô ä! &Ó)ˆŒÜ)GÈÓ)OˆÔ&à�‰Õr1   r2   rI   rJ   r3   Úlabelsrô   r  r   c	           	      ó`  — |�|n| j                   j                  }| j                  |||||||¬«      }
|
d   }| j                  |«      }d}|�«t	        j
                  «       }|�a|j                  d|j                  d   «      dk(  }|j                  d|j                  d   «      |   }||   } |||j                  «       «      }n4 ||j                  d|j                  d   «      |j                  «       «      }|s|f|
dd z   }|�|f|z   S |S t        |||
j                  |
j                  ¬«      S )a"  
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the ELECTRA-style loss. Input should be a sequence of tokens (see `input_ids`
            docstring) Indices should be in `[0, 1]`:

            - 0 indicates the token is an original token,
            - 1 indicates the token was replaced.

        Examples:

        ```python
        >>> from transformers import AutoTokenizer, FunnelForPreTraining
        >>> import torch

        >>> tokenizer = AutoTokenizer.from_pretrained("funnel-transformer/small")
        >>> model = FunnelForPreTraining.from_pretrained("funnel-transformer/small")

        >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
        >>> logits = model(**inputs).logits
        ```N©rI   rJ   r3   rô   r  r   r   ra   r   ©rx  rS  r*  r+  )r   r   rV  r¥  r   r   rö   r‘   r÷   rw  r*  r+  )r.   r2   rI   rJ   r3   r¦  rô   r  r   r“  rP  Údiscriminator_sequence_outputrS  rx  Úloss_fctÚactive_lossÚactive_logitsÚactive_labelsrº   s                      r0   r7   zFunnelForPreTraining.forward”  sj  € ðB &1Ð%<‘kÀ$Ç+Á+×BYÑBYˆà&*§k¡kØØ)Ø)Ø'Ø/Ø!5Ø#ð '2ó '
Ð#ð )DÀAÑ(FÐ%à×/Ñ/Ð0MÓNˆàˆØÐÜ×+Ñ+Ó-ˆHØÐ)Ø,×1Ñ1°"Ð6S×6YÑ6YÐZ[Ñ6\Ó]ÐabÑb�Ø &§¡¨BÐ0M×0SÑ0SÐTUÑ0VÓ WÐXcÑ d�Ø & {Ñ 3�Ù ¨}×/BÑ/BÓ/DÓE‘á §¡¨BÐ0M×0SÑ0SÐTUÑ0VÓ WÐY_×YeÑYeÓYgÓh�áØ�YÐ!<¸Q¸RÐ!@Ñ@ˆFØ)-Ð)9�T�G˜fÑ$ÐE¸vÐEä)ØØØ5×CÑCØ2×=Ñ=ô	
ð 	
r1   r–  )r:   r;   r<   r   r    r   r=   r>   r  r¡   rw  r7   r?   r@   s   @r0   r£  r£  …  sæ   ø„ ð˜|ð °õ ð ð *.Ø.2Ø.2Ø-1Ø&*Ø)-Ø,0Ø#'ñC
à—<‘< $Ñ&ðC
ð Ÿ™ tÑ+ðC
ð Ÿ™ tÑ+ð	C
ð
 —|‘| dÑ*ðC
ð —‘˜tÑ#ðC
ð   $™;ðC
ð # T™kðC
ð ˜D‘[ðC
ð 
Ð+Ñ	+òC
ó ôC
r1   r£  c                   óV  ‡ — e Zd ZddiZdeddfˆ fd„Zdej                  fd„Zdej                  ddfd	„Z
e	 	 	 	 	 	 	 	 dd
ej                  dz  dej                  dz  dej                  dz  dej                  dz  dej                  dz  dedz  dedz  dedz  deez  fd„«       Zˆ xZS )ÚFunnelForMaskedLMzlm_head.weightz(funnel.embeddings.word_embeddings.weightr   r   Nc                 óÂ   •— t         ‰| �  |«       t        |«      | _        t	        j
                  |j                  |j                  «      | _        | j                  «        y r5   )
r   r    r™  rV  r   rÑ   r'   r"   Úlm_headr~  r-   s     €r0   r    zFunnelForMaskedLM.__init__ß  sD   ø€ Ü‰Ñ˜Ô ä! &Ó)ˆŒÜ—y‘y §¡°×1BÑ1BÓCˆŒð 	�‰Õr1   c                 ó   — | j                   S r5   ©r²  r�  s    r0   Úget_output_embeddingsz'FunnelForMaskedLM.get_output_embeddingsè  s   € Ø�|‰|Ðr1   r„  c                 ó   — || _         y r5   r´  r†  s     r0   Úset_output_embeddingsz'FunnelForMaskedLM.set_output_embeddingsë  s	   € Ø%ˆ�r1   r2   rI   rJ   r3   r¦  rô   r  r   c	           	      ó–  — |�|n| j                   j                  }| j                  |||||||¬«      }
|
d   }| j                  |«      }d}|�Ft	        «       } ||j                  d| j                   j                  «      |j                  d«      «      }|s|f|
dd z   }|�|f|z   S |S t        |||
j                  |
j                  ¬«      S )a¢  
        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]`
        Nr¨  r   ra   r   r©  )
r   r   rV  r²  r   rö   r"   r   r*  r+  )r.   r2   rI   rJ   r3   r¦  rô   r  r   r“  r¡  r)  Úprediction_logitsÚmasked_lm_lossr«  rº   s                   r0   r7   zFunnelForMaskedLM.forwardî  sö   € ð& &1Ð%<‘kÀ$Ç+Á+×BYÑBYˆà—+‘+ØØ)Ø)Ø'Ø/Ø!5Ø#ð ó 
ˆð $ A™JÐØ ŸL™LÐ):Ó;ÐàˆØÐÜ'Ó)ˆHÙ%Ð&7×&<Ñ&<¸RÀÇÁ×AWÑAWÓ&XÐZ`×ZeÑZeÐfhÓZiÓjˆNáØ'Ð)¨G°A°B¨KÑ7ˆFØ3AÐ3M�^Ð%¨Ñ.ÐYÐSYÐYäØØ$Ø!×/Ñ/Ø×)Ñ)ô	
ð 	
r1   r–  )r:   r;   r<   Ú_tied_weights_keysr   r    r   rÑ   rµ  r!   r·  r   r=   r>   r  r¡   r   r7   r?   r@   s   @r0   r°  r°  Û  s  ø„ à*Ð,VÐWÐð˜|ð °õ ð r§y¡yó ð&°B·L±Lð &ÀTó &ð ð *.Ø.2Ø.2Ø-1Ø&*Ø)-Ø,0Ø#'ñ/
à—<‘< $Ñ&ð/
ð Ÿ™ tÑ+ð/
ð Ÿ™ tÑ+ð	/
ð
 —|‘| dÑ*ð/
ð —‘˜tÑ#ð/
ð   $™;ð/
ð # T™kð/
ð ˜D‘[ð/
ð 
�Ñ	ò/
ó ô/
r1   r°  zº
    Funnel Transformer Model with a sequence classification/regression head on top (two linear layer on top of the
    first timestep of the last hidden state) e.g. for GLUE tasks.
    c                   ó
  ‡ — e Zd Zdeddfˆ fd„Ze	 	 	 	 	 	 	 	 ddej                  dz  dej                  dz  dej                  dz  dej                  dz  d	ej                  dz  d
edz  dedz  dedz  de	e
z  fd„«       Zˆ xZS )ÚFunnelForSequenceClassificationr   r   Nc                 óÊ   •— t         ‰| �  |«       |j                  | _        || _        t	        |«      | _        t        ||j                  «      | _        | j                  «        y r5   )	r   r    Ú
num_labelsr   r{  rV  rn  Ú
classifierr~  r-   s     €r0   r    z(FunnelForSequenceClassification.__init__(  sN   ø€ Ü‰Ñ˜Ô Ø ×+Ñ+ˆŒØˆŒä% fÓ-ˆŒÜ2°6¸6×;LÑ;LÓMˆŒà�‰Õr1   r2   rI   rJ   r3   r¦  rô   r  r   c	           	      ó,  — |�|n| j                   j                  }| j                  |||||||¬«      }
|
d   }|dd…df   }| j                  |«      }d}|��‡| j                   j                  €�| j
                  dk(  rd| j                   _        nl| j
                  dkD  rL|j                  t        j                  k(  s|j                  t        j                  k(  rd| j                   _        nd| j                   _        | j                   j                  dk(  rIt        «       }| j
                  dk(  r& ||j                  «       |j                  «       «      }nŒ |||«      }n‚| j                   j                  dk(  r=t        «       } ||j                  d| j
                  «      |j                  d«      «      }n,| j                   j                  dk(  rt        «       } |||«      }|s|f|
dd z   }|�|f|z   S |S t        |||
j                   |
j"                  ¬	«      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).
        Nr¨  r   r   Ú
regressionÚsingle_label_classificationÚmulti_label_classificationra   r©  )r   r   rV  rÀ  Úproblem_typer¿  rO   r=   r’   rÀ   r   rR  r   rö   r   r   r*  r+  )r.   r2   rI   rJ   r3   r¦  rô   r  r   r“  r¡  r)  Úpooled_outputrS  rx  r«  rº   s                    r0   r7   z'FunnelForSequenceClassification.forward2  sã  € ð& &1Ð%<‘kÀ$Ç+Á+×BYÑBYˆà—+‘+ØØ)Ø)Ø'Ø/Ø!5Ø#ð ó 
ˆð $ A™JÐØ)ª!¨Q¨$Ñ/ˆØ—‘ Ó/ˆàˆØÑØ�{‰{×'Ñ'Ð/Ø—?‘? aÒ'Ø/;�D—K‘KÕ,Ø—_‘_ qÒ(¨f¯l©l¼e¿j¹jÒ.HÈFÏLÉLÔ\a×\eÑ\eÒLeØ/L�D—K‘KÕ,à/K�D—K‘KÔ,à�{‰{×'Ñ'¨<Ò7Ü"›9�Ø—?‘? aÒ'Ù# F§N¡NÓ$4°f·n±nÓ6FÓG‘Dá# F¨FÓ3‘DØ—‘×)Ñ)Ð-JÒJÜ+Ó-�Ù §¡¨B°·±Ó @À&Ç+Á+ÈbÃ/ÓR‘Ø—‘×)Ñ)Ð-IÒIÜ,Ó.�Ù ¨Ó/�áØ�Y ¨¨ Ñ,ˆFØ)-Ð)9�T�G˜fÑ$ÐE¸vÐEä'ØØØ!×/Ñ/Ø×)Ñ)ô	
ð 	
r1   r–  )r:   r;   r<   r   r    r   r=   r>   r  r¡   r   r7   r?   r@   s   @r0   r½  r½  !  sæ   ø„ ð˜|ð °õ ð ð *.Ø.2Ø.2Ø-1Ø&*Ø)-Ø,0Ø#'ñB
à—<‘< $Ñ&ðB
ð Ÿ™ tÑ+ðB
ð Ÿ™ tÑ+ð	B
ð
 —|‘| dÑ*ðB
ð —‘˜tÑ#ðB
ð   $™;ðB
ð # T™kðB
ð ˜D‘[ðB
ð 
Ð)Ñ	)òB
ó ôB
r1   r½  c                   ó
  ‡ — e Zd Zdeddfˆ fd„Ze	 	 	 	 	 	 	 	 ddej                  dz  dej                  dz  dej                  dz  dej                  dz  d	ej                  dz  d
edz  dedz  dedz  de	e
z  fd„«       Zˆ xZS )ÚFunnelForMultipleChoicer   r   Nc                 ó†   •— t         ‰| �  |«       t        |«      | _        t	        |d«      | _        | j                  «        y r¾   )r   r    r{  rV  rn  rÀ  r~  r-   s     €r0   r    z FunnelForMultipleChoice.__init__z  s4   ø€ Ü‰Ñ˜Ô ä% fÓ-ˆŒÜ2°6¸1Ó=ˆŒà�‰Õr1   r2   rI   rJ   r3   r¦  rô   r  r   c	           	      óî  — |�|n| j                   j                  }|�|j                  d   n|j                  d   }
|�!|j                  d|j	                  d«      «      nd}|�!|j                  d|j	                  d«      «      nd}|�!|j                  d|j	                  d«      «      nd}|�1|j                  d|j	                  d«      |j	                  d«      «      nd}| j                  |||||||¬«      }|d   }|dd…df   }| j                  |«      }|j                  d|
«      }d}|�t        «       } |||«      }|s|f|dd z   }|�|f|z   S |S t        |||j                  |j                  ¬«      S )aJ  
        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)
        Nr   ra   éþÿÿÿr¨  r   r©  )r   r   r‘   rö   rL   rV  rÀ  r   r   r*  r+  )r.   r2   rI   rJ   r3   r¦  rô   r  r   r“  Únum_choicesr¡  r)  rÆ  rS  Úreshaped_logitsrx  r«  rº   s                      r0   r7   zFunnelForMultipleChoice.forward‚  s³  € ð& &1Ð%<‘kÀ$Ç+Á+×BYÑBYˆØ,5Ð,A�i—o‘o aÒ(À}×GZÑGZÐ[\ÑG]ˆà>GÐ>S�I—N‘N 2 y§~¡~°bÓ'9Ô:ÐY]ˆ	ØM[ÐMg˜×,Ñ,¨R°×1DÑ1DÀRÓ1HÔIÐmqˆØM[ÐMg˜×,Ñ,¨R°×1DÑ1DÀRÓ1HÔIÐmqˆð Ð(ð ×Ñ˜r =×#5Ñ#5°bÓ#9¸=×;MÑ;MÈbÓ;QÔRàð 	ð —+‘+ØØ)Ø)Ø'Ø/Ø!5Ø#ð ó 
ˆð $ A™JÐØ)ª!¨Q¨$Ñ/ˆØ—‘ Ó/ˆØ Ÿ+™+ b¨+Ó6ˆàˆØÐÜ'Ó)ˆHÙ˜O¨VÓ4ˆDáØ%Ð'¨'°!°"¨+Ñ5ˆFØ)-Ð)9�T�G˜fÑ$ÐE¸vÐEä(ØØ"Ø!×/Ñ/Ø×)Ñ)ô	
ð 	
r1   r–  )r:   r;   r<   r   r    r   r=   r>   r  r¡   r   r7   r?   r@   s   @r0   rÈ  rÈ  x  sÛ   ø„ ð˜|ð °õ ð ð *.Ø.2Ø.2Ø-1Ø&*Ø)-Ø,0Ø#'ñ;
à—<‘< $Ñ&ð;
ð Ÿ™ tÑ+ð;
ð Ÿ™ tÑ+ð	;
ð
 —|‘| dÑ*ð;
ð —‘˜tÑ#ð;
ð   $™;ð;
ð # T™kð;
ð ˜D‘[ð;
ð 
Ð*Ñ	*ò;
ó ô;
r1   rÈ  c                   ó
  ‡ — e Zd Zdeddfˆ fd„Ze	 	 	 	 	 	 	 	 ddej                  dz  dej                  dz  dej                  dz  dej                  dz  d	ej                  dz  d
edz  dedz  dedz  de	e
z  fd„«       Zˆ xZS )ÚFunnelForTokenClassificationr   r   Nc                 ó,  •— t         ‰| �  |«       |j                  | _        t        |«      | _        t        j                  |j                  «      | _        t        j                  |j                  |j                  «      | _        | j                  «        y r5   )r   r    r¿  r™  rV  r   r*   r+   r,   rÑ   r#   rÀ  r~  r-   s     €r0   r    z%FunnelForTokenClassification.__init__Ã  si   ø€ Ü‰Ñ˜Ô Ø ×+Ñ+ˆŒä! &Ó)ˆŒÜ—z‘z &×"7Ñ"7Ó8ˆŒÜŸ)™) F×$6Ñ$6¸×8IÑ8IÓJˆŒð 	�‰Õr1   r2   rI   rJ   r3   r¦  rô   r  r   c	           	      ó¤  — |�|n| j                   j                  }| j                  |||||||¬«      }
|
d   }| j                  |«      }| j	                  |«      }d}|�<t        «       } ||j                  d| j                  «      |j                  d«      «      }|s|f|
dd z   }|�|f|z   S |S t        |||
j                  |
j                  ¬«      S )zÛ
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
        Nr¨  r   ra   r   r©  )r   r   rV  r,   rÀ  r   rö   r¿  r   r*  r+  )r.   r2   rI   rJ   r3   r¦  rô   r  r   r“  r¡  r)  rS  rx  r«  rº   s                   r0   r7   z$FunnelForTokenClassification.forwardÎ  sô   € ð" &1Ð%<‘kÀ$Ç+Á+×BYÑBYˆà—+‘+ØØ)Ø)Ø'Ø/Ø!5Ø#ð ó 
ˆð $ A™JÐØ ŸL™LÐ):Ó;ÐØ—‘Ð!2Ó3ˆàˆØÐÜ'Ó)ˆHÙ˜FŸK™K¨¨D¯O©OÓ<¸f¿k¹kÈ"»oÓNˆDáØ�Y ¨¨ Ñ,ˆFØ)-Ð)9�T�G˜fÑ$ÐE¸vÐEä$ØØØ!×/Ñ/Ø×)Ñ)ô	
ð 	
r1   r–  )r:   r;   r<   r   r    r   r=   r>   r  r¡   r   r7   r?   r@   s   @r0   rÏ  rÏ  Á  sÛ   ø„ ð	˜|ð 	°õ 	ð ð *.Ø.2Ø.2Ø-1Ø&*Ø)-Ø,0Ø#'ñ.
à—<‘< $Ñ&ð.
ð Ÿ™ tÑ+ð.
ð Ÿ™ tÑ+ð	.
ð
 —|‘| dÑ*ð.
ð —‘˜tÑ#ð.
ð   $™;ð.
ð # T™kð.
ð ˜D‘[ð.
ð 
Ð&Ñ	&ò.
ó ô.
r1   rÏ  c                   ó*  ‡ — e Zd Zdeddfˆ fd„Ze	 	 	 	 	 	 	 	 	 ddej                  dz  dej                  dz  dej                  dz  dej                  dz  d	ej                  dz  d
ej                  dz  dedz  dedz  dedz  de	e
z  fd„«       Zˆ xZS )ÚFunnelForQuestionAnsweringr   r   Nc                 óä   •— t         ‰| �  |«       |j                  | _        t        |«      | _        t        j                  |j                  |j                  «      | _        | j                  «        y r5   )
r   r    r¿  r™  rV  r   rÑ   r#   Ú
qa_outputsr~  r-   s     €r0   r    z#FunnelForQuestionAnswering.__init__  sS   ø€ Ü‰Ñ˜Ô Ø ×+Ñ+ˆŒä! &Ó)ˆŒÜŸ)™) F×$6Ñ$6¸×8IÑ8IÓJˆŒð 	�‰Õr1   r2   rI   rJ   r3   Ústart_positionsÚend_positionsrô   r  r   c
           	      ó$  — |	�|	n| j                   j                  }	| j                  |||||||	¬«      }|d   }| j                  |«      }|j	                  dd¬«      \  }}|j                  d«      j                  «       }|j                  d«      j                  «       }d }|�·|�µt        |j                  «       «      dkD  r|j                  d«      }t        |j                  «       «      dkD  r|j                  d«      }|j                  d«      }|j                  d|«      }|j                  d|«      }t        |¬«      } |||«      } |||«      }||z   dz  }|	s||f|dd  z   }|�|f|z   S |S t        ||||j                  |j                  ¬«      S )	Nr¨  r   r   ra   rb   )Úignore_indexrC   )rx  Ústart_logitsÚ
end_logitsr*  r+  )r   r   rV  rÕ  rê   rR  Ú
contiguousrI  rL   ÚsquezeÚclampr   r   r*  r+  )r.   r2   rI   rJ   r3   rÖ  r×  rô   r  r   r“  r¡  r)  rS  rÚ  rÛ  Ú
total_lossÚignored_indexr«  Ú
start_lossÚend_lossrº   s                         r0   r7   z"FunnelForQuestionAnswering.forward  s¾  € ð &1Ð%<‘kÀ$Ç+Á+×BYÑBYˆà—+‘+ØØ)Ø)Ø'Ø/Ø!5Ø#ð ó 
ˆð $ A™JÐà—‘Ð!2Ó3ˆØ#)§<¡<°°r <Ó#:Ñ ˆ�jØ#×+Ñ+¨BÓ/×:Ñ:Ó<ˆØ×'Ñ'¨Ó+×6Ñ6Ó8ˆ
àˆ
ØÐ&¨=Ð+Dä�?×'Ñ'Ó)Ó*¨QÒ.Ø"1×"8Ñ"8¸Ó"<�Ü�=×%Ñ%Ó'Ó(¨1Ò,Ø -× 5Ñ 5°bÓ 9�à(×-Ñ-¨aÓ0ˆMØ-×3Ñ3°A°}ÓEˆOØ)×/Ñ/°°=ÓAˆMä'°]ÔCˆHÙ! ,°Ó@ˆJÙ 
¨MÓ:ˆHØ$ xÑ/°1Ñ4ˆJáØ" JÐ/°'¸!¸"°+Ñ=ˆFØ/9Ð/E�Z�M FÑ*ÐQÈ6ÐQä+ØØ%Ø!Ø!×/Ñ/Ø×)Ñ)ô
ð 	
r1   )	NNNNNNNNN)r:   r;   r<   r   r    r   r=   r>   r  r¡   r   r7   r?   r@   s   @r0   rÓ  rÓ     sñ   ø„ ð˜|ð °õ ð ð *.Ø.2Ø.2Ø-1Ø/3Ø-1Ø)-Ø,0Ø#'ñ;
à—<‘< $Ñ&ð;
ð Ÿ™ tÑ+ð;
ð Ÿ™ tÑ+ð	;
ð
 —|‘| dÑ*ð;
ð Ÿ™¨Ñ,ð;
ð —|‘| dÑ*ð;
ð   $™;ð;
ð # T™kð;
ð ˜D‘[ð;
ð 
Ð-Ñ	-ò;
ó ô;
r1   rÓ  )	r{  r°  rÈ  r£  rÓ  r½  rÏ  r™  rU  )TF)=r¿   Údataclassesr   Únumpyr_  r=   r   Útorch.nnr   r   r   Ú r	   ra  Úactivationsr
   Úmodeling_outputsr   r   r   r   r   r   Úmodeling_utilsr   Úutilsr   r   r   Úconfiguration_funnelr   Ú
get_loggerr:   Úloggerrø   ÚModuler   rB   r>   rÀ   rÊ   rÌ   r  r  r  r  r?  rA  rL  rU  rn  rw  r{  r™  r£  r°  r½  rÈ  rÏ  rÓ  Ú__all__r#  r1   r0   ú<module>rð     sÒ  ðñ (å !ã Û Ý ß AÑ Aå &Ý !÷÷ õ .ß 9Ñ 9Ý .ð 
ˆ×	Ñ	˜HÓ	%€ð 
€ô�r—y‘yô ô"A ˜rŸy™yô A ðH¨E¯L©Lð Àsð ÐSVð Ð[`×[gÑ[gó ô MG "§)¡)ô MGô`+˜BŸI™Iô +ô&E�"—)‘)ô Eô&<u�B—I‘Iô <uð@ diñØ‡|�|ðØ ðØ.1ðØAEðØ\`ðà
‡\�\óô,.u�B—I‘Iô .uôb R§Y¡Yô ð  ô_˜Oó _ó ð_ô>'˜rŸy™yô 'ñ ðôð
 ô7 ó 7ó óð7ñ ðôô?Ð+ó ?óð?ðD ôZ
Ð'ó Z
ó ðZ
ñz ðôôM
Ð0ó M
óðM
ð` ôB
Ð-ó B
ó ðB
ñJ ðôôN
Ð&;ó N
óðN
ðb ôE
Ð3ó E
ó ðE
ðP ô;
Ð#8ó ;
ó ð;
ð| ôG
Ð!6ó G
ó ðG
òT
�r1   