PyTorch
English
Tevatron
phi3_v
vidore
custom_code
dse-phi3-docmatix-v1 / image_embedding_phi3_v.py
MrLight's picture
Create image_embedding_phi3_v.py
e08889c verified
# coding=utf-8
# Copyright 2024 Microsoft 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.
import torch
from torch import nn
from transformers import CLIPVisionConfig, CLIPVisionModel, PretrainedConfig
from transformers.models.clip.modeling_clip import CLIPAttention
from transformers.utils import logging
try:
from flash_attn import flash_attn_func
except ImportError:
pass
logger = logging.get_logger(__name__)
MAX_INPUT_ID = int(1e9)
CLIP_VIT_LARGE_PATCH14_336_CONFIG = CLIPVisionConfig(
attention_dropout=0.0,
dropout=0.0,
hidden_act="quick_gelu",
hidden_size=1024,
image_size=336,
initializer_factor=1.0,
initializer_range=0.02,
intermediate_size=4096,
layer_norm_eps=1e-05,
num_attention_heads=16,
num_channels=3,
num_hidden_layers=24,
patch_size=14,
projection_dim=768
)
class CLIPAttentionFA2(CLIPAttention):
"""Add flash attention 2 to CLIPAttention. (This is only used in the vision encoder)"""
def forward(self,
hidden_states,
attention_mask=None,
causal_attention_mask=None,
output_attentions=False,
):
"""Input shape: Batch x Time x Channel"""
assert attention_mask is None, "CLIPAttentionFA2 does not support attention_mask"
assert causal_attention_mask is None, "CLIPAttentionFA2 does not support causal_attention_mask"
assert output_attentions is False, "CLIPAttentionFA2 does not support output_attentions"
bsz, tgt_len, embed_dim = hidden_states.size()
query_states = self.q_proj(hidden_states).reshape(bsz, tgt_len, self.num_heads, self.head_dim)
key_states = self.k_proj(hidden_states).reshape(bsz, tgt_len, self.num_heads, self.head_dim)
value_states = self.v_proj(hidden_states).reshape(bsz, tgt_len, self.num_heads, self.head_dim)
attn_output = flash_attn_func(
query_states,
key_states,
value_states,
dropout_p=self.dropout if self.training else 0.0,
softmax_scale=self.scale,
causal=False,
).reshape(bsz, tgt_len, embed_dim)
attn_output = self.out_proj(attn_output)
return attn_output, None
class Phi3ImageEmbedding(nn.Module):
"""Phi3 Image embedding."""
def __init__(self, config: PretrainedConfig, wte=None, **kwargs) -> None:
super().__init__()
# n_embed or hidden_size
hidden_size = config.n_embd if hasattr(config, 'n_embd') else config.hidden_size
if hasattr(config, 'embd_pdrop') or hasattr(config, 'embed_pdrop'):
embd_drop = config.embd_pdrop if hasattr(config, 'embd_pdrop') else config.embed_pdrop
self.drop = nn.Dropout(embd_drop)
else:
self.drop = None
self.wte = wte
if isinstance(config.img_processor, dict) and config.img_processor.get('name', None) == 'clip_vision_model':
assert 'model_name' in config.img_processor, 'model_name must be provided for CLIPVisionModel'
assert 'image_dim_out' in config.img_processor, 'image_dim_out must be provided for CLIPVisionModel'
assert 'num_img_tokens' in config.img_processor, 'num_img_tokens must be provided for CLIPVisionModel'
assert config.img_processor['model_name'] == 'openai/clip-vit-large-patch14-336'
clip_config = CLIP_VIT_LARGE_PATCH14_336_CONFIG
self.img_processor = CLIPVisionModel(clip_config)
image_dim_out = config.img_processor['image_dim_out']
self.num_img_tokens = config.img_processor['num_img_tokens']
# FA2 in CLIP
if config._attn_implementation == 'flash_attention_2':
for layer in self.img_processor.vision_model.encoder.layers:
clip_fa2 = CLIPAttentionFA2(clip_config)
del layer.self_attn
layer.self_attn = clip_fa2
else:
raise NotImplementedError(f'img_processor = {config.img_processor}, not implemented')
self.image_dim_out = image_dim_out
self.img_sizes = None
# global_gn and sub_gn for hd transform, serves as line separator
self.use_hd_transform = kwargs.get('use_hd_transform', False)
self.with_learnable_separator = kwargs.get('with_learnable_separator', False)
self.hd_transform_order = kwargs.get('hd_transform_order', 'glb_sub')
# with_hd_transform and with_learnable_separator should have same value
assert self.use_hd_transform == self.with_learnable_separator, 'use_hd_transform and with_learnable_separator should have same value'
if self.with_learnable_separator:
assert self.use_hd_transform, 'learnable separator is only for hd transform'
# 1024 * 4, merge spatial to channel dimension
self.glb_GN = nn.Parameter(torch.zeros([1, 1, self.image_dim_out * 4]))
self.sub_GN = nn.Parameter(torch.zeros([1, 1, 1, self.image_dim_out * 4]))
logger.info(f'learnable separator enabled for hd transform, hd_transform_order = {self.hd_transform_order}')
projection_cls = kwargs.get('projection_cls', 'linear')
if projection_cls == 'linear':
self.img_projection = nn.Linear(image_dim_out, hidden_size)
elif projection_cls == 'mlp' and self.use_hd_transform:
dim_projection = hidden_size
depth = 2
layers = [nn.Linear(image_dim_out * 4, dim_projection)]
for _ in range(1, depth):
layers.extend([nn.GELU(),
nn.Linear(dim_projection, dim_projection)])
self.img_projection = nn.Sequential(*layers)
elif projection_cls == 'mlp':
dim_projection = hidden_size
depth = 2
layers = [nn.Linear(image_dim_out, dim_projection)]
for _ in range(1, depth):
layers.extend([nn.GELU(),
nn.Linear(dim_projection, dim_projection)])
self.img_projection = nn.Sequential(*layers)
else:
raise NotImplementedError(f'projection_cls = {projection_cls}, not implemented')
self.vocab_size = config.vocab_size
self.img_features = None
if isinstance(config.img_processor, dict):
self.layer_idx = config.img_processor.get('layer_idx', -2)
self.type_feature = config.img_processor.get('type_feature', 'patch')
else:
self.layer_idx = -2
self.type_feature = 'patch'
def set_img_features(self, img_features: torch.FloatTensor) -> None:
self.img_features = img_features
def set_img_sizes(self, img_sizes: torch.LongTensor) -> None:
self.img_sizes = img_sizes
def get_img_features(self, img_embeds: torch.FloatTensor) -> torch.FloatTensor:
LAYER_IDX = self.layer_idx
TYPE_FEATURE = self.type_feature
img_processor_output = self.img_processor(img_embeds, output_hidden_states=True)
img_feature = img_processor_output.hidden_states[LAYER_IDX]
if TYPE_FEATURE == "patch":
patch_feature = img_feature[:, 1:]
return patch_feature
raise NotImplementedError
def forward(
self, input_ids: torch.LongTensor, pixel_values: torch.FloatTensor, image_sizes=None
) -> torch.FloatTensor:
input_shape = input_ids.size()
input_ids = input_ids.view(-1, input_shape[-1])
# positions for image tokens
positions = torch.nonzero((input_ids < 0) & (input_ids > -MAX_INPUT_ID), as_tuple=True)
has_image = len(positions[0].tolist()) > 0
input_ids = input_ids.clamp_min(0).clamp_max(self.vocab_size).detach()
hidden_states = self.wte(input_ids)
if has_image:
assert self.use_hd_transform
num_images, num_crops, c, h, w = pixel_values.shape
assert c == 3 and h == w == 336
img_features = self.get_img_features(pixel_values.flatten(0, 1)).reshape(
num_images, num_crops, -1, self.image_dim_out
)
image_features_proj = self.hd_feature_transform(img_features, image_sizes)
hidden_states = hidden_states.index_put(
positions, image_features_proj, accumulate=False
)
if self.drop is not None:
hidden_states = self.drop(hidden_states)
return hidden_states
def hd_feature_transform(self, image_features, image_sizes):
"""
image_features: (num_images, num_crops+1, 24*24, 1024)
"""
assert (
self.hd_transform_order == 'sub_glb'
), f'hd_transform_order `{self.hd_transform_order}` not implemented'
if isinstance(self.img_projection, nn.Sequential):
target_device = self.img_projection[0].bias.device
target_dtype = self.img_projection[0].bias.dtype
else: # It's a single nn.Linear layer
target_device = self.img_projection.bias.device
target_dtype = self.img_projection.bias.dtype
global_image_features = image_features[:, 0] # (num_images, 24*24, 1024)
# global feature can be viewed as a special HD case with num_crops 1x1
global_image_features_hd = self.reshape_hd_patches_2x2merge(global_image_features, 1, 1)
global_image_features_hd_newline = self.add_image_newline(global_image_features_hd)
all_image_embeddings = []
# need a for loop to process each image because of different image sizes
# (patch arrangement is different for each image)
for i, img_size in enumerate(image_sizes):
h, w = img_size
h_crop = h // 336
w_crop = w // 336
num_crops = h_crop * w_crop
# NOTE: real num_crops is padded
# (num_crops, 24*24, 1024)
sub_image_features = image_features[i, 1 : 1 + num_crops]
sub_image_features_hd = self.reshape_hd_patches_2x2merge(
sub_image_features, h_crop, w_crop
)
sub_image_features_hd_newline = self.add_image_newline(sub_image_features_hd)
# [sub features, separator, global features]
all_image_embeddings.extend(
[
sub_image_features_hd_newline.squeeze(0), # (h_crop*12*(w_crop*12+1), 4096)
self.glb_GN.squeeze(0),
global_image_features_hd_newline[i],
]
)
image_features_proj = self.img_projection(
torch.cat(all_image_embeddings, dim=0).to(target_device).to(target_dtype)
)
return image_features_proj
def reshape_hd_patches_2x2merge(self, image_features, h_crop, w_crop):
"""
image_features: (num_images*num_crops, 24*24, 1024)
output: (num_images, h_crop*12, w_crop*12, 4096), h_crop*w_crop == num_crops
"""
N, L, C = image_features.shape
assert L == 24 * 24 and C == 1024 and N % (h_crop * w_crop) == 0
num_images = N // (h_crop * w_crop)
H = int(L**0.5)
image_features_hd = (
image_features.reshape(N, H, H, C) # N, 24, 24, 1024
.reshape(N, H // 2, 2, H // 2, 2, C) # N, 12, 2, 12, 2, 1024
.permute(0, 1, 3, 2, 4, 5) # N, 12, 12, 2, 2, 1024
.reshape(N, -1, 4 * C) # N, 144, 4096
.reshape(
num_images, h_crop, w_crop, H // 2, H // 2, -1
) # n_img, h_crop, w_crop, 12, 12, 4096
.permute(0, 1, 3, 2, 4, 5) # n_img, h_crop, 12, w_crop, 12, 4096
.reshape(
num_images, h_crop * H // 2, w_crop * H // 2, 4 * C
) # n_img, h_crop*12, w_crop*12, 4096
)
# alternative implementation using einops
# from einops import rearrange
# image_features_nhwc = rearrange(
# image_features,
# 'N (H W) c -> N H W c',
# H=H,
# W=H,
# )
# image_features_2x2merge = rearrange(
# image_features_nhwc,
# 'N (h h_pool) (w w_pool) c -> N h w (h_pool w_pool c)',
# h_pool=2,
# w_pool=2,
# )
# image_features_hd = rearrange(
# image_features_2x2merge,
# '(n_img h_crop w_crop) h w C -> n_img (h_crop h) (w_crop w) C',
# h_crop=h_crop,
# w_crop=w_crop,
# )
return image_features_hd
def add_image_newline(self, image_features_hd):
"""
image_features_hd: (num_images, h_crop*12, w_crop*12, 4096)
output: (num_images, (h_crop*12) * (w_crop*12+1), 4096)
"""
num_images, h, w, hid_dim = image_features_hd.shape
# add the newline token to the HD image feature patches
newline_embeddings = self.sub_GN.expand(num_images, h, -1, -1) # (n_img, h, 1, hid_dim)
image_features_hd_newline = torch.cat(
[image_features_hd, newline_embeddings], dim=2
).reshape(num_images, -1, hid_dim)
return image_features_hd_newline