U 4Afr@sdZddlZddlmZddlmZmZmZmZm Z ddl Z ddl Z ddl m Z ddl mZddlmZdd lmZmZmZmZdd lmZdd lmZmZmZmZmZd d lmZm Z m!Z!e"e#Z$eZ%dEe j&e j'ee(dddZ)dFe j*e j'e j+e(dddZ,dGddZ-dZ.dZ/dZ0dZ1eGdddeZ2eGdddeZ3Gddde j4Z5Gd d!d!e j4Z6Gd"d#d#e j4Z7Gd$d%d%e j4Z8Gd&d'd'e j4Z9Gd(d)d)e j4Z:Gd*d+d+e j4Z;Gd,d-d-e j4ZGd2d3d3e j4Z?Gd4d5d5eZ@Gd6d7d7e@ZAGd8d9d9e@ZBed:e.Gd;d<dd>e j4ZDed?e.Gd@dAdAe@ZEedBe.GdCdDdDe@ZFdS)HzPyTorch KOSMOS-2 model.N) dataclass)AnyListOptionalTupleUnion)nn)CrossEntropyLoss)ACT2FN)BaseModelOutput)BaseModelOutputWithPastAndCrossAttentionsBaseModelOutputWithPooling!CausalLMOutputWithCrossAttentions)PreTrainedModel) ModelOutputadd_start_docstrings%add_start_docstrings_to_model_forwardloggingreplace_return_docstrings) Kosmos2ConfigKosmos2TextConfigKosmos2VisionConfig)maskdtypetgt_lencCsj|\}}|dk r|n|}|ddddddf|d|||}d|}||tjt|jS)z_ Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`. Nr?)sizeexpandtoZ masked_filltorchboolfinfomin)rrrbszsrc_lenZ expanded_maskZ inverted_maskr'P/tmp/pip-unpacked-wheel-zw5xktn0/transformers/models/kosmos2/modeling_kosmos2.py _expand_mask1s  *r))input_ids_shaperdevicepast_key_values_lengthcCs|\}}tj||ft|j|d}tj|d|d}|||d|ddkd||}|dkrtj tj ||||d|gdd}|ddddddf |d|||S)zB Make causal mask used for bi-directional self-attention. )r+rrrr+ZdimN) r!fullr#r$arangerZ masked_fill_viewr catzerosr)r*rr+r,r%rrZ mask_condr'r'r(_make_causal_mask?s"  r5cCs6||}tj|dd|||}||S)a 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`. Args: x: torch.Tensor x: Returns: torch.Tensor rr/)neintr!ZcumsumZtype_aslong) input_ids padding_idxr,rZincremental_indicesr'r'r("create_position_ids_from_input_idsQs r;a@ This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads etc.) This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and behavior. Parameters: config ([`Kosmos2Config`]): Model configuration class with all the parameters of the model. Initializing with a config file does not load the weights associated with the model, only the configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights. a  Args: pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See [`CLIPImageProcessor.__call__`] for details. 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. 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) image_embeds: (`torch.FloatTensor` of shape `(batch_size, latent_query_num, hidden_size)`, *optional*): Sequence of hidden-states at the output of `Kosmos2ImageToTextProjection`. image_embeds_position_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to indicate the location in a sequence to insert the image features . Mask values selected in `[0, 1]`: - 1 for places where to put the image features, - 0 for places that are not for image features (i.e. for text tokens). encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if the model is configured as a decoder. encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): Mask to nullify selected heads of the self-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 cross-attention modules. 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))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`): Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up 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. position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, config.max_position_embeddings - 1]`. [What are position IDs?](../glossary#position-ids) use_cache (`bool`, *optional*): If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see `past_key_values`). 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. a Args: pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See [`CLIPImageProcessor.__call__`] for details. 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) image_embeds_position_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to indicate the location in a sequence to insert the image features . Mask values selected in `[0, 1]`: - 1 for places where to put the image features, - 0 for places that are not for image features (i.e. for text tokens). 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) head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): Mask to nullify selected heads of the self-attention modules. 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))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`): Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up 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)`. image_embeds: (`torch.FloatTensor` of shape `(batch_size, latent_query_num, hidden_size)`, *optional*): Sequence of hidden-states at the output of `Kosmos2ImageToTextProjection`. 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. position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, config.max_position_embeddings - 1]`. [What are position IDs?](../glossary#position-ids) use_cache (`bool`, *optional*): If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see `past_key_values`). 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. c@seZdZUdZdZejed<dZe e e ejed<dZ e e ejed<dZ e e ejed<dZ e ejed<dZe e ejed<dZeed <e ed d d ZdS) Kosmos2ModelOutputa Base class for text model's outputs that also contains a pooling of the last hidden states. Args: last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): Sequence of hidden-states at the output of the last layer of the model. hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. image_embeds (`torch.FloatTensor` of shape `(batch_size, latent_query_num, hidden_size)`, *optional*): Sequence of hidden-states at the output of `Kosmos2ImageToTextProjection`. projection_attentions (`tuple(torch.FloatTensor)`, *optional*): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights given by `Kosmos2ImageToTextProjection`, after the attention softmax, used to compute the weighted average in the self-attention heads. vision_model_output(`BaseModelOutputWithPooling`, *optional*): The output of the [`Kosmos2VisionModel`]. 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 optionally if `config.is_encoder_decoder=True` 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 optionally if `config.is_encoder_decoder=True` in the cross-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. Nlast_hidden_statepast_key_values hidden_states attentions image_embedsprojection_attentionsvision_model_outputreturncstfddDS)Nc3s,|]$}|dkr|n t|VqdS)Ztext_model_outputrCNgetattrto_tuple.0kselfr'r( ;sz.Kosmos2ModelOutput.to_tuple..tuplekeysrMr'rMr(rI:s zKosmos2ModelOutput.to_tuple)__name__ __module__ __qualname____doc__r=r! FloatTensor__annotations__r>rrr?r@rArBrCrrrIr'r'r'r(r< s & r<c@seZdZUdZdZeejed<dZ ejed<dZ ee e ejed<dZ ee ejed<dZ ee ejed<dZeejed<dZee ejed <dZeed <e ed d d ZdS)*Kosmos2ForConditionalGenerationModelOutputaR Model output class for `Kosmos2ForConditionalGeneration`. Args: 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.vocab_size)`): Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. image_embeds (`torch.FloatTensor` of shape `(batch_size, latent_query_num, hidden_size)`, *optional*): Sequence of hidden-states at the output of `Kosmos2ImageToTextProjection`. projection_attentions (`tuple(torch.FloatTensor)`, *optional*): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights given by `Kosmos2ImageToTextProjection`, after the attention softmax, used to compute the weighted average in the self-attention heads. vision_model_output(`BaseModelOutputWithPooling`, *optional*): The output of the [`Kosmos2VisionModel`]. 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 optionally if `config.is_encoder_decoder=True` 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 optionally if `config.is_encoder_decoder=True` in the cross-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. Nlosslogitsr>r?r@rArBrCrDcstfddDS)Nc3s,|]$}|dkr|n t|VqdSrFrGrJrMr'r(rOuszFKosmos2ForConditionalGenerationModelOutput.to_tuple..rPrMr'rMr(rIts z3Kosmos2ForConditionalGenerationModelOutput.to_tuple)rSrTrUrVrZrr!rWrXr[r>rr?r@rArBrCrrrIr'r'r'r(rYAs ( rYcs6eZdZedfdd ZejejdddZZ S)Kosmos2VisionEmbeddingsconfigcst||_|j|_|j|_|j|_tt |j|_ tj |j |j|j|jdd|_|j|jd|_|jd|_t|j|j|_|jdt |jddddS)NF)Z in_channelsZ out_channelsZ kernel_sizeZstridebiasr position_ids)rr- persistent)super__init__r^ hidden_size embed_dimZ image_sizeZ patch_sizer Parameterr!randnclass_embeddingZConv2dZ num_channelspatch_embeddingZ num_patches num_positions Embeddingposition_embeddingregister_bufferr1rrNr^ __class__r'r(re}s"  z Kosmos2VisionEmbeddings.__init__) pixel_valuesrEcCsn|jd}|jjj}||j|d}|ddd}|j|dd}t j ||gdd}|| |j }|S)Nrrr`rr-r/) shaperkweightrr flatten transposerjrr!r3rnra)rNrs batch_sizeZ target_dtypeZ patch_embedsZ class_embeds embeddingsr'r'r(forwards  zKosmos2VisionEmbeddings.forward) rSrTrUrrer!rWTensorr{ __classcell__r'r'rqr(r\|sr\c sneZdZdZfddZejeedddZd eje eje eje e e eje ejfd d d Z Z S) Kosmos2VisionAttention=Multi-headed attention from 'Attention Is All You Need' papercst||_|j|_|j|_|j|j|_|j|j|jkrZtd|jd|jd|jd|_ |j |_ t |j|j|_t |j|j|_t |j|j|_t |j|j|_dS)N;embed_dim must be divisible by num_heads (got `embed_dim`:  and `num_heads`: ).)rdrer^rfrgZnum_attention_heads num_headshead_dim ValueErrorscaleattention_dropoutdropoutrLineark_projv_projq_projout_projrprqr'r(res  zKosmos2VisionAttention.__init__)tensorseq_lenr%cCs ||||j|jddS)Nrr`)r2rrrx contiguous)rNrrr%r'r'r(_shapeszKosmos2VisionAttention._shapeNFr?attention_maskcausal_attention_maskoutput_attentionsrEcCs|\}}}|||j}|||d|} |||d|} ||jd|jf} ||||j| }| j| } | j| } | d} t || dd} | ||j|| fkrt d||j|| fd| |dk rD||d|| fkrt d|d|| fd|| ||j|| |} | ||j|| } |dk r||d|| fkrt d|d|| fd|| ||j|| |} | ||j|| } t jj| dd} |r| ||j|| }|||j|| } nd}t jj| |j|jd }t || }|||j||jfkrRt d ||j||jfd||||j||j}| dd}||||}||}||fS) #Input shape: Batch x Time x Channelr-rr`z$Attention weights should be of size , but is N!Attention mask should be of size r/ptrainingz `attn_output` should be of size )rrrrrrrrr2r!Zbmmrxrr functionalsoftmaxrrZreshaper)rNr?rrrr%rrg query_states key_states value_statesZ proj_shaper& attn_weightsZattn_weights_reshapedZ attn_probs attn_outputr'r'r(r{sX         zKosmos2VisionAttention.forward)NNF)rSrTrUrVrer!r|r7rrr"rr{r}r'r'rqr(r~s r~cs0eZdZfddZejejdddZZS)Kosmos2VisionMLPcsDt||_t|j|_t|j|j |_ t|j |j|_ dSN) rdrer^r Z hidden_act activation_fnrrrfZintermediate_sizefc1fc2rprqr'r(re s   zKosmos2VisionMLP.__init__)r?rEcCs"||}||}||}|Sr)rrrrNr?r'r'r(r{s   zKosmos2VisionMLP.forward)rSrTrUrer!r|r{r}r'r'rqr(r s rcsJeZdZedfdd Zdejejejeee ej dddZ Z S) Kosmos2VisionEncoderLayerr]csRt|j|_t||_tj|j|jd|_ t ||_ tj|j|jd|_ dSNZeps) rdrerfrgr~ self_attnr LayerNormlayer_norm_eps layer_norm1rmlp layer_norm2rprqr'r(res    z"Kosmos2VisionEncoderLayer.__init__FrcCsd|}||}|j||||d\}}||}|}||}||}||}|f}|r`||f7}|S)aI 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. `(config.encoder_attention_heads,)`. output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. )r?rrr)rrrr)rNr?rrrresidualroutputsr'r'r(r{"s"     z!Kosmos2VisionEncoderLayer.forward)F) rSrTrUrrer!r|rr"rrWr{r}r'r'rqr(rs rc s`eZdZdZedfdd Zd eejeejee ee ee e e e fdddZ ZS) Kosmos2VisionEncoderz Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a [`Kosmos2VisionEncoderLayer`]. Args: config: Kosmos2VisionConfig r]cs:t|_tfddtjD|_d|_dS)Ncsg|] }tqSr')rrK_r]r'r( Xsz1Kosmos2VisionEncoder.__init__..F) rdrer^r ModuleListrangenum_hidden_layerslayersgradient_checkpointingrprqr]r(reUs  zKosmos2VisionEncoder.__init__N)rrroutput_hidden_states return_dictrEc Cs|dk r |n|jj}|dk r |n|jj}|dk r4|n|jj}|rDdnd}|rPdnd}|} t|jD]b\} } |rx|| f}|jr|jr|| j | |||} n| | |||d} | d} |rb|| df}qb|r|| f}|st dd| ||fDSt | ||dS) a Args: inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): 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. 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) causal_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): Causal mask for the text model. 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) 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. Nr')rrrcss|]}|dk r|VqdSrr'rKvr'r'r(rOsz/Kosmos2VisionEncoder.forward..)r=r?r@) r^rruse_return_dict enumeraterrr_gradient_checkpointing_func__call__rQr ) rN inputs_embedsrrrrrZencoder_statesZall_attentionsr?idxZ encoder_layer layer_outputsr'r'r(r{[sH&     zKosmos2VisionEncoder.forward)NNNNN)rSrTrUrVrrerr!r|r"rrr r{r}r'r'rqr(rLs  rcsTeZdZedfdd Zdeejeeeeeee e e fdddZ Z S) Kosmos2VisionTransformerr]csRt||_|j}t||_tj||jd|_ t ||_ tj||jd|_ dSr) rdrer^rfr\rzrrr pre_layrnormrencoderpost_layernorm)rNr^rgrqr'r(res   z!Kosmos2VisionTransformer.__init__NrsrrrrEc Cs|dk r |n|jj}|dk r |n|jj}|dk r4|n|jj}|dkrLtd||}||}|j||||d}|d}|dddddf}||}|s||f|ddSt |||j |j dS)Nz You have to specify pixel_values)rrrrrr)r=Z pooler_outputr?r@) r^rrrrrzrrrrr?r@) rNrsrrrr?Zencoder_outputsr=Z pooled_outputr'r'r(r{s2   z Kosmos2VisionTransformer.forward)NNNN)rSrTrUrrerr!rWr"rrrr{r}r'r'rqr(rs  rcseZdZdZdeeeedfdd ZdeeeedddZedeeeedd d Z e de j e j ee j d d dZ ddZZS)(Kosmos2TextSinusoidalPositionalEmbeddingzDThis module produces sinusoidal positional embeddings of any length.Nrl embedding_dimr:cs4td|_||_||_|||j||dS)Nr`)rdreoffsetrr: make_weights)rNrlrr:rqr'r(res  z1Kosmos2TextSinusoidalPositionalEmbedding.__init__)num_embeddingsrr:cCsB||||}t|dr.|j|jj|jjd}|jd|dddS)Nweightsr.Frb) get_embeddinghasattrr rrr+ro)rNrrr:Z emb_weightsr'r'r(rs z5Kosmos2TextSinusoidalPositionalEmbedding.make_weightscCs|d}td|d}ttj|tjd| }tj|tjdd|d}tjt |t |gdd |d}|ddkrtj|t |dgdd}|dk rd||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". r`i'rrtrr/r-N)mathlogr!expr1Zint64float unsqueezer3sincosr2r4r Zget_default_dtype)rrr:Zhalf_dimZembr'r'r(rs  $& z6Kosmos2TextSinusoidalPositionalEmbedding.get_embeddingrr9rr,racCs|dk r4|\}}|dkr\t||j||j}n(|dd\}}|dkr\|||}|jd||}||jdkr|||j|j |j|j d| d |||jj d S)Nr-rr)rr;r:r r+&create_position_ids_from_inputs_embedsrrrr index_selectr2rudetach)rNr9rr,rar%rZmax_posr'r'r(r{s"  z0Kosmos2TextSinusoidalPositionalEmbedding.forwardcCsV|dd}|d}tj|jd||jdtj|jd}|d||S)z We are provided embeddings directly. We cannot infer which are padded so just generate sequential position ids. Args: inputs_embeds: torch.Tensor Returns: torch.Tensor Nr-rr.r) rr!r1r:r8r+rrr)rNrr, input_shapeZsequence_lengthrar'r'r(r*s  zOKosmos2TextSinusoidalPositionalEmbedding.create_position_ids_from_inputs_embeds)N)N)N)NNrN)rSrTrUrVr7rrer staticmethodrr!Zno_gradr|r{rr}r'r'rqr(rs"rc seZdZdZdeeeeeedfdd Zej ej dd d Z dej e ej e e ej e ej e ej ee ej e ej e e ej fd d dZ ZS)KosmosTextAttentionrFT)rgrr is_decoderadd_inner_attn_layernormr_cst||_||_||_|||_|j||jkrNtd|jd|d|jd|_||_t j |||d|_ t j |||d|_ t j |||d|_ t j |||d|_d|_|rt j||jd|_dS)Nrrrr)r_r)rdrergrrrrscalingrrrrrrr inner_attn_lnrr)rNr^rgrrrrr_rqr'r(re@s$   zKosmosTextAttention.__init__) projectionrEcCs6|dd|j|jf}||dddd}|S)Nr-rr`rr )rrrr2permute)rNrZnew_projection_shapeZnew_projectionr'r'r(rbszKosmosTextAttention._shapeN)r?encoder_hidden_statespast_key_valuerlayer_head_maskrrEcCs|dk }|jdd\}} |dk r&|n|} |r\|r\|djd| jdkr\|d} |d} nX||| } ||| } |dk r|stj|d| gdd} tj|d| gdd} ||||j} t| | dd}|j r| | f}| d}|dk r>| |d| |fkr6t d|d| |fd | ||}t jj|dd}|dk r`||}t jj||j|jd }t|| }|dddd || d}|jdk r||}||}|||fS) rNr`rrr/r-rrrr )rurrrr!r3rrmatmulrxrrrrrrrrrrr2rr)rNr?rrrrrZis_cross_attentionry seq_lengthZcurrent_statesrrrrr&Zcontext_statesrr'r'r(r{hs@           zKosmosTextAttention.forward)rFFT)NNNNF)rSrTrUrVr7rr"rer!r|rrrr{r}r'r'rqr(r<s8"  rcs*eZdZedfdd ZddZZS)Kosmos2TextFFNr]csbt|j|_t|j|_|j|_t|j |j |_ t|j |j |_ tj |j |jd|_dSr)rdrerr Zactivation_functionractivation_dropoutrrrgZffn_dimrrrr ffn_layernormrprqr'r(res  zKosmos2TextFFN.__init__cCsT|||}tjj||j|jd}||}||}tjj||j|jd}|S)Nr) rrrrrrrrrrr'r'r(r{s   zKosmos2TextFFN.forward)rSrTrUrrer{r}r'r'rqr(rs rcseZdZedfdd Zd ejeejeejeejeejeejeeejee ee eej eeej ej ffd dd Z Z S) Kosmos2TextBlockr]cst|j|_t||j|j|jddd|_|j|_tj |j|j d|_ |j rt||j|j|jddd|_ tj |j|j d|_t||_tj |j|j d|_dS)NT)rgrrrrrF)rdrergrattention_headsrrrrrrself_attn_layer_normZadd_cross_attention encoder_attnencoder_attn_layer_normrffnfinal_layer_normrprqr'r(res0  zKosmos2TextBlock.__init__NFT) r?rrencoder_attention_maskrcross_attn_layer_head_maskrr use_cacherEc Cs<|} |dk r|ddnd} ||}|j|| |||d\}} } tjj||j|jd}| |}d}d}|dk rt|dstd|d|} ||}|dk r|ddnd}|j ||||||d\}}}tjj||j|jd}| |}| |} |} | |}| |}| |}|f}|r(|| |f7}| r8|| f7}|S) Nr`)r?rrrrrrz'If `encoder_hidden_states` are passed, z` has to be instantiated with cross-attention layers by setting `config.add_cross_attention=True`r)r?rrrrr) rrrrrrrrrrrr)rNr?rrrrrrrrrZself_attn_past_key_valueZself_attn_weightsZpresent_key_valueZcross_attn_present_key_valueZcross_attn_weightsZcross_attn_past_key_valuerr'r'r(r{sV           zKosmos2TextBlock.forward)NNNNNNFT) rSrTrUrrer!r|rrr"rWr{r}r'r'rqr(rs*  rcseZdZdZedfdd ZddZdejejeje ejd d d Z de eje eje eje eje eje eje eje eje e ej e eje eje ee ee ee eeeefd d dZZS)Kosmos2TextTransformerz Transformer decoder consisting of `config.layers` layers. Each layer is a [`Kosmos2TextBlock`]. Args: config: Kosmos2TextConfig r]cst|_j|_j|_jr2tjnd|_ t j j jj d|_tjjj d|_t fddtjD|_t jj|_d|_dS)Nr)r:rcsg|] }tqSr')rrr]r'r(rPsz3Kosmos2TextTransformer.__init__..F)rdrer^r layerdropZscale_embeddingrsqrtrg embed_scalerrm vocab_size pad_token_id embed_tokensrZmax_position_embeddingsembed_positionsrrrrr layer_normrrprqr]r(reAs  zKosmos2TextTransformer.__init__cCs`d}|ddkr$t||j|j|d}|dk r\t||j|dd|j}|dkrT|n||}|S)Nr-r)r+r,r)r5rr+r)r )rNrrrr,Zcombined_attention_maskZexpanded_attn_maskr'r'r(_prepare_decoder_attention_maskUs z6Kosmos2TextTransformer._prepare_decoder_attention_maskNr)rrAimg_input_maskr,rac Cs|dkr||}|dk rB||jd|d||jtjd<||j}|j||||d}||j}||}t j j ||j |j d}|S)Nr-rtrr) rr r+r2rr!r"rrrrrr) rNr9rrAr r,raZ positionsr?r'r'r(forward_embeddingls$     z(Kosmos2TextTransformer.forward_embeddingr9rrAimage_embeds_position_maskrr head_maskcross_attn_head_maskr>rrarrrrrEc Cs| dk r | n|jj} |dk r |n|jj}| dk r4| n|jj} |dk rH|n|jj}|dk rj| dk rjtdnB|dk r|j}|d|d}n"| dk r| dd}ntd| dk r| ddjdnd}|dkrd}d}|j || |||| d}| ||||}|dk r(|dk r(t || j |dd}t jj||j|jd}|jrb|jrb| rbtd d } |rld nd}| rzd nd}| r|dk rd nd}| rd nd}t||gd d gD]V\}}|dk r|dt|jkrtd|dt|jd|ddqt|jD]6\}}|r2||f7}|jrTtg}||jkrTq| dk rf| |nd}|jr|jr||j|||||dk r||nd|dk r||ndd| | }n>||||||dk r||nd|dk r||nd|| | d }|d}| r"||| rdndf7}| r||df7}|dk r||df7}q||}|rl||f7}|stdd|||||fDSt|||||dS)NzDYou cannot specify both input_ids and inputs_embeds at the same timer-z5You have to specify either input_ids or inputs_embedsrr`)r9rrAr r,rarrzZ`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...Fr'rrzThe `z` should be specified for z layers, but it is for .)rrrrrrrrr rcss|]}|dk r|VqdSrr'rr'r'r(rOsz1Kosmos2TextTransformer.forward..)r=r>r?r@cross_attentions) r^rrrrrrur2rr r r)rrrrrrloggerZ warning_onceziplenrrr!ZrandrrrrrQr )rNr9rrAr rrrrr>rrarrrrrr,r?Zall_hidden_statesZall_self_attnsZall_cross_attentionsZpresent_key_value_statesZ attn_maskZ mask_namerZ decoder_layerZdropout_probabilityrrr'r'r(r{s   $         zKosmos2TextTransformer.forward)NNNrN)NNNNNNNNNNNNNNN)rSrTrUrVrrer r!r|r7r rrrWr"rrr r{r}r'r'rqr(r9s` %  rc@s(eZdZdZeZdZddgZddZdS)Kosmos2PreTrainedModelz An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models. TrrcCst|tr|jj}nt|ttfr,|jjj}t|ttfrD|jj }nt|ttfr\|jj j }t|t rt j j|jd|jd|dt j j|jj|jj|dt j j|jj|jj|dn6t|tr|jdd|jjd|}|jd|}t j j|jj|dt j j|jj|dt j j|jj|dt j j|jj|d|jjdk rd|jjj|jjdk r|jjj|jjdk r|jjj|jjdk r|jjjn:t|tr`|jjdd|jjd|}d|jjd|}t j j|j j|dt j j|j!j|d|j jdk r@|j jj|j!jdk r|j!jjnt|t"r|j#jj|j#jj$d|j%jj|j%jj$dnJt|t&r|j'jj|j'jj$d|j(jj|j(jj$dnt|t)rt j j|jj|dt j j|jj|dt j j|jj|dt j j|jj|d|jjdk rp|jjj|jjdk r|jjj|jjdk r|jjj|jjdk r|jjjn.t|t*r6t j j|j j|dt j j|j!j|d|j jdk r|j jj|j!jdk r|j!jjnt|trtt j j|j+j|d|j+jdk r|j+jjnt|t,rt j j|j-j|d|j-jdk r|j-jjnDt|t.r|j/jjjd|d|j/j0dk r|j/jj|j/j0dS)zInitialize the weightsrr)Zmeanstd)rr`Nr)1 isinstanceKosmos2VisionModelr^Zinitializer_factor Kosmos2ModelKosmos2ForConditionalGeneration vision_configKosmos2TextModelKosmos2TextForCausalLMZinit_std text_configr\rinitZnormal_rjrgrkrvZinitializer_rangernr~rrrrrr_dataZzero_rrfrrrrZfill_rrrrrrlm_headKosmos2ImageToTextProjectiondenserrr:)rNmoduleZfactorrZ in_proj_stdZ out_proj_stdZfc_stdr'r'r( _init_weights;s                 z$Kosmos2PreTrainedModel._init_weightsN) rSrTrUrVr config_classZsupports_gradient_checkpointingZ_no_split_modulesr%r'r'r'r(r1s rc seZdZeZdZedfdd ZejdddZ e e e e edd eejeeeeeeeee fd d d ZZS)rrsr]cs"t|t||_|dSr)rdrermodel post_initrprqr'r(res  zKosmos2VisionModel.__init__rDcCs |jjjSr)r'rzrkrMr'r'r(get_input_embeddingssz'Kosmos2VisionModel.get_input_embeddings output_typer&NrcCs|j||||dS) Returns: rsrrrr')rNrsrrrr'r'r(r{s zKosmos2VisionModel.forward)NNNN)rSrTrUrr&main_input_namererModuler)rKOSMOS2_VISION_INPUTS_DOCSTRINGrrrr!rWr"rrr{r}r'r'rqr(rs"  rcseZdZeZedfdd ZejdddZddZ e e e e ed deejeejeejeejeejeejeejeejeeejeejeejeeeeeeeeeee fd d d ZZS)rr]cs"t|t||_|dSr)rdrerr'r(rprqr'r(res  zKosmos2TextModel.__init__rDcCs|jjSrr'rrMr'r'r(r)sz%Kosmos2TextModel.get_input_embeddingscCs ||j_dSrr2rNvaluer'r'r(set_input_embeddingssz%Kosmos2TextModel.set_input_embeddingsr*Nr cCs(|j||||||||| | | | | ||dS)r,r9rrAr rrrrr>rrarrrrr.)rNr9rrAr rrrrr>rrarrrrr'r'r(r{s"zKosmos2TextModel.forward)NNNNNNNNNNNNNNN)rSrTrUrr&rerr0r)r5rKOSMOS2_TEXT_INPUTS_DOCSTRINGrr rr!r|rrWr"rrr{r}r'r'rqr(rsN   rz The text model from KOSMOS-2 with a language modeling head on top (linear layer with weights tied to the input embeddings). cseZdZeZdgZedfdd ZejdddZ dd Z ejdd d Z d d Z e eeeeddeejeejeejeejeejeejeejeejeeejeejeejeejeeeeeeeeeeefdddZdddZeddZZS)rzlm_head.weightr]cs8t|t||_tj|j|jdd|_| dS)NF)Z in_featuresZ out_featuresr_) rdrerr'rrrgrr!r(rprqr'r(res  zKosmos2TextForCausalLM.__init__rDcCs|jjSrr2rMr'r'r(r)sz+Kosmos2TextForCausalLM.get_input_embeddingscCs ||j_dSrr2r3r'r'r(r5sz+Kosmos2TextForCausalLM.set_input_embeddingscCs|jSrr!rMr'r'r(get_output_embeddingssz,Kosmos2TextForCausalLM.get_output_embeddingscCs ||_dSrr8rNZnew_embeddingsr'r'r(set_output_embeddings sz,Kosmos2TextForCausalLM.set_output_embeddingsr*N)r9rrAr rrrrr>rralabelsrrrrrEcCs(|dk r |n|jj}| dk r.| r*tdd} |j||||||||| | | | |||d}||d}d}| dk r| |j} |dddddf}| dddf}|j \}}}t }|| |||| ||}|s |f|dd}|dk r|f|S|St |||j |j|j|jd S) a labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the left-to-right language modeling loss (next word prediction). 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]` Returns: NzJThe `use_cache` argument is changed to `False` since `labels` is provided.Fr6r.r-r)rZr[r>r?r@r)r^rrwarningr'r!r r+rrur r2rr>r?r@r)rNr9rrAr rrrrr>rrar<rrrrrZ lm_logitsrZZ shift_logitsZ shift_labelsryrrZloss_fctoutputr'r'r(r{ sZ    zKosmos2TextForCausalLM.forwardc Ks|j}|dkr||}d} |dk rdt||jjddddddf} |ddddf}d}d}nJ|dk r|\} } |d} tj|tj| | | ftj |j dfdd}|||||| |dS)Nr)r:r,r-)rrr+rr/)r9rAr r>rrar) ruZnew_onesr;r^rrr!r3r4r"r+) rNr9rAr r>rrZ model_kwargsrraryrZmask_lenr'r'r(prepare_inputs_for_generation`s@       z4Kosmos2TextForCausalLM.prepare_inputs_for_generationcs.d}|D] }|tfdd|Df7}q|S)Nr'c3s"|]}|d|jVqdS)rN)rr r+)rKZ past_statebeam_idxr'r(rOsz8Kosmos2TextForCausalLM._reorder_cache..)rQ)r>rAZreordered_pastZ layer_pastr'r@r(_reorder_caches z%Kosmos2TextForCausalLM._reorder_cache)NNNNNNNNNNNNNNNN)NNNNN)rSrTrUrr&_tied_weights_keysrerr0r)r5r9r;rr7rrrr!r|rrW LongTensorr"rrr{r?rrBr}r'r'rqr(rsh    T 3rcs.eZdZdZedfdd ZddZZS)r"zmThe layer that transforms the image model's output to part of the text model's input (namely, image features)r]csbtt|jj|jj|_t t |j |jj|_ t|j|jj|jj|jjddd|_dS)NF)rrr)rdrerrrrfrrgr#rhr!riZlatent_query_num latent_queryrrrx_attnrprqr'r(res z%Kosmos2ImageToTextProjection.__init__cCsZ||}|jd|ddd}tj||gdd}|j||dddd\}}}||fS)Nrr-rr/)r?rrrr)r#rErrrr!r3rF)rNfeaturesr?rEZkey_value_statesrrr'r'r(r{s  z$Kosmos2ImageToTextProjection.forward)rSrTrUrVrrer{r}r'r'rqr(r"sr"z} KOSMOS-2 Model for generating text and image features. The model consists of a vision encoder and a language model. cseZdZeZdZedfdd ZejdddZ dd Z e e e eed deejeejeejeejeejeeejeejeejeejeeeeeeeeeeefd d dZZS)rrsr]cs:t|t|j|_t|j|_t||_ | dSr) rdrerr text_modelrr vision_modelr"image_to_text_projectionr(rprqr'r(res     zKosmos2Model.__init__rDcCs |jjjSrrHr'rrMr'r'r(r)sz!Kosmos2Model.get_input_embeddingscCs||jj_dSrrKr3r'r'r(r5sz!Kosmos2Model.set_input_embeddingsr*N)rsr9r rrr>rArrarrrrrEcCs| dk r | n|jj} | dk r | n|jj} | dk r4| n|jj} d}d}|dkr|dkr\td|j|| | | d}|jj|d}tj j |dd}| |\}}|j |||||||| | | | | d }| s||||f}t dd |DSt|j|j|j|j|||d S) a Returns: Examples: ```python >>> from PIL import Image >>> import requests >>> from transformers import AutoProcessor, Kosmos2Model >>> model = Kosmos2Model.from_pretrained("microsoft/kosmos-2-patch14-224") >>> processor = AutoProcessor.from_pretrained("microsoft/kosmos-2-patch14-224") >>> url = "https://huggingface.co/microsoft/kosmos-2-patch14-224/resolve/main/snowman.jpg" >>> image = Image.open(requests.get(url, stream=True).raw) >>> text = ( ... " An image of a snowman" ... " warming himself by a fire" ... "" ... ) >>> inputs = processor(text=text, images=image, return_tensors="pt", add_eos_token=True) >>> last_hidden_state = model( ... pixel_values=inputs["pixel_values"], ... input_ids=inputs["input_ids"], ... attention_mask=inputs["attention_mask"], ... image_embeds_position_mask=inputs["image_embeds_position_mask"], ... ).last_hidden_state >>> list(last_hidden_state.shape) [1, 91, 2048] ```Nrrarrrrcss|]}|dk r|VqdSrr'rKr>r'r'r(rO7sz'Kosmos2Model.forward..)r=r>r?r@rArBrC)r^rrrrrIr'rrr normalizerJrHrQr<r=r>r?r@)rNrsr9r rrr>rArrarrrrrCrBrr'r'r(r{sX3zKosmos2Model.forward) NNNNNNNNNNNNN)rSrTrUrr&r/rerr0r)r5rKOSMOS2_INPUTS_DOCSTRINGrr<_CONFIG_FOR_DOCrr!r|rrWr"rrr{r}r'r'rqr(rsH    rz KOSMOS-2 Model for generating text and bounding boxes given an image. The model consists of a vision encoder and a language model. cs.eZdZeZdZdgZedfdd Zej dddZ d d Z ej dd d Z d dZ eeeeeddeejeejeejeejeejeeejeejeejeejeejeeeeeeeeeeefdddZdeejeejeejeejeejdddZZS)rrsztext_model.lm_head.weightr]cs:t|t|j|_t|j|_t||_ | dSr) rdrerrrHrrrIr"rJr(rprqr'r(rePs     z(Kosmos2ForConditionalGeneration.__init__rDcCs |jjjSrrKrMr'r'r(r)[sz4Kosmos2ForConditionalGeneration.get_input_embeddingscCs||jj_dSrrKr3r'r'r(r5^sz4Kosmos2ForConditionalGeneration.set_input_embeddingscCs |jSr)rHr9rMr'r'r(r9asz5Kosmos2ForConditionalGeneration.get_output_embeddingscCs|j|dSr)rHr;r:r'r'r(r;dsz5Kosmos2ForConditionalGeneration.set_output_embeddingsr*N)rsr9r rrr>rArrar<rrrrrEcCs| dk r | n|jj} | dk r | n|jj} |dk r4|n|jj}d}d}|dkr|dkr\td|j|| | |d}|jj|d}tj j |dd}| |\}}|j |||||||| | | | | |d }|s||||f}t dd |DSt|j|j|j|j|j|||d S) a labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the left-to-right language modeling loss (next word prediction). 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]` Returns: Examples: ```python >>> from PIL import Image >>> import requests >>> from transformers import AutoProcessor, Kosmos2ForConditionalGeneration >>> model = Kosmos2ForConditionalGeneration.from_pretrained("microsoft/kosmos-2-patch14-224") >>> processor = AutoProcessor.from_pretrained("microsoft/kosmos-2-patch14-224") >>> url = "https://huggingface.co/microsoft/kosmos-2-patch14-224/resolve/main/snowman.jpg" >>> image = Image.open(requests.get(url, stream=True).raw) >>> prompt = " An image of" >>> inputs = processor(text=prompt, images=image, return_tensors="pt") >>> generated_ids = model.generate( ... pixel_values=inputs["pixel_values"], ... input_ids=inputs["input_ids"], ... attention_mask=inputs["attention_mask"], ... image_embeds=None, ... image_embeds_position_mask=inputs["image_embeds_position_mask"], ... use_cache=True, ... max_new_tokens=64, ... ) >>> generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0] >>> processed_text = processor.post_process_generation(generated_text, cleanup_and_extract=False) >>> processed_text ' An image of a snowman warming himself by a fire.' >>> caption, entities = processor.post_process_generation(generated_text) >>> caption 'An image of a snowman warming himself by a fire.' >>> entities [('a snowman', (12, 21), [(0.390625, 0.046875, 0.984375, 0.828125)]), ('a fire', (41, 47), [(0.171875, 0.015625, 0.484375, 0.890625)])] ```NrLr-rr-r/) r9rrAr rr>rrar<rrrrcss|]}|dk r|VqdSrr'rMr'r'r(rOsz:Kosmos2ForConditionalGeneration.forward..)rZr[r>r?r@rArBrC)r^rrrrrIr'rrrrNrJrHrQrYrZr[r>r?r@)rNrsr9r rrr>rArrar<rrrrrCrBZ lm_outputsrr'r'r(r{gs\Az'Kosmos2ForConditionalGeneration.forward)rsr r9rrAc Ks|dd}|dk r,|dk r,td|d|dkr@|dk r@|}|dkr||}|jj|d}tjj|dd}||\}} |j j f||||d|} | S)Ninputsz `inputs`: zp were passed alongside `pixel_values` which is not allowed.Make sure to either pass `inputs` or pixel_values=...rr-r/)r9rrAr ) poprrIr'rrrrNrJrHgenerate) rNrsr r9rrAkwargsrQrCrBr>r'r'r(rSs*   z(Kosmos2ForConditionalGeneration.generate)NNNNNNNNNNNNNN)NNNNN)rSrTrUrr&r/rCrerr0r)r5r9r;rrOrrYrPrr!r|rrWrDr"rrr{rSr}r'r'rqr(rDsj    xr)N)r)r)GrVrZ dataclassesrtypingrrrrrr!Ztorch.utils.checkpointrZtorch.nnr Z activationsr Zmodeling_outputsr r rrZmodeling_utilsrutilsrrrrrZconfiguration_kosmos2rrrZ get_loggerrSrrPr|rr7r)Sizer+r5r;ZKOSMOS2_START_DOCSTRINGr1r7rOr<rYr0r\r~rrrrrrrrrrrrrr"rrr'r'r'r(s~        L?6:$i3b6Xwoy^$:+#}