o Zh@srdZddlZddlZddlmZmZmZddlZddlmZddl m Z ddl m Z ddl mZdd lmZmZdd lmZmZdd lmZdd lmZmZd dlmZeeZGdddejZ GdddejZ!Gdddej"Z#Gdddej"Z$Gdddej"Z%eGdddeZ&Gddde&Z'eddGdd d e&Z(ed!dGd"d#d#e&eZ)d#dgZ*dS)$z/PyTorch TrOCR decoder model (based on RoBERTa).N)OptionalTupleUnion)nn)CrossEntropyLoss)ACT2FN)GenerationMixin)_prepare_4d_attention_mask!_prepare_4d_causal_attention_mask))BaseModelOutputWithPastAndCrossAttentions!CausalLMOutputWithCrossAttentions)PreTrainedModel)auto_docstringlogging) TrOCRConfigcsJeZdZdZdedeffdd Zd dejd ed ejffd d ZZ S)TrOCRLearnedPositionalEmbeddingzN This module learns positional embeddings up to a fixed maximum size. num_embeddings embedding_dimcsd|_t||j|dS)N)offsetsuper__init__)selfrr __class__W/var/www/auris/lib/python3.10/site-packages/transformers/models/trocr/modeling_trocr.pyr.sz(TrOCRLearnedPositionalEmbedding.__init__rN input_idspast_key_values_length position_idscs\|dur |jdd\}}tj|||tj|jjd|d}n|d}t ||j S)z3`input_ids' shape is expected to be [bsz x seqlen].Nr)dtypedevicer) shapetorcharangelongweightr#expand unsqueezerforwardr)rrr r!bszseq_lenrrrr,4s z'TrOCRLearnedPositionalEmbedding.forwardrN) __name__ __module__ __qualname____doc__intrr&Tensorr, __classcell__rrrrr)s(rc sLeZdZdZd dedededeeffdd Zd ej ffd d Z Z S) TrOCRScaledWordEmbeddingz\ This module overrides nn.Embeddings' forward by multiplying with embeddings scale. ?rr padding_idx embed_scalecst|||||_dSN)rrr:)rrrr9r:rrrrHs z!TrOCRScaledWordEmbedding.__init__rcst||jSr;)rr,r:)rrrrrr,Lsz TrOCRScaledWordEmbedding.forward)r8) r0r1r2r3r4rfloatrr&r5r,r6rrrrr7Cs$r7c seZdZdZddededeeffdd Zeddededeefd d Ze dd e j d efddZ dd e j ded eefddZ ZS)"TrOCRSinusoidalPositionalEmbeddingzDThis module produces sinusoidal positional embeddings of any length.N num_positionsrr9csBtd|_||_||_|||||_|dt ddS)Nr _float_tensorr) rrrrr9 get_embeddingweightsZregister_bufferr& FloatTensor)rr>rr9rrrrSs z+TrOCRSinusoidalPositionalEmbedding.__init__rcCs|d}td|d}ttj|tjd| }tj|tjdd|d}tjt |t |gdd |d}|ddkrUtj|t |dgdd}|durad||ddf<| tS) z Build sinusoidal embeddings. This matches the implementation in tensor2tensor, but differs slightly from the description in Section 3.5 of "Attention Is All You Need". ri'r)r"rdimr$N)mathlogr&expr'Zint64r<r+catsincosviewZzerostoZget_default_dtype)rrr9Zhalf_dimZembrrrr@[s $& z0TrOCRSinusoidalPositionalEmbedding.get_embeddingrrr cCs|\}}|||j||j}|jd|}|jdus&||jdkr0|||j|j|_|j|j|_|j d| d ||d }|S)Nrrr$) size"create_position_ids_from_input_idsr9rLr#rAr@rr? index_selectrKdetach)rrr r-r.r!Zmax_posxrrrr,ns "z*TrOCRSinusoidalPositionalEmbedding.forwardcCs6||}tj|dd|||}||S)z Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols are ignored. This is modified from fairseq's `utils.make_positions`. rrC)ner4r&ZcumsumZtype_asr()rrr9r maskZincremental_indicesrrrrNs zETrOCRSinusoidalPositionalEmbedding.create_position_ids_from_input_idsr;)r)r0r1r2r3r4rr staticmethodr@r&Zno_gradr5r,rNr6rrrrr=Ps r=cseZdZdZ      ddededeed eed ed ed ed effdd Zde j dedefddZ     dde j dee j dee e j dee j dee j dede e j ee j ee e j ffddZ ZS) TrOCRAttentionz>Multi-headed attention from 'Attention Is All You Need' paper.NFT embed_dim num_headskdimvdimdropout is_decoderbiasis_cross_attentionc st||_|dur|n||_|dur|n||_||_||_|||_|j||jks9td|jd|d|jd|_ ||_ t j |j||d|_ t j |j||d|_t j |||d|_t j |||d|_dS)Nz;embed_dim must be divisible by num_heads (got `embed_dim`: z and `num_heads`: z).g࿩r])rrrWrYrZrXr[head_dim ValueErrorscalingr\rLineark_projv_projq_projout_proj) rconfigrWrXrYrZr[r\r]r^rrrrs&    zTrOCRAttention.__init__tensorr.r-cCs ||||j|jddS)Nrr)rKrXr` transpose contiguous)rrir.r-rrr_shapes zTrOCRAttention._shape hidden_stateskey_value_statespast_key_valueattention_masklayer_head_maskoutput_attentionsreturncCsV|du}|\}} } |||j} |r"|dur"|d} |d} nZ|r9|||d|} |||d|} nC|durh|||d|} |||d|} tj|d| gdd} tj|d| gdd} n|||d|} |||d|} |jr| | f}||j d|j f}|| | |j |} | j |} | j |} | d}t | | dd}|||j | |fkrtd||j | |fd||dur||d| |fkrtd |d| |fd|| ||j | ||}| ||j | |}tjj|dd}|dur@||j fkr%td |j fd|| dddd| ||j | |}| ||j | |}|rW| ||j | |}| ||j | |}nd}tjj||j|jd }t || }|||j | |j fkrtd ||j | |j fd|| ||j | |j }| dd}||| | }||}|||fS) z#Input shape: Batch x Time x ChannelNrrr$rrCz$Attention weights should be of size z , but is z!Attention mask should be of size z/Head mask for a single layer should be of size ptrainingz `attn_output` should be of size )rMrfrbrlrdrer&rHr\rXr`rKZbmmrjrar functionalZsoftmaxr[rvZreshaperg)rrmrnrorprqrrr^r-tgt_lenrWZ query_statesZ key_statesZ value_statesZ proj_shapeZsrc_lenZ attn_weightsZattn_weights_reshapedZ attn_probsZ attn_outputrrrr,s         "    zTrOCRAttention.forward)NNrVFTF)NNNNF)r0r1r2r3r4rr<boolrr&r5rlrr,r6rrrrrUs\  ! rUcseZdZdeffdd Z        ddejdeejd eejd eejd eejd eejd eeejdee dee fddZ Z S)TrOCRDecoderLayerrhc st|j|_t||j|j|jdd|_|j|_t |j |_ |j |_ t |j|_|jrGt||j|j|j|j|jddd|_t |j|_t |j|j|_t |j|j|_t |j|_dS)NT)rWrXr[r\)rWrXrYrZr[r\r^)rr hidden_sizerWrUZdecoder_attention_headsZattention_dropout self_attnr[rZactivation_function activation_fnactivation_dropoutr LayerNormself_attn_layer_normr\Zcross_attention_hidden_size encoder_attnencoder_attn_layer_normrcZdecoder_ffn_dimfc1fc2final_layer_normrrhrrrr$s8   zTrOCRDecoderLayer.__init__NFTrmrpencoder_hidden_statesencoder_attention_maskrqcross_attn_layer_head_maskrorr use_cachec CsZ|} |dur |ddnd} |j|| |||d\}} } tjj||j|jd}| |}||}d}d}|durk|} |durC|ddnd}|j||||||d\}}}tjj||j|jd}| |}||}| |} |} || |}tjj||j |jd}| |}tjj||j|jd}| |}| |}|f}|r|| |f7}| r|| f7}|S)a Args: hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` attention_mask (`torch.FloatTensor`): attention mask of size `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. encoder_hidden_states (`torch.FloatTensor`): cross attention input to the layer of shape `(batch, seq_len, embed_dim)` encoder_attention_mask (`torch.FloatTensor`): encoder attention mask of size `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size `(encoder_attention_heads,)`. cross_attn_layer_head_mask (`torch.FloatTensor`): mask for cross-attention heads in a given layer of size *(decoder_attention_heads,)*. past_key_value (`Tuple(torch.FloatTensor)`): cached past key and value projection states output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. Nr)rmrorprqrrrt)rmrnrprqrorr) r|rrwr[rvrrrr}rr~rr)rrmrprrrqrrorrrZresidualZself_attn_past_key_valueZself_attn_weightsZpresent_key_valueZcross_attn_present_key_valueZcross_attn_weightsZcross_attn_past_key_valueoutputsrrrr,FsT         zTrOCRDecoderLayer.forward)NNNNNNFT) r0r1r2rrr&r5rrryr,r6rrrrrz#s:%   rzc@s&eZdZeZdZdZdgZddZdS)TrOCRPreTrainedModelmodelTrzcCs|jj}t|tjtjfr%|jjjd|d|j dur#|j j dSdSt|tj rD|jjjd|d|j durF|jj|j  dSdSdS)NrV)meanstd) rhZinit_std isinstancerrcZConv1dr)dataZnormal_r]Zzero_ Embeddingr9)rmodulerrrr _init_weightss   z"TrOCRPreTrainedModel._init_weightsN) r0r1r2rZ config_classZbase_model_prefixZsupports_gradient_checkpointingZ_no_split_modulesrrrrrrs  rcsXeZdZdZdeffdd ZddZddZ d d d ZZ S) TrOCRDecoderz Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`TrOCRDecoderLayer`] Args: config: TrOCRConfig rhcstj|_j|_j|_jrt j nd}t j j |j|d|_ jr5tjj |_ntj|jdj |j|_jrOtj |_nd|_tfddtjD|_d|_|dS)Nr8)r:rcsg|]}tqSr)rz).0_rhrr sz)TrOCRDecoder.__init__..F)rrr[Zdecoder_layerdrop layerdropZ pad_token_idr9Zscale_embeddingrEsqrtr{r7 vocab_size embed_tokensuse_learned_position_embeddingsrZmax_position_embeddingsembed_positionsr=layernorm_embeddingrrZ ModuleListrangeZdecoder_layerslayersgradient_checkpointing post_init)rrhr:rrrrs*   zTrOCRDecoder.__init__cC|jSr;rrrrrget_input_embeddingsz!TrOCRDecoder.get_input_embeddingscC ||_dSr;rrvaluerrrset_input_embeddings z!TrOCRDecoder.set_input_embeddingsNc  Cs| dur| n|jj} | dur| n|jj} | dur| n|jj} | dur$| n|jj} |dur4|dur4td|durD|} |d| jd}n|dur\|dd}|dddddf} ntd|durm|ddjdnd}|durx| |}|jj r|j | |d}n|j ||d}||}|j dur| |}t jj||j|jd}| j}t||||}|dur|durt||j|dd }|jr|jr| rtd d } | rd nd}| rd nd}| r|durd nd}| rd nd}t||gd dgD]*\}}|dur|dt|jkrtd|dt|jd|ddqt|jD]\}}| r0||f7}|jrAtg}||jkrAq$|durJ||nd}|jrv|jrv||j|||||dure||nd|duro||ndd| | }n||||||dur||nd|dur||nd|| | d }|d}| r||| rdndf7}| r||df7}|dur||df7}q$| r||f7}| r|nd}| st dd|||||fDSt!|||||dS)a Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide it. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. [What are attention masks?](../glossary#attention-mask) encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*): Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention of the decoder. encoder_attention_mask (`torch.LongTensor` of shape `(batch_size, encoder_sequence_length)`, *optional*): Mask to avoid performing cross-attention on padding tokens indices of encoder input_ids. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. [What are attention masks?](../glossary#attention-mask) head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*): Mask to nullify selected heads of the attention modules. 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 `(decoder_layers, decoder_attention_heads)`, *optional*): Mask to nullify selected heads of the attention modules in encoder to avoid performing cross-attention on hidden heads. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `decoder_input_ids` of shape `(batch_size, sequence_length)`. inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This is useful if you want more control over how to convert `input_ids` indices into associated vectors than the model's internal embedding lookup matrix. output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for more detail. return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. NzTYou cannot specify both decoder_input_ids and decoder_inputs_embeds at the same timer$zEYou have to specify either decoder_input_ids or decoder_inputs_embedsrr)r rt)rxz^`use_cache = True` is incompatible with gradient checkpointing. Setting `use_cache = False`...Fr head_maskcross_attn_head_maskzThe `z` should be specified for z layers, but it is for .)rprrrqrrorrrrrcss|] }|dur|VqdSr;r)rvrrr sz'TrOCRDecoder.forward..)Zlast_hidden_statepast_key_valuesrm attentionscross_attentions)"rhrroutput_hidden_statesruse_return_dictrarKr%rMrrrrrrwr[rvr r r"rloggerZ warning_onceziplenr enumerater&ZrandrZ_gradient_checkpointing_func__call__tupler )rrrprrrrr inputs_embedsrrrr return_dictinputZ input_shaper Z embed_posrmZall_hidden_statesZall_self_attnsZall_cross_attentionsZnext_decoder_cacheZ attn_maskZ mask_nameidxZ decoder_layerZdropout_probabilityroZ layer_outputsZ next_cacherrrr,sP                  zTrOCRDecoder.forward) NNNNNNNNNNNN) r0r1r2r3rrrrr,r6rrrrrs$ra The TrOCR Model with a language modeling head. Can be used for summarization. This wrapper class is a helper class to correctly load pretrained checkpoints when the causal language model is used in combination with the [`EncoderDecoderModel`] framework. )Z custom_introcs$eZdZfddZddZZS)TrOCRDecoderWrappercst|t||_dSr;)rrrdecoderrrrrrs zTrOCRDecoderWrapper.__init__cOs|j|i|Sr;)r)rargskwargsrrrr,szTrOCRDecoderWrapper.forward)r0r1r2rr,r6rrrrrs rzy The TrOCR Decoder with a language modeling head. Can be used as the decoder part of [`EncoderDecoderModel`] and c seZdZdgZfddZddZddZdd Zd d Zd d Z ddZ e             d#de e jde e jde e jde e jde e jde e jde eee jde e jde e jde ede ede ede edeeeffdd Zed!d"ZZS)$TrOCRForCausalLMzoutput_projection.weightcsNt|}d|_d|_t|t||_tj |j |j dd|_ | dS)NTFr_)copydeepcopyr\Zis_encoder_decoderrrrrrrcr{routput_projectionrrrrrrs    zTrOCRForCausalLM.__init__cCs |jjjSr;rrrrrrrrrz%TrOCRForCausalLM.get_input_embeddingscCs||jj_dSr;rrrrrrsz%TrOCRForCausalLM.set_input_embeddingscCrr;rrrrrget_output_embeddingsrz&TrOCRForCausalLM.get_output_embeddingscCrr;r)rZnew_embeddingsrrrset_output_embeddingsrz&TrOCRForCausalLM.set_output_embeddingscCs ||j_dSr;rr)rrrrr set_decoders zTrOCRForCausalLM.set_decodercCs|jjSr;rrrrr get_decoderszTrOCRForCausalLM.get_decoderNrrprrrrrrlabelsrrrrrrscCs| dur| n|jj} | dur| n|jj} | dur| n|jj} |jj||||||||| | | | d }||d}d}| durNt}||d|jj | d}| sd|f|dd}|durb|f|S|St |||j |j |j |jdS)a cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*): Mask to nullify selected heads of the cross-attention modules. 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, sequence_length)`, *optional*): Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., config.vocab_size]` or -100 (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]`. Example: ```python >>> from transformers import ( ... TrOCRConfig, ... TrOCRProcessor, ... TrOCRForCausalLM, ... ViTConfig, ... ViTModel, ... VisionEncoderDecoderModel, ... ) >>> import requests >>> from PIL import Image >>> # TrOCR is a decoder model and should be used within a VisionEncoderDecoderModel >>> # init vision2text model with random weights >>> encoder = ViTModel(ViTConfig()) >>> decoder = TrOCRForCausalLM(TrOCRConfig()) >>> model = VisionEncoderDecoderModel(encoder=encoder, decoder=decoder) >>> # If you want to start from the pretrained model, load the checkpoint with `VisionEncoderDecoderModel` >>> processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten") >>> model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten") >>> # load image from the IAM dataset >>> url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02.jpg" >>> image = Image.open(requests.get(url, stream=True).raw).convert("RGB") >>> pixel_values = processor(image, return_tensors="pt").pixel_values >>> text = "industry, ' Mr. Brown commented icily. ' Let us have a" >>> # training >>> model.config.decoder_start_token_id = processor.tokenizer.eos_token_id >>> model.config.pad_token_id = processor.tokenizer.pad_token_id >>> model.config.vocab_size = model.config.decoder.vocab_size >>> labels = processor.tokenizer(text, return_tensors="pt").input_ids >>> outputs = model(pixel_values, labels=labels) >>> loss = outputs.loss >>> round(loss.item(), 2) 5.30 >>> # inference >>> generated_ids = model.generate(pixel_values) >>> generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0] >>> generated_text 'industry, " Mr. Brown commented icily. " Let us have a' ```N) rrprrrrrrrrrrrrr$r)losslogitsrrmrr)rhrrrrrrrrrKrr rrmrr)rrrprrrrrrrrrrrrrrrZloss_fctoutputrrrr,sDLzTrOCRForCausalLM.forwardcs.d}|D]}|tfdd|Df7}q|S)Nrc3s$|] }|d|jVqdSr/)rOrLr#)rZ past_statebeam_idxrrrqs"z2TrOCRForCausalLM._reorder_cache..)r)rrZreordered_pastZ layer_pastrrr_reorder_cachels zTrOCRForCausalLM._reorder_cache) NNNNNNNNNNNNN)r0r1r2Z_tied_weights_keysrrrrrrrrrr&Z LongTensorr5rBrryrr r,rTrr6rrrrrsl         ur)+r3rrEtypingrrrr&rZtorch.nnrZ activationsrZ generationr Zmodeling_attn_mask_utilsr r Zmodeling_outputsr r Zmodeling_utilsrutilsrrZconfiguration_trocrrZ get_loggerr0rrrr7Moduler=rUrzrrrr__all__rrrrsJ        >   "