a hl@sjdZddlmZddlmZmZddlZddlZddlmZddl m Z m Z ddl m Z dd lmZdd lmZdd lmZdd lmZdd lmZmZmZmZmZddlmZddlmZe e!Z"eeddGdddeZ#eeddGdddeZ$Gdddej%Z&eGdddeZ'eddGddde'Z(ed dGd!d"d"e'e Z)gd#Z*dS)$zPyTorch PaliGemmamodel.) dataclass)OptionalUnionN)nn)Cache StaticCache)GenerationMixin)FlashAttentionKwargs)BaseModelOutputWithPast)PreTrainedModel)Unpack) ModelOutputTransformersKwargsauto_docstringcan_return_tuplelogging) AutoModel)PaliGemmaConfigzN Base class for Paligemma outputs, with hidden states and attentions. )Z custom_introc@s$eZdZUdZdZeejed<dS)PaligemmaModelOutputWithPasta past_key_values (`Cache`, *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)`) Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. image_hidden_states (`torch.FloatTensor`, *optional*): A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`. image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state. Nimage_hidden_states) __name__ __module__ __qualname____doc__rrtorch FloatTensor__annotations__r r l/var/www/html/assistant/venv/lib/python3.9/site-packages/transformers/models/paligemma/modeling_paligemma.pyr,s  rzU Base class for PaliGemma causal language model (or autoregressive) outputs. c@seZdZUdZdZeejed<dZ eejed<dZ ee e eje fed<dZeeejed<dZeeejed<dZeejed<dS) PaliGemmaCausalLMOutputWithPastaa loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): Language modeling loss (for next-token prediction). logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.text_config.vocab_size)`): Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). past_key_values (`Cache`, *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)`) Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. image_hidden_states (`torch.FloatTensor`, *optional*): A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`. image_hidden_states of the model produced by the vision encoder after projecting last hidden state. Nlosslogitspast_key_values hidden_states attentionsr)rrrrr#rrrrr$r%rlistrr&tupler'rr r r r!r"Bs r"cs*eZdZedfdd ZddZZS)PaliGemmaMultiModalProjectorconfigcs(ttj|jj|jjdd|_dS)NTbias)super__init__rLinear vision_config hidden_sizeZprojection_dimlinearselfr, __class__r r!r0bs z%PaliGemmaMultiModalProjector.__init__cCs||}|SN)r4)r6image_featuresr&r r r!forwardfs z$PaliGemmaMultiModalProjector.forward)rrrrr0r; __classcell__r r r7r!r*asr*c@sDeZdZUeed<dZdZdgZdZdZ dZ dZ dZ dZ ddZd S) PaliGemmaPreTrainedModelr,Tr*r%FcCsNt|jd|jj}t|tjrJ|jjj d|d|j durJ|j j dS)Ninitializer_range)meanstd) getattrr,Zget_text_configr? isinstancerr1weightdataZnormal_r.Zzero_)r6modulerBr r r! _init_weightszs   z&PaliGemmaPreTrainedModel._init_weightsN)rrrrrZbase_model_prefixZsupports_gradient_checkpointingZ_no_split_modulesZ_skip_keys_device_placementZ_can_compile_fullgraphZ_supports_flash_attnZ_supports_sdpaZ_supports_flex_attnZ_supports_attention_backendrHr r r r!r=ls r=z{ The Base Paligemma model which consists of a vision backbone and a language model withou language modeling head., cseZdZddiZdZedfdd ZddZd d Zd d Z d dZ de e dddZ ejdddZejejejdddZeedejeje eje eje eeejefe eje eje eje eje e e e e e e e eeeeefdddZZS)PaliGemmaModelzlanguage_model.modellanguage_modelFr+csjt|tj|jd|_t||_|jj |_ tj|jd}||_ |j j durX|j j nd|_ | dS)Nr+)r/r0r from_configr2 vision_towerr*multi_modal_projector text_config vocab_sizerJr,Z pad_token_id post_init)r6r,rJr7r r!r0s   zPaliGemmaModel.__init__cCs |jSr9)rJget_input_embeddingsr6r r r!rRsz#PaliGemmaModel.get_input_embeddingscCs|j|dSr9)rJset_input_embeddingsr6valuer r r!rTsz#PaliGemmaModel.set_input_embeddingscCs ||_dSr9rJr6decoderr r r! set_decoderszPaliGemmaModel.set_decodercCs|jSr9rWrSr r r! get_decoderszPaliGemmaModel.get_decoderN) is_trainingcCsr|jjjdkr&|dur"d|vr"|SdS|dur2|n|j}t|t}t|jj }|dur\|}|j dd\} } |r|| } n&t|tj r|j dn|d| d} |dur| dkr|Stj| | f||j|jd} | dkr|rtj| dd } nd| ddd| f<| tj| |jd |ddk9} | ddddddf| ddd} |durn| } |j d} |r|dur~td | ddddddd| f|ddddddf| jdkd| ddddddd| f<| ddddddd| f|ddddddf| j}|dk}| ddddddd| f||| ddddddd| f<| S) NZflash_attention_2r@rrKrrZ fill_valuedtypedeviceZdiagonalr`z/Token type ids must be provided during training)r,rOZ_attn_implementationZtrainingrDrrfinfor_minshapeZget_max_cache_shapeTensordimfullr`triuarangereshapeexpandclone ValueError masked_fillto)r6attention_masktoken_type_idsr%cache_position input_tensorr\Zusing_static_cache min_dtypeZinputs_lead_dimsequence_length target_length causal_mask mask_length padding_maskr r r!_update_causal_masksT      $    $ @  z"PaliGemmaModel._update_causal_mask) pixel_valuescCs0||}|j}||}||jjjd}|S)a Obtains image last hidden states from the vision tower and apply multimodal projection. Args: pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`) The tensors corresponding to the input images. Returns: image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). g?)rMlast_hidden_staterNr,rOr3)r6r|Z image_outputsZselected_image_featurer:r r r!get_image_featuress  z!PaliGemmaModel.get_image_features) input_ids inputs_embedsr:cCs|dur8||tj|jjtj|jdk}|d}n ||jjk}|}| d | |j}|j d|j d}|| | krtd|d||S)z Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is equal to the length of multimodal features. If the lengths are different, an error is raised. N)r_r`rKrrz6Image features and image tokens do not match: tokens: z , features )rRrZtensorr,image_token_idlongr`allsum unsqueezeZ expand_asrpreZnumelrn)r6rrr:special_image_maskZn_image_tokensZn_image_featuresr r r!get_placeholder_masks  z#PaliGemmaModel.get_placeholder_mask)rr|rq position_idsr%rrrsrlabels use_cacheoutput_attentionsoutput_hidden_states return_dictkwargsreturnc Ks|du|duArtd| dur$| n|jj} | dur8| n|jj} | durL| n|jj} |duob| du}|dur|jj|jkr||jjk}|}d||<n|}|dur||}|dur|dur| nd}t j |||j d|j d}|dur|dd}|dur>||}||j |j}|j|||d}|||}|||||||}|jf||||| | | d|d |}t|j|j|j|j|dur|ndd S) i 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.text_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.text_config.vocab_size]`. Example: ```python >>> from PIL import Image >>> import requests >>> from transformers import AutoProcessor, PaliGemmaForConditionalGeneration >>> model = PaliGemmaForConditionalGeneration.from_pretrained("google/paligemma2-3b-mix-224") >>> processor = AutoProcessor.from_pretrained("google/paligemma2-3b-mix-224") >>> prompt = "Where is the cat standing?" >>> url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg" >>> image = Image.open(requests.get(url, stream=True).raw) >>> inputs = processor(images=image, text=prompt, return_tensors="pt") >>> # Generate >>> generate_ids = model.generate(**inputs,) >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "Where is the cat standing?\nsnow" ```Nz:You must specify exactly one of input_ids or inputs_embedsrrrb)rr:T) rqrr%rrrrrrs)r}r%r&r'r)rnr,rruse_return_dictrrPrmrRZget_seq_lengthrrjrer`rr~rpr_rZmasked_scatterr{rJrr}r%r&r')r6rr|rqrr%rrrsrrrrrrrr\rZ llm_input_idsZpast_seen_tokensr:rxoutputsr r r!r;sf/         zPaliGemmaModel.forward)NNNNN) NNNNNNNNNNNNN)rrr_checkpoint_conversion_mappingZaccepts_loss_kwargsrr0rRrTrZr[rboolr{rrr~ LongTensorrrrrfrr(rr r r)rr;r<r r r7r!rIsd  B  rIz| The Base Paligemma model which consists of a vision backbone and a language model without language modeling head., csHeZdZdddddZdgZedfdd Zd d Zd d ZddZ ddZ ddZ e ddZ e ddZe ddZeed%ejejeejeejeeeejefeejeejeejeejeeeeeeeeeeejfeeeeefdddZ d&fd d! Z!e"ejeeej#ejed"d#d$Z$Z%S)'!PaliGemmaForConditionalGenerationzmodel.language_modelzmodel.vision_towerzmodel.multi_modal_projectorlm_head)z^language_model.modelz ^vision_towerz^multi_modal_projectorz^language_model.lm_headzlm_head.weightr+cs<t|t||_tj|jj|jjdd|_ | dS)NFr-) r/r0rImodelrr1rOr3rPrrQr5r7r r!r0s  z*PaliGemmaForConditionalGeneration.__init__cCs |jSr9)rrRrSr r r!rRsz6PaliGemmaForConditionalGeneration.get_input_embeddingscCs|j|dSr9)rrTrUr r r!rTsz6PaliGemmaForConditionalGeneration.set_input_embeddingscCs|j|dSr9)rrZrXr r r!rZsz-PaliGemmaForConditionalGeneration.set_decodercCs |jSr9)rr[rSr r r!r[sz-PaliGemmaForConditionalGeneration.get_decodercCs |j|Sr9)rr~)r6r|r r r!r~sz4PaliGemmaForConditionalGeneration.get_image_featurescCs|jjSr9)rrJrSr r r!rJsz0PaliGemmaForConditionalGeneration.language_modelcCs|jjSr9)rrMrSr r r!rMsz.PaliGemmaForConditionalGeneration.vision_towercCs|jjSr9)rrNrSr r r!rNsz7PaliGemmaForConditionalGeneration.multi_modal_projectorNr)rr|rqrr%rrrsrrrrrrlogits_to_keeprrcKs| dur | n|jj} | dur | n|jj} | dur4| n|jj} |jf|||||||| | | | d|d |}|d}t|trt| dn|}||dd|ddf}d}| dur|j f|| |jj j d|}t |||j |j|j|jdS)rNT) rr|rrrqrr%rrrrrrrsr)r$rrP)r#r$r%r&r'r)r,rrrrrDintslicerZ loss_functionrOrPr"r%r&r'r)r6rr|rqrr%rrrsrrrrrrrrrr&Z slice_indicesr$r#r r r!r;sN/ z)PaliGemmaForConditionalGeneration.forwardTc  stj|f|||||| | |d| } | ddurD| dd7<|ddkrX|| d<|duof| du}t|tozt|j}|ddkr|r|dur|n|}|j||||||}|| d<| S)N)r%rrqrrsrrrrrrrr|rq) r/prepare_inputs_for_generationgetrDranyZ is_slidingrr{)r6rr%rrsrr|rqrrrrrrZ model_inputsr\Zis_static_hybrid_cachertrxr7r r!r s6   z?PaliGemmaForConditionalGeneration.prepare_inputs_for_generation)rqrvrwr_rs batch_sizec KsF|dur|dkr|}n&t|j}tj||f|||jd}|dkrVtj|dd}|tj||jd|ddk9}|ddddddf |ddd}|durB| }|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. Nr]r^rrarbrKr)rgrrcrdrhr`rirjrkrlrmrerpro) rqrvrwr_rsrrrxruryrzr r r!5_prepare_4d_causal_attention_mask_with_cache_position<s*  $  6  zWPaliGemmaForConditionalGeneration._prepare_4d_causal_attention_mask_with_cache_position)NNNNNNNNNNNNNr) NNNNNNNTNN)&rrrrZ_tied_weights_keysrr0rRrTrZr[r~propertyrJrMrNrrrrrrrfrr(rrrr rr)r"r;r staticmethodr_rr<r r r7r!rs     [/r)rr=rI)+r dataclassesrtypingrrrZtorch.utils.checkpointrZ cache_utilsrrZ generationr Zmodeling_flash_attention_utilsr Zmodeling_outputsr Zmodeling_utilsr Zprocessing_utilsr utilsrrrrrautorZconfiguration_paligemmarZ get_loggerrloggerrr"Moduler*r=rIr__all__r r r r!sN           zn