o Zh@sdZddlZddlZddlZddlZddlmZmZmZm Z ddl Z ddl m Z ddl m Z mZmZddlmZddlmZmZmZdd lmZdd lmZdd lmZmZmZmZmZm Z m!Z!dd l"m#Z#dd l$m%Z%m&Z&m'Z'ddl(m)Z)m*Z*m+Z+m,Z,m-Z-m.Z.m/Z/m0Z0ddl1m2Z2m3Z3ddl4m5Z5e-rddl6m7Z7ddl8m9Z9e0:e;ZdZ?Gddde j@ZAzddlBmCZCeCZAed>eQZXe,Gd?d@d@eQZYgdAZZdS)BzPyTorch T5 model.N)ListOptionalTupleUnion)nn)BCEWithLogitsLossCrossEntropyLossMSELoss)ACT2FN)Cache DynamicCacheEncoderDecoderCache)GenerationMixin)AttentionMaskConverter)BaseModelOutput)BaseModelOutputWithPastAndCrossAttentionsSeq2SeqLMOutputSeq2SeqModelOutput#Seq2SeqQuestionAnsweringModelOutputSeq2SeqSequenceClassifierOutputTokenClassifierOutput)PreTrainedModel)ALL_LAYERNORM_LAYERS find_pruneable_heads_and_indicesprune_linear_layer) DUMMY_INPUTS DUMMY_MASKadd_start_docstringsauto_docstringis_torch_flex_attn_availableis_torch_fx_proxyis_torchdynamo_compilinglogging)assert_device_mapget_device_map)T5Config) BlockMask)make_flex_block_causal_maskc Cs@zddl}ddl}ddl}Wn tytdwtj|}t d||j |}g}i} |D]\} } t d| d| |j || } | | | | | <q6|D]} | d} tdd | Dr{t d d| | | dqXd | d vrt d d| | | dqX|}| | } | D]}|d |r|d|}n|g}|ddvrt|d}n|ddkrt|d}|d}n|ddkrt|d}|d}n|ddkrt|d}|d}n|ddkrt|drt|d}nt|drt|d}n|ddkrt|d}n||ddks#|ddkr)t|d}nh|ddkr6t|d}n[|dd krE| dd!krEq|dd!krRt|d"}n?|dd#krrt|dkrr|drrt|d$|d}qz t||d}Wntyt d d| Yqwt|dkrt|d}||}q|ddvrt|d}|dd%krt d&| jd'| || } z|j| jkrtd(|jd)| jd*Wnty}z |j|j| jf7_d}~wwt d+| t| |j!|_"| | dqXt d,d-| #d.|S)/z'Load tf checkpoints in a pytorch model.rNzLoading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see https://www.tensorflow.org/install/ for installation instructions.z&Converting TensorFlow checkpoint from zLoading TF weight z with shape /css|]}|dvVqdS))Zadam_vZadam_mZAdamWeightDecayOptimizerZAdamWeightDecayOptimizer_1Z global_stepN).0nr+r+Q/var/www/auris/lib/python3.10/site-packages/transformers/models/t5/modeling_t5.py hs  z(load_tf_weights_in_t5..z Skipping Z_slot_z [A-Za-z]+_\d+z_(\d+))Zkernelscale embeddingweightZself_attentionlayerZenc_dec_attentionr&Zdense_relu_denseZrms_norm layer_normfinal_layer_normr1Z output_biasbetabiasZsquad classifierdecoderlogitslm_headwiZwi_r2z"Transposing numpy weight of shape z for zPointer shape z and array shape mismatchedzInitialize PyTorch weight z%Weights not copied to PyTorch model: z, .)$renumpyZ tensorflow ImportErrorloggererrorospathabspathinfotrainZlist_variablesZ load_variableappendsplitanyjoinpop fullmatchgetattrhasattrlenisdigitAttributeErrorintshape transpose ValueErrorAssertionErrorargstorchZ from_numpyZastypefloat32datakeys)modelconfigZtf_checkpoint_pathrAnptfZtf_pathZ init_varsnamesZ tf_weightsnamerWarrayZtxt_nameZpointerZm_nameZ scope_namesnumer+r+r.load_tf_weights_in_t5Ks                                 *   riaR This is an experimental feature and is a subject to change at a moment's notice. Uses a device map to distribute attention modules of the model across several devices. If no device map is given, it will evenly distribute blocks across all devices. Args: device_map (`Dict[int, list]`, *optional*): A dictionary that maps attention modules to devices. Note that the embedding module and LMHead are always automatically mapped to the first device (for esoteric reasons). That means that the first device should have fewer attention modules mapped to it than other devices. For reference, the t5 models have the following number of attention modules: - google-t5/t5-small: 6 - google-t5/t5-base: 12 - google-t5/t5-large: 24 - google-t5/t5-3b: 24 - google-t5/t5-11b: 24 Example: ```python # Here is an example of a device map on a machine with 4 GPUs using google-t5/t5-3b, which has a total of 24 attention modules: model = T5ForConditionalGeneration.from_pretrained("google-t5/t5-3b") device_map = { 0: [0, 1, 2], 1: [3, 4, 5, 6, 7, 8, 9], 2: [10, 11, 12, 13, 14, 15, 16], 3: [17, 18, 19, 20, 21, 22, 23], } model.parallelize(device_map) ``` a4 Moves the model to cpu from a model parallel state. Example: ```python # On a 4 GPU machine with google-t5/t5-3b: model = T5ForConditionalGeneration.from_pretrained("google-t5/t5-3b") device_map = { 0: [0, 1, 2], 1: [3, 4, 5, 6, 7, 8, 9], 2: [10, 11, 12, 13, 14, 15, 16], 3: [17, 18, 19, 20, 21, 22, 23], } model.parallelize(device_map) # Splits the model across several devices model.deparallelize() # Put the model back on cpu and cleans memory by calling torch.cuda.empty_cache() ``` cs&eZdZdfdd ZddZZS) T5LayerNormư>cs&ttt||_||_dS)zc Construct a layernorm module in the T5 style. No bias and no subtraction of mean. N)super__init__r Parameterr\onesr3variance_epsilon)self hidden_sizeeps __class__r+r.rms  zT5LayerNorm.__init__cCs\|tjdjddd}|t||j}|jjtj tj fvr)||jj}|j|S)Nr5r0T)Zkeepdim) tor\r]powmeanZrsqrtrpr3dtypefloat16Zbfloat16)rq hidden_statesZvariancer+r+r.forwards  zT5LayerNorm.forward)rk)__name__ __module__ __qualname__rmr| __classcell__r+r+rtr.rjsrj) FusedRMSNormzODiscovered apex.normalization.FusedRMSNorm - will use it instead of T5LayerNormzBdiscovered apex but it failed to load, falling back to T5LayerNormc*eZdZdeffdd ZddZZS)T5DenseActDenseracsTttj|j|jdd|_tj|j|jdd|_t|j |_ t |j |_ dSNFr9)rlrmrLineard_modeld_ffr>woDropout dropout_ratedropoutr dense_act_fnactrqrartr+r.rms zT5DenseActDense.__init__cCsl||}||}||}t|jjtjr/|j|jjjkr/|jjjtj kr/| |jjj}||}|SN) r>rr isinstancerr3r\Tensorryint8rvrqr{r+r+r.r|s    zT5DenseActDense.forwardr}r~rr'rmr|rr+r+rtr.rsrcr)T5DenseGatedActDenseracsjttj|j|jdd|_tj|j|jdd|_tj|j|jdd|_t |j |_ t |j |_dSr)rlrmrrrrwi_0wi_1rrrrr rrrrtr+r.rm.s zT5DenseGatedActDense.__init__cCsz|||}||}||}||}t|jjtjr6|j |jjj kr6|jjj tj kr6| |jjj }||}|Sr) rrrrrrr3r\rryrrv)rqr{Z hidden_geluZ hidden_linearr+r+r.r|6s   zT5DenseGatedActDense.forwardrr+r+rtr.r-srcr) T5LayerFFracsJt|jrt||_nt||_t|j|jd|_ t |j |_ dS)Nrs)rlrmZ is_gated_actrDenseReluDenserrjrlayer_norm_epsilonr6rrrrrrtr+r.rmKs   zT5LayerFF.__init__cCs&||}||}|||}|Sr)r6rr)rqr{Zforwarded_statesr+r+r.r|Us  zT5LayerFF.forwardrr+r+rtr.rJs rcsleZdZ  ddedeeffdd ZddZedd d Z dddZ         dddZ Z S) T5AttentionFNra layer_idxcst|j|_||_|j|_|j|_|j|_|j|_|j |_ |j |_ |j |j|_ ||_|dur@|jr@td|jjdtj|j|j dd|_tj|j|j dd|_tj|j|j dd|_tj|j |jdd|_|jrxt|j|j |_t|_d|_dS)NzInstantiating a decoder z without passing `layer_idx` is not recommended and will to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` when creating this class.Fr)rlrm is_decoderhas_relative_attention_biasrelative_attention_num_bucketsrelative_attention_max_distancerd_kvkey_value_proj_dim num_headsn_headsrr inner_dimrrD warning_oncerur}rrqkvo Embeddingrelative_attention_biasset pruned_headsgradient_checkpointingrqrarrrtr+r.rm]s.  zT5Attention.__init__cCst|dkrdSt||j|j|j\}}t|j||_t|j||_t|j||_t|j |dd|_ |jt||_|j|j|_ |j ||_dS)Nrr&dim) rSrrrrrrrrrrunion)rqheadsindexr+r+r. prune_headss zT5Attention.prune_headsT cCsd}|r|d}||dktj|7}t|}n t|t| }|d}||k}|t||t||||tj}t|t ||d}|t |||7}|S)a Adapted from Mesh Tensorflow: https://github.com/tensorflow/mesh/blob/0cb87fe07da627bf0b7e60475d59f95ed6b5be3d/mesh_tensorflow/transformer/transformer_layers.py#L593 Translate relative position to a bucket number for relative attention. The relative position is defined as memory_position - query_position, i.e. the distance in tokens from the attending position to the attended-to position. If bidirectional=False, then positive relative positions are invalid. We use smaller buckets for small absolute relative_position and larger buckets for larger absolute relative_positions. All relative positions >=max_distance map to the same bucket. All relative positions <=-max_distance map to the same bucket. This should allow for more graceful generalization to longer sequences than the model has been trained on Args: relative_position: an int32 Tensor bidirectional: a boolean - whether the attention is bidirectional num_buckets: an integer max_distance: an integer Returns: a Tensor with the same shape as relative_position, containing int32 values in the range [0, num_buckets) rr5r&) rvr\longabsminZ zeros_likelogfloatmathZ full_likewhere)relative_position bidirectional num_buckets max_distanceZrelative_bucketsZ max_exactZis_smallZrelative_position_if_larger+r+r._relative_position_buckets*  z%T5Attention._relative_position_bucketc Cs|dur |jjj}|durtj|tj|ddddf}n |dddf|}tj|tj|ddddf}||}|j||j |j |j d}||} | gd d} | S)z%Compute binned relative position biasN)rydevice)rrr)r5rr&r) rr3rr\arangerrvrrrrZpermuteZ unsqueeze) rq query_length key_lengthrcache_positionZcontext_positionZmemory_positionrZrelative_position_bucketvaluesr+r+r. compute_biass    zT5Attention.compute_biasc Cs|jdd\} } |du} ||}|| d|j|jdd}|dur4|j|j}| r1|j }n|j }| r8|n|}| rO|durO|rO|j |j}|j |j}nE| |}||}|| d|j|jdd}|| d|j|jdd}|dur| s}| nd} ||||jd| i\}}| rd|j|j<t||dd}|dur|jd}|dur|n| dd}|jstjd|j| |f|j|jd }|jr|jrd|_n|j|||j| d }|dddd| dddf}|dur|ddddddd|jdf}||}|jr%t|jd}d |t|j<|dd|f}n|}||7}tjj |!dd "|}tjj#||j#|jd }|durL||}t||}|dd$}|| d|j%}|&|}|||f}| rt||f}|S)z Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states). Nr5r0r&rTr rry)rrrr)ptraining)'rWrviewrrrX is_updatedgetrZcross_attention_cacheself_attention_cacheZ key_cacheZ value_cacherrupdater\matmulrZzerosrryrrZ requires_gradrrrolistboolrZ functionalZsoftmaxrZtype_asr contiguousrr)rqr{maskkey_value_states position_biaspast_key_valuelayer_head_maskr use_cacheoutput_attentionsr batch_size seq_lengthZis_cross_attentionZ query_statesrZcurr_past_key_valueZcurrent_statesZ key_statesZ value_statesZscoresrZreal_seq_length causal_maskZposition_bias_maskedZ attn_weightsZ attn_outputoutputsr+r+r.r|sx           " &     zT5Attention.forwardFN)Trr)NN) NNNNNNFFN) r}r~rr'rrVrmr staticmethodrrr|rr+r+rtr.r\s,#  /rcs@eZdZddeeffdd Z       d ddZZS) T5LayerSelfAttentionFNrcs>tt|||d|_t|j|jd|_t |j |_ dS)Nrrr) rlrmr SelfAttentionrjrrr6rrrrrrtr+r.rmAs zT5LayerSelfAttention.__init__c CsL||} |j| |||||||d} ||| d}|f| dd} | S)N)rrrrrrrrr&)r6rr) rqr{attention_maskrrrrrrnormed_hidden_statesattention_outputrr+r+r.r|Is  zT5LayerSelfAttention.forwardr)NNNNFFNr}r~rrrVrmr|rr+r+rtr.r@s rcsBeZdZddeeffdd Z        d ddZZS) T5LayerCrossAttentionNrcs>tt|d|d|_t|j|jd|_t |j |_ dS)NFrr) rlrmrEncDecAttentionrjrrr6rrrr)rqrarrtr+r.rmes zT5LayerCrossAttention.__init__Fc  CsP||} |j| |||||||| | d } ||| d} | f| dd}|S)N) rrrrrrrrrrr&)r6rr)rqr{rrrrrrrrrrrZ layer_outputrr+r+r.r|ks  zT5LayerCrossAttention.forwardr)NNNNFNFNrr+r+rtr.rds rcsJeZdZd deeffdd Z            d ddZZS) T5BlockFNrcs`t|j|_t|_|jt|||d|jr&|jt||d|jt |dS)Nr)r) rlrmrr ModuleListr4rKrrrrrtr+r.rms   zT5Block.__init__Tc Cs|jd||||| | | | d}|dd\}} |dd}|jtjkrDtt|t|jjdt|jj}tj || |d}|j oJ|du}|r|jd|||||| | dd| | d }|dd\}} |jtjkrtt|t|jjdt|jj}tj || |d}||dd}|jd|}|jtjkrtt|t|jjdt|jj}tj || |d}|f}| r|| f|}|S||}|S) Nr)rrrrrrrr5i)rmaxr&r0)rrrrrrrr) r4ryr\rzrisinfrMfinforclampr)rqr{rrencoder_hidden_statesencoder_attention_maskencoder_decoder_position_biasrcross_attn_layer_head_maskrrr return_dictrZself_attention_outputsZattention_outputsZ clamp_valueZdo_cross_attentionZcross_attention_outputsrr+r+r.r|sn           zT5Block.forwardr) NNNNNNNNFFTNrr+r+rtr.rsrcs<eZdZdZdeffdd ZdejdejfddZZ S) T5ClassificationHeadz-Head for sentence-level classification tasks.racsBtt|j|j|_tj|jd|_t|j|j |_ dS)N)r) rlrmrrrdenserclassifier_dropoutr num_labelsout_projrrtr+r.rms zT5ClassificationHead.__init__r{returncCs6||}||}t|}||}||}|Sr)rrr\tanhrrr+r+r.r|s     zT5ClassificationHead.forward) r}r~r__doc__r'rmr\rr|rr+r+rtr.rsrc@sTeZdZeZeZdZdZdZ dZ dZ dZ dgZ dgZeddZdd Zd d Zd S) T5PreTrainedModel transformerTFrrcCs$tt}tt}|||d}|S)N)decoder_input_ids input_idsdecoder_attention_mask)r\Ztensorrr)rqrZ input_mask dummy_inputsr+r+r.r  s  zT5PreTrainedModel.dummy_inputscCs |jj}t|tr|jj|dd St|ttt t fr^|j jjj d|ddt |dr>|jjs>|jjjj d|ddt |dr\|jjjj d||jjdd|jjjd Sd St|trt |dr}|jjjj d|dd|jjjd Sd St|tr|jjjj d||jjddt |jdr|jjd ur|jjj|jjjj d||jjddt |jdr|jjd ur|jjjd Sd Sd St|tr$|jjjj d||jjddt |jdr|jjd ur|jjj|jjjj d||jjddt |jdr |jjd ur"|jjjd Sd Sd St|tr|jjjj d||jjddt |jdrO|jjd urO|jjj|jjjj d||jjddt |jdrt|jjd urt|jjj|jjjj d||jjddt |jdr|jjd ur|jjjd Sd Sd St|t r|jj}|jj!}|jj"}|j#jjj d|||dd|j$jjj d||dd|j%jjj d||dd|j&jjj d|||dd|j'r|j(jjj d||ddd Sd Sd S) zInitialize the weights?)rxZstdr= qa_outputsr:r9N))raZinitializer_factorrrjr3r^Zfill_T5ModelT5ForConditionalGenerationT5EncoderModelT5ForQuestionAnsweringsharedZnormal_rRtie_word_embeddingsr=r rr9Zzero_T5ForTokenClassificationr:rrrrr>rrrrrrrrrrrrrr)rqmodulefactorrrrr+r+r. _init_weightsst                     zT5PreTrainedModel._init_weightscCs|jj}|jj}|durtdt|r1t|jddd|}tj||dddfgdd}n| |j}|dddf |dddf<||d<|durStd| |d k||S) Nzself.model.config.decoder_start_token_id has to be defined. In T5 it is usually set to the pad_token_id. See T5 docs for more information.r0)r&.rr&).rz1self.model.config.pad_token_id has to be defined.) radecoder_start_token_id pad_token_idrYr!r\fullrWcatZ new_zeroscloneZ masked_fill_)rqrrrZshifted_input_idsr+r+r. _shift_rightUs    zT5PreTrainedModel._shift_rightN)r}r~rr'Z config_classriZload_tf_weightsZbase_model_prefixZis_parallelizableZsupports_gradient_checkpointingZ_supports_quantized_cacheZ_supports_static_cacheZ_supports_cache_classZ_no_split_modulesZ_keep_in_fp32_modulespropertyr rrr+r+r+r.rs  =rc seZdZdfdd ZeedddZeeddZdd Z d d Z             dd d Z dde e jdfde jde jdedef ddZede jdedede jde jdef ddZZS) T5StackNcsxt||_j|_tfddtjD|_t j j d|_ t j|_|d|_d|_d|_dS)Ncs"g|] }tt|dk|dqS)rr)rr)r,irar+r. ys"z$T5Stack.__init__..rF)rlrm embed_tokensrrrrange num_layersblockrjrrr7rrr post_initmodel_parallel device_mapr)rqrar$rtr"r.rmrs  zT5Stack.__init__cCstdt|durtt|jttj n||_ t |j t|jd|_ d|j vr/dn dtt|j |_dtt|j |_|j D]\}}|D]}dt|}|j|||j|<qRqL|j|j|_|j|j|_dS)Na<`T5Stack.parallelize` is deprecated and will be removed in v5 of Transformers, you should load your model with `device_map='balanced'` in the call to `from_pretrained`. You can also provide your own `device_map` but it needs to be a dictionary module_name to device, so for instance {'block.0': 0, 'block.1': 1, ...}Tcpucuda:)warningswarn FutureWarningr%rSr'r%r\cuda device_countr*r$r)r_strr first_devicer last_deviceitemsrvr$r7)rqr*rrr4Z cuda_devicer+r+r. parallelizes"$ * zT5Stack.parallelizecCsxtdtd|_d|_d|_d|_tt|j D] }|j | d|j |<q|j d|_ |j d|_ t jdS)N\Like `parallelize`, `deparallelize` is deprecated and will be removed in v5 of Transformers.Fr+)r-r.r/r)r*r3r4r%rSr'rvr$r7r\r0 empty_cache)rqr!r+r+r. deparallelizeszT5Stack.deparallelizecC|jSrr$rqr+r+r.get_input_embeddingszT5Stack.get_input_embeddingscC ||_dSrr;rqZnew_embeddingsr+r+r.set_input_embeddings zT5Stack.set_input_embeddingsc+Cs0|jrtj|j|j|j|_| dur| n|jj} | dur"| n|jj } | dur,| n|jj } | dur6| n|jj } |durT|durT|j rGdnd}t d|d|d|dure|}|d|d}n|durr|dd}n|j rwdnd}t d|d|d |jr|jr| rtd d } |dur|jdurt d ||}|\}}| d ur|j st d|dd }d }|j r| s|durt|trt|tsd }t|t}n#t|tsd }tdt|}n|durttt}n|j sd}|dur|nd}| durtj||||jd} |dur,ts,||}tj|||jd}|jj rC|||| |dur>|jnd| }n&|durg|ddddddf}|j|j d}d|t!|j j"}nd}|j r|dur|\}}}||f}|durtj||jtj#d}|$|}nd}|%||jj&}|%||jj&}| rdnd}| rdnd}| r|j rdnd}d}d} |'|}!t(|j)D] \}"}#||"}$||"}%|jr0tj|!j|dur||!j}|dur||!j}|dur||!j}|dur||!j}| dur| |!j} |$dur%|$|!j}$|%dur0|%|!j}%| r8||!f}|jrT|jrT|*|#j+|!||||| |$|%d| | | | }&n|#|!||||| |$|%|| | | | d }&| d urx|&ddd|&dd}&|&dd\}!}'|&d}|j r|dur|&| rdnd} | r||&df}|j r||&df}|jr|j,-D]!\}(})|"|)dkrdt.|(|j/kr|!dt.|(d}!qq|0|!}!|'|!}!| r||!f}| r|'nd}*|r|j}*|r|1}*| st2dd |!|*|||fDSt3|!|*|||d!S)"NZdecoder_zYou cannot specify both zinput_ids and zinputs_embeds at the same timer0zYou have to specify either z input_ids or inputs_embedszZ`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...Fz.)last_hidden_statepast_key_valuesr{ attentionscross_attentions)4r)r\r0 set_devicer3r$rvrarroutput_hidden_statesuse_return_dictrrYsizerrrrDrrr rr Zfrom_legacy_cacheget_seq_lengthrrr"ro_update_causal_maskrryrrrZinvert_attention_maskZ get_head_maskr&r enumerater'Z_gradient_checkpointing_funcr|r*r5r2r4r7Zto_legacy_cachetupler)+rqrrrrrD head_maskcross_attn_head_maskrIrrrMrrZerr_msg_prefixZ input_shaperrZreturn_legacy_cacheZreturn_self_attention_cachepast_key_values_lengthZmask_seq_lengthrZencoder_batch_sizeZencoder_sequence_length_Zencoder_hidden_shapeZencoder_extended_attention_maskZall_hidden_statesZall_attentionsZall_cross_attentionsrrr{r!Z layer_modulerrZ layer_outputsZnext_decoder_cacherrZ next_cacher+r+r.r|sl                             "    zT5Stack.forwardFrr( input_tensorrrIrc Cs:|jjdkr|dur|dkr|SdS|jjdkr&t|tjr$t|}|S|dur.|nd}|dur7|jnd}|jjdkrO|sO|sOt j ||||j drOdS|j }|j d} |r^|} nt|tjri|j d n|| d} |j|| | |||j dd } |jjdkr|dur|jjd vr|st|j} t | | } | S) NZflash_attention_2r Zflex_attentionrFZsdpa)rDrVZ is_trainingr&r0)sequence_length target_lengthryrr)r0ZxpuZnpu)raZ_attn_implementationrMrr\rr)rPZis_compileablerZ_ignore_causal_mask_sdparryrWZget_max_cache_shape5_prepare_4d_causal_attention_mask_with_cache_positionrtyperrZ_unmask_unattended) rqrrXrrIrZpast_seen_tokensZusing_compilable_cacheryrYrZr min_dtyper+r+r.rQsT           zT5Stack._update_causal_maskrYrZryrc KsD|dur|dkr|}|St|j}tj||f|||jd}|dkr+tj|dd}|tj||jd|ddk9}|ddddddf |ddd}|dur| }|j d} |ddddddd| f|ddddddf |j} | dk} |ddddddd| f | ||ddddddd| f<|S) aM Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. Args: attention_mask (`torch.Tensor`): A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`. sequence_length (`int`): The sequence length being processed. target_length (`int`): The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet. dtype (`torch.dtype`): The dtype to use for the 4D attention mask. cache_position (`torch.Tensor`): Indices depicting the position of the input sequence tokens in the sequence. batch_size (`torch.Tensor`): Batch size. NrF)Z fill_valueryrr&)ZdiagonalrEr0r)rr\rrrrZtriurZreshapeexpandrrWrvZ masked_fill) rrYrZryrrkwargsrr]Z mask_lengthZ padding_maskr+r+r.r[s,  $ 6  z=T5Stack._prepare_4d_causal_attention_mask_with_cache_positionr) NNNNNNNNNNNNN)F)r}r~rrmrPARALLELIZE_DOCSTRINGr6DEPARALLELIZE_DOCSTRINGr9r=rAr|rr\rr rrQrrVryr[rr+r+rtr.r qsb     Dr a_ The input argument `head_mask` was split into two arguments `head_mask` and `decoder_head_mask`. Currently, `decoder_head_mask` is set to copy `head_mask`, but this feature is deprecated and will be removed in future versions. If you do not want to use any `decoder_head_mask` now, please set `decoder_head_mask = torch.ones(num_layers, num_heads)`. c&sleZdZdgZddgZdeffdd Zeed+dd Z ee d d Z d d Z ddZ ddZddZddZddZe                d,deejdeejdeejdeejdeejdeejdeejdeeeejd eeeejd!eejd"eejd#eed$eed%eed&eed'eejd(eeejeff"d)d*ZZS)-rFdecoder.block.0.layer.1.EncDecAttention.relative_attention_bias.weightencoder.embed_tokens.weightdecoder.embed_tokens.weightracst|t|j|j|_t|}d|_ d|_ d|_ t ||j|_ t|}d|_ d|_ |j|_t ||j|_|d|_d|_dSNFT)rlrmrr vocab_sizerrcopydeepcopyrris_encoder_decoderr encodernum_decoder_layersr&r;r(r)r*rqraencoder_configZdecoder_configrtr+r.rm<s    zT5Model.__init__NcCsptdt|durtt|jjttj n||_ t |j t|jj|j |j |j |j d|_dS)NaL`T5Model.parallelize` is deprecated and will be removed in v5 of Transformers, you should load your model with `device_map='balanced'` in the call to `from_pretrained`. You can also provide your own `device_map` but it needs to be a dictionary module_name to device, so for instance {'encoder.block.0': 0, 'encoder.block.1': 1, ...}T)r-r.r/r%rSrjr'r%r\r0r1r*r$r6r;r)rqr*r+r+r.r6Ss  zT5Model.parallelizecCsVtdt|j|j|jd|_|jd|_d|_d|_t j dSNr7r+F) r-r.r/rjr9r;rvr)r*r\r0r8r<r+r+r.r9fs  zT5Model.deparallelizecCr:rrr<r+r+r.r=tr>zT5Model.get_input_embeddingscC"||_|j||j|dSrrrjrAr;r@r+r+r.rAw zT5Model.set_input_embeddingscC4|jjr||jj|j||jj|jdSdSrrar_tie_or_clone_weightsrjr$rr;r<r+r+r. _tie_weights|zT5Model._tie_weightscCr:rrjr<r+r+r. get_encoderr>zT5Model.get_encodercCr:rr;r<r+r+r. get_decoderr>zT5Model.get_decodercCs*|D]\}}|jj|j|qdS) 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)r5rjr4Z attentionrrqZheads_to_pruner4rr+r+r. _prune_headsszT5Model._prune_headsrrrrrTdecoder_head_maskrUencoder_outputsrIrDdecoder_inputs_embedsrrrMrrrcCs| dur| n|jj} |dur|n|jj}|dur,|dur,|jj|jjkr,ttt|}|dur=|j ||| || ||d}n$|rat |t sat |dt |dkrR|dndt |dkr]|dndd}|d}|j rtj|jj||jj}|dur||jj}|dur||jj}|dur||jj}|j||| | ||||| | |||d }|s||St|j|j|j|j|j|j|j|jdS) a input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. T5 is a model with relative position embeddings so you should be able to pad the inputs on both the right and the left. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for detail. [What are input IDs?](../glossary#input-ids) To know more on how to prepare `input_ids` for pretraining take a look a [T5 Training](./t5#training). decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): Indices of decoder input sequence tokens in the vocabulary. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are decoder input IDs?](../glossary#decoder-input-ids) T5 uses the `pad_token_id` as the starting token for `decoder_input_ids` generation. If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see `past_key_values`). To know more on how to prepare `decoder_input_ids` for pretraining take a look at [T5 Training](./t5#training). decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*): Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also be used by default. decoder_head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): Mask to nullify selected heads of the self-attention modules in the decoder. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. cross_attn_head_mask (`torch.Tensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): Mask to nullify selected heads of the cross-attention modules in the decoder. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. Example: ```python >>> from transformers import AutoTokenizer, T5Model >>> tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-small") >>> model = T5Model.from_pretrained("google-t5/t5-small") >>> input_ids = tokenizer( ... "Studies have been shown that owning a dog is good for you", return_tensors="pt" ... ).input_ids # Batch size 1 >>> decoder_input_ids = tokenizer("Studies show that", return_tensors="pt").input_ids # Batch size 1 >>> # preprocess: Prepend decoder_input_ids with start token which is pad token for T5Model. >>> # This is not needed for torch's T5ForConditionalGeneration as it does this internally using labels arg. >>> decoder_input_ids = model._shift_right(decoder_input_ids) >>> # forward pass >>> outputs = model(input_ids=input_ids, decoder_input_ids=decoder_input_ids) >>> last_hidden_states = outputs.last_hidden_state ```NrrrDrTrrMrrr&r5rHr{rJ rrrDrIrrrTrUrrrMrr)rHrIdecoder_hidden_statesdecoder_attentionsrKencoder_last_hidden_staterencoder_attentions)rarrNr&rkr-r.Z_T5Model__HEAD_MASK_WARNING_MSGr/rjrrrSr)r\r0rLr;r3rvrrHrIr{rJrK)rqrrrrrTrrUrrIrDrrrrMrrr{decoder_outputsr+r+r.r|svQ  zT5Model.forwardrNNNNNNNNNNNNNNNN)r}r~r"_keys_to_ignore_on_load_unexpected_tied_weights_keysr'rmrr`r6rar9r=rArwrzr|rrrr\ LongTensor FloatTensor BoolTensorrrrrrr|rr+r+rtr.r5s        rz: T5 Model with a `language modeling` head on top. )Z custom_introc(seZdZdgZgdZdeffdd Zeed1ddZ ee d d Z d d Z d dZ ddZddZddZddZddZe                 d2deejdeejdeejdeejdeejdeejdeejd eeeejd!eeeejd"eejd#eejd$eejd%eed&eed'eed(eed)eejd*eeejeff$d+d,Zd$ejfd-d.Zd/d0Z Z!S)3rrb)rcrdzlm_head.weightracst||j|_t|j|j|_t |}d|_ d|_ d|_ t ||j|_t |}d|_ d|_ |j|_t ||j|_tj|j|jdd|_|d|_d|_dS)NFTr)rlrmr model_dimrrrfrrgrhrrrir rjrkr&r;rr=r(r)r*rlrtr+r.rm3s"    z#T5ForConditionalGeneration.__init__NcCstdt|durtt|jjttj n||_ t |j t|jj|j |j |j |j |j|jj|_d|_dS)Na_`T5ForConditionalGeneration.parallelize` is deprecated and will be removed in v5 of Transformers, you should load your model with `device_map='balanced'` in the call to `from_pretrained`. You can also provide your own `device_map` but it needs to be a dictionary module_name to device, so for instance {'encoder.block.0': 0, 'encoder.block.1': 1, ...}T)r-r.r/r%rSrjr'r%r\r0r1r*r$r6r;r=rvr3r)rnr+r+r.r6Ns  z&T5ForConditionalGeneration.parallelizecCsdtdt|j|j|jd|_|jd|_|jd|_d|_d|_ t j dSro) r-r.r/rjr9r;rvr=r)r*r\r0r8r<r+r+r.r9bs  z(T5ForConditionalGeneration.deparallelizecCr:rrpr<r+r+r.r=qr>z/T5ForConditionalGeneration.get_input_embeddingscCrqrrrr@r+r+r.rAtrsz/T5ForConditionalGeneration.set_input_embeddingscCrtrrur<r+r+r.rwyrxz'T5ForConditionalGeneration._tie_weightscCr?rr=r@r+r+r.set_output_embeddings~rBz0T5ForConditionalGeneration.set_output_embeddingscCr:rrr<r+r+r.get_output_embeddingsr>z0T5ForConditionalGeneration.get_output_embeddingscCr:rryr<r+r+r.rzr>z&T5ForConditionalGeneration.get_encodercCr:rr{r<r+r+r.r|r>z&T5ForConditionalGeneration.get_decoderrrrrrTrrUrrIrDrlabelsrrrMrrrcCs| dur| n|jj} |dur|n|jj}|dur,|dur,|jj|jjkr,ttt|}|dur=|j ||| ||||d}n$|rat |t sat |dt |dkrR|dndt |dkr]|dndd}|d}|j rptj|jj| dur|dur| dur|| }|j rtj|jj||jj}|dur||jj}|dur||jj}|dur||jj}|j||| | ||||| ||||d }|d}|j rtj|j j|j|j j|_||jjj}|jjr||jd}||}d}| durtd d }| |j} ||d |d | d }|s1|f|dd|}|dur/|f|S|St|||j|j|j |j!|j"|j|j d S) a input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. T5 is a model with relative position embeddings so you should be able to pad the inputs on both the right and the left. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for detail. [What are input IDs?](../glossary#input-ids) To know more on how to prepare `input_ids` for pretraining take a look a [T5 Training](./t5#training). decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): Indices of decoder input sequence tokens in the vocabulary. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are decoder input IDs?](../glossary#decoder-input-ids) T5 uses the `pad_token_id` as the starting token for `decoder_input_ids` generation. If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see `past_key_values`). To know more on how to prepare `decoder_input_ids` for pretraining take a look at [T5 Training](./t5#training). decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*): Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also be used by default. decoder_head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): Mask to nullify selected heads of the self-attention modules in the decoder. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. cross_attn_head_mask (`torch.Tensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): Mask to nullify selected heads of the cross-attention modules in the decoder. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): Labels for computing the sequence classification/regression loss. Indices should be in `[-100, 0, ..., config.vocab_size - 1]`. All labels set to `-100` are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]` Examples: ```python >>> from transformers import AutoTokenizer, T5ForConditionalGeneration >>> tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-small") >>> model = T5ForConditionalGeneration.from_pretrained("google-t5/t5-small") >>> # training >>> input_ids = tokenizer("The walks in park", return_tensors="pt").input_ids >>> labels = tokenizer(" cute dog the ", return_tensors="pt").input_ids >>> outputs = model(input_ids=input_ids, labels=labels) >>> loss = outputs.loss >>> logits = outputs.logits >>> # inference >>> input_ids = tokenizer( ... "summarize: studies have shown that owning a dog is good for you", return_tensors="pt" ... ).input_ids # Batch size 1 >>> outputs = model.generate(input_ids) >>> print(tokenizer.decode(outputs[0], skip_special_tokens=True)) >>> # studies have shown that owning a dog is good for you. ```Nrrr&r5rrr rZ ignore_indexr0 lossr<rIrrrKrrr)#rarrNr&rkr-r.Z2_T5ForConditionalGeneration__HEAD_MASK_WARNING_MSGr/rjrrrSr)r\r0rLr;r3rrvr=r3rrrrrrOrrIr{rJrKrH)rqrrrrrTrrUrrIrDrrrrrMrrr{rsequence_outputZ lm_logitsrloss_fctoutputr+r+r.r|sX       z"T5ForConditionalGeneration.forwardcCs ||Sr)r)rqrr+r+r.%prepare_decoder_input_ids_from_labelsJrBz@T5ForConditionalGeneration.prepare_decoder_input_ids_from_labelsc Cs|dur td|Sd}|D]M}d}|D]}||d||jf}q|dj|djkr@td|djd|djdt|t|krWtdt|dt|d||f}q|S) NzHYou might want to consider setting `use_cache=True` to speed up decodingr+rz%reordered_layer_past_states[0] shape z and layer_past_states[0] shape r?z&length of reordered_layer_past_states z! and length of layer_past_states )rDwarningZ index_selectrvrrWrYrS)rqrIZbeam_idxZreordered_decoder_pastZlayer_past_statesZreordered_layer_past_statesZlayer_past_stater+r+r._reorder_cacheMs(  z)T5ForConditionalGeneration._reorder_cacher)NNNNNNNNNNNNNNNNN)"r}r~rrrr'rmrr`r6rar9r=rArwrrrzr|rrr\rrrrrrrrr|rrrr+r+rtr.r(s        @rcseZdZdgZdgZdeffdd ZeedddZ ee d d Z d d Z d dZ ddZddZddZe       d deejdeejdeejdeejdeedeedeedeeejeffddZZS)!rrcr;racsZt|t|j|j|_t|}d|_ d|_ t ||j|_ | d|_d|_dSNF)rlrmrrrfrrrgrhrrir rjr(r)r*)rqrarmrtr+r.rmqs   zT5EncoderModel.__init__NcCsbtdt|durtt|jjttj n||_ t |j t|jj|j |j d|_dS)NaC`T5EncoderModel.parallelize` is deprecated and will be removed in v5 of Transformers, you should load your model with `device_map='balanced'` in the call to `from_pretrained`. You can also provide your own `device_map` but it needs to be a dictionary module_name to device, so for instance {'block.0': 0, 'block.1': 1, ...}T)r-r.r/r%rSrjr'r%r\r0r1r*r$r6r)rnr+r+r.r6s  zT5EncoderModel.parallelizecCs>tdt|j|jd|_d|_d|_tj dSro) r-r.r/rjr9rvr)r*r\r0r8r<r+r+r.r9s zT5EncoderModel.deparallelizecCr:rrpr<r+r+r.r=r>z#T5EncoderModel.get_input_embeddingscCs||_|j|dSr)rrjrAr@r+r+r.rAsz#T5EncoderModel.set_input_embeddingscCs"|jjr||jj|jdSdSr)rarrvrjr$rr<r+r+r.rwszT5EncoderModel._tie_weightscCr:rryr<r+r+r.rzr>zT5EncoderModel.get_encodercCs0|D]\}}|jj|jdj|qdS)r}rN)r5rjr'r4rrr~r+r+r.rszT5EncoderModel._prune_headsrrrTrDrrMrrc Cs0|dur|n|jj}|j|||||||d}|S)aI input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. T5 is a model with relative position embeddings so you should be able to pad the inputs on both the right and the left. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for detail. To know more on how to prepare `input_ids` for pretraining take a look a [T5 Training](./t5#training). Example: ```python >>> from transformers import AutoTokenizer, T5EncoderModel >>> tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-small") >>> model = T5EncoderModel.from_pretrained("google-t5/t5-small") >>> input_ids = tokenizer( ... "Studies have been shown that owning a dog is good for you", return_tensors="pt" ... ).input_ids # Batch size 1 >>> outputs = model(input_ids=input_ids) >>> last_hidden_states = outputs.last_hidden_state ```Nr)rarNrj) rqrrrTrDrrMrrr+r+r.r|s" zT5EncoderModel.forwardr)NNNNNNN)r}r~rrrr'rmrr`r6rar9r=rArwrzrrrr\rrrrrrr|rr+r+rtr.rlsL    rz T5 model with a sequence classification/head on top (a linear layer on top of the pooled output) e.g. for GLUE tasks. c$seZdZdgZddgZdeffdd Ze               ddee j d ee j d ee j d ee j d ee j d ee j dee j dee e j dee j dee j dee j deedeedeedeedeeeff ddZZS)T5ForSequenceClassificationrbrcrdracs2t|t||_t||_|d|_dSr)rlrmrrrclassification_headr(r)rrtr+r.rms    z$T5ForSequenceClassification.__init__NrrrrrTrrUrrDrrrrrMrrcCsh|dur|n|jj}| durd} |dur!| dur!td|jj|dur6| dur6|dur1td||}|j||||||||| | | | ||d}|d}||jj  |j }t t |ddkrhtd|j\}}}||ddf|d |ddd ddf}||}d}| dur| |j } |jjdur|jjdkrd |j_n|jjdkr| jt jks| jt jkrd |j_nd |j_|jjd krt}|jjdkr||| }n-||| }n'|jjd krt}||d |jj| d }n|jjd krt}||| }|s |f|dd}|dur|f|S|St|||j|j|j|j |j!|j"|j#d S)ag input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. T5 is a model with relative position embeddings so you should be able to pad the inputs on both the right and the left. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for detail. [What are input IDs?](../glossary#input-ids) To know more on how to prepare `input_ids` for pretraining take a look a [T5 Training](./t5#training). decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): Indices of decoder input sequence tokens in the vocabulary. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are decoder input IDs?](../glossary#decoder-input-ids) T5 uses the `pad_token_id` as the starting token for `decoder_input_ids` generation. If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see `past_key_values`). To know more on how to prepare `decoder_input_ids` for pretraining take a look at [T5 Training](./t5#training). decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*): Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also be used by default. decoder_head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): Mask to nullify selected heads of the self-attention modules in the decoder. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. cross_attn_head_mask (`torch.Tensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): Mask to nullify selected heads of the cross-attention modules in the decoder. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. 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 classification loss is computed (Cross-Entropy). NFz8Passing input embeddings is currently not supported for If no `decoder_input_ids` or `decoder_inputs_embeds` are passed, `input_ids` cannot be `None`. Please pass either `input_ids` or `decoder_input_ids` or `decoder_inputs_embeds`.) rrrrTrrUrrDrrrrMrrr&z7All examples must have the same number of tokens.r0Z regressionZsingle_label_classificationZmulti_label_classificationr)$rarNNotImplementedErrorrur}rYrreqZ eos_token_idrvrrSr\Zunique_consecutivesumrWrrZ problem_typerryrrVr squeezerrrrIrrrKrrr)rqrrrrrTrrUrrDrrrrrMrrrZeos_maskrrWrrZsentence_representationr<rrrr+r+r.r|s>   ,      $      z#T5ForSequenceClassification.forward)NNNNNNNNNNNNNNN)r}r~rrrr'rmrrr\rrrrrrrrr|rr+r+rtr.rsj       rcseZdZdgZdeffdd Ze        ddeej deej deej d eej d eej d ee d ee d ee de e ej e ffddZZS)rz'transformer.encoder.embed_tokens.weightracsJt||j|_t||_t|j|_t |j |j|_ | dSr) rlrmrrrrrrrrrrr:r(rrtr+r.rms   z!T5ForTokenClassification.__init__NrrrTrDrrrMrrc  Cs|dur|n|jj}|j|||||||d} | d} || } || } d} |dur:t} | | d|j|d} |sO| | ddf}| durM| f|S|St| | | j | j dS)a< input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. T5 is a model with relative position embeddings so you should be able to pad the inputs on both the right and the left. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for detail. [What are input IDs?](../glossary#input-ids) To know more on how to prepare `input_ids` for pretraining take a look a [T5 Training](./t5#training). 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]`. N)rrTrDrrMrrr0r5)rr<r{rJ) rarNrrr:rrrrr{rJ)rqrrrTrDrrrMrrr{r<rrrr+r+r.r|s4  z T5ForTokenClassification.forward)NNNNNNNN)r}r~rrr'rmrrr\rrrrrr|rr+r+rtr.rs>   rc&s:eZdZdgZddgZdeffdd ZddZd d Zd d Z d dZ ddZ e                d%de ejde ejde ejde ejde ejde ejde ejde eeejde ejde ejde ejde ejde ede ed e ed!e ed"eeejeff"d#d$ZZS)&rrbrcrdracst||j|_t|j|j|_t |}d|_ d|_ d|_ t ||j|_t |}d|_ d|_ |j|_t ||j|_|j|_t|j|j|_|d|_dSre)rlrmrrrrrfrrgrhrrrir rjrkr&r;rrrrr r(r)rlrtr+r.rms"    zT5ForQuestionAnswering.__init__cCr:rrpr<r+r+r.r=r>z+T5ForQuestionAnswering.get_input_embeddingscCrqrrrr@r+r+r.rArsz+T5ForQuestionAnswering.set_input_embeddingscCrtrrur<r+r+r.rwrxz#T5ForQuestionAnswering._tie_weightscCr:rryr<r+r+r.rz r>z"T5ForQuestionAnswering.get_encodercCr:rr{r<r+r+r.r| r>z"T5ForQuestionAnswering.get_decoderNrrrrrTrrUrstart_positions end_positionsrDrrrrMrrcCs|dur|n|jj}| dur| n|jj} | dur| durd} |dur3| dur3|dur.td||}| dur9| n|jj} |durC|n|jj}|dur_|dur_|jj|jjkr_tt t |}|durp|j ||| ||||d}n$|rt |t st |dt|dkr|dndt|dkr|dndd}|d}|j||| d||||| |||d }|d}||}|jdd d \}}|d }|d }d}| dur| durt| dkr| d |j} t| dkr| d |j} |d}| d|} | d|} t|d }||| }||| }||d}|s8||f|dd|}|dur6|f|S|St||||j|j|j|j|j|j|jd S)a= input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. T5 is a model with relative position embeddings so you should be able to pad the inputs on both the right and the left. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for detail. [What are input IDs?](../glossary#input-ids) To know more on how to prepare `input_ids` for pretraining take a look a [T5 Training](./t5#training). decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): Indices of decoder input sequence tokens in the vocabulary. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are decoder input IDs?](../glossary#decoder-input-ids) T5 uses the `pad_token_id` as the starting token for `decoder_input_ids` generation. If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see `past_key_values`). To know more on how to prepare `decoder_input_ids` for pretraining take a look at [T5 Training](./t5#training). decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*): Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also be used by default. decoder_head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): Mask to nullify selected heads of the self-attention modules in the decoder. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. cross_attn_head_mask (`torch.Tensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): Mask to nullify selected heads of the cross-attention modules in the decoder. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. NFrrrr&r5r) rrrDrIrrrTrUrrrMrr0rr) r start_logits end_logitsrIrrrKrrr)rarNrrYrr&rkr-r.Z._T5ForQuestionAnswering__HEAD_MASK_WARNING_MSGr/rjrrrSr;r rLrrrOrvrrrrrIr{rJrKrH)rqrrrrrTrrUrrrrDrrrrMrr{rrr<rrZ total_lossZ ignored_indexrZ start_lossZend_lossrr+r+r.r| s<           zT5ForQuestionAnswering.forwardr)r}r~rrrr'rmr=rArwrzr|rrr\rrrrrrrrr|rr+r+rtr.rsz     r)rrrrrirrr)[rrgrrFr-typingrrrrr\rZtorch.nnrrr Z activationsr Z cache_utilsr r rZ generationrZmodeling_attn_mask_utilsrZmodeling_outputsrrrrrrrZmodeling_utilsrZ pytorch_utilsrrrutilsrrrrr r!r"r#Zutils.model_parallel_utilsr$r%Zconfiguration_t5r'Z!torch.nn.attention.flex_attentionr(Zintegrations.flex_attentionr)Z get_loggerr}rDrir`raModulerjZapex.normalizationrrIrC ExceptionrrKrrrrrrrrrr Z__HEAD_MASK_WARNING_MSGrrrrrr__all__r+r+r+r.s    $ (     n!     e$&dp?sAy#I V