MC-LLaVA-3b / modeling_llava.py

Update modeling file

e9369b2 verified 4 months ago

No virus

70.6 kB

	# coding=utf-8
	import math
	from dataclasses import dataclass
	from typing import List, Optional, Tuple, Union

	import torch
	import torch.nn.functional as F
	import torch.utils.checkpoint
	from torch import nn
	from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
	from transformers import PreTrainedModel, SiglipVisionModel
	from transformers.activations import ACT2FN
	from transformers.cache_utils import Cache, DynamicCache
	from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
	from transformers.modeling_outputs import (
	BaseModelOutputWithPast,
	CausalLMOutputWithPast,
	ModelOutput,
	SequenceClassifierOutputWithPast,
	TokenClassifierOutput,
	)
	from transformers.utils import (
	is_flash_attn_2_available,
	is_flash_attn_greater_or_equal_2_10,
	logging,
	)

	try:
	from flash_attn import flash_attn_func, flash_attn_varlen_func
	from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
	except Exception as exp:
	print(exp)


	from transformers.configuration_utils import PretrainedConfig
	from transformers import SiglipVisionConfig


	logger = logging.get_logger(__name__)


	class PhiConfig(PretrainedConfig):
	model_type = "phi"
	keys_to_ignore_at_inference = ["past_key_values"]

	def __init__(
	self,
	vocab_size=51200,
	hidden_size=2048,
	intermediate_size=8192,
	num_hidden_layers=24,
	num_attention_heads=32,
	num_key_value_heads=None,
	resid_pdrop=0.0,
	embd_pdrop=0.0,
	attention_dropout=0.0,
	hidden_act="gelu_new",
	max_position_embeddings=2048,
	initializer_range=0.02,
	layer_norm_eps=1e-5,
	use_cache=True,
	tie_word_embeddings=False,
	rope_theta=10000.0,
	rope_scaling=None,
	partial_rotary_factor=0.5,
	qk_layernorm=False,
	bos_token_id=1,
	eos_token_id=2,
	**kwargs,
	):
	self.vocab_size = vocab_size
	self.hidden_size = hidden_size
	self.intermediate_size = intermediate_size
	self.num_hidden_layers = num_hidden_layers
	self.num_attention_heads = num_attention_heads

	if num_key_value_heads is None:
	num_key_value_heads = num_attention_heads

	self.num_key_value_heads = num_key_value_heads
	self.resid_pdrop = resid_pdrop
	self.embd_pdrop = embd_pdrop
	self.attention_dropout = attention_dropout
	self.hidden_act = hidden_act
	self.max_position_embeddings = max_position_embeddings
	self.initializer_range = initializer_range
	self.layer_norm_eps = layer_norm_eps
	self.use_cache = use_cache
	self.rope_theta = rope_theta
	self.rope_scaling = rope_scaling
	self.partial_rotary_factor = partial_rotary_factor
	self.qk_layernorm = qk_layernorm
	self._rope_scaling_validation()

	super().__init__(
	bos_token_id=bos_token_id,
	eos_token_id=eos_token_id,
	tie_word_embeddings=tie_word_embeddings,
	**kwargs,
	)

	def _rope_scaling_validation(self):
	"""
	Validate the `rope_scaling` configuration.
	"""
	if self.rope_scaling is None:
	return

	if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2:
	raise ValueError(
	"`rope_scaling` must be a dictionary with with two fields, `type` and `factor`, "
	f"got {self.rope_scaling}"
	)
	rope_scaling_type = self.rope_scaling.get("type", None)
	rope_scaling_factor = self.rope_scaling.get("factor", None)
	if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
	raise ValueError(
	f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"
	)
	if (
	rope_scaling_factor is None
	or not isinstance(rope_scaling_factor, float)
	or rope_scaling_factor <= 1.0
	):
	raise ValueError(
	f"`rope_scaling`'s factor field must be a float > 1, got {rope_scaling_factor}"
	)


	class LlavaConfig(PretrainedConfig):
	model_type = "mc-llava"
	is_composition = False

	def __init__(
	self,
	text_config=None,
	vision_config=None,
	ignore_index=-100,
	image_token_index=50297,
	projector_hidden_act="gelu",
	projector_tokens_num=1,
	vocab_size=51200,
	**kwargs,
	):
	self.ignore_index = ignore_index
	self.image_token_index = image_token_index
	self.projector_hidden_act = projector_hidden_act
	self.projector_tokens_num = projector_tokens_num
	self.vocab_size = vocab_size

	self.text_config = text_config
	if isinstance(self.text_config, dict):
	text_config["model_type"] = (
	text_config["model_type"] if "model_type" in text_config else "phi"
	)
	self.text_config = PhiConfig(**text_config)
	self.vocab_size = self.text_config.vocab_size

	self.vision_config = vision_config
	if isinstance(self.vision_config, dict):
	self.vision_config = SiglipVisionConfig(**vision_config)
	self.vision_embed_dim = self.vision_config.hidden_size

	super().__init__(**kwargs)


	# Copied from transformers.models.llama.modeling_llama._get_unpad_data
	def _get_unpad_data(attention_mask):
	seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
	indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
	max_seqlen_in_batch = seqlens_in_batch.max().item()
	cu_seqlens = F.pad(
	torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0)
	)
	return (
	indices,
	cu_seqlens,
	max_seqlen_in_batch,
	)


	# Copied from transformers.models.llama.modeling_llama.LlamaRotaryEmbedding with Llama->Phi
	class PhiRotaryEmbedding(nn.Module):
	def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
	super().__init__()

	self.dim = dim
	self.max_position_embeddings = max_position_embeddings
	self.base = base
	inv_freq = 1.0 / (
	self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim)
	)
	self.register_buffer("inv_freq", inv_freq, persistent=False)

	# Build here to make `torch.jit.trace` work.
	self._set_cos_sin_cache(
	seq_len=max_position_embeddings,
	device=self.inv_freq.device,
	dtype=torch.get_default_dtype(),
	)

	def _set_cos_sin_cache(self, seq_len, device, dtype):
	self.max_seq_len_cached = seq_len
	t = torch.arange(
	self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype
	)

	freqs = torch.outer(t, self.inv_freq)
	# Different from paper, but it uses a different permutation in order to obtain the same calculation
	emb = torch.cat((freqs, freqs), dim=-1)
	self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
	self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)

	def forward(self, x, seq_len=None):
	# x: [bs, num_attention_heads, seq_len, head_size]
	if seq_len > self.max_seq_len_cached:
	self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)

	return (
	self.cos_cached[:seq_len].to(dtype=x.dtype),
	self.sin_cached[:seq_len].to(dtype=x.dtype),
	)


	# Copied from transformers.models.llama.modeling_llama.LlamaLinearScalingRotaryEmbedding with Llama->Phi
	class PhiLinearScalingRotaryEmbedding(PhiRotaryEmbedding):
	"""PhiRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""

	def __init__(
	self,
	dim,
	max_position_embeddings=2048,
	base=10000,
	device=None,
	scaling_factor=1.0,
	):
	self.scaling_factor = scaling_factor
	super().__init__(dim, max_position_embeddings, base, device)

	def _set_cos_sin_cache(self, seq_len, device, dtype):
	self.max_seq_len_cached = seq_len
	t = torch.arange(
	self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype
	)
	t = t / self.scaling_factor

	freqs = torch.outer(t, self.inv_freq)
	# Different from paper, but it uses a different permutation in order to obtain the same calculation
	emb = torch.cat((freqs, freqs), dim=-1)
	self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
	self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)


	# Copied from transformers.models.llama.modeling_llama.LlamaDynamicNTKScalingRotaryEmbedding with Llama->Phi
	class PhiDynamicNTKScalingRotaryEmbedding(PhiRotaryEmbedding):
	"""PhiRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""

	def __init__(
	self,
	dim,
	max_position_embeddings=2048,
	base=10000,
	device=None,
	scaling_factor=1.0,
	):
	self.scaling_factor = scaling_factor
	super().__init__(dim, max_position_embeddings, base, device)

	def _set_cos_sin_cache(self, seq_len, device, dtype):
	self.max_seq_len_cached = seq_len

	if seq_len > self.max_position_embeddings:
	base = self.base * (
	(self.scaling_factor * seq_len / self.max_position_embeddings)
	- (self.scaling_factor - 1)
	) ** (self.dim / (self.dim - 2))
	inv_freq = 1.0 / (
	base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim)
	)
	self.register_buffer("inv_freq", inv_freq, persistent=False)

	t = torch.arange(
	self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype
	)

	freqs = torch.outer(t, self.inv_freq)
	# Different from paper, but it uses a different permutation in order to obtain the same calculation
	emb = torch.cat((freqs, freqs), dim=-1)
	self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
	self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)


	# Copied from transformers.models.llama.modeling_llama.rotate_half
	def rotate_half(x):
	"""Rotates half the hidden dims of the input."""
	x1 = x[..., : x.shape[-1] // 2]
	x2 = x[..., x.shape[-1] // 2 :]
	return torch.cat((-x2, x1), dim=-1)


	# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
	def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
	cos = cos[position_ids].unsqueeze(unsqueeze_dim)
	sin = sin[position_ids].unsqueeze(unsqueeze_dim)
	q_embed = (q * cos) + (rotate_half(q) * sin)
	k_embed = (k * cos) + (rotate_half(k) * sin)
	return q_embed, k_embed


	# Copied from transformers.models.clip.modeling_clip.CLIPMLP with CLIP->Phi
	class PhiMLP(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.llama.modeling_llama.repeat_kv with llama->phi
	def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
	"""
	This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
	num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
	"""
	batch, num_key_value_heads, slen, head_dim = hidden_states.shape
	if n_rep == 1:
	return hidden_states
	hidden_states = hidden_states[:, :, None, :, :].expand(
	batch, num_key_value_heads, n_rep, slen, head_dim
	)
	return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)


	class PhiAttention(nn.Module):
	"""Multi-headed attention from 'Attention Is All You Need' paper"""

	def __init__(self, config: PhiConfig, layer_idx: Optional[int] = None):
	super().__init__()
	self.config = config
	self.layer_idx = layer_idx
	if layer_idx is None:
	logger.warning_once(
	f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
	"to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
	"when creating this class."
	)

	self.attention_dropout = config.attention_dropout
	self.hidden_size = config.hidden_size
	self.num_heads = config.num_attention_heads
	self.head_dim = self.hidden_size // self.num_heads
	self.num_key_value_heads = config.num_key_value_heads
	self.num_key_value_groups = self.num_heads // self.num_key_value_heads
	self.max_position_embeddings = config.max_position_embeddings
	self.rope_theta = config.rope_theta
	self.partial_rotary_factor = config.partial_rotary_factor
	self.is_causal = True

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

	self.q_proj = nn.Linear(
	self.hidden_size, self.num_heads * self.head_dim, bias=True
	)
	self.k_proj = nn.Linear(
	self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True
	)
	self.v_proj = nn.Linear(
	self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True
	)
	self.dense = nn.Linear(
	self.num_heads * self.head_dim, self.hidden_size, bias=True
	)

	self.qk_layernorm = config.qk_layernorm
	if self.qk_layernorm:
	self.q_layernorm = nn.LayerNorm(
	config.hidden_size // self.num_heads,
	eps=config.layer_norm_eps,
	elementwise_affine=True,
	)
	self.k_layernorm = nn.LayerNorm(
	config.hidden_size // self.num_heads,
	eps=config.layer_norm_eps,
	elementwise_affine=True,
	)

	self._init_rope()

	def _init_rope(self):
	if self.config.rope_scaling is None:
	self.rotary_emb = PhiRotaryEmbedding(
	int(self.partial_rotary_factor * self.head_dim),
	max_position_embeddings=self.max_position_embeddings,
	base=self.rope_theta,
	)
	else:
	scaling_type = self.config.rope_scaling["type"]
	scaling_factor = self.config.rope_scaling["factor"]
	if scaling_type == "linear":
	self.rotary_emb = PhiLinearScalingRotaryEmbedding(
	int(self.partial_rotary_factor * self.head_dim),
	max_position_embeddings=self.max_position_embeddings,
	scaling_factor=scaling_factor,
	base=self.rope_theta,
	)
	elif scaling_type == "dynamic":
	self.rotary_emb = PhiDynamicNTKScalingRotaryEmbedding(
	int(self.partial_rotary_factor * self.head_dim),
	max_position_embeddings=self.max_position_embeddings,
	scaling_factor=scaling_factor,
	base=self.rope_theta,
	)
	else:
	raise ValueError(f"Unknown RoPE scaling type {scaling_type}")

	# Phi-2 has an attention overflow issue (with FP16) and requires autocast to be disabled
	@torch.autocast("cpu", enabled=False)
	@torch.autocast("cuda", enabled=False)
	def forward(
	self,
	hidden_states: torch.Tensor,
	attention_mask: Optional[torch.Tensor] = None,
	position_ids: Optional[torch.LongTensor] = None,
	past_key_value: Optional[Cache] = None,
	output_attentions: bool = False,
	use_cache: bool = False,
	) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
	bsz, q_len, _ = hidden_states.size()

	query_states = self.q_proj(hidden_states)
	key_states = self.k_proj(hidden_states)
	value_states = self.v_proj(hidden_states)

	if self.qk_layernorm:
	query_states = self.q_layernorm(query_states)
	key_states = self.k_layernorm(key_states)

	query_states = query_states.view(
	bsz, q_len, self.num_heads, self.head_dim
	).transpose(1, 2)
	key_states = key_states.view(
	bsz, q_len, self.num_key_value_heads, self.head_dim
	).transpose(1, 2)
	value_states = value_states.view(
	bsz, q_len, self.num_key_value_heads, self.head_dim
	).transpose(1, 2)

	kv_seq_len = key_states.shape[-2]
	if past_key_value is not None:
	if self.layer_idx is None:
	raise ValueError(
	f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
	"for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
	"with a layer index."
	)
	kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
	cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)

	# Partial rotary embedding
	query_rot, query_pass = (
	query_states[..., : self.rotary_emb.dim],
	query_states[..., self.rotary_emb.dim :],
	)
	key_rot, key_pass = (
	key_states[..., : self.rotary_emb.dim],
	key_states[..., self.rotary_emb.dim :],
	)
	# [batch_size, seq_length, num_heads, head_dim // config.partial_rotary_factor]
	query_rot, key_rot = apply_rotary_pos_emb(
	query_rot, key_rot, cos, sin, position_ids
	)

	# [batch_size, seq_length, num_heads, head_dim]
	query_states = torch.cat((query_rot, query_pass), dim=-1)
	key_states = torch.cat((key_rot, key_pass), dim=-1)

	if past_key_value is not None:
	cache_kwargs = {
	"sin": sin,
	"cos": cos,
	"partial_rotation_size": self.rotary_emb.dim,
	}
	key_states, value_states = past_key_value.update(
	key_states, value_states, self.layer_idx, cache_kwargs
	)

	key_states = repeat_kv(key_states, self.num_key_value_groups)
	value_states = repeat_kv(value_states, self.num_key_value_groups)

	# Queries and keys upcast to fp32 is required by Phi-2 to avoid overflow
	attn_weights = torch.matmul(
	query_states.to(torch.float32), key_states.to(torch.float32).transpose(2, 3)
	) / math.sqrt(self.head_dim)

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

	if attention_mask is not None:
	if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
	raise ValueError(
	f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
	)
	attn_weights = attn_weights + attention_mask

	# upcast attention to fp32
	attn_weights = nn.functional.softmax(
	attn_weights, dim=-1, dtype=torch.float32
	).to(value_states.dtype)
	attn_weights = nn.functional.dropout(
	attn_weights, p=self.attention_dropout, training=self.training
	)

	attn_output = torch.matmul(attn_weights, value_states)

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

	attn_output = attn_output.transpose(1, 2).contiguous()
	attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)

	attn_output = self.dense(attn_output)

	if not output_attentions:
	attn_weights = None

	return attn_output, attn_weights, past_key_value


	class PhiFlashAttention2(PhiAttention):
	# Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__
	def __init__(self, args, *kwargs):
	super().__init__(args, *kwargs)

	# TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
	# flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
	# Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
	self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()

	def forward(
	self,
	hidden_states: torch.Tensor,
	attention_mask: Optional[torch.LongTensor] = None,
	position_ids: Optional[torch.LongTensor] = None,
	past_key_value: Optional[Cache] = None,
	output_attentions: bool = False,
	use_cache: bool = False,
	**kwargs,
	) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
	# PhiFlashAttention2 attention does not support output_attentions

	output_attentions = False

	bsz, q_len, _ = hidden_states.size()

	query_states = self.q_proj(hidden_states)
	key_states = self.k_proj(hidden_states)
	value_states = self.v_proj(hidden_states)

	if self.qk_layernorm:
	query_states = self.q_layernorm(query_states)
	key_states = self.k_layernorm(key_states)

	# Flash attention requires the input to have the shape
	# batch_size x seq_length x head_dim x hidden_dim
	# therefore we just need to keep the original shape
	query_states = query_states.view(
	bsz, q_len, self.num_heads, self.head_dim
	).transpose(1, 2)
	key_states = key_states.view(
	bsz, q_len, self.num_key_value_heads, self.head_dim
	).transpose(1, 2)
	value_states = value_states.view(
	bsz, q_len, self.num_key_value_heads, self.head_dim
	).transpose(1, 2)

	kv_seq_len = key_states.shape[-2]
	if past_key_value is not None:
	kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
	cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)

	# Partial rotary embedding
	query_rot, query_pass = (
	query_states[..., : self.rotary_emb.dim],
	query_states[..., self.rotary_emb.dim :],
	)
	key_rot, key_pass = (
	key_states[..., : self.rotary_emb.dim],
	key_states[..., self.rotary_emb.dim :],
	)
	# [batch_size, seq_length, num_heads, head_dim // config.partial_rotary_factor]
	query_rot, key_rot = apply_rotary_pos_emb(
	query_rot, key_rot, cos, sin, position_ids
	)

	# [batch_size, seq_length, num_heads, head_dim]
	query_states = torch.cat((query_rot, query_pass), dim=-1)
	key_states = torch.cat((key_rot, key_pass), dim=-1)

	if past_key_value is not None:
	cache_kwargs = {
	"sin": sin,
	"cos": cos,
	"partial_rotation_size": self.rotary_emb.dim,
	}
	key_states, value_states = past_key_value.update(
	key_states, value_states, self.layer_idx, cache_kwargs
	)

	# TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
	# to be able to avoid many of these transpose/reshape/view.
	query_states = query_states.transpose(1, 2)
	key_states = key_states.transpose(1, 2)
	value_states = value_states.transpose(1, 2)

	attn_dropout = self.attention_dropout if self.training else 0.0

	# In PEFT, usually we cast the layer norms in float32 for training stability reasons
	# therefore the input hidden states gets silently casted in float32. Hence, we need
	# cast them back in the correct dtype just to be sure everything works as expected.
	# This might slowdown training & inference so it is recommended to not cast the LayerNorms
	# in fp32.

	if query_states.dtype == torch.float32:
	if torch.is_autocast_enabled():
	target_dtype = torch.get_autocast_gpu_dtype()
	# Handle the case where the model is quantized
	elif hasattr(self.config, "_pre_quantization_dtype"):
	target_dtype = self.config._pre_quantization_dtype
	else:
	target_dtype = self.q_proj.weight.dtype

	logger.warning_once(
	f"The input hidden states seems to be silently casted in float32, this might be related to"
	f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
	f" {target_dtype}."
	)

	query_states = query_states.to(target_dtype)
	key_states = key_states.to(target_dtype)
	value_states = value_states.to(target_dtype)

	attn_output = self._flash_attention_forward(
	query_states,
	key_states,
	value_states,
	attention_mask,
	q_len,
	dropout=attn_dropout,
	softmax_scale=None,
	)

	attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
	attn_output = self.dense(attn_output)

	if not output_attentions:
	attn_weights = None

	return attn_output, attn_weights, past_key_value

	# Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._flash_attention_forward
	def _flash_attention_forward(
	self,
	query_states,
	key_states,
	value_states,
	attention_mask,
	query_length,
	dropout=0.0,
	softmax_scale=None,
	):
	if not self._flash_attn_uses_top_left_mask:
	causal = self.is_causal
	else:
	# TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.
	causal = self.is_causal and query_length != 1

	# Contains at least one padding token in the sequence
	if attention_mask is not None:
	batch_size = query_states.shape[0]
	(
	query_states,
	key_states,
	value_states,
	indices_q,
	cu_seq_lens,
	max_seq_lens,
	) = self._upad_input(
	query_states, key_states, value_states, attention_mask, query_length
	)

	cu_seqlens_q, cu_seqlens_k = cu_seq_lens
	max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens

	attn_output_unpad = flash_attn_varlen_func(
	query_states,
	key_states,
	value_states,
	cu_seqlens_q=cu_seqlens_q,
	cu_seqlens_k=cu_seqlens_k,
	max_seqlen_q=max_seqlen_in_batch_q,
	max_seqlen_k=max_seqlen_in_batch_k,
	dropout_p=dropout,
	softmax_scale=softmax_scale,
	causal=causal,
	)

	attn_output = pad_input(
	attn_output_unpad, indices_q, batch_size, query_length
	)
	else:
	attn_output = flash_attn_func(
	query_states,
	key_states,
	value_states,
	dropout,
	softmax_scale=softmax_scale,
	causal=causal,
	)

	return attn_output

	# Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._upad_input
	def _upad_input(
	self, query_layer, key_layer, value_layer, attention_mask, query_length
	):
	indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
	batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape

	key_layer = index_first_axis(
	key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim),
	indices_k,
	)
	value_layer = index_first_axis(
	value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim),
	indices_k,
	)
	if query_length == kv_seq_len:
	query_layer = index_first_axis(
	query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim),
	indices_k,
	)
	cu_seqlens_q = cu_seqlens_k
	max_seqlen_in_batch_q = max_seqlen_in_batch_k
	indices_q = indices_k
	elif query_length == 1:
	max_seqlen_in_batch_q = 1
	cu_seqlens_q = torch.arange(
	batch_size + 1, dtype=torch.int32, device=query_layer.device
	) # There is a memcpy here, that is very bad.
	indices_q = cu_seqlens_q[:-1]
	query_layer = query_layer.squeeze(1)
	else:
	# The -q_len: slice assumes left padding.
	attention_mask = attention_mask[:, -query_length:]
	query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(
	query_layer, attention_mask
	)

	return (
	query_layer,
	key_layer,
	value_layer,
	indices_q,
	(cu_seqlens_q, cu_seqlens_k),
	(max_seqlen_in_batch_q, max_seqlen_in_batch_k),
	)


	PHI_ATTENTION_CLASSES = {
	"flash_attention_2": PhiFlashAttention2,
	"eager": PhiAttention,
	}


	class PhiDecoderLayer(nn.Module):
	def __init__(self, config: PhiConfig, layer_idx: int):
	super().__init__()
	if is_flash_attn_2_available():
	config._attn_implementation = "flash_attention_2"
	self.self_attn = PHI_ATTENTION_CLASSES[config._attn_implementation](
	config, layer_idx=layer_idx
	)
	self.mlp = PhiMLP(config)
	self.input_layernorm = nn.LayerNorm(
	config.hidden_size, eps=config.layer_norm_eps
	)
	self.resid_dropout = nn.Dropout(config.resid_pdrop)

	def forward(
	self,
	hidden_states: torch.Tensor,
	attention_mask: Optional[torch.Tensor] = None,
	position_ids: Optional[torch.LongTensor] = None,
	output_attentions: Optional[bool] = False,
	use_cache: Optional[bool] = False,
	past_key_value: Optional[Tuple[torch.Tensor]] = None,
	) -> Tuple[
	torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
	]:
	residual = hidden_states

	hidden_states = self.input_layernorm(hidden_states)

	# Self Attention
	attn_outputs, self_attn_weights, present_key_value = self.self_attn(
	hidden_states=hidden_states,
	attention_mask=attention_mask,
	position_ids=position_ids,
	past_key_value=past_key_value,
	output_attentions=output_attentions,
	use_cache=use_cache,
	)
	attn_outputs = self.resid_dropout(attn_outputs)

	feed_forward_hidden_states = self.resid_dropout(self.mlp(hidden_states))
	hidden_states = attn_outputs + feed_forward_hidden_states + residual
	outputs = (hidden_states,)

	if output_attentions:
	outputs += (self_attn_weights,)

	if use_cache:
	outputs += (present_key_value,)

	return outputs


	class PhiPreTrainedModel(PreTrainedModel):
	config_class = PhiConfig
	base_model_prefix = "model"
	supports_gradient_checkpointing = True
	_no_split_modules = ["PhiDecoderLayer"]
	_skip_keys_device_placement = "past_key_values"
	_supports_flash_attn_2 = True
	_supports_cache_class = True


	class PhiModel(PhiPreTrainedModel):
	"""
	Transformer decoder consisting of config.num_hidden_layers layers. Each layer is a [`PhiDecoderLayer`]

	Args:
	config: PhiConfig
	"""

	def __init__(self, config: PhiConfig):
	super().__init__(config)
	self.padding_idx = config.pad_token_id
	self.vocab_size = config.vocab_size

	self.embed_tokens = nn.Embedding(
	config.vocab_size, config.hidden_size, self.padding_idx
	)
	self.embed_dropout = nn.Dropout(config.embd_pdrop)
	self.layers = nn.ModuleList(
	[
	PhiDecoderLayer(config, layer_idx)
	for layer_idx in range(config.num_hidden_layers)
	]
	)
	self.final_layernorm = nn.LayerNorm(
	config.hidden_size, eps=config.layer_norm_eps
	)
	self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"

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

	def get_input_embeddings(self):
	return self.embed_tokens

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

	def forward(
	self,
	input_ids: torch.LongTensor = None,
	attention_mask: Optional[torch.Tensor] = None,
	position_ids: Optional[torch.LongTensor] = None,
	past_key_values: Optional[List[torch.FloatTensor]] = None,
	inputs_embeds: Optional[torch.FloatTensor] = None,
	use_cache: Optional[bool] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	) -> Union[Tuple, BaseModelOutputWithPast]:
	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
	)

	# retrieve input_ids and inputs_embeds
	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:
	batch_size, seq_length = input_ids.shape[:2]
	elif inputs_embeds is not None:
	batch_size, seq_length = inputs_embeds.shape[:2]
	else:
	raise ValueError("You have to specify either input_ids or inputs_embeds")

	past_key_values_length = 0

	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

	if use_cache:
	use_legacy_cache = not isinstance(past_key_values, Cache)
	if use_legacy_cache:
	past_key_values = DynamicCache.from_legacy_cache(past_key_values)
	past_key_values_length = past_key_values.get_usable_length(seq_length)

	if position_ids is None:
	device = input_ids.device if input_ids is not None else inputs_embeds.device
	position_ids = torch.arange(
	past_key_values_length,
	seq_length + past_key_values_length,
	dtype=torch.long,
	device=device,
	)
	position_ids = position_ids.unsqueeze(0)

	if inputs_embeds is None:
	inputs_embeds = self.embed_tokens(input_ids)

	inputs_embeds = self.embed_dropout(inputs_embeds)

	# Attention mask.
	if self._use_flash_attention_2:
	# 2d mask is passed through the layers
	attention_mask = (
	attention_mask
	if (attention_mask is not None and 0 in attention_mask)
	else None
	)
	else:
	# 4d mask is passed through the layers
	attention_mask = _prepare_4d_causal_attention_mask(
	attention_mask,
	(batch_size, seq_length),
	inputs_embeds,
	past_key_values_length,
	)

	hidden_states = inputs_embeds

	# decoder layers
	all_hidden_states = () if output_hidden_states else None
	all_self_attns = () if output_attentions else None
	next_decoder_cache = None

	for decoder_layer in self.layers:
	if output_hidden_states:
	all_hidden_states += (hidden_states,)

	if self.gradient_checkpointing and self.training:
	layer_outputs = self._gradient_checkpointing_func(
	decoder_layer.__call__,
	hidden_states,
	attention_mask,
	position_ids,
	past_key_values,
	output_attentions,
	)
	else:
	layer_outputs = decoder_layer(
	hidden_states,
	attention_mask=attention_mask,
	position_ids=position_ids,
	past_key_value=past_key_values,
	output_attentions=output_attentions,
	use_cache=use_cache,
	)

	hidden_states = layer_outputs[0]

	if use_cache:
	next_decoder_cache = layer_outputs[2 if output_attentions else 1]

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

	hidden_states = self.final_layernorm(hidden_states)

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

	next_cache = None
	if use_cache:
	next_cache = (
	next_decoder_cache.to_legacy_cache()
	if use_legacy_cache
	else next_decoder_cache
	)
	if not return_dict:
	return tuple(
	v
	for v in [hidden_states, next_cache, all_hidden_states, all_self_attns]
	if v is not None
	)
	return BaseModelOutputWithPast(
	last_hidden_state=hidden_states,
	past_key_values=next_cache,
	hidden_states=all_hidden_states,
	attentions=all_self_attns,
	)


	class PhiForCausalLM(PhiPreTrainedModel):
	_tied_weights_keys = ["lm_head.weight"]

	# Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.__init__ with Llama->Phi,bias=False->bias=True
	def __init__(self, config):
	super().__init__(config)
	self.model = PhiModel(config)
	self.vocab_size = config.vocab_size
	self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=True)

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

	# Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.get_input_embeddings
	def get_input_embeddings(self):
	return self.model.embed_tokens

	# Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.set_input_embeddings
	def set_input_embeddings(self, value):
	self.model.embed_tokens = value

	# Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.get_output_embeddings
	def get_output_embeddings(self):
	return self.lm_head

	# Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.set_output_embeddings
	def set_output_embeddings(self, new_embeddings):
	self.lm_head = new_embeddings

	# Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.set_decoder
	def set_decoder(self, decoder):
	self.model = decoder

	# Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.get_decoder
	def get_decoder(self):
	return self.model

	def forward(
	self,
	input_ids: torch.LongTensor = None,
	attention_mask: Optional[torch.Tensor] = None,
	position_ids: Optional[torch.LongTensor] = None,
	past_key_values: Optional[List[torch.FloatTensor]] = None,
	inputs_embeds: Optional[torch.FloatTensor] = 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, CausalLMOutputWithPast]:
	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
	)

	# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
	outputs = self.model(
	input_ids=input_ids,
	attention_mask=attention_mask,
	position_ids=position_ids,
	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,
	)

	hidden_states = outputs[0]
	logits = self.lm_head(hidden_states)
	logits = logits.float()

	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 CausalLMOutputWithPast(
	loss=loss,
	logits=logits,
	past_key_values=outputs.past_key_values,
	hidden_states=outputs.hidden_states,
	attentions=outputs.attentions,
	)

	@staticmethod
	# Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM._reorder_cache
	def _reorder_cache(past_key_values, beam_idx):
	reordered_past = ()
	for layer_past in past_key_values:
	reordered_past += (
	tuple(
	past_state.index_select(0, beam_idx.to(past_state.device))
	for past_state in layer_past
	),
	)
	return reordered_past


	class PhiForSequenceClassification(PhiPreTrainedModel):
	def __init__(self, config):
	super().__init__(config)
	self.num_labels = config.num_labels
	self.model = PhiModel(config)
	self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)

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

	def get_input_embeddings(self):
	return self.model.embed_tokens

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

	def forward(
	self,
	input_ids: torch.LongTensor = None,
	attention_mask: Optional[torch.Tensor] = None,
	position_ids: Optional[torch.LongTensor] = None,
	past_key_values: Optional[List[torch.FloatTensor]] = None,
	inputs_embeds: Optional[torch.FloatTensor] = 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, SequenceClassifierOutputWithPast]:
	r"""
	labels (`torch.LongTensor` of shape `(batch_size,)`, optional):
	Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
	config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
	`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
	"""
	return_dict = (
	return_dict if return_dict is not None else self.config.use_return_dict
	)

	model_outputs = self.model(
	input_ids,
	attention_mask=attention_mask,
	position_ids=position_ids,
	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,
	)
	hidden_states = model_outputs[0]
	logits = self.score(hidden_states)

	if input_ids is not None:
	batch_size = input_ids.shape[0]
	else:
	batch_size = inputs_embeds.shape[0]

	if self.config.pad_token_id is None and batch_size != 1:
	raise ValueError(
	"Cannot handle batch sizes > 1 if no padding token is defined."
	)
	if self.config.pad_token_id is None:
	sequence_lengths = -1
	else:
	if input_ids is not None:
	# if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
	sequence_lengths = (
	torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
	)
	sequence_lengths = sequence_lengths % input_ids.shape[-1]
	sequence_lengths = sequence_lengths.to(logits.device)
	else:
	sequence_lengths = -1

	pooled_logits = logits[
	torch.arange(batch_size, device=logits.device), sequence_lengths
	]

	loss = None
	if labels is not None:
	labels = labels.to(logits.device)
	if self.config.problem_type is None:
	if self.num_labels == 1:
	self.config.problem_type = "regression"
	elif self.num_labels > 1 and (
	labels.dtype == torch.long or labels.dtype == torch.int
	):
	self.config.problem_type = "single_label_classification"
	else:
	self.config.problem_type = "multi_label_classification"

	if self.config.problem_type == "regression":
	loss_fct = MSELoss()
	if self.num_labels == 1:
	loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
	else:
	loss = loss_fct(pooled_logits, labels)
	elif self.config.problem_type == "single_label_classification":
	loss_fct = CrossEntropyLoss()
	loss = loss_fct(
	pooled_logits.view(-1, self.num_labels), labels.view(-1)
	)
	elif self.config.problem_type == "multi_label_classification":
	loss_fct = BCEWithLogitsLoss()
	loss = loss_fct(pooled_logits, labels)
	if not return_dict:
	output = (pooled_logits,) + model_outputs[1:]
	return ((loss,) + output) if loss is not None else output

	return SequenceClassifierOutputWithPast(
	loss=loss,
	logits=pooled_logits,
	past_key_values=model_outputs.past_key_values,
	hidden_states=model_outputs.hidden_states,
	attentions=model_outputs.attentions,
	)


	class PhiForTokenClassification(PhiPreTrainedModel):
	def __init__(self, config: PhiConfig):
	super().__init__(config)
	self.num_labels = config.num_labels

	self.model = PhiModel(config)
	if (
	hasattr(config, "classifier_dropout")
	and config.classifier_dropout is not None
	):
	classifier_dropout = config.classifier_dropout
	elif hasattr(config, "hidden_dropout") and config.hidden_dropout is not None:
	classifier_dropout = config.hidden_dropout
	else:
	classifier_dropout = 0.1
	self.dropout = nn.Dropout(classifier_dropout)
	self.classifier = nn.Linear(config.hidden_size, config.num_labels)

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

	def forward(
	self,
	input_ids: Optional[torch.LongTensor] = None,
	past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
	attention_mask: Optional[torch.Tensor] = None,
	inputs_embeds: Optional[torch.Tensor] = None,
	labels: 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,
	**deprecated_arguments,
	) -> Union[Tuple[torch.Tensor], TokenClassifierOutput]:
	r"""
	labels (`torch.LongTensor` of shape `(batch_size,)`, optional):
	Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
	config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
	`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
	"""
	return_dict = (
	return_dict if return_dict is not None else self.config.use_return_dict
	)

	model_outputs = self.model(
	input_ids,
	past_key_values=past_key_values,
	attention_mask=attention_mask,
	inputs_embeds=inputs_embeds,
	use_cache=use_cache,
	output_attentions=output_attentions,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	)

	hidden_states = model_outputs[0]
	hidden_states = self.dropout(hidden_states)
	logits = self.classifier(hidden_states)

	loss = None
	if labels is not None:
	# move labels to correct device to enable model parallelism
	labels = labels.to(logits.device)
	batch_size, seq_length = labels.shape
	loss_fct = CrossEntropyLoss()
	loss = loss_fct(
	logits.view(batch_size * seq_length, self.num_labels),
	labels.view(batch_size * seq_length),
	)

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

	return TokenClassifierOutput(
	loss=loss,
	logits=logits,
	hidden_states=model_outputs.hidden_states,
	attentions=model_outputs.attentions,
	)


	@dataclass
	class LlavaCausalLMOutputWithPast(ModelOutput):
	loss: Optional[torch.FloatTensor] = None
	logits: torch.FloatTensor = None
	past_key_values: Optional[List[torch.FloatTensor]] = None
	hidden_states: Optional[Tuple[torch.FloatTensor]] = None
	attentions: Optional[Tuple[torch.FloatTensor]] = None
	image_features: Optional[torch.FloatTensor] = None


	class SiglipVisionEncoder(nn.Module):
	def __init__(self, config: LlavaConfig):
	super().__init__()
	self.vision_tower = SiglipVisionModel(config.vision_config)

	self.coord_embed = nn.Sequential(
	nn.Linear(2, config.vision_embed_dim),
	nn.GELU(),
	nn.Linear(config.vision_embed_dim, config.vision_embed_dim),
	)

	self.num_tokens = 728

	def feature_select(self, image_forward_outs, coord_feature, num_tokens=None):
	image_features = image_forward_outs
	image_features = image_features[:, 1:]
	if num_tokens is None:
	num_tokens = self.num_tokens
	split_size = int(num_tokens / image_features.shape[0])
	sum = 0
	output_list = []
	for i in range(image_features.shape[0]):
	if i == image_features.shape[0] - 1:
	size = num_tokens - sum
	else:
	size = split_size
	sum += size
	chunk_output = image_features[i, -size:, :]
	chunk_output = chunk_output + coord_feature[i]
	output_list.append(chunk_output)
	image_features = torch.cat(output_list)
	return image_features

	def process_image_chunks(self, image_tensor, coord_tensor, num_tokens=None):
	if image_tensor.shape[0] > 50:
	image_forward_out = []
	for i in range(0, image_tensor.shape[0], 50):
	part_forward_out = self.vision_tower(
	image_tensor[i : i + 50], output_hidden_states=True
	).hidden_states[-1]
	image_forward_out.append(part_forward_out)
	image_forward_out = torch.cat(image_forward_out, dim=0)
	else:
	image_forward_out = self.vision_tower(
	image_tensor, output_hidden_states=True
	).hidden_states[-1]
	coord_feature = self.coord_embed(coord_tensor)
	if len(coord_feature.shape) == 1:
	coord_feature = coord_feature.unsqueeze(0)
	image_feature = self.feature_select(
	image_forward_out, coord_feature, num_tokens
	).to(image_tensor.dtype)
	return image_feature

	def forward(
	self, images: List[torch.Tensor], coords: List[torch.Tensor], num_tokens=None
	):
	image_features = []
	for i, image in enumerate(images):
	image_feature = self.process_image_chunks(image, coords[i], num_tokens)
	image_features.append(image_feature)
	image_features = torch.stack(image_features)
	return image_features


	class LlavaMultiModalProjector(nn.Module):
	def __init__(self, config: LlavaConfig):
	super().__init__()

	self.linear_1 = nn.Linear(
	config.vision_embed_dim,
	config.text_config.hidden_size,
	bias=True,
	)
	self.act = nn.GELU()
	self.linear_2 = nn.Linear(
	config.text_config.hidden_size,
	config.text_config.hidden_size,
	bias=True,
	)

	def forward(self, image_features):
	hidden_states = self.linear_1(image_features)
	hidden_states = self.act(hidden_states)
	hidden_states = self.linear_2(hidden_states)
	return hidden_states


	class LlavaPreTrainedModel(PreTrainedModel):
	config_class = LlavaConfig
	base_model_prefix = "model"
	supports_gradient_checkpointing = True
	_no_split_modules = ["LlavaVisionAttention"]
	_skip_keys_device_placement = "past_key_values"
	_supports_flash_attn_2 = True

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

	def _init_weights(self, module):
	return

	@property
	def _supports_sdpa(self):
	"""
	Retrieve language_model's attribute to check whether the model supports
	SDPA or not.
	"""
	return self.language_model._supports_sdpa


	class LlavaForCausalLM(LlavaPreTrainedModel):
	def __init__(self, config: LlavaConfig):
	super().__init__(config)
	self.vision_model = SiglipVisionEncoder(config)

	self.multi_modal_projector = LlavaMultiModalProjector(config)
	self.vocab_size = config.vocab_size
	self.language_model = PhiForCausalLM(config.text_config)
	self.pad_token_id = (
	self.config.pad_token_id if self.config.pad_token_id is not None else -1
	)
	self.post_init()

	def get_input_embeddings(self):
	return self.language_model.get_input_embeddings()

	def set_input_embeddings(self, value):
	self.language_model.set_input_embeddings(value)

	def get_output_embeddings(self):
	return self.language_model.get_output_embeddings()

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

	def set_decoder(self, decoder):
	self.language_model.transformer = decoder

	def get_decoder(self):
	return self.language_model.transformer

	def tie_weights(self):
	return self.language_model.tie_weights()

	def resize_token_embeddings(
	self, new_num_tokens: Optional[int] = None, pad_to_multiple_of=None
	) -> nn.Embedding:
	model_embeds = self.language_model.resize_token_embeddings(
	new_num_tokens, pad_to_multiple_of
	)
	# update vocab size
	self.config.text_config.vocab_size = model_embeds.num_embeddings
	self.config.vocab_size = model_embeds.num_embeddings
	self.vocab_size = model_embeds.num_embeddings
	return model_embeds

	def _merge_input_ids_with_image_features(
	self, image_features, inputs_embeds, input_ids, attention_mask, position_ids
	):
	num_images, num_image_patches, embed_dim = image_features.shape
	batch_size, sequence_length = input_ids.shape
	left_padding = not torch.sum(
	input_ids[:, -1] == torch.tensor(self.pad_token_id)
	)
	# 1. Create a mask to know where special image tokens are
	special_image_token_mask = input_ids == self.config.image_token_index
	num_special_image_tokens = torch.sum(special_image_token_mask, dim=-1)
	# Compute the maximum embed dimension
	max_embed_dim = (
	num_special_image_tokens.max() * (num_image_patches - 1)
	) + sequence_length
	batch_indices, non_image_indices = torch.where(
	input_ids != self.config.image_token_index
	)

	# 2. Compute the positions where text should be written
	# Calculate new positions for text tokens in merged image-text sequence.
	# `special_image_token_mask` identifies image tokens. Each image token will be replaced by `nb_text_tokens_per_images - 1` text tokens.
	# `torch.cumsum` computes how each image token shifts subsequent text token positions.
	# - 1 to adjust for zero-based indexing, as `cumsum` inherently increases indices by one.
	new_token_positions = (
	torch.cumsum((special_image_token_mask * (num_image_patches - 1) + 1), -1)
	- 1
	)
	nb_image_pad = max_embed_dim - 1 - new_token_positions[:, -1]
	if left_padding:
	new_token_positions += nb_image_pad[:, None] # offset for left padding
	text_to_overwrite = new_token_positions[batch_indices, non_image_indices]

	# 3. Create the full embedding, already padded to the maximum position
	final_embedding = torch.zeros(
	batch_size,
	max_embed_dim,
	embed_dim,
	dtype=inputs_embeds.dtype,
	device=inputs_embeds.device,
	)
	final_attention_mask = torch.zeros(
	batch_size,
	max_embed_dim,
	dtype=attention_mask.dtype,
	device=inputs_embeds.device,
	)
	# In case the Vision model or the Language model has been offloaded to CPU, we need to manually
	# set the corresponding tensors into their correct target device.
	target_device = inputs_embeds.device
	batch_indices, non_image_indices, text_to_overwrite = (
	batch_indices.to(target_device),
	non_image_indices.to(target_device),
	text_to_overwrite.to(target_device),
	)
	attention_mask = attention_mask.to(target_device)

	# 4. Fill the embeddings based on the mask. If we have ["hey" "<image>", "how", "are"]
	# we need to index copy on [0, 577, 578, 579] for the text and [1:576] for the image features
	final_embedding[batch_indices, text_to_overwrite] = inputs_embeds[
	batch_indices, non_image_indices
	]
	final_attention_mask[batch_indices, text_to_overwrite] = attention_mask[
	batch_indices, non_image_indices
	]

	# 5. Fill the embeddings corresponding to the images. Anything that is still zeros needs filling
	image_to_overwrite = torch.all(final_embedding == 0, dim=-1)
	image_to_overwrite &= image_to_overwrite.cumsum(-1) - 1 >= nb_image_pad[
	:, None
	].to(target_device)

	if image_to_overwrite.sum() != image_features.shape[:-1].numel():
	raise ValueError(
	f"The input provided to the model are wrong. The number of image tokens is {torch.sum(special_image_token_mask)} while"
	f" the number of image given to the model is {num_images}. This prevents correct indexing and breaks batch generation."
	)

	final_embedding[image_to_overwrite] = (
	image_features.contiguous().reshape(-1, embed_dim).to(target_device)
	)
	final_attention_mask \|= image_to_overwrite
	position_ids = (final_attention_mask.cumsum(-1) - 1).masked_fill_(
	(final_attention_mask == 0), 1
	)
	return final_embedding, final_attention_mask, position_ids

	def forward(
	self,
	input_ids: torch.LongTensor = None,
	image_features: torch.FloatTensor = None,
	attention_mask: Optional[torch.Tensor] = None,
	position_ids: Optional[torch.LongTensor] = None,
	past_key_values: Optional[List[torch.FloatTensor]] = None,
	inputs_embeds: Optional[torch.FloatTensor] = None,
	use_cache: Optional[bool] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	) -> Union[Tuple, LlavaCausalLMOutputWithPast]:
	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 inputs_embeds is None:
	inputs_embeds = self.get_input_embeddings()(input_ids)
	if image_features is not None and input_ids.shape[1] != 1:
	(
	inputs_embeds,
	attention_mask,
	position_ids,
	) = self._merge_input_ids_with_image_features(
	image_features,
	inputs_embeds,
	input_ids,
	attention_mask,
	position_ids,
	)
	else:
	# In case input_ids.shape[1] == 1 & pixel_values==None & past_key_values != None, we are in the case of
	# generation with cache
	if past_key_values is not None and image_features is not None and input_ids.shape[1] == 1:
	# Retrieve the first layer to inspect the logits and mask out the hidden states
	# that are set to 0
	first_layer_past_key_value = past_key_values[0][0][:, :, :, 0]

	# Sum all dimensions of head_dim (-2) to avoid random errors such as: https://github.com/huggingface/transformers/pull/28032#issuecomment-1863691941
	batch_index, non_attended_tokens = torch.where(first_layer_past_key_value.float().sum(-2) == 0)

	# Get the target length
	target_seqlen = first_layer_past_key_value.shape[-1] + 1

	extended_attention_mask = torch.ones(
	(attention_mask.shape[0], target_seqlen - attention_mask.shape[1]),
	dtype=attention_mask.dtype,
	device=attention_mask.device,
	)

	# Filter out only the tokens that can be un-attended, this can happen
	# if one uses Llava + Fused modules where the cache on the
	# first iteration is already big enough, or if one passes custom cache
	valid_indices = non_attended_tokens < extended_attention_mask.size(-1)
	new_batch_index = batch_index[valid_indices]
	new_non_attended_tokens = non_attended_tokens[valid_indices]

	# Zero-out the places where we don't need to attend
	extended_attention_mask[new_batch_index, new_non_attended_tokens] = 0

	attention_mask = torch.cat((attention_mask, extended_attention_mask), dim=1)
	position_ids = torch.sum(attention_mask, dim=1).unsqueeze(-1) - 1

	outputs = self.language_model(
	input_ids=None,
	attention_mask=attention_mask,
	position_ids=position_ids,
	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 = outputs[0]

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

	return LlavaCausalLMOutputWithPast(
	logits=logits,
	past_key_values=outputs.past_key_values,
	hidden_states=outputs.hidden_states,
	attentions=outputs.attentions,
	image_features=image_features,
	)

	def prepare_inputs_for_generation(
	self,
	input_ids,
	past_key_values=None,
	inputs_embeds=None,
	attention_mask=None,
	image_features=None,
	**kwargs,
	):
	if past_key_values is not None:
	if isinstance(past_key_values, Cache):
	cache_length = past_key_values.get_seq_length()
	past_length = past_key_values.seen_tokens
	max_cache_length = past_key_values.get_max_length()
	else:
	cache_length = past_length = past_key_values[0][0].shape[2]
	max_cache_length = None

	# Keep only the unprocessed tokens:
	# 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
	# some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
	# input)
	if (
	attention_mask is not None
	and attention_mask.shape[1] > input_ids.shape[1]
	):
	input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
	# 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
	# input_ids based on the past_length.
	elif past_length < input_ids.shape[1]+image_features.shape[1]-1:
	past_length -= image_features.shape[1]-1
	input_ids = input_ids[:, past_length:]
	attention_mask = attention_mask[:, past_length:]
	# 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.

	# If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
	if (
	max_cache_length is not None
	and attention_mask is not None
	and cache_length + input_ids.shape[1] > max_cache_length
	):
	attention_mask = attention_mask[:, -max_cache_length:]

	position_ids = kwargs.get("position_ids", None)
	if attention_mask is not None and position_ids is None:
	# create position_ids on the fly for batch generation
	position_ids = attention_mask.long().cumsum(-1) - 1
	position_ids.masked_fill_(attention_mask == 0, 1)
	if past_key_values:
	position_ids = position_ids[:, -input_ids.shape[1] :]

	# if `inputs_embeds` are passed, we only want to use them in the 1st generation step
	if inputs_embeds is not None and past_key_values is None:
	model_inputs = {"inputs_embeds": inputs_embeds}
	else:
	model_inputs = {"input_ids": input_ids}

	model_inputs.update(
	{
	"position_ids": position_ids,
	"past_key_values": past_key_values,
	"use_cache": kwargs.get("use_cache"),
	"attention_mask": attention_mask,
	"image_features": image_features,
	}
	)
	return model_inputs

	def _reorder_cache(self, args, *kwargs):
	return self.language_model._reorder_cache(args, *kwargs)