kosmos-2-patch14-224 / modeling_kosmos2.py

Add missing import (#18)

bfb1631 8 months ago

No virus

77 kB

	# coding=utf-8
	# Copyright 2023 Microsoft Research and The HuggingFace Inc. team. All rights reserved.
	#
	# Licensed under the Apache License, Version 2.0 (the "License");
	# you may not use this file except in compliance with the License.
	# You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS,
	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	# See the License for the specific language governing permissions and
	# limitations under the License.
	""" PyTorch KOSMOS-2 model."""


	import math
	from dataclasses import dataclass
	from typing import List, Optional, Tuple, Union

	import torch
	import torch.utils.checkpoint
	from torch import nn
	from torch.nn import CrossEntropyLoss

	from transformers.activations import ACT2FN
	from transformers.modeling_outputs import (
	BaseModelOutput,
	BaseModelOutputWithPastAndCrossAttentions,
	BaseModelOutputWithPooling,
	CausalLMOutputWithCrossAttentions,
	)
	from transformers.modeling_utils import PreTrainedModel
	from transformers.utils import (
	ModelOutput,
	add_start_docstrings,
	add_start_docstrings_to_model_forward,
	logging,
	replace_return_docstrings,
	)
	from .configuration_kosmos2 import Kosmos2Config, Kosmos2TextConfig, Kosmos2VisionConfig


	logger = logging.get_logger(__name__)

	_CHECKPOINT_FOR_DOC = "microsoft/kosmos-2-patch14-224"
	_CONFIG_FOR_DOC = Kosmos2Config
	_EXPECTED_OUTPUT_SHAPE = None


	# Copied from transformers.models.bart.modeling_bart._expand_mask
	def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
	"""
	Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
	"""
	bsz, src_len = mask.size()
	tgt_len = tgt_len if tgt_len is not None else src_len

	expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)

	inverted_mask = 1.0 - expanded_mask

	return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)


	# Copied from transformers.models.bart.modeling_bart._make_causal_mask
	def _make_causal_mask(
	input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0
	):
	"""
	Make causal mask used for bi-directional self-attention.
	"""
	bsz, tgt_len = input_ids_shape
	mask = torch.full((tgt_len, tgt_len), torch.tensor(torch.finfo(dtype).min, device=device), device=device)
	mask_cond = torch.arange(mask.size(-1), device=device)
	mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
	mask = mask.to(dtype)

	if past_key_values_length > 0:
	mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
	return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)


	# Copied from transformers.models.roberta.modeling_roberta.create_position_ids_from_input_ids
	def create_position_ids_from_input_ids(input_ids, padding_idx, past_key_values_length=0):
	"""
	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
	"""
	# The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA.
	mask = input_ids.ne(padding_idx).int()
	incremental_indices = (torch.cumsum(mask, dim=1).type_as(mask) + past_key_values_length) * mask
	return incremental_indices.long() + padding_idx


	KOSMOS2_START_DOCSTRING = r"""Kosmos-2"""
	KOSMOS2_VISION_INPUTS_DOCSTRING = r"""Kosmos-2"""
	KOSMOS2_TEXT_INPUTS_DOCSTRING = r"""Kosmos-2"""
	KOSMOS2_INPUTS_DOCSTRING = r"""Kosmos-2"""


	@dataclass
	class Kosmos2ModelOutput(ModelOutput):
	"""
	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_features (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, optional, returned when being computed by the model):
	Sequence of hidden-states at the output of `Kosmos2ImageToTextConnector`.
	image_connector_attention (`tuple(torch.FloatTensor)`, *optional, returned when being computed by the model):
	Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
	sequence_length)`.

	Attentions weights given by `Kosmos2ImageToTextConnector`, after the attention softmax, used to compute the weighted average in the self-attention
	heads.
	vision_model_output(`BaseModelOutputWithPooling`, optional, returned when being computed by the model):
	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.
	"""

	last_hidden_states: torch.FloatTensor = None
	past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
	hidden_states: Optional[Tuple[torch.FloatTensor]] = None
	attentions: Optional[Tuple[torch.FloatTensor]] = None
	image_features: Optional[torch.FloatTensor] = None
	image_connector_attention: Optional[Tuple[torch.FloatTensor]] = None
	vision_model_output: BaseModelOutputWithPooling = None


	@dataclass
	class Kosmos2ForConditionalGenerationModelOutput(ModelOutput):
	"""
	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_features (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, optional, returned when being computed by the model):
	Sequence of hidden-states at the output of `Kosmos2ImageToTextConnector`.
	image_connector_attention (`tuple(torch.FloatTensor)`, *optional, returned when being computed by the model):
	Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
	sequence_length)`.

	Attentions weights given by `Kosmos2ImageToTextConnector`, after the attention softmax, used to compute the weighted average in the self-attention
	heads.
	vision_model_output(`BaseModelOutputWithPooling`, optional, returned when being computed by the model):
	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.
	"""

	loss: Optional[torch.FloatTensor] = None
	logits: torch.FloatTensor = None
	past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
	hidden_states: Optional[Tuple[torch.FloatTensor]] = None
	attentions: Optional[Tuple[torch.FloatTensor]] = None
	image_features: Optional[torch.FloatTensor] = None
	image_connector_attention: Optional[Tuple[torch.FloatTensor]] = None
	vision_model_output: BaseModelOutputWithPooling = None


	# Copied from transformers.models.clip.modeling_clip.CLIPVisionEmbeddings with CLIP->Kosmos2
	class Kosmos2VisionEmbeddings(nn.Module):
	def __init__(self, config: Kosmos2VisionConfig):
	super().__init__()
	self.config = config
	self.embed_dim = config.hidden_size
	self.image_size = config.image_size
	self.patch_size = config.patch_size

	self.class_embedding = nn.Parameter(torch.randn(self.embed_dim))

	self.patch_embedding = nn.Conv2d(
	in_channels=config.num_channels,
	out_channels=self.embed_dim,
	kernel_size=self.patch_size,
	stride=self.patch_size,
	bias=False,
	)

	self.num_patches = (self.image_size // self.patch_size) ** 2
	self.num_positions = self.num_patches + 1
	self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
	self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False)

	def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
	batch_size = pixel_values.shape[0]
	patch_embeds = self.patch_embedding(pixel_values) # shape = [*, width, grid, grid]
	patch_embeds = patch_embeds.flatten(2).transpose(1, 2)

	class_embeds = self.class_embedding.expand(batch_size, 1, -1)
	embeddings = torch.cat([class_embeds, patch_embeds], dim=1)
	embeddings = embeddings + self.position_embedding(self.position_ids)
	return embeddings


	# Copied from transformers.models.clip.modeling_clip.CLIPAttention with CLIP->Kosmos2Vision
	class Kosmos2VisionAttention(nn.Module):
	"""Multi-headed attention from 'Attention Is All You Need' paper"""

	def __init__(self, config):
	super().__init__()
	self.config = config
	self.embed_dim = config.hidden_size
	self.num_heads = config.num_attention_heads
	self.head_dim = self.embed_dim // self.num_heads
	if self.head_dim * self.num_heads != self.embed_dim:
	raise ValueError(
	f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
	f" {self.num_heads})."
	)
	self.scale = self.head_dim**-0.5
	self.dropout = config.attention_dropout

	self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
	self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
	self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
	self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)

	def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
	return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()

	def forward(
	self,
	hidden_states: torch.Tensor,
	attention_mask: Optional[torch.Tensor] = None,
	causal_attention_mask: Optional[torch.Tensor] = None,
	output_attentions: Optional[bool] = False,
	) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
	"""Input shape: Batch x Time x Channel"""

	bsz, tgt_len, embed_dim = hidden_states.size()

	# get query proj
	query_states = self.q_proj(hidden_states) * self.scale
	key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
	value_states = self._shape(self.v_proj(hidden_states), -1, bsz)

	proj_shape = (bsz * self.num_heads, -1, self.head_dim)
	query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape)
	key_states = key_states.view(*proj_shape)
	value_states = value_states.view(*proj_shape)

	src_len = key_states.size(1)
	attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))

	if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):
	raise ValueError(
	f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"
	f" {attn_weights.size()}"
	)

	# apply the causal_attention_mask first
	if causal_attention_mask is not None:
	if causal_attention_mask.size() != (bsz, 1, tgt_len, src_len):
	raise ValueError(
	f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is"
	f" {causal_attention_mask.size()}"
	)
	attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + causal_attention_mask
	attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)

	if attention_mask is not None:
	if attention_mask.size() != (bsz, 1, tgt_len, src_len):
	raise ValueError(
	f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"
	)
	attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask
	attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)

	attn_weights = nn.functional.softmax(attn_weights, dim=-1)

	if output_attentions:
	# this operation is a bit akward, but it's required to
	# make sure that attn_weights keeps its gradient.
	# In order to do so, attn_weights have to reshaped
	# twice and have to be reused in the following
	attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
	attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len)
	else:
	attn_weights_reshaped = None

	attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)

	attn_output = torch.bmm(attn_probs, value_states)

	if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):
	raise ValueError(
	f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is"
	f" {attn_output.size()}"
	)

	attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)
	attn_output = attn_output.transpose(1, 2)
	attn_output = attn_output.reshape(bsz, tgt_len, embed_dim)

	attn_output = self.out_proj(attn_output)

	return attn_output, attn_weights_reshaped


	# Copied from transformers.models.clip.modeling_clip.CLIPMLP with CLIP->Kosmos2Vision
	class Kosmos2VisionMLP(nn.Module):
	def __init__(self, config):
	super().__init__()
	self.config = config
	self.activation_fn = ACT2FN[config.hidden_act]
	self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
	self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)

	def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
	hidden_states = self.fc1(hidden_states)
	hidden_states = self.activation_fn(hidden_states)
	hidden_states = self.fc2(hidden_states)
	return hidden_states


	# Copied from transformers.models.clip.modeling_clip.CLIPEncoderLayer with CLIP->Kosmos2Vision
	class Kosmos2VisionEncoderLayer(nn.Module):
	def __init__(self, config: Kosmos2VisionConfig):
	super().__init__()
	self.embed_dim = config.hidden_size
	self.self_attn = Kosmos2VisionAttention(config)
	self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
	self.mlp = Kosmos2VisionMLP(config)
	self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)

	def forward(
	self,
	hidden_states: torch.Tensor,
	attention_mask: torch.Tensor,
	causal_attention_mask: torch.Tensor,
	output_attentions: Optional[bool] = False,
	) -> Tuple[torch.FloatTensor]:
	"""
	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.
	"""
	residual = hidden_states

	hidden_states = self.layer_norm1(hidden_states)
	hidden_states, attn_weights = self.self_attn(
	hidden_states=hidden_states,
	attention_mask=attention_mask,
	causal_attention_mask=causal_attention_mask,
	output_attentions=output_attentions,
	)
	hidden_states = residual + hidden_states

	residual = hidden_states
	hidden_states = self.layer_norm2(hidden_states)
	hidden_states = self.mlp(hidden_states)
	hidden_states = residual + hidden_states

	outputs = (hidden_states,)

	if output_attentions:
	outputs += (attn_weights,)

	return outputs


	# Copied from transformers.models.clip.modeling_clip.CLIPEncoder with CLIP->Kosmos2Vision
	class Kosmos2VisionEncoder(nn.Module):
	"""
	Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
	[`Kosmos2VisionEncoderLayer`].

	Args:
	config: Kosmos2VisionConfig
	"""

	def __init__(self, config: Kosmos2VisionConfig):
	super().__init__()
	self.config = config
	self.layers = nn.ModuleList([Kosmos2VisionEncoderLayer(config) for _ in range(config.num_hidden_layers)])
	self.gradient_checkpointing = False

	def forward(
	self,
	inputs_embeds,
	attention_mask: Optional[torch.Tensor] = None,
	causal_attention_mask: Optional[torch.Tensor] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	) -> Union[Tuple, BaseModelOutput]:
	r"""
	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.
	"""
	output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
	output_hidden_states = (
	output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
	)
	return_dict = return_dict if return_dict is not None else self.config.use_return_dict

	encoder_states = () if output_hidden_states else None
	all_attentions = () if output_attentions else None

	hidden_states = inputs_embeds
	for idx, encoder_layer in enumerate(self.layers):
	if output_hidden_states:
	encoder_states = encoder_states + (hidden_states,)
	if self.gradient_checkpointing and self.training:

	def create_custom_forward(module):
	def custom_forward(*inputs):
	return module(*inputs, output_attentions)

	return custom_forward

	layer_outputs = torch.utils.checkpoint.checkpoint(
	create_custom_forward(encoder_layer),
	hidden_states,
	attention_mask,
	causal_attention_mask,
	)
	else:
	layer_outputs = encoder_layer(
	hidden_states,
	attention_mask,
	causal_attention_mask,
	output_attentions=output_attentions,
	)

	hidden_states = layer_outputs[0]

	if output_attentions:
	all_attentions = all_attentions + (layer_outputs[1],)

	if output_hidden_states:
	encoder_states = encoder_states + (hidden_states,)

	if not return_dict:
	return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
	return BaseModelOutput(
	last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions
	)


	# Copied from transformers.models.clip.modeling_clip.CLIPVisionTransformer with CLIPVision->Kosmos2Vision,CLIP_VISION->KOSMOS2_VISION,CLIP->Kosmos2Vision
	class Kosmos2VisionTransformer(nn.Module):
	def __init__(self, config: Kosmos2VisionConfig):
	super().__init__()
	self.config = config
	embed_dim = config.hidden_size

	self.embeddings = Kosmos2VisionEmbeddings(config)
	self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
	self.encoder = Kosmos2VisionEncoder(config)
	self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)

	@add_start_docstrings_to_model_forward(KOSMOS2_VISION_INPUTS_DOCSTRING)
	def forward(
	self,
	pixel_values: Optional[torch.FloatTensor] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	) -> Union[Tuple, BaseModelOutputWithPooling]:
	output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
	output_hidden_states = (
	output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
	)
	return_dict = return_dict if return_dict is not None else self.config.use_return_dict

	if pixel_values is None:
	raise ValueError("You have to specify pixel_values")

	hidden_states = self.embeddings(pixel_values)
	hidden_states = self.pre_layrnorm(hidden_states)

	encoder_outputs = self.encoder(
	inputs_embeds=hidden_states,
	output_attentions=output_attentions,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	)

	last_hidden_state = encoder_outputs[0]
	pooled_output = last_hidden_state[:, 0, :]
	pooled_output = self.post_layernorm(pooled_output)

	if not return_dict:
	return (last_hidden_state, pooled_output) + encoder_outputs[1:]

	return BaseModelOutputWithPooling(
	last_hidden_state=last_hidden_state,
	pooler_output=pooled_output,
	hidden_states=encoder_outputs.hidden_states,
	attentions=encoder_outputs.attentions,
	)


	# Copied from transformers.models.m2m_100.modeling_m2m_100.M2M100SinusoidalPositionalEmbedding with M2M100->Kosmos2
	class Kosmos2TextSinusoidalPositionalEmbedding(nn.Module):
	"""This module produces sinusoidal positional embeddings of any length."""

	def __init__(self, num_positions: int, embedding_dim: int, padding_idx: Optional[int] = None):
	super().__init__()
	self.offset = 2
	self.embedding_dim = embedding_dim
	self.padding_idx = padding_idx
	self.make_weights(num_positions + self.offset, embedding_dim, padding_idx)

	def make_weights(self, num_embeddings: int, embedding_dim: int, padding_idx: Optional[int] = None):
	emb_weights = self.get_embedding(num_embeddings, embedding_dim, padding_idx)
	if hasattr(self, "weights"):
	# in forward put the weights on the correct dtype and device of the param
	emb_weights = emb_weights.to(dtype=self.weights.dtype, device=self.weights.device)

	self.register_buffer("weights", emb_weights, persistent=False)

	@staticmethod
	def get_embedding(num_embeddings: int, embedding_dim: int, padding_idx: Optional[int] = None):
	"""
	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".
	"""
	half_dim = embedding_dim // 2
	emb = math.log(10000) / (half_dim - 1)
	emb = torch.exp(torch.arange(half_dim, dtype=torch.float) * -emb)
	emb = torch.arange(num_embeddings, dtype=torch.float).unsqueeze(1) * emb.unsqueeze(0)
	emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(num_embeddings, -1)
	if embedding_dim % 2 == 1:
	# zero pad
	emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1)
	if padding_idx is not None:
	emb[padding_idx, :] = 0

	return emb.to(torch.get_default_dtype())

	@torch.no_grad()
	def forward(
	self, input_ids: torch.Tensor = None, inputs_embeds: torch.Tensor = None, past_key_values_length: int = 0
	):
	if input_ids is not None:
	bsz, seq_len = input_ids.size()
	# Create the position ids from the input token ids. Any padded tokens remain padded.
	position_ids = create_position_ids_from_input_ids(input_ids, self.padding_idx, past_key_values_length).to(
	input_ids.device
	)
	else:
	bsz, seq_len = inputs_embeds.size()[:-1]
	position_ids = self.create_position_ids_from_inputs_embeds(inputs_embeds, past_key_values_length)

	# expand embeddings if needed
	max_pos = self.padding_idx + 1 + seq_len + past_key_values_length
	if max_pos > self.weights.size(0):
	self.make_weights(max_pos + self.offset, self.embedding_dim, self.padding_idx)

	return self.weights.index_select(0, position_ids.view(-1)).view(bsz, seq_len, self.weights.shape[-1]).detach()

	def create_position_ids_from_inputs_embeds(self, inputs_embeds, past_key_values_length):
	"""
	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
	"""
	input_shape = inputs_embeds.size()[:-1]
	sequence_length = input_shape[1]

	position_ids = torch.arange(
	self.padding_idx + 1, sequence_length + self.padding_idx + 1, dtype=torch.long, device=inputs_embeds.device
	)
	return position_ids.unsqueeze(0).expand(input_shape).contiguous() + past_key_values_length


	# Similar to transformers.models.bart.modeling_bart.BartAttention with an additional `inner_attn_ln`.
	class KosmosTextAttention(nn.Module):
	"""Multi-headed attention from 'Attention Is All You Need' paper"""

	def __init__(
	self,
	config,
	embed_dim: int,
	num_heads: int,
	dropout: float = 0.0,
	is_decoder: bool = False,
	add_inner_attn_layernorm: bool = False,
	bias: bool = True,
	):
	super().__init__()
	self.embed_dim = embed_dim
	self.num_heads = num_heads
	self.dropout = dropout
	self.head_dim = embed_dim // num_heads

	if (self.head_dim * num_heads) != self.embed_dim:
	raise ValueError(
	f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
	f" and `num_heads`: {num_heads})."
	)
	self.scaling = self.head_dim**-0.5
	self.is_decoder = is_decoder

	self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
	self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
	self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
	self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)

	self.inner_attn_ln = None
	if add_inner_attn_layernorm:
	self.inner_attn_ln = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)

	def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
	return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()

	def forward(
	self,
	hidden_states: torch.Tensor,
	key_value_states: Optional[torch.Tensor] = None,
	past_key_value: Optional[Tuple[torch.Tensor]] = None,
	attention_mask: Optional[torch.Tensor] = None,
	layer_head_mask: Optional[torch.Tensor] = None,
	output_attentions: bool = False,
	) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
	"""Input shape: Batch x Time x Channel"""

	# if key_value_states are provided this layer is used as a cross-attention layer
	# for the decoder
	is_cross_attention = key_value_states is not None

	bsz, tgt_len, _ = hidden_states.size()

	# get query proj
	query_states = self.q_proj(hidden_states) * self.scaling
	# get key, value proj
	# `past_key_value[0].shape[2] == key_value_states.shape[1]`
	# is checking that the `sequence_length` of the `past_key_value` is the same as
	# the provided `key_value_states` to support prefix tuning
	if (
	is_cross_attention
	and past_key_value is not None
	and past_key_value[0].shape[2] == key_value_states.shape[1]
	):
	# reuse k,v, cross_attentions
	key_states = past_key_value[0]
	value_states = past_key_value[1]
	elif is_cross_attention:
	# cross_attentions
	key_states = self._shape(self.k_proj(key_value_states), -1, bsz)
	value_states = self._shape(self.v_proj(key_value_states), -1, bsz)
	elif past_key_value is not None:
	# reuse k, v, self_attention
	key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
	value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
	key_states = torch.cat([past_key_value[0], key_states], dim=2)
	value_states = torch.cat([past_key_value[1], value_states], dim=2)
	else:
	# self_attention
	key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
	value_states = self._shape(self.v_proj(hidden_states), -1, bsz)

	if self.is_decoder:
	# if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
	# Further calls to cross_attention layer can then reuse all cross-attention
	# key/value_states (first "if" case)
	# if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
	# all previous decoder key/value_states. Further calls to uni-directional self-attention
	# can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
	# if encoder bi-directional self-attention `past_key_value` is always `None`
	past_key_value = (key_states, value_states)

	proj_shape = (bsz * self.num_heads, -1, self.head_dim)
	query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape)
	key_states = key_states.reshape(*proj_shape)
	value_states = value_states.reshape(*proj_shape)

	src_len = key_states.size(1)
	attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))

	if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):
	raise ValueError(
	f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"
	f" {attn_weights.size()}"
	)

	if attention_mask is not None:
	if attention_mask.size() != (bsz, 1, tgt_len, src_len):
	raise ValueError(
	f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"
	)
	attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask
	attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)

	attn_weights = nn.functional.softmax(attn_weights, dim=-1)

	if layer_head_mask is not None:
	if layer_head_mask.size() != (self.num_heads,):
	raise ValueError(
	f"Head mask for a single layer should be of size {(self.num_heads,)}, but is"
	f" {layer_head_mask.size()}"
	)
	attn_weights = layer_head_mask.view(1, -1, 1, 1) * attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
	attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)

	if output_attentions:
	# this operation is a bit awkward, but it's required to
	# make sure that attn_weights keeps its gradient.
	# In order to do so, attn_weights have to be reshaped
	# twice and have to be reused in the following
	attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
	attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len)
	else:
	attn_weights_reshaped = None

	attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)

	attn_output = torch.bmm(attn_probs, value_states)

	if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):
	raise ValueError(
	f"`attn_output` should be of size {(bsz * self.num_heads, tgt_len, self.head_dim)}, but is"
	f" {attn_output.size()}"
	)

	attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)
	attn_output = attn_output.transpose(1, 2)

	# Use the `embed_dim` from the config (stored in the class) rather than `hidden_state` because `attn_output` can be
	# partitioned across GPUs when using tensor-parallelism.
	attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim)

	if self.inner_attn_ln is not None:
	attn_output = self.inner_attn_ln(attn_output)

	attn_output = self.out_proj(attn_output)

	return attn_output, attn_weights_reshaped, past_key_value


	class Kosmos2TextFFN(nn.Module):
	def __init__(self, config: Kosmos2TextConfig):
	super().__init__()

	self.dropout = config.dropout
	self.activation_fn = ACT2FN[config.activation_function]
	self.activation_dropout = config.activation_dropout

	self.fc1 = nn.Linear(config.embed_dim, config.ffn_dim)
	self.fc2 = nn.Linear(config.ffn_dim, config.embed_dim)

	self.ffn_layernorm = nn.LayerNorm(config.ffn_dim, eps=config.layer_norm_eps)

	def forward(self, hidden_states):
	hidden_states = self.activation_fn(self.fc1(hidden_states))
	hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
	hidden_states = self.ffn_layernorm(hidden_states)
	hidden_states = self.fc2(hidden_states)
	hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)

	return hidden_states


	class Kosmos2TextBlock(nn.Module):
	def __init__(self, config: Kosmos2TextConfig):
	super().__init__()
	self.embed_dim = config.embed_dim

	self.self_attn = KosmosTextAttention(
	config,
	embed_dim=self.embed_dim,
	num_heads=config.attention_heads,
	dropout=config.attention_dropout,
	is_decoder=True,
	add_inner_attn_layernorm=True,
	)
	self.dropout = config.dropout
	self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)

	if config.add_cross_attention:
	self.encoder_attn = KosmosTextAttention(
	config,
	embed_dim=self.embed_dim,
	num_heads=config.attention_heads,
	dropout=config.attention_dropout,
	is_decoder=True,
	add_inner_attn_layernorm=False,
	)
	self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)

	self.ffn = Kosmos2TextFFN(config)
	self.final_layer_norm = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)

	def forward(
	self,
	hidden_states: torch.Tensor,
	attention_mask: Optional[torch.Tensor] = None,
	encoder_hidden_states: Optional[torch.Tensor] = None,
	encoder_attention_mask: Optional[torch.Tensor] = None,
	layer_head_mask: Optional[torch.Tensor] = None,
	cross_attn_layer_head_mask: Optional[torch.Tensor] = None,
	past_key_value: Optional[Tuple[torch.Tensor]] = None,
	output_attentions: Optional[bool] = False,
	use_cache: Optional[bool] = True,
	) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
	residual = hidden_states

	# Self Attention
	# decoder uni-directional self-attention cached key/values tuple is at positions 1,2
	self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None

	hidden_states = self.self_attn_layer_norm(hidden_states)

	# add present self-attn cache to positions 1,2 of present_key_value tuple
	hidden_states, self_attn_weights, present_key_value = self.self_attn(
	hidden_states=hidden_states,
	past_key_value=self_attn_past_key_value,
	attention_mask=attention_mask,
	layer_head_mask=layer_head_mask,
	output_attentions=output_attentions,
	)
	hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
	hidden_states = residual + hidden_states

	# Cross-Attention Block
	cross_attn_present_key_value = None
	cross_attn_weights = None
	if encoder_hidden_states is not None:
	if not hasattr(self, "encoder_attn"):
	raise ValueError(
	f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers"
	" by setting `config.add_cross_attention=True`"
	)

	residual = hidden_states

	hidden_states = self.encoder_attn_layer_norm(hidden_states)

	# cross_attn cached key/values tuple is at positions 3,4 of present_key_value tuple
	cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
	hidden_states, cross_attn_weights, cross_attn_present_key_value = self.encoder_attn(
	hidden_states=hidden_states,
	key_value_states=encoder_hidden_states,
	attention_mask=encoder_attention_mask,
	layer_head_mask=cross_attn_layer_head_mask,
	past_key_value=cross_attn_past_key_value,
	output_attentions=output_attentions,
	)
	hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
	hidden_states = residual + hidden_states

	# add cross-attn to positions 3,4 of present_key_value tuple
	present_key_value = present_key_value + cross_attn_present_key_value

	# Fully Connected
	residual = hidden_states

	hidden_states = self.final_layer_norm(hidden_states)

	# FFN
	hidden_states = self.ffn(hidden_states)
	hidden_states = residual + hidden_states

	outputs = (hidden_states,)

	if output_attentions:
	outputs += (self_attn_weights, cross_attn_weights)

	if use_cache:
	outputs += (present_key_value,)

	return outputs


	class Kosmos2TextTransformer(nn.Module):
	"""
	Transformer decoder consisting of `config.layers` layers. Each layer is a [`Kosmos2TextBlock`].

	Args:
	config: Kosmos2TextConfig
	"""

	def __init__(self, config: Kosmos2TextConfig):
	super().__init__()
	self.config = config
	self.dropout = config.dropout
	self.layerdrop = config.layerdrop

	self.embed_scale = math.sqrt(config.embed_dim) if config.scale_embedding else 1.0
	self.embed_tokens = nn.Embedding(config.vocab_size, config.embed_dim, padding_idx=config.pad_token_id)

	self.embed_positions = Kosmos2TextSinusoidalPositionalEmbedding(
	num_positions=config.max_position_embeddings,
	embedding_dim=config.embed_dim,
	padding_idx=config.pad_token_id,
	)

	self.layers = nn.ModuleList([Kosmos2TextBlock(config) for _ in range(config.layers)])
	self.layer_norm = nn.LayerNorm(config.embed_dim, config.layer_norm_eps)

	self.gradient_checkpointing = False

	# Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask
	def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):
	# create causal mask
	# [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
	combined_attention_mask = None
	if input_shape[-1] > 1:
	combined_attention_mask = _make_causal_mask(
	input_shape,
	inputs_embeds.dtype,
	device=inputs_embeds.device,
	past_key_values_length=past_key_values_length,
	)

	if attention_mask is not None:
	# [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
	expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(
	inputs_embeds.device
	)
	combined_attention_mask = (
	expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
	)

	return combined_attention_mask

	def forward_embedding(
	self, input_ids, inputs_embeds=None, img_features=None, img_input_mask=None, past_key_values_length: int = 0
	):
	# The argument `inputs_embeds` should be the one without being multiplied by `self.embed_scale`.
	if inputs_embeds is None:
	inputs_embeds = self.embed_tokens(input_ids)

	if img_features is not None:
	inputs_embeds[img_input_mask.to(dtype=torch.bool)] = img_features.view(-1, img_features.size(-1))

	inputs_embeds = inputs_embeds * self.embed_scale

	# embed positions
	positions = self.embed_positions(
	input_ids=input_ids, inputs_embeds=inputs_embeds, past_key_values_length=past_key_values_length
	)
	positions = positions.to(inputs_embeds.device)

	hidden_states = inputs_embeds + positions

	hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)

	return hidden_states

	def forward(
	self,
	input_ids: Optional[torch.Tensor] = None,
	attention_mask: Optional[torch.Tensor] = None,
	img_features: Optional[torch.Tensor] = None,
	img_attn_mask: Optional[torch.Tensor] = None,
	encoder_hidden_states: Optional[torch.Tensor] = None,
	encoder_attention_mask: Optional[torch.Tensor] = None,
	head_mask: Optional[torch.Tensor] = None,
	cross_attn_head_mask: Optional[torch.Tensor] = None,
	past_key_values: Optional[List[torch.FloatTensor]] = None,
	inputs_embeds: Optional[torch.Tensor] = None,
	use_cache: Optional[bool] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	) -> Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]:
	output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
	output_hidden_states = (
	output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
	)
	use_cache = use_cache if use_cache is not None else self.config.use_cache
	return_dict = return_dict if return_dict is not None else self.config.use_return_dict

	if input_ids is not None and inputs_embeds is not None:
	raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
	elif input_ids is not None:
	input_shape = input_ids.shape
	input_ids = input_ids.view(-1, input_shape[-1])
	elif inputs_embeds is not None:
	input_shape = inputs_embeds.size()[:-1]
	else:
	raise ValueError("You have to specify either input_ids or inputs_embeds")

	# past_key_values_length
	past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0

	# We don't need img info. when `past_key_values_length` > 0
	if past_key_values_length > 0:
	img_features = None
	img_attn_mask = None

	hidden_states = self.forward_embedding(
	input_ids=input_ids,
	inputs_embeds=inputs_embeds,
	img_features=img_features,
	img_input_mask=img_attn_mask,
	past_key_values_length=past_key_values_length,
	)

	attention_mask = self._prepare_decoder_attention_mask(
	attention_mask, input_shape, hidden_states, past_key_values_length
	)

	# expand encoder attention mask
	if encoder_hidden_states is not None and encoder_attention_mask is not None:
	# [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
	encoder_attention_mask = _expand_mask(encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1])

	hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)

	if self.gradient_checkpointing and self.training:
	if use_cache:
	logger.warning_once(
	"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
	)
	use_cache = False

	# decoder layers
	all_hidden_states = () if output_hidden_states else None
	all_self_attns = () if output_attentions else None
	all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None
	next_decoder_cache = () if use_cache else None

	# check if head_mask/cross_attn_head_mask has a correct number of layers specified if desired
	for attn_mask, mask_name in zip([head_mask, cross_attn_head_mask], ["head_mask", "cross_attn_head_mask"]):
	if attn_mask is not None:
	if attn_mask.size()[0] != (len(self.layers)):
	raise ValueError(
	f"The `{mask_name}` should be specified for {len(self.layers)} layers, but it is for"
	f" {head_mask.size()[0]}."
	)

	for idx, decoder_layer in enumerate(self.layers):
	# add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
	if output_hidden_states:
	all_hidden_states += (hidden_states,)
	if self.training:
	dropout_probability = torch.rand([])
	if dropout_probability < self.layerdrop:
	continue

	past_key_value = past_key_values[idx] if past_key_values is not None else None

	if self.gradient_checkpointing and self.training:

	def create_custom_forward(module):
	def custom_forward(*inputs):
	# None for past_key_value
	return module(*inputs, output_attentions, use_cache)

	return custom_forward

	layer_outputs = torch.utils.checkpoint.checkpoint(
	create_custom_forward(decoder_layer),
	hidden_states,
	attention_mask,
	encoder_hidden_states,
	encoder_attention_mask,
	head_mask[idx] if head_mask is not None else None,
	cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None,
	None,
	)
	else:
	layer_outputs = decoder_layer(
	hidden_states,
	attention_mask=attention_mask,
	encoder_hidden_states=encoder_hidden_states,
	encoder_attention_mask=encoder_attention_mask,
	layer_head_mask=(head_mask[idx] if head_mask is not None else None),
	cross_attn_layer_head_mask=(
	cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None
	),
	past_key_value=past_key_value,
	output_attentions=output_attentions,
	use_cache=use_cache,
	)
	hidden_states = layer_outputs[0]

	if use_cache:
	next_decoder_cache += (layer_outputs[3 if output_attentions else 1],)

	if output_attentions:
	all_self_attns += (layer_outputs[1],)

	if encoder_hidden_states is not None:
	all_cross_attentions += (layer_outputs[2],)

	# add final layer norm
	hidden_states = self.layer_norm(hidden_states)

	# add hidden states from the last decoder layer
	if output_hidden_states:
	all_hidden_states += (hidden_states,)

	next_cache = next_decoder_cache if use_cache else None
	if not return_dict:
	return tuple(
	v
	for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_cross_attentions]
	if v is not None
	)
	return BaseModelOutputWithPastAndCrossAttentions(
	last_hidden_state=hidden_states,
	past_key_values=next_cache,
	hidden_states=all_hidden_states,
	attentions=all_self_attns,
	cross_attentions=all_cross_attentions,
	)


	class Kosmos2PreTrainedModel(PreTrainedModel):
	"""
	An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
	models.
	"""

	config_class = Kosmos2Config
	supports_gradient_checkpointing = True


	@add_start_docstrings(
	"""The vision model from KOSMOS-2 without any head or projection on top.""",
	KOSMOS2_START_DOCSTRING,
	)
	class Kosmos2VisionModel(Kosmos2PreTrainedModel):
	config_class = Kosmos2VisionConfig
	main_input_name = "pixel_values"

	# Copied from transformers.models.clip.modeling_clip.CLIPVisionModel.__init__ with CLIP_VISION->KOSMOS2_VISION,CLIP->Kosmos2
	def __init__(self, config: Kosmos2VisionConfig):
	super().__init__(config)
	self.model = Kosmos2VisionTransformer(config)
	# Initialize weights and apply final processing
	self.post_init()

	# Copied from transformers.models.clip.modeling_clip.CLIPVisionModel.get_input_embeddings with CLIP_VISION->KOSMOS2_VISION,CLIP->Kosmos2
	def get_input_embeddings(self) -> nn.Module:
	return self.model.embeddings.patch_embedding

	@add_start_docstrings_to_model_forward(KOSMOS2_VISION_INPUTS_DOCSTRING)
	@replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=Kosmos2VisionConfig)
	def forward(
	self,
	pixel_values: Optional[torch.FloatTensor] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	) -> Union[Tuple, BaseModelOutputWithPooling]:
	r"""
	Returns:

	"""
	return self.model(
	pixel_values=pixel_values,
	output_attentions=output_attentions,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	)


	@add_start_docstrings(
	"""The text model from KOSMOS-2 without any head or projection on top.""",
	KOSMOS2_START_DOCSTRING,
	)
	class Kosmos2TextModel(Kosmos2PreTrainedModel):
	config_class = Kosmos2TextConfig

	_no_split_modules = ["Kosmos2TextBlock"]

	def __init__(self, config: Kosmos2TextConfig):
	super().__init__(config)
	self.model = Kosmos2TextTransformer(config)
	# Initialize weights and apply final processing
	self.post_init()

	def get_input_embeddings(self) -> nn.Module:
	return self.model.embed_tokens

	def set_input_embeddings(self, value):
	self.model.embed_tokens = value

	@add_start_docstrings_to_model_forward(KOSMOS2_TEXT_INPUTS_DOCSTRING)
	@replace_return_docstrings(output_type=BaseModelOutputWithPastAndCrossAttentions, config_class=Kosmos2TextConfig)
	def forward(
	self,
	input_ids: Optional[torch.Tensor] = None,
	attention_mask: Optional[torch.Tensor] = None,
	img_features: Optional[torch.Tensor] = None,
	img_attn_mask: Optional[torch.Tensor] = None,
	encoder_hidden_states: Optional[torch.Tensor] = None,
	encoder_attention_mask: Optional[torch.Tensor] = None,
	head_mask: Optional[torch.Tensor] = None,
	cross_attn_head_mask: Optional[torch.Tensor] = None,
	past_key_values: Optional[List[torch.FloatTensor]] = None,
	inputs_embeds: Optional[torch.Tensor] = None,
	use_cache: Optional[bool] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	) -> Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]:
	r"""
	Returns:

	"""
	return self.model(
	input_ids=input_ids,
	attention_mask=attention_mask,
	img_features=img_features,
	img_attn_mask=img_attn_mask,
	encoder_hidden_states=encoder_hidden_states,
	encoder_attention_mask=encoder_attention_mask,
	head_mask=head_mask,
	cross_attn_head_mask=cross_attn_head_mask,
	past_key_values=past_key_values,
	inputs_embeds=inputs_embeds,
	use_cache=use_cache,
	output_attentions=output_attentions,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	)


	@add_start_docstrings(
	"""
	The text model from KOSMOS-2 with a language modeling head on top (linear layer with weights tied to the input
	embeddings).
	""",
	KOSMOS2_START_DOCSTRING,
	)
	class Kosmos2TextForCausalLM(Kosmos2PreTrainedModel):
	config_class = Kosmos2TextConfig
	_tied_weights_keys = ["lm_head.weight"]

	def __init__(self, config: Kosmos2TextConfig):
	super().__init__(config)

	self.model = Kosmos2TextTransformer(config)
	self.lm_head = nn.Linear(in_features=config.embed_dim, out_features=config.vocab_size, bias=False)

	# Initialize weights and apply final processing
	self.post_init()

	def get_input_embeddings(self) -> nn.Module:
	return self.model.embed_tokens

	def set_input_embeddings(self, value):
	self.model.embed_tokens = value

	def get_output_embeddings(self) -> nn.Module:
	return self.lm_head

	def set_output_embeddings(self, new_embeddings):
	self.lm_head = new_embeddings

	@add_start_docstrings_to_model_forward(KOSMOS2_TEXT_INPUTS_DOCSTRING)
	@replace_return_docstrings(output_type=CausalLMOutputWithCrossAttentions, config_class=Kosmos2TextConfig)
	def forward(
	self,
	input_ids: Optional[torch.Tensor] = None,
	attention_mask: Optional[torch.Tensor] = None,
	img_features: Optional[torch.Tensor] = None,
	img_attn_mask: Optional[torch.Tensor] = None,
	encoder_hidden_states: Optional[torch.Tensor] = None,
	encoder_attention_mask: Optional[torch.Tensor] = None,
	head_mask: Optional[torch.Tensor] = None,
	cross_attn_head_mask: Optional[torch.Tensor] = None,
	past_key_values: Optional[List[torch.FloatTensor]] = None,
	inputs_embeds: Optional[torch.Tensor] = None,
	labels: Optional[torch.LongTensor] = None,
	use_cache: Optional[bool] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:
	r"""
	Returns:

	"""
	return_dict = return_dict if return_dict is not None else self.config.use_return_dict

	if labels is not None:
	if use_cache:
	logger.warning("The `use_cache` argument is changed to `False` since `labels` is provided.")
	use_cache = False

	outputs = self.model(
	input_ids=input_ids,
	attention_mask=attention_mask,
	img_features=img_features,
	img_attn_mask=img_attn_mask,
	encoder_hidden_states=encoder_hidden_states,
	encoder_attention_mask=encoder_attention_mask,
	head_mask=head_mask,
	cross_attn_head_mask=cross_attn_head_mask,
	past_key_values=past_key_values,
	inputs_embeds=inputs_embeds,
	use_cache=use_cache,
	output_attentions=output_attentions,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	)
	logits = self.lm_head(outputs[0])

	loss = None
	if labels is not None:
	# Shift so that tokens < n predict n
	shift_logits = logits[..., :-1, :].contiguous()
	shift_labels = labels[..., 1:].contiguous()
	# Flatten the tokens
	loss_fct = CrossEntropyLoss()
	shift_logits = shift_logits.view(-1, self.config.vocab_size)
	shift_labels = shift_labels.view(-1)
	# Enable model parallelism
	shift_labels = shift_labels.to(shift_logits.device)
	loss = loss_fct(shift_logits, shift_labels)

	if not return_dict:
	output = (logits,) + outputs[1:]
	return (loss,) + output if loss is not None else output

	return CausalLMOutputWithCrossAttentions(
	loss=loss,
	logits=logits,
	past_key_values=outputs.past_key_values,
	hidden_states=outputs.hidden_states,
	attentions=outputs.attentions,
	cross_attentions=outputs.cross_attentions,
	)

	def prepare_inputs_for_generation(
	self,
	input_ids,
	img_features,
	img_attn_mask,
	past_key_values=None,
	attention_mask=None,
	use_cache=None,
	**model_kwargs,
	):
	input_shape = input_ids.shape
	# if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly
	if attention_mask is None:
	attention_mask = input_ids.new_ones(input_shape)

	# cut input_ids if past_key_values is used
	if past_key_values is not None:
	input_ids = input_ids[:, -1:]
	# the image info. is already encoded into the past keys/values
	img_features = None
	img_attn_mask = None
	elif img_attn_mask is not None:
	# appending `False` to `img_attn_mask` (because `input_ids` grows during generation)
	batch_size, seq_len = input_ids.size()
	mask_len = img_attn_mask.size()[-1]
	img_attn_mask = torch.cat(
	(img_attn_mask, torch.zeros(size=(batch_size, seq_len - mask_len), dtype=torch.bool, device=input_ids.device)), dim=1
	)

	return {
	"input_ids": input_ids,
	"img_features": img_features,
	"img_attn_mask": img_attn_mask,
	"past_key_values": past_key_values,
	"attention_mask": attention_mask,
	"use_cache": use_cache,
	}


	class Kosmos2ImageToTextConnector(nn.Module):
	"""The layer that transforms the image model's output to part of the text model's input (namely, image features)"""

	def __init__(self, config: Kosmos2Config):
	super().__init__()
	self.dense = nn.Linear(config.vision_config.hidden_size, config.text_config.embed_dim)
	self.latent_query = nn.Parameter(torch.randn(config.latent_query_num, config.text_config.embed_dim))

	self.x_attn = KosmosTextAttention(
	config.text_config,
	config.text_config.embed_dim,
	config.text_config.attention_heads,
	dropout=config.text_config.attention_dropout,
	is_decoder=False,
	add_inner_attn_layernorm=False,
	)

	def forward(self, features):
	hidden_states = self.dense(features)

	# shape = [batch, latent_query_num, h_dim]
	latent_query = self.latent_query.unsqueeze(0).expand(hidden_states.size(0), -1, -1)
	key_value_states = torch.cat([hidden_states, latent_query], dim=1)

	hidden_states, attn_weights, _ = self.x_attn(
	hidden_states=latent_query,
	key_value_states=key_value_states,
	past_key_value=None,
	attention_mask=None,
	output_attentions=None,
	)

	return hidden_states, attn_weights


	@add_start_docstrings(
	"""
	KOSMOS-2 Model for generating text and image features. The model consists of a vision encoder (CLIP) and a language
	model.
	""",
	KOSMOS2_START_DOCSTRING,
	)
	class Kosmos2Model(Kosmos2PreTrainedModel):
	config_class = Kosmos2Config

	def __init__(self, config: Kosmos2Config):
	super().__init__(config)

	self.text_model = Kosmos2TextModel(config.text_config)
	self.vision_model = Kosmos2VisionModel(config.vision_config)
	self.image_to_text_connector = Kosmos2ImageToTextConnector(config)

	# Initialize weights and apply final processing
	self.post_init()

	def get_input_embeddings(self) -> nn.Module:
	return self.text_model.model.embed_tokens

	def set_input_embeddings(self, value):
	self.text_model.model.embed_tokens = value

	@add_start_docstrings_to_model_forward(KOSMOS2_INPUTS_DOCSTRING)
	@replace_return_docstrings(output_type=Kosmos2ModelOutput, config_class=Kosmos2Config)
	def forward(
	self,
	pixel_values: Optional[torch.Tensor] = None,
	input_ids: Optional[torch.Tensor] = None,
	attention_mask: Optional[torch.Tensor] = None,
	img_attn_mask: Optional[torch.Tensor] = None,
	head_mask: Optional[torch.Tensor] = None,
	img_features: Optional[torch.Tensor] = None,
	past_key_values: Optional[List[torch.FloatTensor]] = None,
	inputs_embeds: Optional[torch.Tensor] = None,
	use_cache: Optional[bool] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	) -> Union[Tuple, Kosmos2ModelOutput]:
	# TODO: Add this
	r"""
	Returns:

	```"""
	output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
	output_hidden_states = (
	output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
	)
	return_dict = return_dict if return_dict is not None else self.config.use_return_dict

	vision_model_output = None
	image_connector_attention = None
	if img_features is None:
	if pixel_values is None:
	raise ValueError("You have to specify either `pixel_values` or `img_features`.")

	vision_model_output = self.vision_model(pixel_values)
	# HF's CLIP has `last_hidden_state` without going through `post_layernorm`.
	# Here we need the whole `last_hidden_state` through `post_layernorm` instead of just `pooled_output`.
	img_features = self.vision_model.model.post_layernorm(vision_model_output.last_hidden_state)
	# normalized features
	img_features = nn.functional.normalize(img_features, dim=-1)
	img_features, image_connector_attention = self.image_to_text_connector(img_features)

	outputs = self.text_model(
	input_ids=input_ids,
	attention_mask=attention_mask,
	img_features=img_features,
	img_attn_mask=img_attn_mask,
	head_mask=head_mask,
	past_key_values=past_key_values,
	inputs_embeds=inputs_embeds,
	use_cache=use_cache,
	output_attentions=output_attentions,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	)

	if not return_dict:
	outputs = outputs + (img_features, image_connector_attention, vision_model_output)
	return tuple(output for output in outputs if output is not None)

	return Kosmos2ModelOutput(
	last_hidden_states=outputs.last_hidden_state,
	past_key_values=outputs.past_key_values,
	hidden_states=outputs.hidden_states,
	attentions=outputs.attentions,
	image_features=img_features,
	image_connector_attention=image_connector_attention,
	vision_model_output=vision_model_output,
	)


	@add_start_docstrings(
	"""
	KOSMOS-2 Model for generating text and bounding boxes given an image. The model consists of a vision encoder (CLIP)
	and a language model.
	""",
	KOSMOS2_START_DOCSTRING,
	)
	class Kosmos2ForConditionalGeneration(Kosmos2PreTrainedModel):
	config_class = Kosmos2Config
	_tied_weights_keys = ["text_model.lm_head.weight"]

	def __init__(self, config: Kosmos2Config):
	super().__init__(config)

	self.text_model = Kosmos2TextForCausalLM(config.text_config)
	self.vision_model = Kosmos2VisionModel(config.vision_config)

	self.image_to_text_connector = Kosmos2ImageToTextConnector(config)

	# Initialize weights and apply final processing
	self.post_init()

	def get_input_embeddings(self) -> nn.Module:
	return self.text_model.model.embed_tokens

	def set_input_embeddings(self, value):
	self.text_model.model.embed_tokens = value

	def get_output_embeddings(self) -> nn.Module:
	return self.text_model.get_output_embeddings()

	def set_output_embeddings(self, new_embeddings):
	self.text_model.set_output_embeddings(new_embeddings)

	@add_start_docstrings_to_model_forward(KOSMOS2_INPUTS_DOCSTRING)
	@replace_return_docstrings(output_type=Kosmos2ForConditionalGenerationModelOutput, config_class=Kosmos2Config)
	def forward(
	self,
	pixel_values: Optional[torch.Tensor] = None,
	img_attn_mask=None,
	input_ids: Optional[torch.Tensor] = None,
	attention_mask=None,
	head_mask: Optional[torch.Tensor] = None,
	img_features: Optional[List[torch.FloatTensor]] = None,
	past_key_values: Optional[List[torch.FloatTensor]] = None,
	inputs_embeds: Optional[torch.Tensor] = None,
	labels: Optional[torch.LongTensor] = None,
	use_cache: Optional[bool] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	) -> Union[Tuple, Kosmos2ForConditionalGenerationModelOutput]:
	r"""
	Returns:

	Examples:

	```python
	>>> from PIL import Image
	>>> from transformers import AutoProcessor, Kosmos2ForConditionalGeneration

	>>> model = Kosmos2ForConditionalGeneration.from_pretrained("ydshieh/kosmos-2-patch14-224")
	>>> processor = AutoProcessor.from_pretrained("ydshieh/kosmos-2-patch14-224")

	>>> prompt = "<grounding> An image of"
	>>> image = Image.open("snowman.jpg")

	>>> inputs = processor(text=prompt, images=image, return_tensors="pt")

	>>> generated_ids = model.generate(
	... pixel_values=inputs["pixel_values"],
	... input_ids=inputs["input_ids"][:, :-1],
	... attention_mask=inputs["attention_mask"][:, :-1],
	... img_features=None,
	... img_attn_mask=inputs["img_attn_mask"][:, :-1],
	... use_cache=True,
	... max_new_tokens=64,
	... )

	>>> generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
	>>> result = processor.post_processor_generation(generated_text)
	>>> result
	<grounding> An image of<phrase> a snowman</phrase><object><patch_index_0044><patch_index_0863></object> warming himself by<phrase> a fire</phrase><object><patch_index_0005><patch_index_0911></object>.
	```"""
	output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
	output_hidden_states = (
	output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
	)
	return_dict = return_dict if return_dict is not None else self.config.use_return_dict

	vision_model_output = None
	image_connector_attention = None
	if img_features is None:
	if pixel_values is None:
	raise ValueError("You have to specify either `pixel_values` or `img_features`.")

	vision_model_output = self.vision_model(pixel_values)
	# HF's CLIP has `last_hidden_state` without going through `post_layernorm`.
	# Here we need the whole `last_hidden_state` through `post_layernorm` instead of just `pooled_output`.
	img_features = self.vision_model.model.post_layernorm(vision_model_output.last_hidden_state)
	# normalized features
	img_features = nn.functional.normalize(img_features, dim=-1)
	img_features, image_connector_attention = self.image_to_text_connector(img_features)

	lm_outputs = self.text_model(
	input_ids=input_ids,
	attention_mask=attention_mask,
	img_features=img_features,
	img_attn_mask=img_attn_mask,
	head_mask=head_mask,
	past_key_values=past_key_values,
	inputs_embeds=inputs_embeds,
	labels=labels,
	use_cache=use_cache,
	output_attentions=output_attentions,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	)

	if not return_dict:
	outputs = lm_outputs + (img_features, image_connector_attention, vision_model_output)
	return tuple(output for output in outputs if output is not None)

	return Kosmos2ForConditionalGenerationModelOutput(
	loss=lm_outputs.loss,
	logits=lm_outputs.logits,
	past_key_values=lm_outputs.past_key_values,
	hidden_states=lm_outputs.hidden_states,
	attentions=lm_outputs.attentions,
	image_features=img_features,
	image_connector_attention=image_connector_attention,
	vision_model_output=vision_model_output,
	)

	def generate(
	self,
	input_ids=None,
	attention_mask=None,
	img_features=None,
	inputs_embeds=None,
	pixel_values=None,
	**kwargs,
	):
	# in order to allow `inputs` argument (as in `GenerationMixin`)
	inputs = kwargs.pop("inputs", None)
	if pixel_values is not None and inputs is not None:
	raise ValueError(
	f"`inputs`: {inputs} were passed alongside `pixel_values` which is not allowed."
	f"Make sure to either pass `inputs` or pixel_values=..."
	)
	if pixel_values is None and inputs is not None:
	pixel_values = inputs

	if img_features is None:
	vision_model_output = self.vision_model(pixel_values)
	# HF's CLIP has `last_hidden_state` without going through `post_layernorm`.
	# Here we need the whole `last_hidden_state` through `post_layernorm` instead of just `pooled_output`.
	img_features = self.vision_model.model.post_layernorm(vision_model_output.last_hidden_state)
	# normalized features
	img_features = nn.functional.normalize(img_features, dim=-1)
	img_features, image_connector_attention = self.image_to_text_connector(img_features)

	output = self.text_model.generate(
	input_ids=input_ids,
	attention_mask=attention_mask,
	img_features=img_features,
	input_embeds=inputs_embeds,
	**kwargs,
	)

	return output