@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0a995d5434

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 21:00 UTC

[Model] New model support for Phi-4-multimodal-instruct (#14119)

Congcong Chen committed Mar 4, 2025 at 20:57 UTC 0a995d543464582522edcc88f8bf70af715c0129

10 files changed +7159 -3

docs/source/models/supported_models.md

+8 -1

@@ -410,7 +410,7 @@ See [this page](#generative-models) for more information on how to use generativ

410 * ✅︎

411 - * `Phi3ForCausalLM`

412 * Phi-4, Phi-3

413 - * `microsoft/Phi-4`, `microsoft/Phi-3-mini-4k-instruct`, `microsoft/Phi-3-mini-128k-instruct`, `microsoft/Phi-3-medium-128k-instruct`, etc.

413 + * `microsoft/Phi-4-mini-instruct`, `microsoft/Phi-4`, `microsoft/Phi-3-mini-4k-instruct`, `microsoft/Phi-3-mini-128k-instruct`, `microsoft/Phi-3-medium-128k-instruct`, etc.

414 * ✅︎

415 * ✅︎

416 - * `Phi3SmallForCausalLM`

@@ -856,6 +856,13 @@ See [this page](#generative-models) for more information on how to use generativ

856 *

857 * ✅︎

858 * ✅︎

859 +- * `Phi4MMForCausalLM`

860 + * Phi-4-multimodal

861 + * T + I+ / T + A+ / I+ + A+

862 + * `microsoft/Phi-4-multimodal-instruct`, etc.

863 + * ✅︎

864 + *

865 + *

866 - * `PixtralForConditionalGeneration`

867 * Pixtral

868 * T + I+

requirements-common.txt

@@ -37,3 +37,4 @@ depyf==0.18.0 # required for profiling and debugging with compilation config

37 cloudpickle # allows pickling lambda functions in model_executor/models/registry.py

38 watchfiles # required for http server to monitor the updates of TLS files

39 python-json-logger # Used by logging as per examples/other/logging_configuration.md

40 +scipy # Required for phi-4-multimodal-instruct

\ No newline at end of file

tests/models/registry.py

@@ -272,6 +272,8 @@ _MULTIMODAL_EXAMPLE_MODELS = {

272 extras={"v2": "google/paligemma2-3b-ft-docci-448"}), # noqa: E501

273 "Phi3VForCausalLM": _HfExamplesInfo("microsoft/Phi-3-vision-128k-instruct",

274 trust_remote_code=True),

275 + "Phi4MMForCausalLM": _HfExamplesInfo("microsoft/Phi-4-multimodal-instruct",

276 + trust_remote_code=True),

277 "PixtralForConditionalGeneration": _HfExamplesInfo("mistralai/Pixtral-12B-2409", # noqa: E501

278 tokenizer_mode="mistral"),

279 "QwenVLForConditionalGeneration": _HfExamplesInfo("Qwen/Qwen-VL",

vllm/config.py

+2 -2

@@ -2284,9 +2284,9 @@ class LoRAConfig:

2284 return hash_str

2285

2286 def __post_init__(self):

2287 - # Setting the maximum rank to 256 should be able to satisfy the vast

2287 + # Setting the maximum rank to 512 should be able to satisfy the vast

2288 # majority of applications.

2289 - possible_max_ranks = (8, 16, 32, 64, 128, 256)

2289 + possible_max_ranks = (8, 16, 32, 64, 128, 256, 320, 512)

2290 possible_lora_extra_vocab_size = (0, 256, 512)

2291 if self.max_lora_rank not in possible_max_ranks:

2292 raise ValueError(

vllm/entrypoints/chat_utils.py

@@ -395,6 +395,8 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]):

395 if model_type == "phi3_v":

396 # Workaround since this token is not defined in the tokenizer

397 return f"<|image_{current_count}|>"

398 + if model_type == "phi4mm":

399 + return "<|endoftext10|>" # 200010 (see vocab.json in hf model)

400 if model_type in ("minicpmo", "minicpmv"):

401 return "(<image>./</image>)"

402 if model_type in ("blip-2", "chatglm", "fuyu", "paligemma",

@@ -424,6 +426,8 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]):

426 elif modality == "audio":

427 if model_type == "ultravox":

428 return "<|audio|>"

429 + if model_type == "phi4mm":

430 + return "<|endoftext11|>" # 200011 (see vocab.json in hf model)

431 if model_type == "qwen2_audio":

432 return (f"Audio {current_count}: "

433 f"<|audio_bos|><|AUDIO|><|audio_eos|>")

vllm/model_executor/models/phi4mm.py new

+1803

@@ -0,0 +1,1803 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +import math

3 +import re

4 +from functools import lru_cache

5 +from typing import (Dict, Iterable, List, Literal, Mapping, Optional, Tuple,

6 + TypedDict, Union)

7 +

8 +import numpy as np

9 +import scipy.signal

10 +import torch

11 +import torch.nn as nn

12 +import torchvision.transforms as T

13 +from PIL import Image

14 +from transformers import PretrainedConfig

15 +from transformers.utils import logging

16 +

17 +from vllm.config import VllmConfig

18 +from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

19 +from vllm.inputs import (INPUT_REGISTRY, DecoderOnlyInputs, DummyData,

20 + InputContext)

21 +from vllm.inputs.data import TokenInputs, token_inputs

22 +from vllm.model_executor.layers.logits_processor import LogitsProcessor

23 +from vllm.model_executor.layers.quantization import QuantizationConfig

24 +from vllm.model_executor.layers.sampler import Sampler, SamplerOutput

25 +from vllm.model_executor.layers.vocab_parallel_embedding import (

26 + DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead)

27 +from vllm.model_executor.models.llama import LlamaModel

28 +from vllm.model_executor.sampling_metadata import SamplingMetadata

29 +from vllm.multimodal import MULTIMODAL_REGISTRY

30 +from vllm.multimodal.inputs import MultiModalInputs, NestedTensors

31 +from vllm.sequence import IntermediateTensors, SequenceData

32 +from vllm.transformers_utils.tokenizer import cached_tokenizer_from_config

33 +

34 +from .interfaces import SupportsLoRA, SupportsMultiModal

35 +from .phi4mm_audio import AudioEmbedding

36 +from .utils import maybe_prefix

37 +from .vision_siglip_navit import get_siglip_vision_model

38 +

39 +# <|endoftext10|> (see vocab.json in hf model)

40 +_IMAGE_PLACEHOLDER_TOKEN_ID = 200010

41 +# <|endoftext11|>

42 +_AUDIO_PLACEHOLDER_TOKEN_ID = 200011

43 +

44 +_AUDIO_MAX_SOUNDFILE_SIZE = 241_000

45 +DUMMY_SAMPLING_FREQUENCY = 16_000 # kHz

46 +

47 +DYNAMIC_HD = 16

48 +AUDIO_TOKEN_PATTERN = r"<\|audio_(\d+)\|>"

49 +IMAGE_TOKEN_PATTERN = r"<\|image_(\d+)\|>"

50 +

51 +SIGLIP_NAME = "siglip-so400m-patch14-448"

52 +VISION_ENCODER_TO_PROCESSING_CONFIG = {

53 + 'siglip-so400m-patch14-448': {

54 + 'dynamic_hd': 16,

55 + 'vit_image_size': 448,

56 + 'vit_patch_size': 14,

57 + 'token_compression_factor': 2,

58 + },

59 +}

60 +logger = logging.get_logger(__name__)

61 +# This is a workaround to prevent text (user input) + audio + image

62 +# from being used in the same prompt.

63 +# It includes token ids for "/n" and tokens in added_tokens_decoder

64 +# from the tokenizer_confg.json file.

65 +NON_USER_INPUT_TOKENS = {

66 + 198, 200010, 200011, 199999, 200018, 200019, 200020, 200021, 200022,

67 + 200023, 200024, 200025, 200026, 200027, 200028

68 +}

69 +

70 +

71 +def get_max_dummy_image(ctx: InputContext):

72 + hf_config = ctx.get_hf_config()

73 + vision_encoder_name = hf_config.img_processor

74 + if vision_encoder_name is None:

75 + vision_encoder_name = SIGLIP_NAME

76 + prepro_config = VISION_ENCODER_TO_PROCESSING_CONFIG[vision_encoder_name]

77 + dynamic_hd_size = prepro_config['dynamic_hd']

78 + vit_image_size = prepro_config['vit_image_size']

79 +

80 + max_side = vit_image_size * dynamic_hd_size

81 + dummy_image = dummy_image_for_phi4mm(vit_image_size, max_side)

82 + return dummy_image

83 +

84 +

85 +# image token length

86 +def get_max_phi4mm_image_tokens(ctx: InputContext):

87 + dummy_image = get_max_dummy_image(ctx)

88 +

89 + hf_config = ctx.get_hf_config()

90 + vision_encoder_name = hf_config.img_processor

91 + if vision_encoder_name is None:

92 + vision_encoder_name = SIGLIP_NAME

93 + prepro_config = VISION_ENCODER_TO_PROCESSING_CONFIG[vision_encoder_name]

94 + dynamic_hd_size = prepro_config['dynamic_hd']

95 + vit_image_size = prepro_config['vit_image_size']

96 + vit_patch_size = prepro_config['vit_patch_size']

97 + token_compression_factor = prepro_config['token_compression_factor']

98 +

99 + image_num_tokens = _compute_num_image_tokens(dummy_image, dynamic_hd_size,

100 + vit_image_size,

101 + vit_patch_size,

102 + token_compression_factor)

103 + return image_num_tokens

104 +

105 +

106 +def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height,

107 + image_size):

108 + best_ratio_diff = float('inf')

109 + best_ratio = (1, 1)

110 + area = width * height

111 + for ratio in target_ratios:

112 + target_aspect_ratio = ratio[0] / ratio[1]

113 + ratio_diff = abs(aspect_ratio - target_aspect_ratio)

114 + if ratio_diff < best_ratio_diff:

115 + best_ratio_diff = ratio_diff

116 + best_ratio = ratio

117 + elif ratio_diff == best_ratio_diff:

118 + if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:

119 + best_ratio = ratio

120 + return best_ratio

121 +

122 +

123 +def _find_target_aspect_ratio(image, image_size, max_num, min_num):

124 + orig_width, orig_height = image.size

125 +

126 + w_crop_num = math.ceil(orig_width / float(image_size))

127 + h_crop_num = math.ceil(orig_height / float(image_size))

128 + if w_crop_num * h_crop_num > max_num:

129 + aspect_ratio = orig_width / orig_height

130 +

131 + # calculate the existing image aspect ratio

132 + target_ratios = set((i, j) for i in range(1, max_num + 1)

133 + for j in range(1, max_num + 1)

134 + if i * j <= max_num and i * j >= min_num)

135 + target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])

136 +

137 + # find the closest aspect ratio to the target

138 + target_aspect_ratio = find_closest_aspect_ratio(

139 + aspect_ratio, target_ratios, orig_width, orig_height, image_size)

140 +

141 + # calculate the target width and height

142 + target_width = image_size * target_aspect_ratio[0]

143 + target_height = image_size * target_aspect_ratio[1]

144 + logger.debug("target_aspect_ratio: %s", target_aspect_ratio)

145 + else:

146 + target_width = image_size * w_crop_num

147 + target_height = image_size * h_crop_num

148 + target_aspect_ratio = (w_crop_num, h_crop_num)

149 + return target_aspect_ratio, target_height, target_width

150 +

151 +

152 +def _get_padding_size(image, target_height, target_width):

153 + orig_width, orig_height = image.size

154 + ratio_width = target_width / orig_width

155 + ratio_height = target_height / orig_height

156 +

157 + if ratio_width < ratio_height:

158 + padding_width = 0

159 + padding_height = target_height - int(orig_height * ratio_width)

160 + else:

161 + padding_width = target_width - int(orig_width * ratio_height)

162 + padding_height = 0

163 + return padding_height, padding_width

164 +

165 +

166 +def dynamic_preprocess(image,

167 + min_num=1,

168 + max_num=12,

169 + image_size=384,

170 + mask_size=27):

171 + target_aspect_ratio, target_height, target_width =\

172 + _find_target_aspect_ratio(

173 + image, image_size, max_num, min_num)

174 + padding_height, padding_width = _get_padding_size(image, target_height,

175 + target_width)

176 +

177 + # Calculate the ratio

178 + orig_width, orig_height = image.size

179 + ratio_width = target_width / orig_width

180 + ratio_height = target_height / orig_height

181 + if ratio_width < ratio_height:

182 + new_size = (target_width, int(orig_height * ratio_width))

183 + else:

184 + new_size = (int(orig_width * ratio_height), target_height)

185 +

186 + attention_mask = torch.ones((int(mask_size * target_aspect_ratio[1]),

187 + int(mask_size * target_aspect_ratio[0])))

188 + if padding_width >= 14:

189 + attention_mask[:, -math.floor(padding_width / 14):] = 0

190 + if padding_height >= 14:

191 + attention_mask[-math.floor(padding_height / 14):, :] = 0

192 + assert attention_mask.sum(

193 + ) > 0, f'attention mask is empty {attention_mask}'

194 +

195 + if min(new_size[1], target_height) < 10 or min(new_size[0],

196 + target_width) < 10:

197 + raise ValueError(f'the aspect ratio is very extreme {new_size}')

198 +

199 + image = T.functional.resize(

200 + image,

201 + [new_size[1], new_size[0]],

202 + )

203 +

204 + resized_img = T.functional.pad(image,

205 + [0, 0, padding_width, padding_height],

206 + fill=[255, 255, 255])

207 +

208 + return resized_img, attention_mask

209 +

210 +

211 +def pad_to_max_num_crops(images, max_crops=5):

212 + """

213 + images: B x 3 x H x W, B<=max_crops

214 + """

215 + B, _, H, W = images.shape

216 + if max_crops > B:

217 + pad = torch.zeros(max_crops - B,

218 + 3,

219 + H,

220 + W,

221 + dtype=images.dtype,

222 + device=images.device)

223 + images = torch.cat([images, pad], dim=0)

224 + return images

225 +

226 +

227 +def pad_mask_to_max_num_crops(masks, max_crops=5):

228 + B, H, W = masks.shape

229 + if max_crops > B:

230 + pad = torch.ones(max_crops - B,

231 + H,

232 + W,

233 + dtype=masks.dtype,

234 + device=masks.device)

235 + masks = torch.cat([masks, pad], dim=0)

236 + return masks

237 +

238 +

239 +def preprocess(images, dynamic_hd_size, vit_resolution, vit_patch_size):

240 +

241 + # Basic settings.

242 + img_processor = T.Compose([

243 + T.ToTensor(),

244 + T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),

245 + ])

246 + # Dynamic HD

247 + base_resolution = vit_resolution

248 + images = [image.convert('RGB') for image in images]

249 + # cover 384 and 448 resolution

250 + mask_resolution = base_resolution // vit_patch_size

251 + elems, image_attention_masks = [], []

252 + for im in images:

253 + elem, attention_mask = dynamic_preprocess(im,

254 + max_num=dynamic_hd_size,

255 + image_size=base_resolution,

256 + mask_size=mask_resolution)

257 + elems.append(elem)

258 + image_attention_masks.append(attention_mask)

259 + hd_images = [img_processor(im) for im in elems]

260 + global_image = [

261 + torch.nn.functional.interpolate(

262 + im.unsqueeze(0).float(),

263 + size=(base_resolution, base_resolution),

264 + mode='bicubic',

265 + ).to(im.dtype) for im in hd_images

266 + ]

267 + shapes = [[im.size(1), im.size(2)] for im in hd_images]

268 + mask_shapes = [[mask.size(0), mask.size(1)]

269 + for mask in image_attention_masks]

270 + global_attention_mask = [

271 + torch.ones((1, mask_resolution, mask_resolution)) for _ in hd_images

272 + ]

273 + hd_images_reshape = [

274 + im.reshape(1, 3, h // base_resolution, base_resolution,

275 + w // base_resolution, base_resolution).permute(

276 + 0, 2, 4, 1, 3, 5).reshape(-1, 3, base_resolution,

277 + base_resolution).contiguous()

278 + for im, (h, w) in zip(hd_images, shapes)

279 + ]

280 + attention_masks_reshape = [

281 + mask.reshape(1, h // mask_resolution, mask_resolution,

282 + w // mask_resolution, mask_resolution).permute(

283 + 0, 1, 3, 2, 4).reshape(-1, mask_resolution,

284 + mask_resolution).contiguous()

285 + for mask, (h, w) in zip(image_attention_masks, mask_shapes)

286 + ]

287 + # NOTE token compression is hard coded here, and odd numbers seems to fail

288 + downsample_attention_masks = [

289 + mask[:, 0::2,

290 + 0::2].reshape(1, h // mask_resolution, w // mask_resolution,

291 + mask_resolution // 2 + mask_resolution % 2,

292 + mask_resolution // 2 + mask_resolution % 2).permute(

293 + 0, 1, 3, 2, 4)

294 + for mask, (h, w) in zip(attention_masks_reshape, mask_shapes)

295 + ]

296 + downsample_attention_masks = [

297 + mask.reshape(mask.size(1) * mask.size(2),

298 + mask.size(3) * mask.size(4))

299 + for mask in downsample_attention_masks

300 + ]

301 + # NOTE hard coded number of tokens

302 + num_img_tokens = [

303 + 256 + 1 + int(mask.sum().item()) + int(mask[:, 0].sum().item()) + 16

304 + for mask in downsample_attention_masks

305 + ]

306 +

307 + hd_images_reshape = [

308 + torch.cat([_global_image] + [_im], dim=0)

309 + for _global_image, _im in zip(global_image, hd_images_reshape)

310 + ]

311 + hd_masks_reshape = [

312 + torch.cat([_global_mask] + [_mask],

313 + dim=0) for _global_mask, _mask in zip(

314 + global_attention_mask, attention_masks_reshape)

315 + ]

316 + max_crops = max([img.size(0) for img in hd_images_reshape])

317 + image_transformed = [

318 + pad_to_max_num_crops(im, max_crops) for im in hd_images_reshape

319 + ]

320 + image_transformed = torch.stack(image_transformed, dim=0)

321 + mask_transformed = [

322 + pad_mask_to_max_num_crops(mask, max_crops) \

323 + for mask in hd_masks_reshape

324 + ]

325 + mask_transformed = torch.stack(mask_transformed, dim=0)

326 +

327 + returned_input_image_embeds = image_transformed

328 + returned_image_sizes = torch.tensor(shapes, dtype=torch.long)

329 + returned_image_attention_mask = mask_transformed

330 + returned_num_img_tokens = num_img_tokens

331 +

332 + data = {

333 + "pixel_values": returned_input_image_embeds,

334 + "image_sizes": returned_image_sizes,

335 + "image_attention_mask": returned_image_attention_mask,

336 + "num_img_tokens": returned_num_img_tokens,

337 + }

338 + return data

339 +

340 +

341 +class Phi4MMImageEncoder(nn.Module):

342 + """Image embedding."""

343 +

344 + def __init__(self,

345 + config: PretrainedConfig,

346 + quant_config: Optional[QuantizationConfig],

347 + prefix: str = "",

348 + model_dir: str = "") -> None:

349 + super().__init__()

350 +

351 + # n_embed or hidden_size

352 + hidden_size = config.n_embd if hasattr(

353 + config, 'n_embd') else config.hidden_size

354 + if hasattr(config, 'embd_pdrop') or hasattr(config, 'embed_pdrop'):

355 + embd_drop = config.embd_pdrop if hasattr(

356 + config, 'embd_pdrop') else config.embed_pdrop

357 + self.drop = nn.Dropout(embd_drop)

358 + else:

359 + self.drop = None

360 +

361 + # layer_idx to output the img features

362 + if isinstance(config.img_processor, dict):

363 + self.layer_idx = config.img_processor.get('layer_idx', -2)

364 + self.type_feature = config.img_processor.get(

365 + 'type_feature', 'patch')

366 + else:

367 + self.layer_idx = -2

368 + self.type_feature = 'patch'

369 +

370 + self.img_processor = get_siglip_vision_model(

371 + _flash_attn_2_enabled=True)

372 +

373 + pe_weight = self.img_processor.embeddings.position_embedding.weight

374 + L, D = pe_weight.size()

375 + H = int(math.sqrt(L))

376 + assert H**2 == L, f'position embedding size {L} is not square'

377 + if H % 2 != 0:

378 + self.img_processor_padding = nn.ReflectionPad2d((0, 1, 0, 1))

379 + H += 1

380 + image_dim_out = D

381 + # ((448/14)//2)**2

382 + self.num_img_tokens = (H // 2)**2

383 + self.base_feat_height_target = H

384 +

385 + self.image_dim_out = image_dim_out

386 + self.img_sizes = None

387 + self.image_attention_mask = None

388 +

389 + # global_gn and sub_gn for hd transform, serves as line separator

390 + self.use_hd_transform = True

391 + self.with_learnable_separator = True

392 + self.hd_transform_order = "sub_glb"

393 + self.freeze_img_processor = False

394 + self.crop_size = 448

395 +

396 + # image token compression

397 + self.image_token_compression_cls = 'avg_pool_2d'

398 + self.image_token_compression = nn.AvgPool2d(kernel_size=2, stride=2)

399 + self.base_feat_height_reduction = 1

400 + self.base_feat_height_target = self.base_feat_height_target // 2

401 +

402 + # with_hd_transform and with_learnable_separator should have same value

403 + assert self.use_hd_transform == self.with_learnable_separator, \

404 + 'use_hd_transform and with_learnable_separator should have same value'

405 + assert self.use_hd_transform, \

406 + 'learnable separator is only for hd transform'

407 + # 1024 * 4, merge spatial to channel dimension

408 + self.glb_GN = nn.Parameter(

409 + torch.zeros([

410 + 1, 1, self.image_dim_out * self.base_feat_height_reduction**2

411 + ]))

412 + self.sub_GN = nn.Parameter(

413 + torch.zeros([

414 + 1, 1, 1,

415 + self.image_dim_out * self.base_feat_height_reduction**2

416 + ]))

417 +

418 + dim_projection = hidden_size

419 + depth = 2

420 + layers = [

421 + nn.Linear(image_dim_out * self.base_feat_height_reduction**2,

422 + dim_projection)

423 + ]

424 + for _ in range(1, depth):

425 + layers.extend(

426 + [nn.GELU(),

427 + nn.Linear(dim_projection, dim_projection)])

428 + self.img_projection = nn.Sequential(*layers)

429 +

430 + self.vocab_size = config.vocab_size

431 + self.img_features = None

432 +

433 + self.use_out_place_operations = False

434 +

435 + def get_img_features(self,

436 + img_embeds: torch.FloatTensor,

437 + attention_mask=None) -> torch.FloatTensor:

438 + LAYER_IDX = self.layer_idx

439 + TYPE_FEATURE = self.type_feature

440 +

441 + img_processor_output = self.img_processor(

442 + img_embeds,

443 + output_hidden_states=True,

444 + patch_attention_mask=attention_mask)

445 + img_feature = img_processor_output.hidden_states[LAYER_IDX]

446 +

447 + if TYPE_FEATURE == "patch":

448 + patch_feature = img_feature

449 +

450 + use_token_compression = self.image_token_compression is not None

451 + use_padding = getattr(self, 'img_processor_padding',

452 + None) is not None

453 + if use_token_compression or use_padding:

454 + # reshape to 2D tensor

455 + width = int(math.sqrt(patch_feature.size(1)))

456 + patch_feature = patch_feature.view(-1, width, width,

457 + patch_feature.size(-1))

458 + # convert to NCHW

459 + patch_feature = patch_feature.permute(0, 3, 1, 2)

460 +

461 + if use_padding:

462 + patch_feature = self.img_processor_padding(patch_feature)

463 + if use_token_compression:

464 + patch_feature = self.image_token_compression(patch_feature)

465 +

466 + # convert to NHWC

467 + patch_feature = patch_feature.permute(0, 2, 3, 1)

468 + patch_feature = patch_feature.view(

469 + -1,

470 + patch_feature.size(1) * patch_feature.size(2),

471 + patch_feature.size(-1))

472 +

473 + return patch_feature

474 +

475 + raise NotImplementedError

476 +

477 + def forward(self, pixel_values: torch.FloatTensor,

478 + image_sizes: torch.Tensor,

479 + image_attention_mask: torch.Tensor) -> torch.FloatTensor:

480 + """

481 + process image and return vision embeddings.

482 +

483 + pixel_values: (num_images, num_crops, c, h, w)

484 + image_sizes: [[h1, w1], [h2, w2]]

485 + image_attention_mask: num_images x num_crops x 32 x 32

486 + output: (num_images, num_img_tokens, hidden_size)

487 + """

488 +

489 + # eg

490 + # pixel_values: torch.Size([1, 7, 3, 448, 448])

491 + # image_sizes: tensor([[ 896, 1344]], device='cuda:0')

492 + # output: torch.Size([1, 1841, 3072])

493 +

494 + if isinstance(self.img_projection, nn.Sequential):

495 + target_device = self.img_projection[0].bias.device

496 + target_dtype = self.img_projection[0].bias.dtype

497 + else: # It's a single nn.Linear layer

498 + target_device = self.img_projection.bias.device

499 + target_dtype = self.img_projection.bias.dtype

500 +

501 + img_sizes = image_sizes

502 + num_images, num_crops, c, h, w = pixel_values.shape

503 + bs = num_images

504 + pixel_values = pixel_values.flatten(0, 1)

505 +

506 + img_features = self.get_img_features(

507 + pixel_values,

508 + image_attention_mask.type(torch.BoolTensor).flatten(

509 + 0, 1).to(target_device))

510 +

511 + base_feat_height_target = self.base_feat_height_target

512 + base_resolution = self.crop_size

513 + base_feat_height_reduction = self.base_feat_height_reduction

514 +

515 + base_feat_height = base_feat_width = int(np.sqrt(

516 + img_features.shape[1]))

517 + assert base_feat_height == base_feat_height_target \

518 + and base_feat_width == base_feat_height_target, \

519 + f'base_feat_height: {base_feat_height},"\

520 + f" base_feat_width: {base_feat_width}, "\

521 + f"expect {base_feat_height_target} features for hd transform'

522 +

523 + # bs x max_num_crops x (24x24) x C

524 + img_features = img_features.view(bs, -1,

525 + base_feat_height * base_feat_width,

526 + self.image_dim_out)

527 + C = self.image_dim_out

528 + H = base_feat_height

529 +

530 + output_imgs = []

531 + output_len = []

532 + # training is tensor, inference is list

533 + if isinstance(img_sizes, torch.Tensor):

534 + img_sizes = img_sizes.view(-1, 2)

535 + for _bs in range(bs):

536 + h, w = img_sizes[_bs]

537 + h = h // base_resolution

538 + w = w // base_resolution

539 + B_ = h * w

540 +

541 + # 1 x (24x24) x 1024

542 + global_img_feature = img_features[_bs, :1]

543 +

544 + # 1 x 12 x 12 x 4096

545 + glb_img = global_img_feature.reshape(1, H, H, C).reshape(

546 + 1, H // base_feat_height_reduction, base_feat_height_reduction,

547 + H // base_feat_height_reduction, base_feat_height_reduction,

548 + C).contiguous().permute(0, 1, 3, 2, 4, 5).reshape(

549 + 1, H // base_feat_height_reduction,

550 + H // base_feat_height_reduction,

551 + base_feat_height_reduction * base_feat_height_reduction *

552 + C).contiguous()

553 + temp_glb_GN = self.sub_GN.repeat(1,

554 + H // base_feat_height_reduction,

555 + 1, 1)

556 +

557 + # 1 x 156 x 4096

558 + glb_img = torch.cat([glb_img, temp_glb_GN], dim=2).reshape(

559 + 1, -1,

560 + base_feat_height_reduction * base_feat_height_reduction * C)

561 +

562 + # (max_num_crops-1) x (12x12) x C

563 + sub_img = img_features[_bs, 1:]

564 + # 16x574x1024

565 + # get rid of padding sub_img

566 + sub_img = sub_img[:B_]

567 +

568 + # (num_crops, 12, 2, 12, 2, 1024) ->

569 + # (num_crops, 12, 12, 2, 2, 1024) -> (num_crops, 12*12, 4*1024)

570 + sub_img = sub_img.reshape(B_, H, H, C).reshape(

571 + B_, H // base_feat_height_reduction,

572 + base_feat_height_reduction, H // base_feat_height_reduction,

573 + base_feat_height_reduction,

574 + C).contiguous().permute(0, 1, 3, 2, 4, 5).reshape(

575 + B_, -1, base_feat_height_reduction *

576 + base_feat_height_reduction * C).contiguous()

577 + sub_img = sub_img.reshape(

578 + 1, h, w, base_feat_height // base_feat_height_reduction,

579 + base_feat_width // base_feat_height_reduction,

580 + -1).permute(0, 1, 3, 2, 4, 5).reshape(

581 + 1, h * base_feat_height // base_feat_height_reduction,

582 + w * base_feat_width // base_feat_height_reduction,

583 + base_feat_height_reduction * base_feat_height_reduction *

584 + C)

585 +

586 + if image_attention_mask is not None and len(

587 + image_attention_mask) > 0:

588 + reshaped_image_attention_mask = image_attention_mask[

589 + _bs, 1:B_ + 1, 0::2, 0::2].reshape(

590 + 1, h, w,

591 + base_feat_height // base_feat_height_reduction,

592 + base_feat_width // base_feat_height_reduction).permute(

593 + 0, 1, 3, 2, 4).reshape(

594 + 1, h * base_feat_height //

595 + base_feat_height_reduction, w *

596 + base_feat_width // base_feat_height_reduction)

597 + useful_height = int(

598 + reshaped_image_attention_mask[0, :, 0].sum().item())

599 + useful_width = int(

600 + reshaped_image_attention_mask[0, 0, :].sum().item())

601 + sub_img = sub_img[:, :useful_height, :useful_width]

602 + temp_sub_GN = self.sub_GN.repeat(1, useful_height, 1, 1)

603 + temp_len = int(

604 + image_attention_mask[_bs, :B_ + 1, 0::2, 0::2].sum().item(

605 + )) + (useful_height +

606 + 1) + base_feat_height // base_feat_height_reduction

607 + else:

608 + temp_sub_GN = self.sub_GN.repeat(

609 + 1, h * base_feat_height // base_feat_height_reduction, 1,

610 + 1)

611 + temp_len = int((h * w + 1) * self.num_img_tokens + 1 +

612 + (h + 1) * base_feat_height //

613 + base_feat_height_reduction)

614 +

615 + sub_img = torch.cat([sub_img, temp_sub_GN], dim=2).reshape(

616 + 1, -1,

617 + base_feat_height_reduction * base_feat_height_reduction * C)

618 + # (1, num_img_tokens, 1024*4)

619 +

620 + # glb + sub

621 + if self.hd_transform_order == 'glb_sub':

622 + output_imgs.append(

623 + torch.cat([glb_img, self.glb_GN, sub_img], dim=1))

624 + elif self.hd_transform_order == 'sub_glb':

625 + output_imgs.append(

626 + torch.cat([sub_img, self.glb_GN, glb_img], dim=1))

627 + else:

628 + raise NotImplementedError(

629 + f'hd_transform_order = {self.hd_transform_order}, "\

630 + "not implemented')

631 +

632 + #temp_len = int((h*w+1)*144 + 1 + (h+1)*12)

633 + assert temp_len == output_imgs[-1].shape[

634 + 1], f'temp_len: {temp_len}, output_imgs[-1].shape[1]: "\

635 + "{output_imgs[-1].shape[1]}'

636 +

637 + output_len.append(temp_len)

638 +

639 + img_set_tensor = []

640 + for _output_img in output_imgs:

641 + img_feature_proj = self.img_projection(

642 + _output_img.to(target_device).to(target_dtype))

643 + img_set_tensor.append(img_feature_proj)

644 +

645 + return img_set_tensor

646 +

647 +

648 +class Phi4MMAudioFeatureInputs(TypedDict):

649 + type: Literal["audio_features"]

650 + data: Tuple[NestedTensors]

651 + """Shape: `((batch_size, num_audios, 80, M), )"""

652 +

653 +

654 +class Phi4MMAudioEmbeddingInputs(TypedDict):

655 + type: Literal["audio_embeds"]

656 + data: NestedTensors

657 + """Shape: `(batch_size, num_audios, audio_feature_size, hidden_size)"""

658 +

659 +

660 +Phi4MMAudioInputs = Union[Phi4MMAudioFeatureInputs, Phi4MMAudioEmbeddingInputs]

661 +

662 +

663 +def speechlib_mel(sample_rate, n_fft, n_mels, fmin=None, fmax=None):

664 + """Create a Mel filter-bank the same as SpeechLib FbankFC.

665 +

666 + Args:

667 + sample_rate (int): Sample rate in Hz. number > 0 [scalar]

668 + n_fft (int): FFT size. int > 0 [scalar]

669 + n_mel (int): Mel filter size. int > 0 [scalar]

670 + fmin (float): lowest frequency (in Hz). If None use 0.0.

671 + float >= 0 [scalar]

672 + fmax: highest frequency (in Hz). If None use sample_rate / 2.

673 + float >= 0 [scalar]

674 +

675 + Returns

676 + out (numpy.ndarray): Mel transform matrix

677 + [shape=(n_mels, 1 + n_fft/2)]

678 + """

679 +

680 + bank_width = int(n_fft // 2 + 1)

681 + if fmax is None:

682 + fmax = sample_rate / 2

683 + if fmin is None:

684 + fmin = 0

685 + assert fmin >= 0, "fmin cannot be negative"

686 + assert (fmin < fmax <=

687 + sample_rate / 2), "fmax must be between (fmin, samplerate / 2]"

688 +

689 + def mel(f):

690 + return 1127.0 * np.log(1.0 + f / 700.0)

691 +

692 + def bin2mel(fft_bin):

693 + return 1127.0 * np.log(1.0 + fft_bin * sample_rate / (n_fft * 700.0))

694 +

695 + def f2bin(f):

696 + return int((f * n_fft / sample_rate) + 0.5)

697 +

698 + # Spec 1: FFT bin range [f2bin(fmin) + 1, f2bin(fmax) - 1]

699 + klo = f2bin(fmin) + 1

700 + khi = f2bin(fmax)

701 +

702 + khi = max(khi, klo)

703 +

704 + # Spec 2: SpeechLib uses triangles in Mel space

705 + mlo = mel(fmin)

706 + mhi = mel(fmax)

707 + m_centers = np.linspace(mlo, mhi, n_mels + 2)

708 + ms = (mhi - mlo) / (n_mels + 1)

709 +

710 + matrix = np.zeros((n_mels, bank_width), dtype=np.float32)

711 + for m in range(0, n_mels):

712 + left = m_centers[m]

713 + center = m_centers[m + 1]

714 + right = m_centers[m + 2]

715 + for fft_bin in range(klo, khi):

716 + mbin = bin2mel(fft_bin)

717 + if left < mbin < right:

718 + matrix[m, fft_bin] = 1.0 - abs(center - mbin) / ms

719 +

720 + return matrix

721 +

722 +

723 +class LogFbankProcessor:

724 +

725 + def __init__(self):

726 +

727 + self._eightk_method = "fillzero"

728 + self._mel = speechlib_mel(16000, 512, 80, fmin=None, fmax=7690).T

729 +

730 + self._hamming400 = np.hamming(400) # for 16k audio

731 + self._hamming200 = np.hamming(200) # for 8k audio

732 +

733 + def extract_spectrogram(self, wav, fs):

734 + """Extract spectrogram features from waveform.

735 + Args:

736 + wav (1D array): waveform of the input

737 + fs (int): sampling rate of the waveform, 16000 or 8000.

738 + If fs=8000, the waveform will be resampled to 16000Hz.

739 + Output:

740 + log_fbank (2D array): a TxD matrix of log Mel filterbank features.

741 + D=80, and T is the number of frames.

742 + """

743 + if wav.ndim > 1:

744 + wav = np.squeeze(wav)

745 +

746 + # by default, we extract the mean if stereo

747 + if len(wav.shape) == 2:

748 + wav = wav.mean(1)

749 +

750 + # Resample to 16000 or 8000 if needed

751 + if fs > 16000:

752 + wav = scipy.signal.resample_poly(wav, 1, fs // 16000)

753 + fs = 16000

754 + elif 8000 < fs < 16000:

755 + wav = scipy.signal.resample_poly(wav, 1, fs // 8000)

756 + fs = 8000

757 + elif fs < 8000:

758 + raise RuntimeError(f"Unsupported sample rate {fs}")

759 +

760 + if fs == 8000:

761 + if self._eightk_method == "resample":

762 + # Input audio is 8 kHz. Convert to 16 kHz before feature

763 + # extraction

764 + wav = scipy.signal.resample_poly(wav, 2, 1)

765 + fs = 16000

766 + # Do nothing here for fillzero method

767 + elif fs != 16000:

768 + # Input audio is not a supported sample rate.

769 + raise RuntimeError(

770 + f"Input data using an unsupported sample rate: {fs}")

771 +

772 + preemphasis = 0.97

773 +

774 + if fs == 8000:

775 + n_fft = 256

776 + win_length = 200

777 + hop_length = 80

778 + fft_window = self._hamming200

779 + elif fs == 16000:

780 + n_fft = 512

781 + win_length = 400

782 + hop_length = 160

783 + fft_window = self._hamming400

784 +

785 + # Spec 1: SpeechLib cut remaining sample insufficient for a hop

786 + n_batch = (wav.shape[0] - win_length) // hop_length + 1

787 + # Here we don't use stride_tricks since the input array may not satisfy

788 + # memory layout requirement and we need writeable output

789 + # Here we only use list of views before copy to destination

790 + # so it is more efficient than broadcasting

791 + y_frames = np.array(

792 + [

793 + wav[_stride:_stride + win_length]

794 + for _stride in range(0, hop_length * n_batch, hop_length)

795 + ],

796 + dtype=np.float32,

797 + )

798 +

799 + # Spec 2: SpeechLib applies preemphasis within each batch

800 + y_frames_prev = np.roll(y_frames, 1, axis=1)

801 + y_frames_prev[:, 0] = y_frames_prev[:, 1]

802 + y_frames = (y_frames - preemphasis * y_frames_prev) * 32768

803 +

804 + S = np.fft.rfft(fft_window * y_frames, n=n_fft,

805 + axis=1).astype(np.complex64)

806 +

807 + if fs == 8000:

808 + # Need to pad the output to look like 16 kHz data but with zeros in

809 + # the 4 to 8 kHz bins.

810 + frames, bins = S.shape

811 + padarray = np.zeros((frames, bins))

812 + S = np.concatenate((S[:, 0:-1], padarray),

813 + axis=1) # Nyquist bin gets set to zero

814 +

815 + spec = np.abs(S).astype(np.float32)

816 + return spec

817 +

818 + def extract_features(self, wav, fs):

819 + """Extract log filterbank features from waveform.

820 + Args:

821 + wav (1D array): waveform of the input

822 + fs (int): sampling rate of the waveform, 16000 or 8000.

823 + If fs=8000, the waveform will be resampled to 16000Hz.

824 + Output:

825 + log_fbank (2D array): a TxD matrix of log Mel filterbank features.

826 + D=80, and T is the number of frames.

827 + """

828 + spec = self.extract_spectrogram(wav, fs)

829 + spec_power = spec**2

830 +

831 + fbank_power = np.clip(spec_power.dot(self._mel), 1.0, None)

832 + log_fbank = np.log(fbank_power).astype(np.float32)

833 +

834 + return log_fbank

835 +

836 +

837 +@lru_cache

838 +def audio_feature_extractor() -> LogFbankProcessor:

839 + # Creates an instance of the audio processor, needed to extract the

840 + # the audio features from the sound file

841 + # LRU cache ensures that we only make one copy

842 + return LogFbankProcessor()

843 +

844 +

845 +def _compute_num_image_tokens(image, dynamic_hd_size, vit_image_size,

846 + vit_patch_size, token_compression_factor):

847 + """

848 + compute the number of tokens an image is expected to take up considering

849 + the image encoder architecture and exclude output features containing

850 + only padding pixels

851 +

852 + for siglip, vit_image_size=448, vit_patch_size=14, so output will be

853 + 32x32 feature map

854 + NOTE right now, Phi4MM uses hard-coded token_compression_factor=2

855 + """

856 + assert vit_image_size % vit_patch_size == 0, \

857 + "vit_image_size must be divisible by vit_patch_size"

858 + assert vit_image_size // vit_patch_size % token_compression_factor == 0, \

859 + "vit_image_size // vit_patch_size must be divisible by "\

860 + "token_compression_factor"

861 +

862 + target_aspect_ratio, target_height, target_width = (

863 + _find_target_aspect_ratio(image,

864 + vit_image_size,

865 + dynamic_hd_size,

866 + min_num=1))

867 + assert target_aspect_ratio[

868 + 0] * vit_image_size == target_width, \

869 + f"{target_aspect_ratio[0]} * {vit_image_size} != {target_width}"

870 + assert target_aspect_ratio[

871 + 1] * vit_image_size == target_height, \

872 + f"{target_aspect_ratio[1]} * {vit_image_size} != {target_height}"

873 + assert (target_height % vit_image_size == 0

874 + and target_width % vit_image_size == 0)

875 +

876 + padding_height, padding_width = _get_padding_size(image, target_height,

877 + target_width)

878 + assert padding_width == 0 or padding_height == 0, \

879 + "padding_width or padding_height must be 0"

880 +

881 + target_feat_width = target_width // vit_patch_size

882 + target_feat_height = target_height // vit_patch_size

883 + if padding_width >= vit_patch_size:

884 + assert padding_height == 0, "padding_height not 0"

885 + non_pad_feat_width = target_feat_width - math.floor(

886 + padding_width / vit_patch_size)

887 + non_pad_feat_height = target_feat_height

888 + elif padding_height >= vit_patch_size:

889 + assert padding_width == 0, "padding_width not 0"

890 + non_pad_feat_height = target_feat_height - math.floor(

891 + padding_height / vit_patch_size)

892 + non_pad_feat_width = target_feat_width

893 + else:

894 + # small padding shorter than a vit patch

895 + non_pad_feat_width = target_feat_width

896 + non_pad_feat_height = target_feat_height

897 +

898 + feat_width = non_pad_feat_width // token_compression_factor

899 + feat_height = non_pad_feat_height // token_compression_factor

900 + # NOTE it's possible that the non-padding feature is not divisible

901 + if non_pad_feat_width % token_compression_factor != 0:

902 + feat_width += 1

903 + if non_pad_feat_height % token_compression_factor != 0:

904 + feat_height += 1

905 + num_hd_patch_tokens = feat_width * feat_height

906 + num_hd_newline_tokens = feat_height

907 + vit_feature_size = vit_image_size // vit_patch_size

908 + num_global_image_tokens = (vit_feature_size // token_compression_factor)**2

909 + num_sep_tokens = 1

910 + num_global_image_newline_tokens = \

911 + vit_feature_size // token_compression_factor

912 +

913 + return (num_global_image_tokens + num_sep_tokens + num_hd_patch_tokens +

914 + num_hd_newline_tokens + num_global_image_newline_tokens)

915 +

916 +

917 +def compute_logfbank_output_size(wav_length: int, fs: int) -> Tuple[int, int]:

918 + """

919 + Compute the output size of the `extract_features` method.

920 +

921 + Args:

922 + wav_length (int): Length of the input waveform in samples.

923 + fs (int): Sampling rate of the waveform, either 16000 or 8000.

924 +

925 + Returns:

926 + tuple (int, int): Output size as (T, D), where:

927 + T: Number of time frames.

928 + D: Number of Mel filterbank bins (80).

929 + """

930 +

931 + # Resample to 16000 or 8000 if needed

932 + if fs > 16000:

933 + wav_length //= fs // 16000

934 + fs = 16000

935 + elif 8000 <= fs < 16000:

936 + # We'll resample to 16K from 8K

937 + wav_length *= 2

938 + fs = 16000

939 + elif fs < 8000:

940 + raise RuntimeError(f"Unsupported sample rate {fs}")

941 +

942 + # Spectrogram parameters for 16 kHz

943 + win_length = 400 # Frame length in samples

944 + hop_length = 160 # Frame shift in samples

945 + mel_bins = 80 # Number of mel filterbank bins

946 +

947 + # Calculate number of frames (T)

948 + T = (wav_length - win_length) // hop_length + 1

949 + if T < 1:

950 + raise ValueError("Waveform too short for given parameters.")

951 +

952 + # Return time frames (T) and mel bins (D)

953 + return T, mel_bins

954 +

955 +

956 +def _get_audio_embed_sizes(audios, ctx: InputContext):

957 + """

958 + Get the audio embedding sizes for each audio file.

959 +

960 + Args:

961 + audios (List[Tuple[np.ndarray, int]]): List of audio files as tuples of

962 + waveform and sample rate.

963 + ctx (InputContext): Input context.

964 +

965 + Returns:

966 + List[int]: List of audio embedding sizes.

967 + """

968 + audio_embed_sizes = []

969 + for audio in audios:

970 + audio_data, sf = audio

971 + audio_frames, _ = compute_logfbank_output_size(len(audio_data), sf)

972 + audio_embed_size = _compute_audio_embed_size(ctx.get_hf_config(),

973 + audio_frames)

974 + audio_embed_sizes.append(audio_embed_size)

975 + return audio_embed_sizes

976 +

977 +

978 +def _get_audio_id_to_input_ids(audios, ctx: InputContext, prompt_str=""):

979 + """

980 + The following will search for `<|audio_{idx}|>` tokens and

981 + return a mapping of audio placeholder tokens to audio placeholder token ids

982 + based on the size of the audio embeddings.

983 +

984 + Args:

985 + audios (List[Tuple[np.ndarray, int]]): List of audio files as tuples of

986 + waveform and sample rate.

987 + ctx (InputContext): Input context.

988 + prompt_str (str): The prompt string.

989 +

990 + Returns:

991 + Dict[str, List[int]]: Mapping of audio placeholder tokens to audio

992 + placeholder token ids.

993 +

994 + """

995 + if len(audios) == 0:

996 + return {}

997 +

998 + audio_embed_sizes = _get_audio_embed_sizes(audios, ctx)

999 + audio_ids = re.findall(AUDIO_TOKEN_PATTERN, prompt_str)

1000 + audio_ids = [int(audio_id) for audio_id in audio_ids]

1001 + assert len(audio_ids) == len(

1002 + audio_embed_sizes

1003 + ), "Number of audio tokens and audio features do not match"

1004 + assert tuple(audio_ids) == tuple(range(1,

1005 + len(audio_ids) +

1006 + 1)), "Audio ids are not in order!"

1007 + audio_id_to_input_ids = {

1008 + f"<|audio_{audio_id}|>":

1009 + [_AUDIO_PLACEHOLDER_TOKEN_ID] * audio_embed_size

1010 + for audio_id, audio_embed_size in zip(audio_ids, audio_embed_sizes)

1011 + }

1012 +

1013 + return audio_id_to_input_ids

1014 +

1015 +

1016 +def _count_image_tokens(images, ctx: InputContext):

1017 + hf_config = ctx.get_hf_config()

1018 + vision_encoder_name = hf_config.img_processor

1019 + if vision_encoder_name is None:

1020 + vision_encoder_name = SIGLIP_NAME

1021 + prepro_config = VISION_ENCODER_TO_PROCESSING_CONFIG[vision_encoder_name]

1022 + dynamic_hd_size = prepro_config['dynamic_hd']

1023 + vit_image_size = prepro_config['vit_image_size']

1024 + vit_patch_size = prepro_config['vit_patch_size']

1025 + token_compression_factor = prepro_config['token_compression_factor']

1026 +

1027 + image_token_counts = [

1028 + _compute_num_image_tokens(image, dynamic_hd_size, vit_image_size,

1029 + vit_patch_size, token_compression_factor)

1030 + for image in images

1031 + ]

1032 + return image_token_counts

1033 +

1034 +

1035 +def _get_image_id_to_input_ids(images, prompt, ctx: InputContext):

1036 + if len(images) == 0:

1037 + return {}

1038 +

1039 + image_ids = re.findall(IMAGE_TOKEN_PATTERN, prompt)

1040 + image_ids = [int(image_id) for image_id in image_ids]

1041 + assert len(image_ids) == len(

1042 + set(image_ids)), "Duplicate image tokens in prompt"

1043 + assert len(images) == len(

1044 + image_ids), "Number of images and image tokens in prompt do not match"

1045 +

1046 + # NOTE the following assertion is not strictly necessary

1047 + assert tuple(image_ids) == tuple(range(1,

1048 + len(image_ids) +

1049 + 1)), "Image ids are not in order"

1050 +

1051 + image_token_counts = _count_image_tokens(images, ctx)

1052 + image_id_to_input_ids = {

1053 + f"<|image_{image_id}|>": [_IMAGE_PLACEHOLDER_TOKEN_ID] * num_tokens

1054 + for image_id, num_tokens in zip(image_ids, image_token_counts)

1055 + }

1056 + return image_id_to_input_ids

1057 +

1058 +

1059 +def input_processor_for_phi4mm(ctx: InputContext,

1060 + inputs: DecoderOnlyInputs) -> TokenInputs:

1061 + """

1062 + Implements the input processor, which transforms the input prompt ids

1063 + to include the audio placeholder token. This will become the `input_ids`

1064 + in `forward` for the model.

1065 +

1066 + Args:

1067 + ctx (InputContext): Input context.

1068 + inputs (DecoderOnlyInputs): The inputs (e.g. prompt, prompt_token_ids)

1069 + to process.

1070 +

1071 + Returns:

1072 + TokenInputs: Processed inputs

1073 + """

1074 + multi_modal_data = inputs.get("multi_modal_data")

1075 + if (multi_modal_data is None or

1076 + ("audio" not in multi_modal_data and "image" not in multi_modal_data)):

1077 + # pure text input, so no need to do pre-processing

1078 + return inputs

1079 +

1080 + prompt_str = inputs.get("prompt")

1081 + prompt_token_ids = inputs.get("prompt_token_ids")

1082 + # for offline_inference, we will get str input and we parse MM special

1083 + # tokens from it

1084 + # (ignore prompt_token_ids)

1085 + # for OAI server, we will get prompt_token_ids, where MM special tokens

1086 + # are already parsed

1087 +

1088 + if 'audio' in multi_modal_data:

1089 + audios = multi_modal_data["audio"]

1090 +

1091 + if not isinstance(audios, list):

1092 + audios = [audios]

1093 + if prompt_str is not None:

1094 + audio_id_to_input_ids = _get_audio_id_to_input_ids(

1095 + audios, ctx, prompt_str=prompt_str)

1096 + audio_embed_sizes = []

1097 + elif prompt_token_ids is not None:

1098 + audio_id_to_input_ids = {}

1099 + audio_embed_sizes = _get_audio_embed_sizes(audios, ctx)

1100 + else:

1101 + audio_id_to_input_ids = {}

1102 + audio_embed_sizes = []

1103 +

1104 + if 'image' in multi_modal_data:

1105 + # PIL Image or list of PIL Images

1106 + images = multi_modal_data["image"]

1107 + if not isinstance(images, list):

1108 + images = [images]

1109 + if prompt_str is not None:

1110 + image_id_to_input_ids = _get_image_id_to_input_ids(

1111 + images, prompt_str, ctx)

1112 + image_token_counts = []

1113 + elif prompt_token_ids is not None:

1114 + image_id_to_input_ids = {}

1115 + image_token_counts = _count_image_tokens(images, ctx)

1116 + else:

1117 + image_id_to_input_ids = {}

1118 + image_token_counts = []

1119 +

1120 + # Handle the case where the prompt is a string and we need to manually

1121 + # tokenize it.

1122 + # In this case, the `audio_id_to_input_ids` dict will be mapping from

1123 + # an audio placeholder

1124 + # string (e.g. `<|audio_1|>`) to the audio placeholder tokens for the

1125 + # given audio length.

1126 + if prompt_str:

1127 + pattern = r"(<\|image_\d+\|>|<\|audio_\d+\|>)"

1128 + prompt_chunk_strings = re.split(pattern, prompt_str)

1129 + prompt_chunk_strings = [s for s in prompt_chunk_strings if s != ""]

1130 +

1131 + # Create the new input_ids with the placeholder image and audio

1132 + # tokens inserted

1133 + tokenizer = cached_tokenizer_from_config(ctx.model_config)

1134 + input_ids = []

1135 + has_imag, has_audio, has_user_text_input = False, False, False

1136 + for prompt_chunk_string in prompt_chunk_strings:

1137 + if re.match(IMAGE_TOKEN_PATTERN, prompt_chunk_string):

1138 + input_ids.extend(image_id_to_input_ids[prompt_chunk_string])

1139 + has_imag = True

1140 + elif re.match(AUDIO_TOKEN_PATTERN, prompt_chunk_string):

1141 + input_ids.extend(audio_id_to_input_ids[prompt_chunk_string])

1142 + has_audio = True

1143 + else:

1144 + curr_token_ids = tokenizer(prompt_chunk_string).input_ids

1145 + if not has_user_text_input:

1146 + for token_id in curr_token_ids:

1147 + if token_id not in NON_USER_INPUT_TOKENS:

1148 + has_user_text_input = True

1149 + break

1150 + input_ids.extend(curr_token_ids)

1151 + if has_audio and has_imag and has_user_text_input:

1152 + raise ValueError(

1153 + "Phi4MMForCausalLM does not support text + audio + image" +

1154 + " inputs in the same prompt")

1155 + # Handle the case where the prompt is already tokenized

1156 + else:

1157 + assert prompt_token_ids is not None, \

1158 + "If string prompt isn't provided, prompt_token_ids must be"

1159 +

1160 + i = 0

1161 + input_ids = prompt_token_ids

1162 + # only needed for later assertion

1163 + img_cnt, audio_cnt, user_text_input_cnt = 0, 0, 0

1164 + image_token_count_iter = iter(image_token_counts)

1165 + audio_embed_size_iter = iter(audio_embed_sizes)

1166 + while i < len(input_ids):

1167 + token_id = input_ids[i]

1168 + if token_id == _AUDIO_PLACEHOLDER_TOKEN_ID:

1169 + token_count = next(audio_embed_size_iter)

1170 + audio_cnt += 1

1171 + elif token_id == _IMAGE_PLACEHOLDER_TOKEN_ID:

1172 + token_count = next(image_token_count_iter)

1173 + img_cnt += 1

1174 + else:

1175 + user_text_input_cnt += 1 if token_id not in \

1176 + NON_USER_INPUT_TOKENS else 0

1177 + i += 1

1178 + continue

1179 + tokens = [token_id] * token_count

1180 + input_ids = input_ids[:i] + tokens + input_ids[i + 1:]

1181 + i += token_count

1182 +

1183 + if audio_cnt > 0 and img_cnt > 0 and user_text_input_cnt > 0:

1184 + raise ValueError(

1185 + "Phi4MMForCausalLM does not support text + audio + image" +

1186 + " inputs in the same prompt")

1187 + # If the below assertion fails, it might be that input pure-text

1188 + # messages contain image/audio special tokens literally

1189 + # (<|endoftext10|>, <|endoftext11|>).

1190 + assert (img_cnt == len(image_token_counts)), (

1191 + f"Number of image tokens in prompt_token_ids ({img_cnt}) "

1192 + f"does not match number of images ({len(image_token_counts)})")

1193 + assert (audio_cnt == len(audio_embed_sizes)), (

1194 + f"Number of audio tokens in prompt_token_ids ({audio_cnt}) "

1195 + f"does not match number of audios ({len(audio_embed_sizes)})")

1196 +

1197 + # NOTE: Create a defensive copy of the original inputs

1198 + return token_inputs(

1199 + prompt_token_ids=input_ids,

1200 + prompt=prompt_str,

1201 + multi_modal_data=multi_modal_data,

1202 + )

1203 +

1204 +

1205 +def _compute_audio_embed_size(hf_config, audio_frames):

1206 + """

1207 + Compute the audio embedding size based on the audio frames and

1208 + compression rate.

1209 + """

1210 + compression_rate = hf_config.embd_layer['audio_embd_layer'][

1211 + 'compression_rate']

1212 + # NOTE: this is a hard-coded value but might be configurable in the future

1213 + qformer_compression_rate = 1

1214 + integer = audio_frames // compression_rate

1215 + remainder = audio_frames % compression_rate

1216 +

1217 + result = integer if remainder == 0 else integer + 1

1218 +

1219 + integer = result // qformer_compression_rate

1220 + remainder = result % qformer_compression_rate

1221 + result = integer if remainder == 0 else integer + 1 # qformer compression

1222 +

1223 + return result

1224 +

1225 +

1226 +def get_max_phi4mm_audio_tokens(ctx: InputContext) -> int:

1227 + return 10000

1228 +

1229 +

1230 +def dummy_audio_for_phi4mm(audio_count: int) -> dict:

1231 + """

1232 + Create dummy audio data for the Phi4MM model, which is used for profiling.

1233 +

1234 + Args:

1235 + audio_count (int): Number of audio samples.

1236 +

1237 + Returns:

1238 + dict: Dummy audio data.

1239 + """

1240 + dummy_audio = np.full((_AUDIO_MAX_SOUNDFILE_SIZE, ), 0.0)

1241 + return [(dummy_audio, DUMMY_SAMPLING_FREQUENCY)] * audio_count

1242 +

1243 +

1244 +def dummy_image_for_phi4mm(width: int, height: int):

1245 + image = Image.new('RGB', (width, height), color='black')

1246 + return image

1247 +

1248 +

1249 +def dummy_data_for_phi4mm(ctx: InputContext, seq_len: int,

1250 + mm_counts: Mapping[str, int]) -> DummyData:

1251 + """

1252 + Create dummy sequence (input_ids) and audio data for the Phi4MM model,

1253 + which is used for profiling.

1254 +

1255 + In this case, the sequence data is a bunch of 0s with a number of audio

1256 + tokens that correspond to the audio embed size of the

1257 + _AUDIO_MAX_SOUNDFILE_SIZE.

1258 +

1259 + Args:

1260 + ctx (InputContext): Input context.

1261 + seq_len (int): Length of the sequence.

1262 + mm_counts (Mapping[str, int]): Multi-modal counts.

1263 +

1264 + Returns:

1265 + Tuple: Dummy sequence data and dummy audio data.

1266 + """

1267 + audio_count = mm_counts["audio"]

1268 + audio_frames, _ = compute_logfbank_output_size(_AUDIO_MAX_SOUNDFILE_SIZE,

1269 + DUMMY_SAMPLING_FREQUENCY)

1270 + audio_feature_size = _compute_audio_embed_size(ctx.get_hf_config(),

1271 + audio_frames)

1272 +

1273 + image_count = mm_counts["image"]

1274 + dummy_image = get_max_dummy_image(ctx)

1275 + max_image_tokens = get_max_phi4mm_image_tokens(ctx)

1276 + total_image_tokens = image_count * max_image_tokens

1277 +

1278 + if seq_len - audio_feature_size * audio_count - total_image_tokens < 0:

1279 + raise RuntimeError(

1280 + f"Phi4MM cannot process {audio_count} audios and {image_count}"

1281 + f"images in a prompt, please increase max_model_len to be at"

1282 + f" larger than "

1283 + f"{audio_feature_size * audio_count + total_image_tokens}"

1284 + " or reduce audio/image limit by --limit-mm-per-prompt.")

1285 +

1286 + if audio_feature_size * audio_count > total_image_tokens:

1287 + seq_data = SequenceData.from_prompt_token_counts(

1288 + (_AUDIO_PLACEHOLDER_TOKEN_ID, audio_feature_size * audio_count),

1289 + (0, seq_len - audio_feature_size * audio_count),

1290 + )

1291 + mm_data = {

1292 + "audio": dummy_audio_for_phi4mm(audio_count),

1293 + }

1294 + else:

1295 + seq_data = SequenceData.from_prompt_token_counts(

1296 + (_IMAGE_PLACEHOLDER_TOKEN_ID, total_image_tokens),

1297 + (0, seq_len - total_image_tokens),

1298 + )

1299 + mm_data = {

1300 + "image": [dummy_image] * image_count,

1301 + }

1302 + return DummyData(seq_data, mm_data)

1303 +

1304 +

1305 +def input_mapper_for_phi4mm_audio(ctx: InputContext,

1306 + data: object) -> MultiModalInputs:

1307 + """

1308 + This function is used to create the MultiModalInputs for the Phi4MM

1309 + (audio) model.

1310 + Specifically, for audio, we extract the audio features from the sound

1311 + file and create pairs of audio features and audio embed lengths (the

1312 + latter of which is used to repeat the audio placeholder token in the

1313 + input prompt IDs).

1314 + These pairs are used, downstream, in `_audio_features_to_embeddings`

1315 + (via `_process_audio_input`).

1316 +

1317 + Note that the incoming audio data (each entry in `data`) is a tuple of

1318 + the audio data and the sampling frequency (e.g. from soundfile.read).

1319 +

1320 + Args:

1321 + ctx (InputContext): Input context.

1322 + data (object): Audio data.

1323 +

1324 + Returns:

1325 + MultiModalInputs: Multi-modal inputs.

1326 + """

1327 + if not isinstance(data, list):

1328 + data = [data]

1329 +

1330 + if len(data) == 0:

1331 + return MultiModalInputs()

1332 +

1333 + audio_features = []

1334 + for audio_input in data:

1335 + if not isinstance(audio_input, tuple):

1336 + raise NotImplementedError(

1337 + f"Unsupported data type: {type(audio_input)}")

1338 +

1339 + audio, sf = audio_input

1340 + feature_extractor = audio_feature_extractor()

1341 + single_audio_features = feature_extractor.extract_features(audio, sf)

1342 + feat_stride = (1 if not hasattr(feature_extractor, "stride") else

1343 + feature_extractor.stride)

1344 + audio_frames = len(single_audio_features) * feat_stride

1345 + single_audio_embed_size = _compute_audio_embed_size(

1346 + ctx.get_hf_config(), audio_frames)

1347 + single_audio_feature_audio_len_pair = (

1348 + single_audio_features,

1349 + [single_audio_embed_size],

1350 + )

1351 + audio_features.append(single_audio_feature_audio_len_pair)

1352 + return MultiModalInputs({"audio_features": audio_features})

1353 +

1354 +

1355 +def input_mapper_for_phi4mm_image(ctx: InputContext, data: object):

1356 + if not isinstance(data, list):

1357 + data = [data]

1358 + # data: list of PIL images

1359 + if len(data) == 0:

1360 + return MultiModalInputs()

1361 + hf_config = ctx.get_hf_config()

1362 + vision_encoder_name = hf_config.img_processor

1363 + if vision_encoder_name is None:

1364 + vision_encoder_name = SIGLIP_NAME

1365 + prepro_config = VISION_ENCODER_TO_PROCESSING_CONFIG[vision_encoder_name]

1366 + dynamic_hd_size = prepro_config['dynamic_hd']

1367 + vit_image_size = prepro_config['vit_image_size']

1368 + vit_patch_size = prepro_config['vit_patch_size']

1369 +

1370 + image_input_dict = preprocess(data, dynamic_hd_size, vit_image_size,

1371 + vit_patch_size)

1372 + return MultiModalInputs({

1373 + "pixel_values":

1374 + image_input_dict["pixel_values"],

1375 + "image_sizes":

1376 + image_input_dict["image_sizes"],

1377 + "image_attention_mask":

1378 + image_input_dict["image_attention_mask"],

1379 + "num_img_tokens":

1380 + image_input_dict["num_img_tokens"],

1381 + })

1382 +

1383 +

1384 +def cat_with_pad(tensors, dim, padding_value=0):

1385 + """

1386 + cat along dim, while pad to max for all other dims

1387 + """

1388 + ndim = tensors[0].dim()

1389 + assert all(

1390 + t.dim() == ndim for t in

1391 + tensors[1:]), "All tensors must have the same number of dimensions"

1392 +

1393 + out_size = [max(t.shape[i] for t in tensors) for i in range(ndim)]

1394 + out_size[dim] = sum(t.shape[dim] for t in tensors)

1395 + output = tensors[0].new_full(out_size, padding_value)

1396 +

1397 + index = 0

1398 + for t in tensors:

1399 + # Create a slice list where every dimension except dim is full slice

1400 + slices = [slice(0, t.shape[d]) for d in range(ndim)]

1401 + # Update only the concat dimension slice

1402 + slices[dim] = slice(index, index + t.shape[dim])

1403 +

1404 + output[slices] = t

1405 + index += t.shape[dim]

1406 +

1407 + return output

1408 +

1409 +

1410 +@MULTIMODAL_REGISTRY.register_input_mapper("audio",

1411 + input_mapper_for_phi4mm_audio)

1412 +@MULTIMODAL_REGISTRY.register_input_mapper("image",

1413 + input_mapper_for_phi4mm_image)

1414 +@MULTIMODAL_REGISTRY.register_max_multimodal_tokens(

1415 + "audio", get_max_phi4mm_audio_tokens)

1416 +@MULTIMODAL_REGISTRY.register_max_multimodal_tokens(

1417 + "image", get_max_phi4mm_image_tokens)

1418 +@INPUT_REGISTRY.register_dummy_data(dummy_data_for_phi4mm)

1419 +@INPUT_REGISTRY.register_input_processor(input_processor_for_phi4mm)

1420 +class Phi4MMForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):

1421 + """

1422 + Implements the Phi-4-multimodal-instruct model in VLLM.

1423 + """

1424 + # LoRA specific attributes

1425 + packed_modules_mapping = {

1426 + "qkv_proj": [

1427 + "qkv_proj",

1428 + ],

1429 + "gate_up_proj": [

1430 + "gate_up_proj",

1431 + ],

1432 + }

1433 + supported_lora_modules = [

1434 + "qkv_proj", "o_proj", "gate_up_proj", "down_proj"

1435 + ]

1436 + # Phi4MMForCausalLM does not apply LoRA to the embedding layer.

1437 + embedding_modules = {}

1438 + embedding_padding_modules = []

1439 +

1440 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

1441 + super().__init__()

1442 + config = vllm_config.model_config.hf_config

1443 + multimodal_config = vllm_config.model_config.multimodal_config

1444 + assert multimodal_config, "multimodal_config is required"

1445 + quant_config = vllm_config.quant_config

1446 + lora_config = vllm_config.lora_config

1447 +

1448 + self.config = config

1449 + self.multimodal_config = multimodal_config

1450 + self.quant_config = quant_config

1451 + self.lora_config = lora_config

1452 +

1453 + # Tensor/Pipeline parallel not supported for now.

1454 + assert get_tensor_model_parallel_world_size(

1455 + ) == 1, "tensor parallel is not supported"

1456 + assert get_pp_group(

1457 + ).world_size == 1, "pipeline parallel is not supported"

1458 +

1459 + self.vision_encoder = Phi4MMImageEncoder(

1460 + config,

1461 + quant_config,

1462 + prefix="model.vision_embed_tokens",

1463 + model_dir=config._name_or_path)

1464 +

1465 + if isinstance(config.embd_layer["audio_embd_layer"], dict):

1466 + embedding_config = {

1467 + "embedding_cls":

1468 + config.embd_layer["audio_embd_layer"]["embedding_cls"],

1469 + **config.embd_layer["audio_embd_layer"],

1470 + }

1471 + else:

1472 + embedding_config = {

1473 + "embedding_cls": self.config.embd_layer["embedding_cls"]

1474 + }

1475 +

1476 + self.embed_tokens_extend = AudioEmbedding(config, **embedding_config)

1477 + self.model = LlamaModel(vllm_config=vllm_config,

1478 + prefix=maybe_prefix(prefix, "model"))

1479 +

1480 + self.unpadded_vocab_size = config.vocab_size

1481 + if lora_config:

1482 + self.unpadded_vocab_size += lora_config.lora_extra_vocab_size

1483 + self.lm_head = ParallelLMHead(

1484 + self.unpadded_vocab_size,

1485 + config.hidden_size,

1486 + org_num_embeddings=config.vocab_size,

1487 + padding_size=(

1488 + DEFAULT_VOCAB_PADDING_SIZE

1489 + # We need bigger padding if using lora for kernel

1490 + # compatibility

1491 + if not lora_config else lora_config.lora_vocab_padding_size),

1492 + quant_config=quant_config,

1493 + )

1494 + if config.tie_word_embeddings:

1495 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

1496 + logit_scale = getattr(config, "logit_scale", 1.0)

1497 + self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,

1498 + config.vocab_size, logit_scale)

1499 + self.sampler = Sampler()

1500 +

1501 + def _audio_features_to_embeddings(

1502 + self,

1503 + input_ids: torch.Tensor,

1504 + input_features: List[torch.Tensor],

1505 + audio_input_sizes: torch.Tensor,

1506 + audio_projection_mode: str,

1507 + ) -> torch.Tensor:

1508 + """

1509 + Convert audio features to embeddings, which are used as input to the

1510 + model (via `inputs_embeds`).

1511 +

1512 + Args:

1513 + input_ids (torch.Tensor): Input IDs (the prompt in this case).

1514 + input_features (list[torch.Tensor]): Input features (the audio

1515 + embeddings).

1516 + audio_input_sizes (list[torch.Tensor]): Audio input sizes (the

1517 + audio embed lengths to use for padding the audio placeholder token

1518 + in the input prompt IDs).

1519 + """

1520 + # The audio projection can either be a single linear or Sequential,

1521 + # so handle both cases

1522 + if isinstance(self.embed_tokens_extend.audio_projection,

1523 + nn.Sequential):

1524 + target_dtype = self.embed_tokens_extend.audio_projection[

1525 + 0].bias.dtype

1526 + else:

1527 + target_dtype = self.embed_tokens_extend.audio_projection.bias.dtype

1528 +

1529 + audio_input = [

1530 + input.unsqueeze(0).to(target_dtype) for input in input_features

1531 + ]

1532 + kwargs = {

1533 + "wte": self.model.embed_tokens,

1534 + 'audio_projection_mode': audio_projection_mode

1535 + }

1536 + audio_embeddings = self.embed_tokens_extend(input_ids, audio_input,

1537 + audio_input_sizes,

1538 + **kwargs)

1539 + audio_embeddings = audio_embeddings.to(target_dtype)

1540 + return audio_embeddings

1541 +

1542 + def _parse_and_validate_audio_input(

1543 + self, **kwargs: object) -> Optional[Phi4MMAudioInputs]:

1544 + """

1545 + Parse and validate the audio input to the model. This handles both

1546 + audio features and audio embeddings, but only the former is used for

1547 + now.

1548 +

1549 + Args:

1550 + kwargs (object): Keyword arguments.

1551 +

1552 + Returns:

1553 + Optional[Phi4MMAudioInputs]: Parsed and validated audio inputs.

1554 + """

1555 + audio_features = kwargs.pop("audio_features", None)

1556 + audio_embeds = kwargs.pop("audio_embeds", None)

1557 +

1558 + if audio_features is None and audio_embeds is None:

1559 + return None

1560 +

1561 + if audio_features is not None:

1562 + if not isinstance(audio_features, (torch.Tensor, list)):

1563 + raise ValueError("Incorrect type of audio features. "

1564 + f"Got type: {type(audio_features)}")

1565 +

1566 + return Phi4MMAudioFeatureInputs(type="audio_features",

1567 + data=audio_features)

1568 +

1569 + if audio_embeds is not None:

1570 + if not isinstance(audio_embeds, (torch.Tensor, list)):

1571 + raise ValueError("Incorrect type of audio embeds. "

1572 + f"Got type: {type(audio_embeds)}")

1573 +

1574 + return Phi4MMAudioEmbeddingInputs(type="audio_embeds",

1575 + data=audio_embeds)

1576 +

1577 + raise AssertionError("This line should be unreachable.")

1578 +

1579 + def _process_audio_input(self, input_ids: torch.Tensor,

1580 + audio_input: Phi4MMAudioInputs,

1581 + audio_projection_mode: str) -> NestedTensors:

1582 + """

1583 + Create the audio embeddings from the audio input, where the audio input

1584 + is pairs of audio features and audio embed lengths. The audio input is

1585 + created by `input_mapper_for_phi4mm_audio`.

1586 +

1587 + Args:

1588 + input_ids (torch.Tensor): Input IDs (the prompt in this case,

1589 + before the audio token replication).

1590 + audio_input (Phi4MMAudioInputs): Audio input.

1591 +

1592 + Returns:

1593 + NestedTensors: Audio embeddings

1594 + """

1595 + if audio_input["type"] == "audio_embeds":

1596 + return audio_input["data"]

1597 +

1598 + audio_features = audio_input["data"]

1599 + # (e.g. multiple examples) and the second dim is the multi-audio dim

1600 + # (e.g. multiple audios in the same example)

1601 + audio_feature = [i[0] for j in audio_features for i in j]

1602 + audio_feature_len = [i[1].item() for j in audio_features for i in j]

1603 + # Add the batch dim via `squeeze`

1604 +

1605 + return self._audio_features_to_embeddings(

1606 + input_ids.unsqueeze(0),

1607 + audio_feature,

1608 + audio_feature_len,

1609 + audio_projection_mode,

1610 + ).squeeze(0)

1611 +

1612 + def _parse_and_validate_image_input(self,

1613 + **kwargs: object) -> Optional[Dict]:

1614 + pixel_values: Optional[Dict] = kwargs.get("pixel_values")

1615 + if pixel_values is None:

1616 + return None

1617 +

1618 + image_sizes = kwargs.get("image_sizes")

1619 + image_attention_mask = kwargs.get("image_attention_mask")

1620 + num_img_tokens = kwargs.get("num_img_tokens")

1621 + assert image_sizes is not None and image_attention_mask is not None\

1622 + and num_img_tokens is not None, "Missing image inputs"

1623 +

1624 + if isinstance(pixel_values, list):

1625 + assert pixel_values[0].dim() == 5, "Incorrect image inputs"

1626 + # list len is batch_size.

1627 + # each tensor has dimension: num_img_per_example, num_hd_patches,

1628 + # channels, height, width.

1629 + # need to pad along num_hd_patches.

1630 + # mask size num_img_per_prompt, num_hd_patches, feat_h, heat_w.

1631 + pixel_values = cat_with_pad(pixel_values, dim=0)

1632 + elif isinstance(pixel_values, torch.Tensor):

1633 + # dimension: batch_size, num_img_per_example, num_hd_patches,

1634 + # channels, height, width.

1635 + # we flatten first 2 dims to make it a single large batch for

1636 + # SigLIP Encoder.

1637 + assert pixel_values.dim() == 6, "Incorrect image inputs"

1638 + pixel_values = pixel_values.flatten(0, 1)

1639 + else:

1640 + raise ValueError("Incorrect pixel_values inputs")

1641 +

1642 + if isinstance(image_attention_mask, list):

1643 + image_attention_mask = cat_with_pad(image_attention_mask, dim=0)

1644 + elif isinstance(image_attention_mask, torch.Tensor):

1645 + image_attention_mask = image_attention_mask.flatten(0, 1)

1646 + else:

1647 + raise ValueError("Incorrect image_attention_mask inputs")

1648 +

1649 + if isinstance(image_sizes, list):

1650 + image_sizes = torch.cat(image_sizes, dim=0)

1651 + elif isinstance(image_sizes, torch.Tensor):

1652 + image_sizes = image_sizes.flatten(0, 1)

1653 + else:

1654 + raise ValueError("Incorrect image_attention_mask inputs")

1655 +

1656 + if isinstance(num_img_tokens, list):

1657 + num_img_tokens = [

1658 + n for num_tensor in num_img_tokens

1659 + for n in num_tensor.tolist()

1660 + ]

1661 + elif isinstance(num_img_tokens, torch.Tensor):

1662 + num_img_tokens = num_img_tokens.flatten(0, 1).tolist()

1663 + else:

1664 + raise ValueError("Incorrect image_attention_mask inputs")

1665 +

1666 + return {

1667 + 'pixel_values': pixel_values,

1668 + 'image_sizes': image_sizes,

1669 + 'image_attention_mask': image_attention_mask,

1670 + 'num_img_tokens': num_img_tokens,

1671 + }

1672 +

1673 + def merge_image_features_to_inputs_embeds(

1674 + self,

1675 + input_ids: torch.Tensor,

1676 + inputs_embeds: torch.Tensor,

1677 + image_set_tensors: List[torch.Tensor],

1678 + ):

1679 + position_tuple = (input_ids == _IMAGE_PLACEHOLDER_TOKEN_ID).nonzero(

1680 + as_tuple=True)

1681 +

1682 + assert all([t.shape[0] == 1 for t in image_set_tensors

1683 + ]), 'img_set_tensor should have shape (1, N_tokens, C)'

1684 + # Shape: (merged_N_tokens, C)

1685 + image_set_tensor = torch.cat(image_set_tensors, dim=1).squeeze(0)

1686 + image_set_tensor = image_set_tensor.to(inputs_embeds.dtype).to(

1687 + inputs_embeds.device)

1688 + merged_embeds = inputs_embeds.index_put(

1689 + indices=position_tuple,

1690 + values=image_set_tensor,

1691 + accumulate=False,

1692 + )

1693 + return merged_embeds

1694 +

1695 + def load_weights(self, weights: Iterable[Tuple[str,

1696 + torch.Tensor]]) -> None:

1697 + weights = {name: weight for name, weight in weights}

1698 + adjusted_weights = {}

1699 +

1700 + for name, weight in weights.items():

1701 + # NOTE vision-speech tasks use a separate projection layer

1702 + audio_proj_4v = \

1703 + "model.embed_tokens_extend.audio_embed.audio_projection.vision"

1704 + if name.startswith(audio_proj_4v):

1705 + name = name.replace(

1706 + audio_proj_4v,

1707 + "embed_tokens_extend.audio_projection_for_vision")

1708 +

1709 + name = (name.replace(

1710 + "model.embed_tokens_extend.audio_embed."\

1711 + "audio_projection.speech.",

1712 + "embed_tokens_extend.audio_projection.",

1713 + ).replace(

1714 + "model.embed_tokens_extend.audio_embed.",

1715 + "embed_tokens_extend.",

1716 + ).replace("model.embed_tokens_extend.image_embed.",

1717 + "vision_encoder."))

1718 + # NOTE: this is deal with LoRA injection, where `base_layer`

1719 + # remains as the original layer in the model

1720 + if name.endswith(".base_layer.weight"):

1721 + name = name.replace(".base_layer.weight", ".weight")

1722 + adjusted_weights[name] = weight

1723 +

1724 + missing_keys, unexpected_keys = self.load_state_dict(adjusted_weights,

1725 + strict=False)

1726 + logger.debug("*** missing keys:")

1727 + for key in missing_keys:

1728 + logger.debug(key)

1729 + logger.debug("**** unexpected keys:")

1730 + for key in unexpected_keys:

1731 + logger.debug(key)

1732 +

1733 + def forward(

1734 + self,

1735 + input_ids: torch.Tensor,

1736 + positions: torch.Tensor,

1737 + intermediate_tensors: Optional[IntermediateTensors] = None,

1738 + **kwargs: object,

1739 + ) -> torch.Tensor:

1740 + if intermediate_tensors is not None:

1741 + input_ids = None

1742 + inputs_embeds = None

1743 + else:

1744 + # Each entry in this is a pair of audio_features and audio_embed

1745 + # lengths

1746 + audio_input = self._parse_and_validate_audio_input(**kwargs)

1747 + image_inputs = self._parse_and_validate_image_input(**kwargs)

1748 +

1749 + has_audio = audio_input is not None

1750 + has_image = image_inputs is not None

1751 +

1752 + if has_audio:

1753 + audio_projection_mode = 'vision' if has_image else 'speech'

1754 + inputs_embeds = self._process_audio_input(

1755 + input_ids, audio_input, audio_projection_mode)

1756 +

1757 + if has_image:

1758 + dtype = self.vision_encoder.img_processor.embeddings.\

1759 + patch_embedding.weight.dtype

1760 + pixel_values = image_inputs['pixel_values'].to(dtype)

1761 + image_sizes = image_inputs['image_sizes']

1762 + image_attention_mask = image_inputs['image_attention_mask']

1763 + image_set_tensors = self.vision_encoder(

1764 + pixel_values, image_sizes, image_attention_mask)

1765 + if not has_audio:

1766 + inputs_embeds = self.model.embed_tokens(input_ids)

1767 +

1768 + inputs_embeds = self.merge_image_features_to_inputs_embeds(

1769 + input_ids, inputs_embeds, image_set_tensors)

1770 +

1771 + if has_image or has_audio:

1772 + # multi-modal input, we have set inputs_embeds properly in

1773 + # previous steps

1774 + input_ids = None

1775 + else:

1776 + # text-only, we keep using original input_ids

1777 + inputs_embeds = None

1778 +

1779 + hidden_states = self.model(

1780 + input_ids,

1781 + positions,

1782 + intermediate_tensors,

1783 + inputs_embeds=inputs_embeds,

1784 + )

1785 +

1786 + return hidden_states

1787 +

1788 + def compute_logits(

1789 + self,

1790 + hidden_states: torch.Tensor,

1791 + sampling_metadata: SamplingMetadata,

1792 + ) -> Optional[torch.Tensor]:

1793 + logits = self.logits_processor(self.lm_head, hidden_states,

1794 + sampling_metadata)

1795 + return logits

1796 +

1797 + def sample(

1798 + self,

1799 + logits: torch.Tensor,

1800 + sampling_metadata: SamplingMetadata,

1801 + ) -> Optional[SamplerOutput]:

1802 + next_tokens = self.sampler(logits, sampling_metadata)

1803 + return next_tokens

vllm/model_executor/models/phi4mm_audio.py new

+1403

@@ -0,0 +1,1403 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# Copyright (c) Microsoft Corporation.

3 +# Licensed under the MIT license.

4 +# Code copied from Microsoft/MoE by Jacob Platin (jacobplatin@microsoft.com)

5 +# but implemented by the Phi-Speech team

6 +#!/usr/bin/env python3

7 +import abc

8 +import math

9 +from functools import partial

10 +from typing import Callable, Dict, List, Literal, Optional, Union

11 +

12 +import numpy as np

13 +import torch

14 +import torch.nn.functional as F

15 +from torch import Tensor, nn

16 +from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (

17 + CheckpointImpl, CheckpointWrapper, checkpoint_wrapper, offload_wrapper)

18 +from torch.distributed.fsdp.fully_sharded_data_parallel import (

19 + FullyShardedDataParallel)

20 +from torch.utils.checkpoint import checkpoint

21 +from transformers import PretrainedConfig

22 +

23 +from vllm.model_executor.models.phi4mm_utils import (

24 + AbsolutePositionalEncoding, ConvModule, FeedForward, MeanVarianceNormLayer,

25 + MultiHeadedAttention, NemoConvSubsampling, T5RelativeAttentionLogitBias,

26 + adaptive_enc_mask, attn_checkpointing, embedding_checkpoint_wrapper,

27 + get_offset, repeat, unfold_tensor, validate_checkpointing_config)

28 +

29 +_AUDIO_PLACEHOLDER_TOKEN_ID = 200011 # <|endoftext11|>

30 +

31 +

32 +def encoder_checkpoint_wrapper(

33 + activation_checkpointing: Union[str, Dict],

34 + layer_cls: type,

35 + idx: int = 0,

36 +) -> Callable:

37 + """return encoder activation checkpoint wrapper"""

38 + validate_checkpointing_config(activation_checkpointing)

39 +

40 + if isinstance(activation_checkpointing, str):

41 + if activation_checkpointing:

42 + if activation_checkpointing == "offload":

43 + return offload_wrapper

44 + return partial(checkpoint_wrapper)

45 + return lambda x: x

46 +

47 + if isinstance(activation_checkpointing, dict):

48 + target_layer_cls = activation_checkpointing.get(

49 + "module", "transformer")

50 + if target_layer_cls.lower() == "transformer":

51 + target_layer_cls = (

52 + "EncoderLayer",

53 + "ConformerEncoderLayer",

54 + )

55 + elif target_layer_cls.lower() == "attention":

56 + target_layer_cls = ("MultiHeadedAttention", "MultiHeadAttention")

57 + checkpointing_interval = activation_checkpointing.get("interval", 1)

58 + offloading = activation_checkpointing.get("offload", False)

59 + impl = (CheckpointImpl.REENTRANT if activation_checkpointing.get(

60 + "reentrant", True) else CheckpointImpl.NO_REENTRANT)

61 +

62 + if (idx % checkpointing_interval == 0

63 + and layer_cls.__name__ in target_layer_cls):

64 + if offloading:

65 + return offload_wrapper

66 + return partial(checkpoint_wrapper, checkpoint_impl=impl)

67 + return lambda x: x

68 +

69 + raise ValueError("Invalid activation_checkpointing config")

70 +

71 +

72 +class ConformerEncoderLayer(nn.Module):

73 + """ConformerEncoder Layer module.

74 + for more details see conformer paper:

75 + https://arxiv.org/abs/2005.08100

76 + This module implement the Conformer block layer.

77 +

78 + Args:

79 + d_model: int

80 + attention dim.

81 + ext_pw_out_channel: int

82 + if > 0, ext_pw_out_channel is a dim channel size

83 + for the last pointwise conv after swish activation.

84 + depthwise_seperable_out_channel: int

85 + if set different to 0, the number of

86 + depthwise_seperable_out_channel will be used as a

87 + channel_out of the second conv1d layer.

88 + otherwise, it equal to 0, the second conv1d layer is skipped.

89 + depthwise_multiplier: int

90 + number of input_dim channels duplication. this value

91 + will be used to compute the hidden channels of the Conv1D.

92 + n_head: int

93 + the number of heads for multihead attention module.

94 + d_ffn: int

95 + output size of the feed_forward blocks.

96 + ext_pw_kernel_size: int

97 + kernel size of the conv pointwise of the conformer.

98 + kernel_size: int

99 + kernel size.

100 + dropout_rate: float

101 + dropout rate.

102 + causal: bool, optional

103 + if set to True, convolution have no access

104 + to future frames. default False.

105 + batch_norm: bool, optional

106 + if set to True, apply batchnorm before activation

107 + in ConvModule layer of the conformer.

108 + default False

109 + activation: str, optional

110 + activation function name,

111 + one of ["relu", "swish", "sigmoid"],

112 + sigmoid activation is only used with "glu_in_fnn=True",

113 + default "relu".

114 + chunk_se: int, optional

115 + 0 for offline SE.

116 + 1 for streaming SE, where mean is computed

117 + by accumulated history until current chunk_se.

118 + 2 for streaming SE, where mean is computed

119 + by only the current chunk.

120 + default 0.

121 + chunk_size: int, optional

122 + chunk_size for cnn. default 18

123 + conv_activation: str, optional

124 + activation function used in ConvModule part

125 + of the conformer, default "relu".

126 + conv_glu_type: str, optional

127 + activation function used for the glu inside

128 + the ConvModule part of the conformer.

129 + default: "sigmoid".

130 + bias_in_glu: bool, optional

131 + if set to True, use additive bias in the weight module

132 + before GLU.

133 + linear_glu_in_convm: bool, optional

134 + if set to True, use GLULinear module,

135 + otherwise, used GLUPointWiseConv module.

136 + default to False.

137 + attention_innner_dim: int, optional

138 + if equal to -1, attention dim for linears k/q/v is

139 + equal to d_model. otherwise attention_innner_dim is used.

140 + default -1.

141 + attention_glu_type: str, optional

142 + activation function for glu used in the multihead attention,

143 + default "swish".

144 + activation_checkpointing: str, optional

145 + a dictionarry of {"module","interval","offload"}, where

146 + "module": str

147 + accept ["transformer", "attention"] to select

148 + which module should do activation checkpointing.

149 + "interval": int, default 1,

150 + interval of applying activation checkpointing,

151 + interval = 1 means that we apply checkpointing

152 + on every layer (if activation), otherwise,

153 + we apply it every x interval.

154 + "offload": bool, default False,

155 + if set to True, we offload activation to cpu and

156 + reload it during backward, otherwise,

157 + we recalculate activation in backward.

158 + default "".

159 + export: bool, optional

160 + if set to True, it remove the padding from convolutional layers

161 + and allow the onnx conversion for inference.

162 + default False.

163 + use_pt_scaled_dot_product_attention: bool, optional

164 + if set to True, use pytorch's scaled dot product attention

165 + implementation in training.

166 + attn_group_sizes: int, optional

167 + the number of groups to use for attention, default 1

168 + (Multi-Head Attention),

169 + 1 = typical Multi-Head Attention,

170 + 1 < attn_group_sizes < attention_heads = Grouped-Query Attention

171 + attn_group_sizes = attenion_heads = Multi-Query Attention

172 + """

173 +

174 + def __init__(

175 + self,

176 + d_model=512,

177 + ext_pw_out_channel=0,

178 + depthwise_seperable_out_channel=256,

179 + depthwise_multiplier=1,

180 + n_head=4,

181 + d_ffn=2048,

182 + ext_pw_kernel_size=1,

183 + kernel_size=3,

184 + dropout_rate=0.1,

185 + causal=False,

186 + batch_norm=False,

187 + activation="relu",

188 + chunk_se=0,

189 + chunk_size=18,

190 + conv_activation="relu",

191 + conv_glu_type="sigmoid",

192 + bias_in_glu=True,

193 + linear_glu_in_convm=False,

194 + attention_innner_dim=-1,

195 + attention_glu_type="swish",

196 + activation_checkpointing="",

197 + export=False,

198 + use_pt_scaled_dot_product_attention=False,

199 + attn_group_sizes: int = 1,

200 + ):

201 + super().__init__()

202 +

203 + self.feed_forward_in = FeedForward(

204 + d_model=d_model,

205 + d_inner=d_ffn,

206 + dropout_rate=dropout_rate,

207 + activation=activation,

208 + bias_in_glu=bias_in_glu,

209 + )

210 +

211 + self.self_attn = encoder_checkpoint_wrapper(

212 + activation_checkpointing,

213 + MultiHeadedAttention,

214 + )(MultiHeadedAttention(

215 + n_head,

216 + d_model,

217 + dropout_rate,

218 + attention_innner_dim,

219 + attention_glu_type,

220 + bias_in_glu,

221 + use_pt_scaled_dot_product_attention=

222 + use_pt_scaled_dot_product_attention,

223 + group_size=attn_group_sizes,

224 + ))

225 + self.conv = ConvModule(

226 + d_model,

227 + ext_pw_out_channel,

228 + depthwise_seperable_out_channel,

229 + ext_pw_kernel_size,

230 + kernel_size,

231 + depthwise_multiplier,

232 + dropout_rate,

233 + causal,

234 + batch_norm,

235 + chunk_se,

236 + chunk_size,

237 + conv_activation,

238 + conv_glu_type,

239 + bias_in_glu,

240 + linear_glu_in_convm,

241 + export=export,

242 + )

243 +

244 + self.feed_forward_out = FeedForward(

245 + d_model=d_model,

246 + d_inner=d_ffn,

247 + dropout_rate=dropout_rate,

248 + activation=activation,

249 + bias_in_glu=bias_in_glu,

250 + )

251 +

252 + self.layer_norm_att = nn.LayerNorm(d_model)

253 + self.layer_norm = nn.LayerNorm(d_model)

254 +

255 + def forward(

256 + self,

257 + x,

258 + pos_k,

259 + pos_v,

260 + mask,

261 + relative_attention_bias: Optional[Tensor] = None,

262 + ):

263 + """ConformerEncoder forward.

264 +

265 + Args:

266 + x: torch.Tensor

267 + input feature of shape (batch, max_time_in, size)

268 + pos_k: torch.Tensor

269 + positional key embedding.

270 + mask: torch.Tensor

271 + mask for x (batch, max_time_in)

272 + relative_attention_bias: Optional[torch.Tensor]

273 + bias added to attention logits w.r.t. relative positions

274 + (1, n_head, time1, time2)

275 + """

276 + x = x + 0.5 * self.feed_forward_in(x)

277 + norm_x = self.layer_norm_att(x)

278 +

279 + x = x + self.self_attn(

280 + norm_x,

281 + norm_x,

282 + norm_x,

283 + pos_k,

284 + pos_v,

285 + mask,

286 + relative_attention_bias=relative_attention_bias,

287 + )

288 + x = x + self.conv(x)

289 + x = x + 0.5 * self.feed_forward_out(x)

290 +

291 + out = self.layer_norm(x)

292 +

293 + return out, pos_k, pos_v, mask

294 +

295 +

296 +class TransformerEncoderBase(abc.ABC, nn.Module):

297 + """The Base class for Transformer based encoders

298 +

299 + Please set causal = True in streaming model

300 + Args:

301 + input_size: int

302 + input feature dimension.

303 + chunk_size: int, list(int)

304 + Number of frames for each chunk

305 + This variable can take 2 forms:

306 + int: Used for inference, or single chunk size training

307 + list(int) : Used only for variable chunk size training

308 + Some examples for the 2 cases:

309 + chunk_size = 12

310 + chunk_size = [6, 8, 12, 24]

311 + left_chunk: int, list(int)

312 + Number of chunks used for masking in streaming mode.

313 + This variable can take 2 forms:

314 + int: Used for inference, or single chunk size training

315 + list(int) : Used only for variable chunk size training. When

316 + chunk_size is a list, left_chunk must be a list with same length.

317 + Some examples for the 2 cases:

318 + left_chunk = 6

319 + left_chunk = [12, 9, 6, 3]

320 + attention_dim: int, optional

321 + attention dimension. default 256.

322 + attention_heads: int, optional

323 + the number of heads. default 4

324 + input_layer: str, optional

325 + input layer type before Conformer,

326 + one of ["linear", "conv2d", "custom", "vgg2l", "embed"],

327 + default "conv2d"

328 + cnn_out: int, optional

329 + the number of CNN channels before Conformer.

330 + default -1.

331 + cnn_layer_norm: bool, optional

332 + layer norm between Conformer and the first CNN.

333 + default False.

334 + time_reduction: int, optional

335 + time reduction factor

336 + default 4

337 + dropout_rate: float, optional

338 + dropout rate. default 0.1

339 + padding_idx: int, optional

340 + padding index for input_layer=embed

341 + default -1

342 + relative_attention_bias_args: dict, optional

343 + use more efficient scalar bias-based relative multihead attention

344 + (Q*K^T + B) implemented in cmb.basics.embedding.

345 + [T5/ALiBi]RelativeAttentionLogitBias

346 + usage: relative_attention_bias_args={"type": t5/alibi}

347 + additional method-specific arguments can be provided (see

348 + transformer_base.py)

349 + positional_dropout_rate: float, optional

350 + dropout rate after positional encoding. default 0.0

351 + nemo_conv_settings: dict, optional

352 + A dictionary of settings for NeMo Subsampling.

353 + default None

354 + conv2d_extra_padding: str, optional

355 + Add extra padding in conv2d subsampling layers. Choices are

356 + (feat, feat_time, none, True).

357 + if True or feat_time, the extra padding is added into non full

358 + supraframe utts in batch.

359 + Default: none

360 + attention_group_size: int, optional

361 + the number of groups to use for attention, default 1

362 + (Multi-Head Attention),

363 + 1 = typical Multi-Head Attention,

364 + 1 < attention_group_size < attention_heads = Grouped-Query

365 + Attention

366 + attention_group_size = attenion_heads = Multi-Query Attention

367 + """

368 +

369 + def __init__(

370 + self,

371 + input_size,

372 + chunk_size,

373 + left_chunk,

374 + attention_dim=256,

375 + attention_heads=4,

376 + input_layer="nemo_conv",

377 + cnn_out=-1,

378 + cnn_layer_norm=False,

379 + time_reduction=4,

380 + dropout_rate=0.0,

381 + padding_idx=-1,

382 + relative_attention_bias_args=None,

383 + positional_dropout_rate=0.0,

384 + nemo_conv_settings=None,

385 + conv2d_extra_padding: Literal["feat", "feat_time", "none",

386 + True] = "none",

387 + attention_group_size=1,

388 + encoder_embedding_config=None,

389 + ):

390 + super().__init__()

391 + self.input_size = input_size

392 + self.input_layer = input_layer

393 + self.chunk_size = chunk_size

394 + self.left_chunk = left_chunk

395 + self.attention_dim = attention_dim

396 + self.num_heads = attention_heads

397 + self.attention_group_size = attention_group_size

398 + self.time_reduction = time_reduction

399 + self.nemo_conv_settings = nemo_conv_settings

400 + self.encoder_embedding_config = encoder_embedding_config

401 +

402 + if self.input_layer == "nemo_conv":

403 + default_nemo_conv_settings = {

404 + "subsampling": "dw_striding",

405 + "subsampling_factor": self.time_reduction,

406 + "feat_in": input_size,

407 + "feat_out": attention_dim,

408 + "conv_channels": 256,

409 + "subsampling_conv_chunking_factor": 1,

410 + "activation": nn.ReLU(),

411 + "is_causal": False,

412 + }

413 + # Override any of the defaults with the incoming, user settings

414 + if nemo_conv_settings:

415 + default_nemo_conv_settings.update(nemo_conv_settings)

416 + for i in ["subsampling_factor", "feat_in", "feat_out"]:

417 + assert (

418 + i not in nemo_conv_settings

419 + ), "{i} should be specified outside of the NeMo dictionary"

420 +

421 + self.embed = NemoConvSubsampling(**default_nemo_conv_settings, )

422 + else:

423 + raise ValueError("unknown input_layer: " + input_layer)

424 +

425 + self.pos_emb = AbsolutePositionalEncoding(attention_dim,

426 + positional_dropout_rate)

427 +

428 + self.relative_attention_bias_type = (

429 + relative_attention_bias_args.get("type")

430 + if relative_attention_bias_args else None)

431 + if self.relative_attention_bias_type == "t5":

432 + assert (self.num_heads % self.attention_group_size == 0

433 + ), "attention_group_size must divide n_head"

434 + self.relative_attention_bias_layer = T5RelativeAttentionLogitBias(

435 + self.num_heads // self.attention_group_size,

436 + max_distance=relative_attention_bias_args.get(

437 + "t5_bias_max_distance", 1000),

438 + symmetric=relative_attention_bias_args.get(

439 + "t5_bias_symmetric", False),

440 + )

441 + else:

442 + raise NotImplementedError

443 +

444 + def post_init(self, init_model_config):

445 +

446 + pretrained_speech_encoder_path = init_model_config.get(

447 + "pretrained_speech_encoder_path", None)

448 + if pretrained_speech_encoder_path:

449 + model_state = torch.load(pretrained_speech_encoder_path,

450 + map_location="cpu")

451 + encoder_state_dict = {}

452 + for k, v in model_state.items():

453 + if "encoder." in k:

454 + tmp_k = k.replace("encoder.", "")

455 + encoder_state_dict[tmp_k] = v

456 +

457 + if hasattr(self, "encoder_embedding"):

458 + del self.encoder_embedding

459 + self.load_state_dict(encoder_state_dict)

460 +

461 + if not hasattr(self, "encoder_embedding"):

462 + self.encoder_embedding = MeanVarianceNormLayer(

463 + self.encoder_embedding_config["input_size"])

464 +

465 + def compute_lens_change(self, feature_lens):

466 + """feature_lens: int

467 + return updated feature lens.

468 +

469 + This used to return a different lambda function for each case that

470 + computed the right thing. That does not work within Torchscript.

471 + If you really need this to be faster, create nn.Module()-s for all

472 + the cases and return one of them. Torchscript does support that.

473 + """

474 + if self.input_layer == "nemo_conv":

475 + # Handle the special causal case

476 + subsampling_causal_cond = self.nemo_conv_settings.get(

477 + "subsampling", "dw_striding") in [

478 + "dw_striding",

479 + "striding",

480 + "striding_conv1d",

481 + ]

482 + is_causal = self.nemo_conv_settings.get("is_causal", False)

483 + if is_causal and subsampling_causal_cond:

484 + lens_change = (torch.ceil(feature_lens /

485 + self.time_reduction).long()

486 + if isinstance(feature_lens, Tensor) else

487 + math.ceil(feature_lens / self.time_reduction))

488 + feature_lens_remainder = feature_lens % self.time_reduction

489 + if isinstance(feature_lens, Tensor):

490 + lens_change[feature_lens_remainder != 1] += 1

491 + elif feature_lens_remainder != 1:

492 + lens_change += 1

493 + return lens_change

494 + ceil_func = (math.ceil

495 + if isinstance(feature_lens, int) else torch.ceil)

496 + return ceil_func(feature_lens / self.time_reduction)

497 +

498 + @abc.abstractmethod

499 + def forward(self):

500 + """Abstract forward method implementation."""

501 +

502 + def _chunk_size_selection(self, chunk_size=None, left_chunk=None):

503 + """If chunk size is a list, we will randomly select a chunk size."""

504 +

505 + if chunk_size is None:

506 + chunk_size = self.chunk_size

507 + if left_chunk is None:

508 + left_chunk = self.left_chunk

509 + if isinstance(chunk_size, list):

510 + # Variable chunk size during training

511 + chunk_size_index = int(

512 + torch.randint(low=0, high=len(chunk_size), size=(1, )))

513 + chunk_size_train_eff = chunk_size[chunk_size_index]

514 + if not isinstance(left_chunk, list):

515 + raise ValueError(

516 + "Since chunk_size is a list, left_chunk must be a list")

517 + if len(left_chunk) != len(chunk_size):

518 + raise ValueError(

519 + "The length of left_chunk must be the same as length of "\

520 + "chunk_size."

521 + )

522 + left_chunk_train_eff = left_chunk[chunk_size_index]

523 + else:

524 + chunk_size_train_eff = chunk_size

525 + left_chunk_train_eff = left_chunk

526 +

527 + return chunk_size_train_eff, left_chunk_train_eff

528 +

529 + def _get_embed_class(self, embed):

530 + # pylint: disable=protected-access

531 + is_embed_using_act_chkpt = isinstance(embed, CheckpointWrapper)

532 + is_embed_fsdp_wrapped = isinstance(embed, FullyShardedDataParallel)

533 + embed_class = embed

534 + if is_embed_using_act_chkpt:

535 + embed_class = embed._checkpoint_wrapped_module

536 + if is_embed_fsdp_wrapped:

537 + embed_class = embed.module

538 + return embed_class

539 +

540 + def _forward_embeddings_core(self, input_tensor, masks):

541 + embed_class = self._get_embed_class(self.embed)

542 + assert isinstance(embed_class, NemoConvSubsampling)

543 + input_tensor, masks = self.embed(input_tensor, masks)

544 + return input_tensor, masks

545 +

546 + def _position_embedding(self, input_tensor):

547 + pos_k = None

548 + pos_v = None

549 + if self.relative_attention_bias_layer is None:

550 + input_tensor = self.pos_emb(

551 + input_tensor) # default to add abs sinusoid embedding

552 + return pos_k, pos_v

553 +

554 + def _streaming_mask(self, seq_len, batch_size, chunk_size, left_chunk):

555 + chunk_size_train_eff, left_chunk_train_eff = \

556 + self._chunk_size_selection(chunk_size, left_chunk)

557 +

558 + # Create mask matrix for streaming

559 + # S stores start index. if chunksize is 18, s is [0,18,36,....]

560 + chunk_start_idx = np.arange(0, seq_len, chunk_size_train_eff)

561 + # avoid randomness when run evaluation or decoding

562 + if self.training and np.random.rand() > 0.5:

563 + # Either first or last chunk is not complete.

564 + # If only the last one is not complete, EOS is not effective

565 + chunk_start_idx = seq_len - chunk_start_idx

566 + chunk_start_idx = chunk_start_idx[::-1]

567 + chunk_start_idx = chunk_start_idx[:-1]

568 + chunk_start_idx = np.insert(chunk_start_idx, 0, 0)

569 +

570 + enc_streaming_mask = (adaptive_enc_mask(

571 + seq_len, chunk_start_idx,

572 + left_window=left_chunk_train_eff).unsqueeze(0).expand(

573 + [batch_size, -1, -1]))

574 + return enc_streaming_mask

575 +

576 + def forward_embeddings(self,

577 + xs_pad,

578 + masks,

579 + chunk_size_nc=None,

580 + left_chunk_nc=None):

581 + """Forwarding the inputs through the top embedding layers

582 +

583 + Args:

584 + xs_pad: torch.Tensor

585 + input tensor

586 + masks: torch.Tensor

587 + input mask

588 + chunk_size_nc: (optional, default is None) chunk size for

589 + non-causal layers

590 + left_chunk_nc: (optional, default is None) # of left chunks for

591 + non-causal layers

592 + """

593 + # pylint: disable=R0915

594 + # get new lens.

595 + seq_len = int(self.compute_lens_change(xs_pad.shape[1]))

596 + if seq_len <= 0:

597 + raise ValueError(

598 + f"""The sequence length after time reduction is invalid:

599 + {seq_len}. Your input feature is too short. Consider

600 + filtering out the very short sentence from data

601 + loader""", )

602 +

603 + batch_size = xs_pad.shape[0]

604 +

605 + enc_streaming_mask = self._streaming_mask(seq_len, batch_size,

606 + self.chunk_size,

607 + self.left_chunk)

608 +

609 + if xs_pad.is_cuda:

610 + enc_streaming_mask = enc_streaming_mask.cuda()

611 + xs_pad = xs_pad.cuda()

612 +

613 + input_tensor = xs_pad

614 + input_tensor, masks = self._forward_embeddings_core(

615 + input_tensor, masks)

616 +

617 + streaming_mask = enc_streaming_mask

618 + if streaming_mask is not None and masks is not None:

619 + hs_mask = masks & streaming_mask

620 + elif masks is not None:

621 + hs_mask = masks

622 + else:

623 + hs_mask = streaming_mask

624 +

625 + if chunk_size_nc is not None:

626 + enc_streaming_mask_nc = self._streaming_mask(

627 + seq_len, batch_size, chunk_size_nc, left_chunk_nc)

628 + if xs_pad.is_cuda:

629 + enc_streaming_mask_nc = enc_streaming_mask_nc.cuda()

630 + if masks is not None:

631 + hs_mask_nc = masks & enc_streaming_mask_nc

632 + else:

633 + hs_mask_nc = enc_streaming_mask_nc

634 + else:

635 + hs_mask_nc = None

636 +

637 + pos_k, pos_v = self._position_embedding(input_tensor)

638 +

639 + if chunk_size_nc is None:

640 + return input_tensor, pos_k, pos_v, hs_mask, masks

641 + return input_tensor, pos_k, pos_v, hs_mask, masks, hs_mask_nc

642 +

643 + def get_offset(self):

644 + """Returns offset used when retaining inputs for decoding.

645 +

646 + This is essentially, how many additional frames have to be added to

647 + the front-end CNN input to ensure it can produce a single output.

648 + So if the "padding" parameter is 0, typically offset will be > 0.

649 + """

650 + return get_offset(self.input_layer, self.time_reduction)

651 +

652 +

653 +class ConformerEncoder(TransformerEncoderBase):

654 + """ConformerEncoder module.

655 + see original paper for more details:

656 + https://arxiv.org/abs/2005.08100

657 +

658 + Please set causal = True in streaming model

659 + Args:

660 + input_size: int

661 + input feature dimension.

662 + chunk_size: int, list(int)

663 + Number of frames for each chunk

664 + This variable can take 2 forms:

665 + int: Used for inference, or single chunk size training

666 + list(int) : Used only for variable chunk size training

667 + Some examples for the 2 cases:

668 + chunk_size = 12

669 + chunk_size = [6, 8, 12, 24]

670 + left_chunk: int, list(int)

671 + Number of chunks used for masking in streaming mode.

672 + This variable can take 2 forms:

673 + int: Used for inference, or single chunk size training

674 + list(int) : Used only for variable chunk size training. When

675 + chunk_size is a list, left_chunk must be a list with same length.

676 + Some examples for the 2 cases:

677 + left_chunk = 6

678 + left_chunk = [12, 9, 6, 3]

679 + left_chunk: int

680 + number of chunks used for masking in streaming mode.

681 + num_lang: int

682 + This parameter is used to store the number of languages in the

683 + lang_dict, only used for multiseed/multilingual models.

684 + default None.

685 + attention_dim: int, optional

686 + attention dimension. default 256.

687 + attention_heads: int, optional

688 + the number of heads. default 4

689 + linear_units:

690 + the number of units of position-wise feed forward.

691 + default 2048

692 + num_block:

693 + number of Transformer layer. default 6

694 + dropout_rate: float, optional

695 + dropout rate. default 0.1

696 + input_layer: str, optional

697 + input layer type before Conformer,

698 + one of ["linear", "conv2d", "custom", "vgg2l", "embed"],

699 + default "conv2d"

700 + causal: bool, optional

701 + if set to True, convolution have no access

702 + to future frames. default False.

703 + batch_norm: bool, optional

704 + if set to True, apply batchnorm before activation

705 + in ConvModule layer of the conformer.

706 + default False

707 + cnn_out: int, optional

708 + the number of CNN channels before Conformer.

709 + default -1.

710 + cnn_layer_norm: bool, optional

711 + layer norm between Conformer and the first CNN.

712 + default False.

713 + ext_pw_out_channel: int, optional

714 + the number of channel for CNN

715 + before depthwise_seperable_CNN.

716 + If 0 then use linear. default 0.

717 + ext_pw_kernel_size: int, optional

718 + kernel size of N before depthwise_seperable_CNN.

719 + only work for ext_pw_out_channel > 0.

720 + default 1

721 + depthwise_seperable_out_channel: int, optional

722 + the number of channel for

723 + depthwise_seperable_CNN.

724 + default 256.

725 + depthwise_multiplier: int, optional

726 + the number of multiplier for

727 + depthwise_seperable_CNN.

728 + default 1.

729 + chunk_se: int, optional

730 + 0 for offline SE.

731 + 1 for streaming SE, where mean is computed

732 + by accumulated history until current chunk_se.

733 + 2 for streaming SE, where mean is computed

734 + by only the current chunk.

735 + default 0.

736 + kernel_size: int, optional

737 + the number of kernels for depthwise_seperable_CNN.

738 + default 3.

739 + activation: str, optional

740 + FeedForward block activation.

741 + one of ["relu", "swish", "sigmoid"]

742 + default "relu".

743 + conv_activation: str, optional

744 + activation function used in ConvModule part

745 + of the conformer, default "relu".

746 + conv_glu_type: str, optional

747 + activation used use glu in depthwise_seperable_CNN,

748 + default "sigmoid"

749 + bias_in_glu: bool, optional

750 + if set to True, use additive bias in the weight module

751 + before GLU. default True

752 + linear_glu_in_convm: bool, optional

753 + if set to True, use GLULinear module,

754 + otherwise, used GLUPointWiseConv module.

755 + default to False.

756 + attention_glu_type: str

757 + only work for glu_in_attention !=0

758 + default "swish".

759 + export: bool, optional

760 + if set to True, it remove the padding from convolutional layers

761 + and allow the onnx conversion for inference.

762 + default False.

763 + activation_checkpointing: str, optional

764 + a dictionarry of {"module","interval","offload"}, where

765 + "module": str

766 + accept ["transformer", "attention"] to select

767 + which module should do activation checkpointing.

768 + "interval": int, default 1,

769 + interval of applying activation checkpointing,

770 + interval = 1 means that we apply checkpointing

771 + on every layer (if activation), otherwise,

772 + we apply it every x interval.

773 + "offload": bool, default False,

774 + if set to True, we offload activation to cpu and

775 + reload it during backward, otherwise,

776 + we recalculate activation in backward.

777 + default "".

778 + extra_layer_output_idx: int

779 + the layer index to be exposed.

780 + relative_attention_bias_args: dict, optional

781 + use more efficient scalar bias-based relative multihead attention

782 + (Q*K^T + B) implemented in cmb.basics.embedding.

783 + [T5/ALiBi]RelativeAttentionLogitBias

784 + usage: relative_attention_bias_args={"type": t5/alibi}

785 + additional method-specific arguments can be provided (see

786 + transformer_base.py)

787 + time_reduction: int optional

788 + time reduction factor

789 + default 4

790 + use_pt_scaled_dot_product_attention: whether to use pytorch scaled

791 + dot product attention in training.

792 + Default: False

793 + nemo_conv_settings: dict, optional

794 + A dictionary of settings for NeMo Subsampling.

795 + default: None

796 + usage: nemo_conv_settings=

797 + {

798 + "subsampling":

799 + dw_striding/striding/dw_striding_conv1d/striding_conv1d,

800 + "conv_channels": int,

801 + "subsampling_conv_chunking_factor": int,

802 + "is_causal": True/False

803 + }

804 + conv2d_extra_padding: str, optional

805 + Add extra padding in conv2d subsampling layers. Choices are

806 + (feat, feat_time, none, True)

807 + Default: none

808 + replication_pad_for_subsample_embedding: For batched-streaming

809 + decoding, use "replication" padding for the cache at start of

810 + utterance.

811 + Default: False

812 + attention_group_size: int, optional

813 + the number of groups to use for attention, default 1

814 + (Multi-Head Attention),

815 + 1 = typical Multi-Head Attention,

816 + 1 < attention_group_size < attention_heads = Grouped-Query

817 + Attention

818 + attention_group_size = attenion_heads = Multi-Query Attention

819 + """

820 +

821 + extra_multi_layer_output_idxs: List[int]

822 +

823 + def __init__( # pylint: disable-all

824 + self,

825 + input_size,

826 + chunk_size,

827 + left_chunk,

828 + num_lang=None,

829 + attention_dim=256,

830 + attention_heads=4,

831 + linear_units=2048,

832 + num_blocks=6,

833 + dropout_rate=0.1,

834 + input_layer="nemo_conv",

835 + causal=True,

836 + batch_norm=False,

837 + cnn_out=-1,

838 + cnn_layer_norm=False,

839 + ext_pw_out_channel=0,

840 + ext_pw_kernel_size=1,

841 + depthwise_seperable_out_channel=256,

842 + depthwise_multiplier=1,

843 + chunk_se=0,

844 + kernel_size=3,

845 + activation="relu",

846 + conv_activation="relu",

847 + conv_glu_type="sigmoid",

848 + bias_in_glu=True,

849 + linear_glu_in_convm=False,

850 + attention_glu_type="swish",

851 + export=False,

852 + extra_layer_output_idx=-1,

853 + extra_multi_layer_output_idxs=[], # noqa

854 + activation_checkpointing="",

855 + relative_attention_bias_args=None,

856 + time_reduction=4,

857 + use_pt_scaled_dot_product_attention=False,

858 + nemo_conv_settings=None,

859 + conv2d_extra_padding: Literal["feat", "feat_time", "none",

860 + True] = "none",

861 + replication_pad_for_subsample_embedding=False,

862 + attention_group_size=1,

863 + encoder_embedding_config=None,

864 + ):

865 + super().__init__(

866 + input_size,

867 + chunk_size,

868 + left_chunk,

869 + attention_dim,

870 + attention_heads,

871 + input_layer,

872 + cnn_out,

873 + cnn_layer_norm,

874 + time_reduction,

875 + dropout_rate=dropout_rate,

876 + relative_attention_bias_args=relative_attention_bias_args,

877 + positional_dropout_rate=0.0,

878 + nemo_conv_settings=nemo_conv_settings,

879 + conv2d_extra_padding=conv2d_extra_padding,

880 + attention_group_size=attention_group_size,

881 + encoder_embedding_config=encoder_embedding_config,

882 + )

883 + self.num_blocks = num_blocks

884 + self.num_lang = num_lang

885 + self.kernel_size = kernel_size

886 + self.embed = embedding_checkpoint_wrapper(activation_checkpointing)(

887 + self.embed)

888 + self.replication_pad_for_subsample_embedding: bool = (

889 + replication_pad_for_subsample_embedding)

890 + assert (self.num_heads % attention_group_size == 0

891 + ), "attention_group_size must divide n_head"

892 + self.num_heads_k = self.num_heads // attention_group_size

893 +

894 + self.encoders = repeat(

895 + num_blocks,

896 + lambda i: encoder_checkpoint_wrapper(activation_checkpointing,

897 + ConformerEncoderLayer, i)

898 + (ConformerEncoderLayer(

899 + d_model=attention_dim,

900 + ext_pw_out_channel=ext_pw_out_channel,

901 + depthwise_seperable_out_channel=

902 + depthwise_seperable_out_channel,

903 + depthwise_multiplier=depthwise_multiplier,

904 + n_head=attention_heads,

905 + d_ffn=linear_units,

906 + ext_pw_kernel_size=ext_pw_kernel_size,

907 + kernel_size=kernel_size,

908 + dropout_rate=dropout_rate,

909 + causal=causal,

910 + batch_norm=batch_norm,

911 + activation=activation,

912 + chunk_se=chunk_se,

913 + chunk_size=chunk_size,

914 + conv_activation=conv_activation,

915 + conv_glu_type=conv_glu_type,

916 + bias_in_glu=bias_in_glu,

917 + linear_glu_in_convm=linear_glu_in_convm,

918 + attention_glu_type=attention_glu_type,

919 + activation_checkpointing=attn_checkpointing(

920 + activation_checkpointing, i),

921 + export=export,

922 + use_pt_scaled_dot_product_attention=

923 + use_pt_scaled_dot_product_attention,

924 + attn_group_sizes=attention_group_size,

925 + )),

926 + )

927 + self.extra_layer_output_idx = extra_layer_output_idx

928 + self.extra_multi_layer_output_idxs = extra_multi_layer_output_idxs

929 + # Make a zeros scalar we can use in get_initial_state to determine

930 + # the device and the needed dtype:

931 + self.register_buffer("dev_type", torch.zeros(()), persistent=False)

932 +

933 + def init_relative_attention_bias(self, input_tensor):

934 + if self.relative_attention_bias_layer:

935 + return self.relative_attention_bias_layer(input_tensor)

936 +

937 + def calculate_hs_mask(self, xs_pad, device, mask):

938 + max_audio_length = xs_pad.shape[1]

939 + batch_size = xs_pad.shape[0]

940 + enc_streaming_mask = self._streaming_mask(max_audio_length, batch_size,

941 + self.chunk_size,

942 + self.left_chunk)

943 + enc_streaming_mask = enc_streaming_mask.to(device)

944 + if mask is None:

945 + return enc_streaming_mask

946 +

947 + feature_lens = mask.sum(1)

948 + padding_length = feature_lens

949 + pad_mask = (torch.arange(0, max_audio_length,

950 + device=device).expand(padding_length.size(0),

951 + -1)

952 + < padding_length.unsqueeze(1))

953 + pad_mask = pad_mask.unsqueeze(1)

954 + pad_mask = pad_mask & enc_streaming_mask

955 + return pad_mask

956 +

957 + @torch.jit.ignore

958 + def forward(self, xs_pad, masks):

959 + """Conformer Forward function

960 +

961 + Args:

962 + xs_pad: torch.Tensor

963 + input tensor

964 + masks: torch.Tensor

965 + post-embedding input lengths

966 + """

967 + xs_pad = self.encoder_embedding(xs_pad)

968 + input_tensor, pos_k, pos_v, hs_mask, masks = self.forward_embeddings(

969 + xs_pad, masks)

970 +

971 + unfolded = False

972 + ori_bz, seq_len, D = input_tensor.shape

973 + max_seq_len = 500 #maximum position for absolute positional encoding

974 + if seq_len > max_seq_len:

975 + # audio sequence is longer than max_seq_len, unfold it into chunks

976 + # of max_seq_len

977 + unfolded = True

978 + # the unfold op will drop residual frames, pad it to the multiple

979 + # of max_seq_len

980 + if seq_len % max_seq_len > 0:

981 + chunk_pad_size = max_seq_len - (seq_len % max_seq_len)

982 + else:

983 + chunk_pad_size = 0

984 + if chunk_pad_size > 0:

985 + input_tensor_pad = F.pad(input_tensor,

986 + (0, 0, 0, chunk_pad_size), "constant",

987 + 0)

988 + input_tensor = input_tensor_pad.to(input_tensor.device)

989 + input_tensor = unfold_tensor(input_tensor, max_seq_len)

990 + if masks is not None:

991 + # revise hs_mask here because the previous calculated hs_mask

992 + # did not consider extra pad

993 + subsampled_pad_mask = masks.squeeze(

994 + 1) # [bz, subsampled_unmask_seq_len]

995 + extra_padded_subsamlped_pad_mask = F.pad(

996 + subsampled_pad_mask, (0, chunk_pad_size), "constant",

997 + False) # extra padding to the pad mask

998 + extra_padded_subsamlped_pad_mask = \

999 + extra_padded_subsamlped_pad_mask.unsqueeze(-1).float()

1000 + masks_unfold = unfold_tensor(

1001 + extra_padded_subsamlped_pad_mask, max_seq_len

1002 + ) # unfold the pad mask like we did to the input tensor

1003 + masks_unfold = masks_unfold.squeeze(

1004 + -1).bool() # unfold op does not support bool tensor

1005 + else:

1006 + masks_unfold = None

1007 + hs_mask = self.calculate_hs_mask(

1008 + input_tensor, input_tensor.device, masks_unfold

1009 + ) # calculate hs_mask based on the unfolded pad mask

1010 +

1011 + # layer_emb = None

1012 +

1013 + relative_attention_bias = self.init_relative_attention_bias(

1014 + input_tensor)

1015 +

1016 + _simplified_path = (self.extra_layer_output_idx == -1

1017 + and relative_attention_bias is None)

1018 +

1019 + if _simplified_path:

1020 + input_tensor, *_ = self.encoders(input_tensor, pos_k, pos_v,

1021 + hs_mask)

1022 + else:

1023 + for i, layer in enumerate(self.encoders):

1024 + input_tensor, _, _, _ = layer(

1025 + input_tensor,

1026 + pos_k,

1027 + pos_v,

1028 + hs_mask,

1029 + relative_attention_bias=relative_attention_bias,

1030 + )

1031 +

1032 + # if i == self.extra_layer_output_idx:

1033 + # layer_emb = input_tensor

1034 +

1035 + if unfolded:

1036 + embed_dim = input_tensor.shape[-1]

1037 + input_tensor = input_tensor.reshape(ori_bz, -1, embed_dim)

1038 + # if we ever padded before unfolding, we need to remove the padding

1039 + if chunk_pad_size > 0:

1040 + input_tensor = input_tensor[:, :-chunk_pad_size, :]

1041 +

1042 + return input_tensor, masks # , layer_emb

1043 +

1044 + def gradient_checkpointing_enable(self):

1045 + pass

1046 +

1047 +

1048 +class WindowQformer(nn.Module):

1049 + """Window-level Qformer"""

1050 +

1051 + def __init__(

1052 + self,

1053 + window_size: int = 8,

1054 + num_queries: int = 1,

1055 + num_blocks: int = 2,

1056 + attention_dim: int = 512,

1057 + attention_heads: int = 8,

1058 + linear_units: int = 2048,

1059 + dropout_rate: float = 0.0,

1060 + normalize_before: bool = True,

1061 + ):

1062 + super().__init__()

1063 +

1064 + self.decoders = nn.ModuleList([

1065 + nn.TransformerDecoderLayer(

1066 + d_model=attention_dim,

1067 + nhead=attention_heads,

1068 + dim_feedforward=linear_units,

1069 + dropout=dropout_rate,

1070 + activation="relu",

1071 + batch_first=True,

1072 + norm_first=normalize_before, # TODO need to verify

1073 + ) for _ in range(num_blocks)

1074 + ])

1075 +

1076 + self.queries = nn.Parameter(torch.zeros(1, num_queries, attention_dim))

1077 + self.after_norm = (nn.LayerNorm(attention_dim, eps=1e-12)

1078 + if normalize_before else None)

1079 + self.window_size = window_size

1080 + self.gradient_checkpointing_enable = False

1081 +

1082 + def enable_gradient_checkpointing(self):

1083 + self.gradient_checkpointing_enable = True

1084 +

1085 + def disable_gradient_checkpointing(self):

1086 + self.gradient_checkpointing_enable = False

1087 +

1088 + def forward(self, audio_embed, mask, embed_len=None):

1089 + """forward decoder"""

1090 + # audio_embed: N x T x D => N x D x T

1091 +

1092 + audio_embed = audio_embed.transpose(1, 2)

1093 + # audio_embed: N x D x 1 x T => N x DK x T'

1094 + padding = audio_embed.shape[-1] % self.window_size

1095 + if padding > 0:

1096 + audio_embed = F.pad(audio_embed, (0, self.window_size - padding),

1097 + "constant", 0)

1098 +

1099 + embed_chunk = F.unfold(

1100 + audio_embed[..., None, :],

1101 + kernel_size=(1, self.window_size),

1102 + stride=(1, self.window_size),

1103 + )

1104 + bsz, _, slen = embed_chunk.shape

1105 + # N x D x K x T'

1106 + embed_chunk = embed_chunk.view(bsz, -1, self.window_size, slen)

1107 + # N x T' x K x D

1108 + embed_chunk = embed_chunk.transpose(1, 3).contiguous()

1109 + # NT' x K x D

1110 + embed_chunk = embed_chunk.view(bsz * slen, self.window_size, -1)

1111 + # NT' x 1 x D

1112 + q = self.queries.expand(bsz * slen, -1, -1)

1113 + for layer in self.decoders:

1114 + if self.gradient_checkpointing_enable and self.training:

1115 + q = checkpoint(

1116 + layer.__call__,

1117 + q,

1118 + embed_chunk,

1119 + None,

1120 + mask,

1121 + use_reentrant=True,

1122 + )

1123 + else:

1124 + q = layer(tgt=q,

1125 + memory=embed_chunk,

1126 + tgt_mask=None,

1127 + memory_mask=mask)

1128 +

1129 + if self.after_norm is not None:

1130 + q = self.after_norm(q)

1131 +

1132 + if embed_len is not None:

1133 + embed_len = embed_len // self.window_size

1134 + # N x T' x D

1135 + out = q.view(bsz, slen, -1)

1136 +

1137 + return out, embed_len

1138 +

1139 +

1140 +class AudioEmbedding(nn.Module):

1141 + """Image embedding."""

1142 +

1143 + def __init__(self, config: PretrainedConfig, **kwargs) -> None:

1144 + super().__init__()

1145 + self.config = config

1146 + # n_embed or hidden_size for text LM

1147 + hidden_size = (config.n_embd

1148 + if hasattr(config, "n_embd") else config.hidden_size)

1149 +

1150 + if hasattr(config, "embd_pdrop") or hasattr(config, "embed_pdrop"):

1151 + embd_drop = (config.embd_pdrop if hasattr(config, "embd_pdrop")

1152 + else config.embed_pdrop)

1153 + self.drop = nn.Dropout(embd_drop)

1154 + else:

1155 + self.drop = None

1156 +

1157 + # self.wte = nn.Embedding(config.vocab_size, hidden_size)

1158 +

1159 + audio_dim_out = (

1160 + None # Set this variable according to the actual audio processor

1161 + )

1162 + self.layer_idx = -2

1163 +

1164 + if (isinstance(config.audio_processor, dict)

1165 + and config.audio_processor.get("name", None) == "cascades"):

1166 + encoder_config = config.audio_processor.get("config", None)

1167 + assert encoder_config is not None

1168 + self.encoder = ConformerEncoder(**encoder_config)

1169 +

1170 + # fake initialization, create encoder_embedding layer only so that

1171 + # in decoding, all parameters can be loaded in

1172 + # from_pretrained_function in training, we do post init after

1173 + # from_pretrained function to make sure the correct initialization

1174 + self.encoder.post_init({})

1175 +

1176 + audio_dim_out = encoder_config["attention_dim"]

1177 + n_mels = encoder_config["input_size"]

1178 + else:

1179 + raise NotImplementedError("")

1180 +

1181 + assert (audio_dim_out

1182 + is not None), "Remember to set values for audio_dim_out"

1183 + self.audio_dim_out = audio_dim_out

1184 + self.audio_dim_in = n_mels

1185 +

1186 + self.freeze_audio_processor = kwargs.get("freeze_audio_processor",

1187 + False)

1188 +

1189 + self.downsample_rate = kwargs.get("downsample_rate", 1)

1190 +

1191 + if kwargs.get("use_qformer", False):

1192 + qformer_config = kwargs.get("qformer_config", {})

1193 + qformer_config["attention_dim"] = audio_dim_out

1194 + self.qformer = WindowQformer(**qformer_config)

1195 + else:

1196 + self.qformer = None

1197 +

1198 + if kwargs.get("use_conv_downsample", False):

1199 + assert (self.qformer is None

1200 + ), "don't support use qformer and conv downsample together"

1201 + nemo_conv_settings = kwargs.get("nemo_conv_settings", {})

1202 + default_nemo_conv_settings = {

1203 + "subsampling": "dw_striding",

1204 + "subsampling_factor": self.downsample_rate,

1205 + "feat_in": audio_dim_out,

1206 + "feat_out": audio_dim_out,

1207 + "conv_channels": 256,

1208 + "subsampling_conv_chunking_factor": 1,

1209 + "activation": nn.ReLU(),

1210 + "is_causal": False,

1211 + }

1212 + # Override any of the defaults with the incoming, user settings

1213 + if nemo_conv_settings:

1214 + default_nemo_conv_settings.update(nemo_conv_settings)

1215 + for i in ["subsampling_factor", "feat_in", "feat_out"]:

1216 + assert (

1217 + i not in nemo_conv_settings

1218 + ), "{i} should be specified outside of the NeMo dictionary"

1219 +

1220 + self.conv_ds = NemoConvSubsampling(**default_nemo_conv_settings, )

1221 + else:

1222 + self.conv_ds = None

1223 +

1224 + enable_gradient_checkpointing = kwargs.get(

1225 + "enable_gradient_checkpointing", False)

1226 + if enable_gradient_checkpointing:

1227 + self.encoder.gradient_checkpointing_enable()

1228 +

1229 + if self.qformer:

1230 + self.qformer.enable_gradient_checkpointing()

1231 +

1232 + projection_cls = kwargs.get("projection_cls", "linear")

1233 + if projection_cls == "linear":

1234 + self.audio_projection = nn.Linear(audio_dim_out, hidden_size)

1235 + elif projection_cls == "mlp":

1236 + # follow llava-v1.5's implementation

1237 + # (do not use image_projection and image_proj_norm)

1238 + dim_projection = hidden_size

1239 + depth = 2

1240 + self.linear_downsample_rate = (1 if (self.qformer or self.conv_ds)

1241 + else self.downsample_rate)

1242 + layers = [

1243 + nn.Linear(audio_dim_out * self.linear_downsample_rate,

1244 + dim_projection)

1245 + ]

1246 + for _ in range(1, depth):

1247 + layers.extend(

1248 + [nn.GELU(),

1249 + nn.Linear(dim_projection, dim_projection)])

1250 + self.audio_projection = nn.Sequential(*layers)

1251 + # NOTE vision-speech tasks use a separate projection layer

1252 + layers = [

1253 + nn.Linear(audio_dim_out * self.linear_downsample_rate,

1254 + dim_projection)

1255 + ]

1256 + for _ in range(1, depth):

1257 + layers.extend(

1258 + [nn.GELU(),

1259 + nn.Linear(dim_projection, dim_projection)])

1260 + self.audio_projection_for_vision = nn.Sequential(*layers)

1261 + else:

1262 + raise NotImplementedError(

1263 + f"projection_cls = {projection_cls}, not implemented")

1264 +

1265 + # TODO: audio sequence compression - Qformer

1266 + self.vocab_size = config.vocab_size

1267 + self.input_embeds = None

1268 + self.audio_embed_sizes = None

1269 +

1270 + def set_audio_embeds(self, input_embeds: torch.FloatTensor) -> None:

1271 + self.input_embeds = input_embeds

1272 +

1273 + def set_audio_embed_sizes(self,

1274 + audio_embed_sizes: torch.LongTensor) -> None:

1275 + self.audio_embed_sizes = audio_embed_sizes

1276 +

1277 + def get_audio_features(

1278 + self,

1279 + input_embeds: torch.FloatTensor,

1280 + audio_attention_mask: torch.Tensor = None,

1281 + audio_projection_mode: str = "speech",

1282 + ):

1283 +

1284 + if self.freeze_audio_processor:

1285 + with torch.no_grad():

1286 + audio_features, masks = self.encoder(input_embeds,

1287 + audio_attention_mask)

1288 + else:

1289 + audio_features, masks = self.encoder(input_embeds,

1290 + audio_attention_mask)

1291 +

1292 + if self.qformer is not None:

1293 + audio_features, _ = self.qformer(audio_features, mask=None)

1294 +

1295 + if self.conv_ds is not None:

1296 + if masks is not None:

1297 + masks = masks.squeeze(1)

1298 +

1299 + audio_features, masks = self.conv_ds(audio_features, mask=masks)

1300 +

1301 + if self.linear_downsample_rate != 1:

1302 + bs, seq_len, feat_dim = audio_features.size()

1303 + padding = seq_len % self.linear_downsample_rate

1304 + if padding > 0:

1305 + audio_features = F.pad(

1306 + audio_features,

1307 + (0, 0, 0, self.linear_downsample_rate - padding),

1308 + "constant",

1309 + 0,

1310 + )

1311 +

1312 + seq_len = audio_features.size(1)

1313 + audio_features = audio_features.view(

1314 + bs,

1315 + seq_len // self.linear_downsample_rate,

1316 + feat_dim * self.linear_downsample_rate,

1317 + )

1318 +

1319 + if audio_projection_mode == 'speech':

1320 + audio_set_tensor = self.audio_projection(audio_features)

1321 + elif audio_projection_mode == 'vision':

1322 + audio_set_tensor = self.audio_projection_for_vision(audio_features)

1323 + else:

1324 + raise ValueError(

1325 + f"audio_projection_mode = {audio_projection_mode} not "\

1326 + "implemented"

1327 + )

1328 +

1329 + return audio_set_tensor

1330 +

1331 + def forward(

1332 + self,

1333 + input_ids: torch.LongTensor,

1334 + input_embeds: torch.FloatTensor,

1335 + audio_embed_sizes,

1336 + **kwargs,

1337 + ) -> torch.FloatTensor:

1338 + """

1339 + arguments:

1340 + input_ids: input text ids (B, U)

1341 + input_embeds: audio features (B, T, D) B: num audios in a sequence

1342 + """

1343 + assert input_embeds is not None and len(input_embeds) == len(

1344 + audio_embed_sizes)

1345 +

1346 + input_shape = input_ids.size()

1347 + input_ids = input_ids.view(-1, input_shape[-1])

1348 +

1349 + with torch.no_grad():

1350 + positions = (input_ids == _AUDIO_PLACEHOLDER_TOKEN_ID).nonzero(

1351 + as_tuple=False)

1352 +

1353 + if not isinstance(input_embeds, list):

1354 + input_embeds = [input_embeds]

1355 +

1356 + audio_projection_mode = kwargs.get("audio_projection_mode", "speech")

1357 + audio_set_tensor = [

1358 + self.get_audio_features(

1359 + input_embed, audio_projection_mode=audio_projection_mode)

1360 + for input_embed in input_embeds

1361 + ]

1362 +

1363 + with torch.no_grad():

1364 + input_ids.clamp_min_(0).clamp_max_(self.vocab_size)

1365 +

1366 + if "wte" in kwargs:

1367 + # we use the token embedding layer from the huggingface model, this

1368 + # is REQUIRED to make sure we are using the loaded weights.

1369 + hidden_states = kwargs["wte"](input_ids)

1370 + else:

1371 + # otherwise, we use token embedding in pretrained mixformer from

1372 + # phi team

1373 + hidden_states = self.wte(input_ids)

1374 +

1375 + if len(positions.tolist()) > 0:

1376 + assert sum(audio_embed_sizes) == len(

1377 + positions

1378 + ), "please ensure the encoder outputs have the same length as"\

1379 + " defined in input_ids!"

1380 + idx = 0

1381 + for i in range(len(audio_embed_sizes)):

1382 + cnt = audio_embed_sizes[i]

1383 + assert audio_set_tensor[i].shape[0] == 1

1384 + hidden_states[

1385 + positions[idx, 0],

1386 + positions[idx, 1]:positions[idx, 1] + cnt,

1387 + ] = (audio_set_tensor[i][0, :audio_embed_sizes[i], :].to(

1388 + hidden_states.dtype).to(hidden_states.device))

1389 + idx += cnt

1390 +

1391 + else:

1392 + if self.training:

1393 + # hidden_states[:, 0:img_set_tensor.shape[0]] =

1394 + # hidden_states[:, 0:img_set_tensor.shape[0]] +

1395 + # 0 * img_set_tensor.to(hidden_states.dtype)

1396 + # .to(hidden_states.device)

1397 + hidden_states[:, 0:1] = hidden_states[:, 0:1] + \

1398 + 0 * audio_set_tensor[:, 0:1].to(hidden_states.dtype)\

1399 + .to(hidden_states.device)

1400 +

1401 + if self.drop is not None:

1402 + hidden_states = self.drop(hidden_states)

1403 + return hidden_states

vllm/model_executor/models/phi4mm_utils.py new

+1969

@@ -0,0 +1,1969 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# Copyright (c) Microsoft Corporation.

3 +# Licensed under the MIT license.

4 +# Code copied from Microsoft/MoE by Jacob Platin (jacobplatin@microsoft.com)

5 +# but implemented by the Phi-Speech team

6 +#!/usr/bin/env python3

7 +import math

8 +from functools import partial

9 +from typing import Callable, Dict, Optional, Tuple, Union

10 +

11 +import torch

12 +import torch.nn.functional as F

13 +from torch import Tensor, nn

14 +from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (

15 + CheckpointImpl, checkpoint_wrapper, offload_wrapper)

16 +

17 +

18 +class Block(nn.Module):

19 + """Block abstract module"""

20 +

21 + def __init__(self, input_size, output_size):

22 + super().__init__()

23 + self.input_size = input_size

24 + self.output_size = output_size

25 +

26 +

27 +def get_activation(name="relu"):

28 + """Select an activation function by name

29 +

30 + Args:

31 + name: str

32 + activation function name,

33 + one of ["relu", "gelu", "swish", "sigmoid"],

34 + default "relu".

35 + """

36 + name = name.lower()

37 + if name == "relu":

38 + return nn.ReLU(inplace=True)

39 + if name == "gelu":

40 + return nn.GELU()

41 + if name == "swish":

42 + return Swish()

43 + if name == "sigmoid":

44 + return torch.nn.Sigmoid()

45 + return nn.Identity()

46 +

47 +

48 +def adaptive_enc_mask(x_len, chunk_start_idx, left_window=0, right_window=0):

49 + """

50 + The function is very important for Transformer Transducer Streaming mode

51 + Args:

52 + xs_len (int): sequence length

53 + chunk_start_idx (list): first idx of each chunk, such as [0,18,36,48].

54 + It also supports adaptive chunk size [0,10,15,45]

55 + left_window (int): how many left chunks can be seen

56 + right_window (int): how many right chunks can be seen. It is used for

57 + chunk overlap model.

58 + Returns:

59 + mask (torch.Tensor): a mask tensor for streaming model

60 + Torch 1.0.1

61 + tensor([[1., 1., 0., 0.],

62 + [0., 1., 1., 0.],

63 + [0., 0., 1., 1.]])

64 + Torch 1.4.1

65 + tensor([[True., True., False., False.],

66 + [False., True., True., False.],

67 + [False., False., True., True.]])

68 + """

69 + chunk_start_idx = torch.Tensor(chunk_start_idx).long(

70 + ) # first idx of each chunk, such as [0,18,36,48].

71 + start_pad = torch.nn.functional.pad(

72 + chunk_start_idx,

73 + (1, 0)) # append 0 to the beginning, so it becomes [0, 0, 18, 36, 48]

74 + end_pad = torch.nn.functional.pad(

75 + chunk_start_idx, (0, 1), value=x_len

76 + ) # append x_len to the end, so it becomes [0,18,36,48, x_len]

77 + seq_range = torch.arange(0,

78 + x_len).unsqueeze(-1) # seq_range size: [x_len, 1]

79 + idx = ((seq_range < end_pad) &

80 + (seq_range >= start_pad)).nonzero()[:, 1] # idx size: [x_len]

81 + # boundary = end_pad[idx] # boundary size: [x_len]

82 + seq_range_expand = (torch.arange(0, x_len).unsqueeze(0).expand(x_len, -1)

83 + ) # seq_range_expand size [x_len, x_len]

84 + idx_left = idx - left_window

85 + idx_left[idx_left < 0] = 0

86 + boundary_left = start_pad[idx_left]

87 + mask_left = seq_range_expand >= boundary_left.unsqueeze(-1)

88 + idx_right = idx + right_window

89 + idx_right[idx_right > len(chunk_start_idx)] = len(chunk_start_idx)

90 + boundary_right = end_pad[idx_right]

91 + mask_right = seq_range_expand < boundary_right.unsqueeze(-1)

92 + return mask_left & mask_right

93 +

94 +

95 +class Swish(nn.Module):

96 + """Implement Swish activation module.

97 + From https://arxiv.org/pdf/2005.03191.pdf

98 +

99 + """

100 +

101 + def __init__(self) -> None:

102 + super().__init__()

103 + self.act_fn = nn.Sigmoid()

104 +

105 + def forward(self, x: Tensor) -> Tensor:

106 + """Apply Swish function

107 +

108 + Args:

109 + x: torch.Tensor

110 + Input.

111 + """

112 + return x * self.act_fn(x)

113 +

114 +

115 +class GLU(nn.Module):

116 + """Implement Gated Linear Unit (GLU) module"""

117 +

118 + def __init__(self, dim: int = -1, act_name: str = "sigmoid") -> None:

119 + super().__init__()

120 + self.dim = dim

121 + self.act_name = act_name.lower()

122 +

123 + if self.act_name == "relu":

124 + self.act_fn = nn.ReLU(inplace=True)

125 + elif self.act_name == "gelu":

126 + self.act_fn = nn.GELU()

127 + elif self.act_name == "swish":

128 + self.act_fn = Swish()

129 + elif self.act_name == "sigmoid":

130 + self.act_fn = nn.Sigmoid()

131 + else:

132 + self.act_fn = nn.Identity()

133 +

134 + def forward(self, x: Tensor) -> Tensor:

135 + """GLU forward

136 + Apply Swish function on the first half of input matrices

137 + with sigmoid of the second half.

138 +

139 + Args:

140 + x: torch.Tensor

141 + Input.

142 +

143 + """

144 + half_x, gate = x.chunk(2, dim=self.dim)

145 + return half_x * self.act_fn(gate)

146 +

147 +

148 +# TODO: Abdel, this can be improved using GLU module

149 +class GLUPointWiseConv(nn.Module):

150 + """GLUPointWiseConv module

151 + used for conformer architecture,

152 + for more details see:

153 + https://arxiv.org/pdf/2005.08100v1.pdf

154 +

155 + Args:

156 + input_dim: int

157 + input channel size.

158 + output_dim: int

159 + output channel size.

160 + kernel_size: int

161 + kernel size

162 + glu_type: str, optional

163 + activation function one of

164 + ["sigmoid", "relu", "gelu"]

165 + default "sigmoid".

166 + bias_in_glu: bool, optional

167 + use addtive bias in glu

168 + causal: bool, optional

169 + if set to True, padding is set to the half of

170 + kernel size, ie, convolution can't see future frames.

171 + default False.

172 +

173 + """

174 +

175 + def __init__(

176 + self,

177 + input_dim,

178 + output_dim,

179 + kernel_size,

180 + glu_type="sigmoid",

181 + bias_in_glu=True,

182 + causal=False,

183 + ):

184 + super().__init__()

185 +

186 + self.glu_type = glu_type

187 + self.output_dim = output_dim

188 + self.bias_in_glu = bias_in_glu

189 + if causal:

190 + self.ext_pw_conv_1d = nn.Conv1d(

191 + input_dim,

192 + output_dim * 2,

193 + kernel_size,

194 + 1,

195 + padding=(kernel_size - 1),

196 + )

197 + else:

198 + self.ext_pw_conv_1d = nn.Conv1d(

199 + input_dim,

200 + output_dim * 2,

201 + kernel_size,

202 + 1,

203 + padding=(kernel_size - 1) // 2,

204 + )

205 +

206 + if glu_type == "sigmoid":

207 + self.glu_act = nn.Sigmoid()

208 + elif glu_type == "relu":

209 + self.glu_act = nn.ReLU()

210 + elif glu_type == "gelu":

211 + self.glu_act = nn.GELU()

212 + elif glu_type == "swish":

213 + self.glu_act = Swish()

214 + else:

215 + raise ValueError(f"Unsupported activation type {self.glu_act}")

216 +

217 + if bias_in_glu:

218 + self.b1 = nn.Parameter(torch.zeros(1, output_dim, 1))

219 + self.b2 = nn.Parameter(torch.zeros(1, output_dim, 1))

220 +

221 + def forward(self, x):

222 + """

223 + Args:

224 + x: torch.Tensor

225 + input tensor

226 + """

227 + # to be consistent with GLULinear, we assume the input always has the

228 + # #channel (#dim) in the last dimension of the tensor, so need to

229 + # switch the dimension first for 1D-Conv case

230 + x = x.permute([0, 2, 1])

231 + x = self.ext_pw_conv_1d(x)

232 + if self.glu_type == "bilinear":

233 + if self.bias_in_glu:

234 + x = (x[:, 0:self.output_dim, :] + self.b1) * (

235 + x[:, self.output_dim:self.output_dim * 2, :] + self.b2)

236 + else:

237 + x = (x[:, 0:self.output_dim, :]) * (

238 + x[:, self.output_dim:self.output_dim * 2, :])

239 + else:

240 + if self.bias_in_glu:

241 + x = (x[:, 0:self.output_dim, :] + self.b1) * self.glu_act(

242 + x[:, self.output_dim:self.output_dim * 2, :] + self.b2)

243 + else:

244 + x = (x[:, 0:self.output_dim, :]) * self.glu_act(

245 + x[:, self.output_dim:self.output_dim * 2, :])

246 +

247 + x = x.permute([0, 2, 1])

248 + return x

249 +

250 +

251 +class DepthWiseSeperableConv1d(nn.Module):

252 + """DepthWiseSeperableConv1d module used in Convnet module

253 + for the conformer, for more details see:

254 + https://arxiv.org/pdf/2005.08100v1.pdf

255 +

256 + Args:

257 + input_dim: int

258 + input channel size.

259 + depthwise_seperable_out_channel: int

260 + if set different to 0, the number of

261 + depthwise_seperable_out_channel will be used as a channel_out

262 + of the second conv1d layer.

263 + otherwise, it equal to 0, the second conv1d layer is skipped.

264 + kernel_size: int

265 + kernel_size

266 + depthwise_multiplier: int

267 + number of input_dim channels duplication. this value

268 + will be used to compute the hidden channels of the Conv1D.

269 + padding: int, optional

270 + padding for the conv1d,

271 + default: 0.

272 +

273 + """

274 +

275 + def __init__(

276 + self,

277 + input_dim,

278 + depthwise_seperable_out_channel,

279 + kernel_size,

280 + depthwise_multiplier,

281 + padding=0,

282 + ):

283 + super().__init__()

284 +

285 + self.dw_conv = nn.Conv1d(

286 + input_dim,

287 + input_dim * depthwise_multiplier,

288 + kernel_size,

289 + 1,

290 + padding=padding,

291 + groups=input_dim,

292 + )

293 +

294 + if depthwise_seperable_out_channel != 0:

295 + self.pw_conv = nn.Conv1d(

296 + input_dim * depthwise_multiplier,

297 + depthwise_seperable_out_channel,

298 + 1,

299 + 1,

300 + 0,

301 + )

302 + else:

303 + self.pw_conv = nn.Identity()

304 + self.depthwise_seperable_out_channel = depthwise_seperable_out_channel

305 +

306 + def forward(self, x):

307 + """

308 +

309 + Args:

310 + x: torch.Tensor

311 + input tensor

312 + """

313 + x = self.dw_conv(x)

314 + if self.depthwise_seperable_out_channel != 0:

315 + x = self.pw_conv(x)

316 + return x

317 +

318 +

319 +class ConvModule(nn.Module):

320 + """ConvModule Module for the conformer block.

321 + for more details see:

322 + https://arxiv.org/pdf/2005.08100v1.pdf

323 +

324 + Args:

325 + input_dim: int

326 + input channel size.

327 + ext_pw_out_channel: int

328 + if > 0, ext_pw_out_channel is a dim channel size

329 + for the last pointwise conv after swish activation.

330 + depthwise_seperable_out_channel: int

331 + if set different to 0, the number of

332 + depthwise_seperable_out_channel

333 + will be used as a channel_out of the second conv1d layer.

334 + otherwise, it equal to 0, the second conv1d layer is skipped.

335 + ext_pw_kernel_size: int

336 + kernel size of the conv pointwise of the conformer.

337 + kernel_size: int

338 + kernel size.

339 + depthwise_multiplier: int

340 + number of input_dim channels duplication. this value

341 + will be used to compute the hidden channels of the Conv1D.

342 + dropout_rate: float

343 + dropout rate.

344 + causal: bool, optional

345 + if set to True, convolution have no access

346 + to future frames. default False.

347 + batch_norm: bool, optional

348 + if set to True, apply batchnorm before activation.

349 + default False

350 + chunk_se: int, optional

351 + 0 for offline SE.

352 + 1 for streaming SE, where mean is computed

353 + by accumulated history until current chunk_se.

354 + 2 for streaming SE, where mean is computed

355 + by only the current chunk.

356 + chunk_size: int, optional

357 + chunk size for cnn. default 18

358 + activation: str, optional

359 + activation function used in ConvModule,

360 + default: "relu".

361 + glu_type: str, optional

362 + activation function used for the glu,

363 + default: "sigmoid".

364 + bias_in_glu: bool, optional

365 + if set to True, use additive bias in the weight module

366 + before GLU.

367 + linear_glu_in_convm: bool, optional

368 + if set to True, use GLULinear module,

369 + otherwise, used GLUPointWiseConv module.

370 + default to False.

371 + export: bool, optional,

372 + if set to True, padding is equal to 0. This is for inference,

373 + or onnx export. Typically this is set by the export program or

374 + the decoder program, and it isn't present in your config file.

375 + default False

376 + """

377 +

378 + def __init__(

379 + self,

380 + input_dim,

381 + ext_pw_out_channel,

382 + depthwise_seperable_out_channel,

383 + ext_pw_kernel_size,

384 + kernel_size,

385 + depthwise_multiplier,

386 + dropout_rate,

387 + causal=False,

388 + batch_norm=False,

389 + chunk_se=0,

390 + chunk_size=18,

391 + activation="relu",

392 + glu_type="sigmoid",

393 + bias_in_glu=True,

394 + linear_glu_in_convm=False,

395 + export=False,

396 + ):

397 + super().__init__()

398 + self.layer_norm = nn.LayerNorm(input_dim)

399 + self.input_dim = input_dim

400 + self.ext_pw_out_channel = ext_pw_out_channel

401 + self.ext_pw_kernel_size = ext_pw_kernel_size

402 + self.depthwise_seperable_out_channel = depthwise_seperable_out_channel

403 + self.glu_type = glu_type

404 + self.bias_in_glu = bias_in_glu

405 + self.linear_glu_in_convm = linear_glu_in_convm

406 + self.causal = causal

407 +

408 + self._add_ext_pw_layer()

409 +

410 + self.batch_norm = batch_norm

411 + self.kernel_size = kernel_size

412 +

413 + if batch_norm:

414 + self.bn_layer = nn.BatchNorm1d(input_dim)

415 +

416 + self.act = get_activation(activation)

417 + self.dropout = nn.Dropout(dropout_rate)

418 + self.export = export

419 +

420 + if causal:

421 + padding = 0 if export else kernel_size - 1

422 + else:

423 + padding = (kernel_size - 1) // 2

424 +

425 + self.dw_sep_conv_1d = DepthWiseSeperableConv1d(

426 + input_dim,

427 + depthwise_seperable_out_channel,

428 + kernel_size,

429 + depthwise_multiplier,

430 + padding=padding,

431 + )

432 +

433 + if depthwise_seperable_out_channel != 0:

434 + if input_dim != depthwise_seperable_out_channel:

435 + self.ln2 = nn.Linear(depthwise_seperable_out_channel,

436 + input_dim)

437 + else:

438 + if depthwise_multiplier != 1:

439 + self.ln2 = nn.Linear(input_dim * depthwise_multiplier,

440 + input_dim)

441 +

442 + def _add_ext_pw_layer(self):

443 + """

444 + This function is an extension of __init__ function

445 + and dedicated to the convolution module creation

446 + of the conformer.

447 + """

448 + self.ln1 = self.glu = self.bn_layer = self.ext_pw_conv_1d = (

449 + nn.Identity()) # jit hacks.

450 + self.squeeze_excitation = nn.Identity() # jit.

451 + self.apply_ln1 = self.fix_len1 = False # jit.

452 +

453 + if self.ext_pw_out_channel != 0:

454 + if self.causal:

455 + self.ext_pw_conv_1d = nn.Conv1d(

456 + self.input_dim,

457 + self.ext_pw_out_channel,

458 + self.ext_pw_kernel_size,

459 + 1,

460 + padding=(self.ext_pw_kernel_size - 1),

461 + )

462 + if self.ext_pw_kernel_size > 1:

463 + self.fix_len1 = True

464 + else:

465 + self.fix_len1 = False

466 + else:

467 + self.ext_pw_conv_1d = nn.Conv1d(

468 + self.input_dim,

469 + self.ext_pw_out_channel,

470 + self.ext_pw_kernel_size,

471 + 1,

472 + padding=(self.ext_pw_kernel_size - 1) // 2,

473 + )

474 + self.fix_len1 = False

475 +

476 + if self.linear_glu_in_convm:

477 + self.glu = GLULinear(

478 + self.input_dim,

479 + self.ext_pw_out_channel,

480 + self.glu_type,

481 + self.bias_in_glu,

482 + )

483 + else:

484 + self.glu = GLUPointWiseConv(

485 + self.input_dim,

486 + self.ext_pw_out_channel,

487 + self.ext_pw_kernel_size,

488 + self.glu_type,

489 + self.bias_in_glu,

490 + self.causal,

491 + )

492 +

493 + if self.input_dim != self.ext_pw_out_channel:

494 + self.apply_ln1 = True

495 + self.ln1 = nn.Linear(self.ext_pw_out_channel, self.input_dim)

496 + else:

497 + self.apply_ln1 = False

498 + else:

499 + self.pw_conv_simplify_w = torch.nn.Parameter(torch.ones(3))

500 + self.pw_conv_simplify_b = torch.nn.Parameter(torch.zeros(3))

501 +

502 + def forward(self, x):

503 + """ConvModule Forward.

504 +

505 + Args:

506 + x: torch.Tensor

507 + input tensor.

508 + """

509 + x = self.layer_norm(x)

510 +

511 + if self.ext_pw_out_channel != 0:

512 + x = self.glu(x)

513 + if self.causal and self.ext_pw_kernel_size > 1:

514 + x = x[:, :-(self.ext_pw_kernel_size - 1), :]

515 + if self.apply_ln1:

516 + x = self.ln1(x)

517 + else:

518 + x_0 = x * self.pw_conv_simplify_w[0] + self.pw_conv_simplify_b[0]

519 + x_1 = x * self.pw_conv_simplify_w[1] + self.pw_conv_simplify_b[1]

520 + x = x_0 + x_1

521 +

522 + x = x.permute([0, 2, 1])

523 +

524 + x = self.dw_sep_conv_1d(x)

525 + if self.causal and self.kernel_size > 1:

526 + x = x[:, :, :-(self.kernel_size - 1)]

527 + if hasattr(self, "ln2"):

528 + x = x.permute([0, 2, 1])

529 + x = self.ln2(x)

530 + x = x.permute([0, 2, 1])

531 + if self.batch_norm:

532 + x = self.bn_layer(x)

533 + x = self.act(x)

534 +

535 + if self.ext_pw_out_channel != 0:

536 + x = self.ext_pw_conv_1d(x)

537 + if self.fix_len1:

538 + x = x[:, :, :-(self.ext_pw_kernel_size - 1)]

539 +

540 + if self.apply_ln1:

541 + x = x.permute([0, 2, 1])

542 + x = self.ln1(x)

543 + x = x.permute([0, 2, 1])

544 +

545 + x = x.permute([0, 2, 1])

546 + else:

547 + x = x.unsqueeze(1).permute([0, 1, 3, 2])

548 + x = x * self.pw_conv_simplify_w[2] + self.pw_conv_simplify_b[2]

549 + x = x.squeeze(1)

550 +

551 + x = self.dropout(x)

552 + return x

553 +

554 +

555 +class GLULinear(nn.Module):

556 + """Linear + GLU module

557 +

558 + Args:

559 + input_dim: int

560 + input size

561 + output_dim: int

562 + output size.

563 + glu_type:

564 + activation function name used in glu module.

565 + default "sigmoid" (swish function).

566 + bias_in_glu: bool, optional

567 + If True, the addtive bias is added. Default False.

568 + """

569 +

570 + def __init__(

571 + self,

572 + input_dim,

573 + output_dim,

574 + glu_type="sigmoid",

575 + bias_in_glu=True,

576 + ):

577 + super().__init__()

578 + self.linear = nn.Linear(input_dim, output_dim * 2, bias_in_glu)

579 + self.glu_act = GLU(-1, glu_type)

580 +

581 + def forward(self, x):

582 + """GLULinear forward

583 +

584 + Args:

585 + x: torch.Tensor

586 + inpute tensor.

587 + """

588 + x = self.linear(x)

589 + return self.glu_act(x)

590 +

591 +

592 +class FeedForward(nn.Module):

593 + """FeedForward Module.

594 + For more details see Conformer paper:

595 + https://arxiv.org/pdf/2005.08100.pdf

596 +

597 + Args:

598 + d_model: int

599 + input size.

600 + d_inner: int

601 + output size.

602 + dropout_rate: float,

603 + dropout rate.

604 + activation: str,

605 + activation function name,

606 + one of ["relu", "swish", "sigmoid"],

607 + sigmoid activation is only used with "glu_in_fnn=True",

608 + default "sigmoid".

609 + bias_in_glu: bool, optional

610 + """

611 +

612 + def __init__(

613 + self,

614 + d_model,

615 + d_inner,

616 + dropout_rate,

617 + activation="sigmoid",

618 + bias_in_glu=True,

619 + ):

620 + super().__init__()

621 + self.d_model = d_model

622 + self.d_inner = d_inner

623 +

624 + self.layer_norm = nn.LayerNorm(d_model)

625 + module = GLULinear(d_model, d_inner, activation, bias_in_glu)

626 + self.net = nn.Sequential(

627 + module,

628 + nn.Dropout(dropout_rate),

629 + nn.Linear(d_inner, d_model),

630 + nn.Dropout(dropout_rate),

631 + )

632 +

633 + def forward(self, x):

634 + """FeedForward forward function.

635 +

636 + Args:

637 + x: torch.Tensor

638 + input tensor.

639 + """

640 + out = self.net(self.layer_norm(x))

641 +

642 + return out

643 +

644 +

645 +#### positional encoding starts here

646 +def _pre_hook(

647 + state_dict,

648 + prefix,

649 + local_metadata,

650 + strict,

651 + missing_keys,

652 + unexpected_keys,

653 + error_msgs,

654 +):

655 + """Perform pre-hook in load_state_dict for backward compatibility.

656 +

657 + Note:

658 + We saved self.pe until v.0.5.2 but we have omitted it later.

659 + Therefore, we remove the item "pe" from `state_dict` for backward

660 + compatibility.

661 +

662 + """

663 + k = prefix + "pe"

664 + if k in state_dict:

665 + state_dict.pop(k)

666 +

667 +

668 +class T5RelativeAttentionLogitBias(nn.Module):

669 + """

670 + This module implements the relative position bias described in Section

671 + 2.1 of the T5 paper: https://arxiv.org/pdf/1910.10683.pdf

672 +

673 + The Huggingface implementation is used as a reference

674 + https://github.com/huggingface/transformers/blob/v4.30.0/src/

675 + transformers/models/t5/modeling_t5.py#L435

676 +

677 + Modifies attention as Q*K^T + B, where B is a learned scalar bias based

678 + on relative position of the query and key. It is HxNxN, where H is the

679 + number of heads, N is the sequence length.

680 +

681 + I've made these modifications to the original T5 bias:

682 + - Skipping of the bucketing step. Original T5 bias converted rel

683 + position distances into logarithmically increasing buckets. This is

684 + supposed to help with length generalization.

685 + - I just directly use rel position index as bias values, as we don't

686 + need length generalization (40s max is good enough for ASR encoder),

687 + and it keeps ONNX export simple.

688 + - I've also extended it so that biases can be asymmetric, the default

689 + implementation treats L->R and R->L the same. Asymmetric was found to

690 + yield better results in my experiments.

691 +

692 + Args:

693 + num_heads: int

694 + Number of attention heads

695 + num_buckets: int

696 + Number of buckets to use for relative attention bias. This is the

697 + size of the learnable bias parameter. Bucketing is not yet

698 + supported, so this defaults to -1 which means no bucketing is

699 + used (max_distance determines size of bias param).

700 + max_distance: int

701 + Maximum distance to use for relative attention bias. With

702 + num_buckets=-1, this directly controls the max size of the bias

703 + parameter. When num_buckets > 0 is supported, this will control

704 + the maximum distance for logarithmic bucketing after which all

705 + positions are in the same bucket.

706 + symmetric: bool

707 + Whether to use symmetric or asymmetric biases. symmetric=False uses

708 + 2x number of bias params to distinguish L->R from R->L. This was

709 + found to be better for the encoder.

710 + """

711 +

712 + def __init__(self,

713 + num_heads,

714 + num_buckets=-1,

715 + max_distance=1000,

716 + symmetric=False):

717 + super().__init__()

718 + self.num_heads = num_heads

719 + self.num_buckets = num_buckets

720 + self.max_distance = max_distance

721 + self.symmetric = symmetric

722 + self._skip_bucketing = self.num_buckets < 0

723 + if self._skip_bucketing:

724 + self.num_buckets = max_distance

725 + else:

726 + raise NotImplementedError(

727 + "T5 attention bias with bucketed positions is not yet tested")

728 + if not self.symmetric:

729 + self.num_buckets *= 2

730 + self.bias_values = nn.Embedding(self.num_buckets, self.num_heads)

731 +

732 + def forward(self, x):

733 + # instantiate bias compatible with shape of x

734 + maxpos = x.size(1)

735 + context_position = torch.arange(maxpos,

736 + device=x.device,

737 + dtype=torch.long)[:, None]

738 + memory_position = torch.arange(maxpos,

739 + device=x.device,

740 + dtype=torch.long)[None, :]

741 + relative_position = memory_position - context_position

742 + # clipping to a maximum distance using ops that play well with ONNX

743 + # export

744 + relative_position = relative_position.masked_fill(

745 + relative_position < -self.max_distance, -self.max_distance)

746 + relative_position = relative_position.masked_fill(

747 + relative_position > self.max_distance - 1, self.max_distance - 1)

748 +

749 + # mapping from relative position to index in the bias parameter

750 + if self._skip_bucketing:

751 + bias_idx = relative_position

752 + else:

753 + bias_idx = self._bucket_relative_position(relative_position)

754 + if self.symmetric:

755 + bias_idx = bias_idx.abs()

756 + else:

757 + bias_idx += self.num_buckets // 2

758 +

759 + t5_rel_att_bias = self.bias_values(bias_idx) # [L, L, H]

760 + t5_rel_att_bias = t5_rel_att_bias.permute(2, 0, 1).unsqueeze(

761 + 0) # [1, H, L, L]

762 +

763 + return t5_rel_att_bias

764 +

765 + def _bucket_relative_position(self, relative_position):

766 + # this is a placeholder (isn't tested, likely buggy) using HuggingFace

767 + # implem as a reference this also needs to be extended to support

768 + # asymmetric +/- ve positions

769 + relative_buckets = 0

770 + if not self.causal:

771 + self.num_buckets //= 2

772 + relative_buckets += (relative_position > 0).to(

773 + torch.long) * self.num_buckets

774 + relative_position = torch.abs(relative_position)

775 + else:

776 + relative_position = -torch.min(relative_position,

777 + torch.zeros_like(relative_position))

778 + # now relative_position is in the range [0, inf)

779 +

780 + # half of the buckets are for exact increments in positions

781 + max_exact = self.num_buckets // 2

782 + is_small = relative_position < max_exact

783 +

784 + # The other half of the buckets are for logarithmically bigger bins in

785 + # positions up to max_distance

786 + relative_position_if_large = max_exact + (

787 + torch.log(relative_position.float() / max_exact) /

788 + math.log(self.max_distance / max_exact) *

789 + (self.num_buckets - max_exact)).to(torch.long)

790 + relative_position_if_large = torch.min(

791 + relative_position_if_large,

792 + torch.full_like(relative_position_if_large, self.num_buckets - 1),

793 + )

794 +

795 + relative_buckets += torch.where(is_small, relative_position,

796 + relative_position_if_large)

797 + return relative_buckets

798 +

799 +

800 +class AbsolutePositionalEncoding(nn.Module):

801 + """Absolute Positional encoding module.

802 + This module implement Absolute sinusoidal positional encoding

803 + from: https://arxiv.org/pdf/1706.03762.pdf

804 +

805 + Args:

806 + d_model: int

807 + Input embedding size.

808 + dropout_rate: float

809 + dropout rate

810 + max_len: int, optional

811 + Maximum input length sequence, Default 5000

812 +

813 + """

814 +

815 + def __init__(self, d_model, dropout_rate, max_len=5000):

816 + """Construct an PositionalEncoding object."""

817 + super().__init__()

818 + self.d_model = d_model

819 + self.xscale = math.sqrt(self.d_model)

820 + self.dropout = torch.nn.Dropout(p=dropout_rate)

821 + self.pe = None

822 + self.extend_pe(torch.tensor(0.0).expand(1, max_len))

823 + self._register_load_state_dict_pre_hook(_pre_hook)

824 +

825 + def extend_pe(self, x):

826 + """Reset the positional encodings.

827 +

828 + Args:

829 + x: torch.Tensor

830 + """

831 + if self.pe is not None and self.pe.size(1) >= x.size(1):

832 + if self.pe.dtype != x.dtype or self.pe.device != x.device:

833 + self.pe = self.pe.to(dtype=x.dtype, device=x.device)

834 + return

835 + pe = torch.zeros(x.size(1), self.d_model)

836 + position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1)

837 + div_term = torch.exp(

838 + torch.arange(0, self.d_model, 2, dtype=torch.float32) *

839 + -(math.log(10000.0) / self.d_model))

840 + pe[:, 0::2] = torch.sin(position * div_term)

841 + pe[:, 1::2] = torch.cos(position * div_term)

842 + pe = pe.unsqueeze(0)

843 + self.pe = pe.to(device=x.device, dtype=x.dtype)

844 +

845 + def forward(self, x: torch.Tensor):

846 + """Add positional encoding.

847 +

848 + Args:

849 + x: torch.Tensor

850 + Input tensor. shape is (batch, time, ...)

851 +

852 + Returns:

853 + torch.Tensor: Encoded tensor. Its shape is (batch, time, ...)

854 +

855 + """

856 + self.extend_pe(x)

857 + x = x * self.xscale + self.pe[:, :x.size(1)]

858 + return self.dropout(x)

859 +

860 +

861 +#### forward embedding layers starts here

862 +class MeanVarianceNormLayer(nn.Module):

863 + """Mean/variance normalization layer.

864 +

865 + Will subtract mean and multiply input by inverted standard deviation.

866 + Typically used as a very first layer in a model.

867 +

868 + Args:

869 + input_size: int

870 + layer input size.

871 + """

872 +

873 + def __init__(self, input_size):

874 + super().__init__()

875 + self.input_size = input_size

876 + self.register_buffer("global_mean", torch.zeros(input_size))

877 + self.register_buffer("global_invstd", torch.ones(input_size))

878 + self.global_mean: Optional[Tensor]

879 + self.global_invstd: Optional[Tensor]

880 +

881 + def forward(self, input_: Tensor) -> Tensor:

882 + """MeanVarianceNormLayer Forward

883 +

884 + Args:

885 + input_: torch.Tensor

886 + input tensor.

887 + """

888 + return (input_ - self.global_mean) * self.global_invstd

889 +

890 +

891 +class CausalConv1D(nn.Conv1d):

892 + """

893 + A causal version of nn.Conv1d where each step would have limited access to

894 + locations on its right or left

895 + All arguments are the same as nn.Conv1d except padding.

896 +

897 + If padding is set None, then paddings are set automatically to make it a

898 + causal convolution where each location would not see any steps on its right.

899 +

900 + If padding is set as a list (size of 2), then padding[0] would be used as

901 + left padding and padding[1] as right padding.

902 + It would make it possible to control the number of steps to be accessible

903 + on the right and left.

904 + This mode is not supported when stride > 1. padding[0]+padding[1] should

905 + be equal to (kernel_size - 1).

906 + """

907 +

908 + def __init__(

909 + self,

910 + in_channels: int,

911 + out_channels: int,

912 + kernel_size: int,

913 + stride: int = 1,

914 + padding: Union[str, int] = 0,

915 + dilation: int = 1,

916 + groups: int = 1,

917 + bias: bool = True,

918 + padding_mode: str = "zeros",

919 + device=None,

920 + dtype=None,

921 + ) -> None:

922 + self.cache_drop_size = None

923 + if padding is None:

924 + self._left_padding = kernel_size - 1

925 + self._right_padding = stride - 1

926 + else:

927 + if stride != 1 and padding != kernel_size - 1:

928 + raise ValueError(

929 + "No striding allowed for non-symmetric convolutions!")

930 + if isinstance(padding, int):

931 + self._left_padding = padding

932 + self._right_padding = padding

933 + elif (isinstance(padding, list) and len(padding) == 2

934 + and padding[0] + padding[1] == kernel_size - 1):

935 + self._left_padding = padding[0]

936 + self._right_padding = padding[1]

937 + else:

938 + raise ValueError(f"Invalid padding param: {padding}!")

939 +

940 + self._max_cache_len = self._left_padding

941 +

942 + super().__init__(

943 + in_channels=in_channels,

944 + out_channels=out_channels,

945 + kernel_size=kernel_size,

946 + stride=stride,

947 + padding=0,

948 + dilation=dilation,

949 + groups=groups,

950 + bias=bias,

951 + padding_mode=padding_mode,

952 + device=device,

953 + dtype=dtype,

954 + )

955 +

956 + def update_cache(self, x, cache=None):

957 + if cache is None:

958 + new_x = F.pad(x, pad=(self._left_padding, self._right_padding))

959 + next_cache = cache

960 + else:

961 + new_x = F.pad(x, pad=(0, self._right_padding))

962 + new_x = torch.cat([cache, new_x], dim=-1)

963 + if self.cache_drop_size > 0:

964 + next_cache = new_x[:, :, :-self.cache_drop_size]

965 + else:

966 + next_cache = new_x

967 + next_cache = next_cache[:, :, -cache.size(-1):]

968 + return new_x, next_cache

969 +

970 + def forward(self, x, cache=None):

971 + x, cache = self.update_cache(x, cache=cache)

972 + x = super().forward(x)

973 + if cache is None:

974 + return x

975 + else:

976 + return x, cache

977 +

978 +

979 +class CausalConv2D(nn.Conv2d):

980 + """

981 + A causal version of nn.Conv2d where each location in the 2D matrix would

982 + have no access to locations on its right or down

983 + All arguments are the same as nn.Conv2d except padding which should be

984 + set as None

985 + """

986 +

987 + def __init__(

988 + self,

989 + in_channels: int,

990 + out_channels: int,

991 + kernel_size: int,

992 + stride: int = 1,

993 + padding: Union[str, int] = 0,

994 + dilation: int = 1,

995 + groups: int = 1,

996 + bias: bool = True,

997 + padding_mode: str = "zeros",

998 + device=None,

999 + dtype=None,

1000 + ) -> None:

1001 + if padding is not None:

1002 + raise ValueError(

1003 + "Argument padding should be set to None for CausalConv2D.")

1004 + self._left_padding = kernel_size - 1

1005 + self._right_padding = stride - 1

1006 +

1007 + padding = 0

1008 + super().__init__(

1009 + in_channels,

1010 + out_channels,

1011 + kernel_size,

1012 + stride,

1013 + padding,

1014 + dilation,

1015 + groups,

1016 + bias,

1017 + padding_mode,

1018 + device,

1019 + dtype,

1020 + )

1021 +

1022 + def forward(

1023 + self,

1024 + x,

1025 + ):

1026 + if self.training:

1027 + x = F.pad(

1028 + x,

1029 + pad=(

1030 + self._left_padding,

1031 + self._right_padding,

1032 + self._left_padding,

1033 + self._right_padding,

1034 + ),

1035 + )

1036 + else:

1037 + x = F.pad(

1038 + x,

1039 + pad=(self._left_padding, self._right_padding, 0, 0),

1040 + )

1041 + x = super().forward(x)

1042 + return x

1043 +

1044 +

1045 +class NemoConvSubsampling(torch.nn.Module):

1046 + """Convlutional subsampling module, taken from NeMo ASR

1047 + (https://github.com/NVIDIA/NeMo/blob/b367413645d5c72db3c2c96e46e95a

1048 + 34501479cf/nemo/collections/asr/parts/submodules/subsampling.py)

1049 +

1050 + Striding Subsampling: "Speech-Transformer: A No-Recurrence

1051 + Sequence-to-Sequence Model for Speech Recognition" by Linhao Dong

1052 + et al. (https://ieeexplore.ieee.org/document/8462506)

1053 +

1054 +

1055 + Compared with the EncoderConv2D (`input_layer: custom`), this is a

1056 + much simplified approach, and uses no LayerNorm and far fewer Conv2Ds.

1057 + Moreover, depthwise convolutions are used to reduce FLOPs, but the first

1058 + layer is kept as a regular convolution so as not to degrade accuracy.

1059 +

1060 + `Striding` and `dw_striding` are the same except that the latter uses

1061 + depthwise convolutions after the first layer, whereas the former does not.

1062 +

1063 + Args:

1064 + subsampling_factor (int): Time reduction factor

1065 + feat_in (int): size of the input features

1066 + feat_out (int): size of the output features

1067 + subsampling (str): The subsampling technique, choose from

1068 + {"striding", "dw-striding", "striding_conv1d",

1069 + "dw_striding_conv1d"}

1070 + conv_channels (int): Number of channels for the convolution layers,

1071 + default is 256.

1072 + subsampling_conv_chunking_factor (int): Input chunking factor which

1073 + can be -1 (no chunking) 1 (auto) or a power of 2. Default is 1

1074 + activation (Module): activation function, default is nn.ReLU()

1075 + is_causal (bool): whether to use causal Conv1/2D, where each step will

1076 + have limited access to locations on its right or left

1077 + """

1078 +

1079 + def __init__(

1080 + self,

1081 + feat_in,

1082 + feat_out,

1083 + subsampling_factor=4,

1084 + subsampling="dw_striding",

1085 + conv_channels=256,

1086 + subsampling_conv_chunking_factor=1,

1087 + activation=nn.ReLU(), # noqa: B008

1088 + is_causal=False,

1089 + ):

1090 + super().__init__()

1091 + self._subsampling = subsampling

1092 + self._conv_channels = conv_channels

1093 + self._feat_in = feat_in

1094 + self._feat_out = feat_out

1095 +

1096 + if subsampling_factor % 2 != 0:

1097 + raise ValueError("Sampling factor should be a multiply of 2!")

1098 + self._sampling_num = int(math.log(subsampling_factor, 2))

1099 + self.subsampling_factor = subsampling_factor

1100 + self.is_causal = is_causal

1101 + self.subsampling_causal_cond = subsampling in (

1102 + "dw_striding",

1103 + "striding",

1104 + "striding_conv1d",

1105 + )

1106 +

1107 + if (subsampling_conv_chunking_factor != -1

1108 + and subsampling_conv_chunking_factor != 1

1109 + and subsampling_conv_chunking_factor % 2 != 0):

1110 + raise ValueError(

1111 + "subsampling_conv_chunking_factor should be -1, 1, or a "\

1112 + "power of 2"

1113 + )

1114 + self.subsampling_conv_chunking_factor = \

1115 + subsampling_conv_chunking_factor

1116 +

1117 + in_channels = 1

1118 + layers = []

1119 +

1120 + if subsampling == "dw_striding":

1121 + self._stride = 2

1122 + self._kernel_size = 3

1123 + self._ceil_mode = False

1124 +

1125 + if self.is_causal:

1126 + self._left_padding = self._kernel_size - 1

1127 + self._right_padding = self._stride - 1

1128 + self._max_cache_len = subsampling_factor + 1

1129 + else:

1130 + self._left_padding = (self._kernel_size - 1) // 2

1131 + self._right_padding = (self._kernel_size - 1) // 2

1132 + self._max_cache_len = 0

1133 +

1134 + # Layer 1

1135 + if self.is_causal:

1136 + layers.append(

1137 + CausalConv2D(

1138 + in_channels=in_channels,

1139 + out_channels=conv_channels,

1140 + kernel_size=self._kernel_size,

1141 + stride=self._stride,

1142 + padding=None,

1143 + ))

1144 + else:

1145 + layers.append(

1146 + torch.nn.Conv2d(

1147 + in_channels=in_channels,

1148 + out_channels=conv_channels,

1149 + kernel_size=self._kernel_size,

1150 + stride=self._stride,

1151 + padding=self._left_padding,

1152 + ))

1153 + in_channels = conv_channels

1154 + layers.append(activation)

1155 +

1156 + for i in range(self._sampling_num - 1):

1157 + if self.is_causal:

1158 + layers.append(

1159 + CausalConv2D(

1160 + in_channels=in_channels,

1161 + out_channels=in_channels,

1162 + kernel_size=self._kernel_size,

1163 + stride=self._stride,

1164 + padding=None,

1165 + groups=in_channels,

1166 + ))

1167 + else:

1168 + layers.append(

1169 + torch.nn.Conv2d(

1170 + in_channels=in_channels,

1171 + out_channels=in_channels,

1172 + kernel_size=self._kernel_size,

1173 + stride=self._stride,

1174 + padding=self._left_padding,

1175 + groups=in_channels,

1176 + ))

1177 +

1178 + layers.append(

1179 + torch.nn.Conv2d(

1180 + in_channels=in_channels,

1181 + out_channels=conv_channels,

1182 + kernel_size=1,

1183 + stride=1,

1184 + padding=0,

1185 + groups=1,

1186 + ))

1187 + layers.append(activation)

1188 + in_channels = conv_channels

1189 +

1190 + elif subsampling == "striding":

1191 + self._stride = 2

1192 + self._kernel_size = 3

1193 + self._ceil_mode = False

1194 +

1195 + if self.is_causal:

1196 + self._left_padding = self._kernel_size - 1

1197 + self._right_padding = self._stride - 1

1198 + self._max_cache_len = subsampling_factor + 1

1199 + else:

1200 + self._left_padding = (self._kernel_size - 1) // 2

1201 + self._right_padding = (self._kernel_size - 1) // 2

1202 + self._max_cache_len = 0

1203 +

1204 + for i in range(self._sampling_num):

1205 + if self.is_causal:

1206 + layers.append(

1207 + CausalConv2D(

1208 + in_channels=in_channels,

1209 + out_channels=conv_channels,

1210 + kernel_size=self._kernel_size,

1211 + stride=self._stride,

1212 + padding=None,

1213 + ))

1214 + else:

1215 + layers.append(

1216 + torch.nn.Conv2d(

1217 + in_channels=in_channels,

1218 + out_channels=conv_channels,

1219 + kernel_size=self._kernel_size,

1220 + stride=self._stride,

1221 + padding=self._left_padding,

1222 + ))

1223 + layers.append(activation)

1224 + in_channels = conv_channels

1225 +

1226 + elif subsampling == "striding_conv1d":

1227 + in_channels = feat_in

1228 +

1229 + self._stride = 2

1230 + self._kernel_size = 5

1231 + self._ceil_mode = False

1232 +

1233 + if self.is_causal:

1234 + self._left_padding = self._kernel_size - 1

1235 + self._right_padding = self._stride - 1

1236 + self._max_cache_len = subsampling_factor + 1

1237 + else:

1238 + self._left_padding = (self._kernel_size - 1) // 2

1239 + self._right_padding = (self._kernel_size - 1) // 2

1240 + self._max_cache_len = 0

1241 +

1242 + for i in range(self._sampling_num):

1243 + if self.is_causal:

1244 + layers.append(

1245 + CausalConv1D(

1246 + in_channels=in_channels,

1247 + out_channels=(feat_out if self._sampling_num == i +

1248 + 1 else conv_channels),

1249 + kernel_size=self._kernel_size,

1250 + stride=self._stride,

1251 + padding=None,

1252 + ))

1253 + else:

1254 + layers.append(

1255 + torch.nn.Conv1d(

1256 + in_channels=in_channels,

1257 + out_channels=(feat_out if self._sampling_num == i +

1258 + 1 else conv_channels),

1259 + kernel_size=self._kernel_size,

1260 + stride=self._stride,

1261 + padding=self._left_padding,

1262 + ))

1263 + layers.append(activation)

1264 + in_channels = conv_channels

1265 +

1266 + elif subsampling == "dw_striding_conv1d":

1267 + in_channels = feat_in

1268 +

1269 + self._stride = 2

1270 + self._kernel_size = 5

1271 + self._ceil_mode = False

1272 +

1273 + self._left_padding = (self._kernel_size - 1) // 2

1274 + self._right_padding = (self._kernel_size - 1) // 2

1275 +

1276 + # Layer 1

1277 + layers.extend([

1278 + torch.nn.Conv1d(

1279 + in_channels=in_channels,

1280 + out_channels=in_channels,

1281 + kernel_size=self._kernel_size,

1282 + stride=self._stride,

1283 + padding=self._left_padding,

1284 + groups=in_channels,

1285 + ),

1286 + torch.nn.Conv1d(

1287 + in_channels=in_channels,

1288 + out_channels=(feat_out if self._sampling_num == 1 else

1289 + conv_channels),

1290 + kernel_size=1,

1291 + stride=1,

1292 + padding=0,

1293 + groups=1,

1294 + ),

1295 + ])

1296 + in_channels = conv_channels

1297 + layers.append(activation)

1298 +

1299 + for i in range(self._sampling_num - 1):

1300 + layers.extend([

1301 + torch.nn.Conv1d(

1302 + in_channels=in_channels,

1303 + out_channels=in_channels,

1304 + kernel_size=self._kernel_size,

1305 + stride=self._stride,

1306 + padding=self._left_padding,

1307 + groups=in_channels,

1308 + ),

1309 + torch.nn.Conv1d(

1310 + in_channels=in_channels,

1311 + out_channels=(feat_out if self._sampling_num == i +

1312 + 2 else conv_channels),

1313 + kernel_size=1,

1314 + stride=1,

1315 + padding=0,

1316 + groups=1,

1317 + ),

1318 + ])

1319 + layers.append(activation)

1320 + in_channels = conv_channels

1321 +

1322 + else:

1323 + raise ValueError(f"Not valid sub-sampling: {subsampling}!")

1324 +

1325 + if subsampling in ["dw_striding", "striding"]:

1326 + in_length = torch.tensor(feat_in, dtype=torch.float)

1327 + out_length = calc_length(

1328 + lengths=in_length,

1329 + all_paddings=self._left_padding + self._right_padding,

1330 + kernel_size=self._kernel_size,

1331 + stride=self._stride,

1332 + ceil_mode=self._ceil_mode,

1333 + repeat_num=self._sampling_num,

1334 + )

1335 + self.out = torch.nn.Linear(conv_channels * int(out_length),

1336 + feat_out)

1337 + self.conv2d_subsampling = True

1338 + elif subsampling in ["striding_conv1d", "dw_striding_conv1d"]:

1339 + self.out = None

1340 + self.conv2d_subsampling = False

1341 + else:

1342 + raise ValueError(f"Not valid sub-sampling: {subsampling}!")

1343 +

1344 + self.conv = torch.nn.Sequential(*layers)

1345 +

1346 + def get_sampling_frames(self):

1347 + return [1, self.subsampling_factor]

1348 +

1349 + def get_streaming_cache_size(self):

1350 + return [0, self.subsampling_factor + 1]

1351 +

1352 + def forward(self, x, mask):

1353 + """

1354 + Forward method for NeMo subsampling.

1355 +

1356 + Args:

1357 + x[Batch, Time, Filters]: torch.Tensor

1358 + input tensor

1359 + x_mask: torch.Tensor

1360 + input mask

1361 +

1362 + Returns:

1363 + x: torch.Tensor

1364 + Resulting tensor from subsampling (B, T //

1365 + time_reduction_factor, feat_out)

1366 + pad_mask: torch.Tensor

1367 + tensor of padded hidden state sequences (B, 1, T //

1368 + time_reduction_factor)

1369 + """

1370 + x = x.unsqueeze(1) if self.conv2d_subsampling else x.transpose(1, 2)

1371 +

1372 + # split inputs if chunking_factor is set

1373 + if (self.subsampling_conv_chunking_factor != -1

1374 + and self.conv2d_subsampling):

1375 + if self.subsampling_conv_chunking_factor == 1:

1376 + # if subsampling_conv_chunking_factor is 1, we split only

1377 + # if needed.

1378 + # avoiding a bug / feature limiting indexing of tensors

1379 + # to 2**31.

1380 + # see https://github.com/pytorch/pytorch/issues/80020

1381 + x_ceil = (2**31 / self._conv_channels * self._stride *

1382 + self._stride)

1383 + need_to_split = torch.numel(x) > x_ceil

1384 + else:

1385 + # if subsampling_conv_chunking_factor > 1 we always split

1386 + need_to_split = True

1387 +

1388 + if need_to_split:

1389 + x, success = self.conv_split_by_batch(x)

1390 + if not success: # if unable to split by batch, try by channel

1391 + if self._subsampling == "dw_striding":

1392 + x = self.conv_split_by_channel(x)

1393 + else:

1394 + x = self.conv(x) # try anyway

1395 + else:

1396 + x = self.conv(x)

1397 + else:

1398 + x = self.conv(x)

1399 +

1400 + # Flatten Channel and Frequency Axes

1401 + if self.conv2d_subsampling:

1402 + b, c, t, f = x.size()

1403 + x = self.out(x.transpose(1, 2).reshape(b, t, -1))

1404 + # Transpose to Channel Last mode

1405 + else:

1406 + x = x.transpose(1, 2)

1407 +

1408 + if mask is None:

1409 + return x, None

1410 +

1411 + max_audio_length = x.shape[1]

1412 + feature_lens = mask.sum(1)

1413 + padding_length = torch.ceil(feature_lens / self.subsampling_factor)

1414 + if self.is_causal and self.subsampling_causal_cond:

1415 + feature_lens_remainder = feature_lens % self.subsampling_factor

1416 + padding_length[feature_lens_remainder != 1] += 1

1417 + pad_mask = torch.arange(0, max_audio_length, device=x.device).expand(

1418 + padding_length.size(0), -1) < padding_length.unsqueeze(1)

1419 + return x, pad_mask.unsqueeze(1)

1420 +

1421 + def reset_parameters(self):

1422 + # initialize weights

1423 + if self._subsampling == "dw_striding":

1424 + with torch.no_grad():

1425 + # init conv

1426 + scale = 1.0 / self._kernel_size

1427 + dw_max = (self._kernel_size**2)**-0.5

1428 + pw_max = self._conv_channels**-0.5

1429 +

1430 + torch.nn.init.uniform_(self.conv[0].weight, -scale, scale)

1431 + torch.nn.init.uniform_(self.conv[0].bias, -scale, scale)

1432 +

1433 + for idx in range(2, len(self.conv), 3):

1434 + torch.nn.init.uniform_(self.conv[idx].weight, -dw_max,

1435 + dw_max)

1436 + torch.nn.init.uniform_(self.conv[idx].bias, -dw_max,

1437 + dw_max)

1438 + torch.nn.init.uniform_(self.conv[idx + 1].weight, -pw_max,

1439 + pw_max)

1440 + torch.nn.init.uniform_(self.conv[idx + 1].bias, -pw_max,

1441 + pw_max)

1442 +

1443 + # init fc (80 * 64 = 5120 from https://github.com/kssteven418/

1444 + # Squeezeformer/blob/13c97d6cf92f2844d2cb3142b4c5bfa9ad1a8951/

1445 + # src/models/conformer_encoder.py#L487

1446 + fc_scale = (self._feat_out * self._feat_in /

1447 + self._sampling_num)**-0.5

1448 + torch.nn.init.uniform_(self.out.weight, -fc_scale, fc_scale)

1449 + torch.nn.init.uniform_(self.out.bias, -fc_scale, fc_scale)

1450 +

1451 + def conv_split_by_batch(self, x):

1452 + """Tries to split input by batch, run conv and concat results"""

1453 + b, _, _, _ = x.size()

1454 + if b == 1: # can't split if batch size is 1

1455 + return x, False

1456 +

1457 + if self.subsampling_conv_chunking_factor > 1:

1458 + cf = self.subsampling_conv_chunking_factor

1459 + else:

1460 + # avoiding a bug / feature limiting indexing of tensors to 2**31

1461 + # see https://github.com/pytorch/pytorch/issues/80020

1462 + x_ceil = 2**31 / self._conv_channels * self._stride * self._stride

1463 + p = math.ceil(math.log(torch.numel(x) / x_ceil, 2))

1464 + cf = 2**p

1465 +

1466 + new_batch_size = b // cf

1467 + if new_batch_size == 0: # input is too big

1468 + return x, False

1469 +

1470 + return (

1471 + torch.cat([

1472 + self.conv(chunk)

1473 + for chunk in torch.split(x, new_batch_size, 0)

1474 + ]),

1475 + True,

1476 + )

1477 +

1478 + def conv_split_by_channel(self, x):

1479 + """For dw convs, tries to split input by time, run conv and concat

1480 + results"""

1481 + x = self.conv[0](x) # full conv2D

1482 + x = self.conv[1](x) # activation

1483 +

1484 + for i in range(self._sampling_num - 1):

1485 + _, c, t, _ = x.size()

1486 +

1487 + if self.subsampling_conv_chunking_factor > 1:

1488 + cf = self.subsampling_conv_chunking_factor

1489 + else:

1490 + # avoiding a bug / feature limiting indexing of tensors

1491 + # to 2**31

1492 + # see https://github.com/pytorch/pytorch/issues/80020

1493 + p = math.ceil(math.log(torch.numel(x) / 2**31, 2))

1494 + cf = 2**p

1495 +

1496 + new_c = int(c // cf)

1497 + if new_c == 0:

1498 + new_c = 1

1499 +

1500 + new_t = int(t // cf)

1501 + if new_t == 0:

1502 + new_t = 1

1503 +

1504 + x = self.channel_chunked_conv(self.conv[i * 3 + 2], new_c,

1505 + x) # conv2D, depthwise

1506 +

1507 + # splitting pointwise convs by time

1508 + x = torch.cat(

1509 + [

1510 + self.conv[i * 3 + 3](chunk)

1511 + for chunk in torch.split(x, new_t, 2)

1512 + ],

1513 + 2,

1514 + ) # conv2D, pointwise

1515 + x = self.conv[i * 3 + 4](x) # activation

1516 + return x

1517 +

1518 + def channel_chunked_conv(self, conv, chunk_size, x):

1519 + """Performs channel chunked convolution"""

1520 +

1521 + ind = 0

1522 + out_chunks = []

1523 + for chunk in torch.split(x, chunk_size, 1):

1524 + step = chunk.size()[1]

1525 +

1526 + if self.is_causal:

1527 + chunk = nn.functional.pad(

1528 + chunk,

1529 + pad=(

1530 + self._kernel_size - 1,

1531 + self._stride - 1,

1532 + self._kernel_size - 1,

1533 + self._stride - 1,

1534 + ),

1535 + )

1536 + ch_out = nn.functional.conv2d(

1537 + chunk,

1538 + conv.weight[ind:ind + step, :, :, :],

1539 + bias=conv.bias[ind:ind + step],

1540 + stride=self._stride,

1541 + padding=0,

1542 + groups=step,

1543 + )

1544 + else:

1545 + ch_out = nn.functional.conv2d(

1546 + chunk,

1547 + conv.weight[ind:ind + step, :, :, :],

1548 + bias=conv.bias[ind:ind + step],

1549 + stride=self._stride,

1550 + padding=self._left_padding,

1551 + groups=step,

1552 + )

1553 + out_chunks.append(ch_out)

1554 + ind += step

1555 +

1556 + return torch.cat(out_chunks, 1)

1557 +

1558 + def change_subsampling_conv_chunking_factor(

1559 + self, subsampling_conv_chunking_factor: int):

1560 + if (subsampling_conv_chunking_factor != -1

1561 + and subsampling_conv_chunking_factor != 1

1562 + and subsampling_conv_chunking_factor % 2 != 0):

1563 + raise ValueError(

1564 + "subsampling_conv_chunking_factor should be -1, 1, or a "\

1565 + "power of 2"

1566 + )

1567 + self.subsampling_conv_chunking_factor = subsampling_conv_chunking_factor

1568 +

1569 +

1570 +def calc_length(lengths,

1571 + all_paddings,

1572 + kernel_size,

1573 + stride,

1574 + ceil_mode,

1575 + repeat_num=1):

1576 + """Calculates the output length of a Tensor passed through a convolution or

1577 + max pooling layer"""

1578 + add_pad: float = all_paddings - kernel_size

1579 + one: float = 1.0

1580 + for i in range(repeat_num):

1581 + lengths = (torch.div(lengths.to(dtype=torch.float) + add_pad, stride) +

1582 + one)

1583 + lengths = torch.ceil(lengths) if ceil_mode else torch.floor(lengths)

1584 + return lengths.to(dtype=torch.int)

1585 +

1586 +

1587 +#### multihead attention starts here

1588 +class AttModule(nn.Module):

1589 + """Attention abstraction module"""

1590 +

1591 + def __init__(self):

1592 + super().__init__()

1593 + self.export_mode = False

1594 +

1595 + def set_export(self, mode=True):

1596 + """set the export mode"""

1597 + self.export_mode = mode

1598 +

1599 + def forward(

1600 + self,

1601 + x: Tensor,

1602 + memory: Optional[Tensor] = None,

1603 + pos_emb: Optional[Tensor] = None,

1604 + att_mask: Optional[Tensor] = None,

1605 + ) -> Tuple[Tensor, Tensor, Optional[Tensor], Optional[Tensor]]:

1606 + """AttModule forward

1607 +

1608 + Args:

1609 + x: torch.Tensor

1610 + input tensor.

1611 + memory: torch.Tensor, optional

1612 + memory tensor.

1613 + pos_emb: torch.Tensor, optional

1614 + positional encoder embedding.

1615 + att_mask: torch.Tensor, optional

1616 + attention mask tensor.

1617 + """

1618 + return x, memory, pos_emb, att_mask

1619 +

1620 +

1621 +class AttBlock(Block, AttModule):

1622 + """Attention Block module to support both Attention and Block module."""

1623 +

1624 + def memory_dims(self, max_len=False):

1625 + """memory dimensions"""

1626 + return (1, self.input_size)

1627 +

1628 +

1629 +def masked_softmax(

1630 + scores,

1631 + mask: Optional[Tensor],

1632 +):

1633 + if mask is not None:

1634 + mask = mask.unsqueeze(1).eq(0) # (batch, 1, time1, time2)

1635 + scores = scores.masked_fill(mask, -torch.inf)

1636 + attn = torch.softmax(scores, dim=-1).masked_fill(

1637 + mask, 0.0) # (batch, head, time1, time2)

1638 + else:

1639 + attn = torch.softmax(scores, dim=-1) # (batch, head, time1, time2)

1640 + return attn

1641 +

1642 +

1643 +class MultiHeadedAttention(nn.Module):

1644 + """Multi-Head Attention layer with optional relative position embedding

1645 + and GLU.

1646 +

1647 + Args:

1648 + n_head: int

1649 + the number of heads.

1650 + n_feat: int

1651 + input size features.

1652 + dropout_rate: float

1653 + dropout rate.

1654 + use_LN: bool

1655 + apply layer norm or not

1656 + dropout_at_output: bool

1657 + whether to apply dropout at output

1658 + attention_inner_dim: int, optional

1659 + the attention dimension used in the class,

1660 + it can be different from the input dimension n_feat.

1661 + default: -1 (equal to n_feat).

1662 + use_pt_scaled_dot_product_attention: bool, optional

1663 + if set True, use pytorch scaled dot product attention in training.

1664 + NOTE: this will NOT be used in ONNX decoding due to a lack of

1665 + support. In that case, we use the original attention

1666 + implementation, which shows no regression.

1667 + default: False.

1668 + n_value: int, optional

1669 + if set to values other than -1, use a different dimension for

1670 + value. With the default value (i.e. -1), it is backward compatible.

1671 + group_size: int, optional. must divide `n_head`

1672 + if group_size > 1: GQA

1673 + if group_size = 1: MHA

1674 + if group_size = n_head: MQA

1675 + """

1676 +

1677 + inv_sqrt_d_k: torch.jit.Final[float]

1678 + h: torch.jit.Final[int]

1679 + h_k: torch.jit.Final[int]

1680 + g: torch.jit.Final[int]

1681 +

1682 + def __init__(

1683 + self,

1684 + n_head,

1685 + n_feat,

1686 + dropout_rate,

1687 + attention_inner_dim=-1,

1688 + glu_type="swish",

1689 + bias_in_glu=True,

1690 + use_pt_scaled_dot_product_attention=False,

1691 + n_value=-1,

1692 + group_size: int = 1,

1693 + ):

1694 + super().__init__()

1695 + if n_value == -1:

1696 + n_value = n_feat

1697 + if attention_inner_dim == -1:

1698 + attention_inner_dim = n_feat

1699 + assert attention_inner_dim % n_head == 0

1700 +

1701 + # We assume d_v always equals d_k

1702 + self.d_k = attention_inner_dim // n_head

1703 + self.inv_sqrt_d_k = 1.0 / math.sqrt(self.d_k)

1704 + self.h = n_head

1705 + assert n_head % group_size == 0, "group_size must divide n_head"

1706 + self.g = group_size

1707 + self.h_k = n_head // group_size

1708 +

1709 + self.linear_q = nn.Linear(n_feat, attention_inner_dim)

1710 + self.linear_k = nn.Linear(n_feat, attention_inner_dim // group_size)

1711 + self.linear_v = nn.Linear(n_value, attention_inner_dim // group_size)

1712 + self.linear_out = nn.Linear(attention_inner_dim // group_size, n_value)

1713 +

1714 + self.attn = torch.jit.Attribute(None, Optional[Tensor])

1715 + self.dropout = nn.Dropout(p=dropout_rate)

1716 + self.dropout_rate = dropout_rate

1717 + self.use_pt_scaled_dot_product_attention = (

1718 + use_pt_scaled_dot_product_attention)

1719 +

1720 + if use_pt_scaled_dot_product_attention and group_size > 1:

1721 + raise ValueError("Cannot use PT Scaled Attention with GQA")

1722 +

1723 + # Torchscript eager quantization. Note that these functions below are

1724 + # NOOPs and have very little impact on performance unless quantization

1725 + # is enabled.

1726 + self.quant_q = torch.ao.quantization.QuantStub()

1727 + self.quant_x = torch.ao.quantization.QuantStub()

1728 + self.dequant = torch.ao.quantization.DeQuantStub()

1729 + self.ffunc = torch.ao.nn.quantized.FloatFunctional()

1730 +

1731 + def forward(

1732 + self,

1733 + query: Tensor,

1734 + key: Tensor,

1735 + value: Tensor,

1736 + pos_k: Tensor,

1737 + pos_v: Tensor,

1738 + mask: Optional[Tensor],

1739 + relative_attention_bias: Optional[Tensor] = None,

1740 + ):

1741 + """Compute 'Scaled Dot Product Attention'.

1742 +

1743 + Args:

1744 + query: torch.Tensor

1745 + query tensor (batch, time1, size)

1746 + key: torch.Tensor

1747 + key tensor (batch, time2, size)

1748 + value: torch.Tensor

1749 + value tensor (batch, time1, size)

1750 + pos_k: torch.Tensor

1751 + key tensor used for relative positional embedding.

1752 + pos_v: torch.Tensor

1753 + value tensor used for relative positional embedding.

1754 + mask: torch.Tensor

1755 + mask tensor (batch, time1, time2)

1756 + relative_attention_bias: torch.Tensor

1757 + bias added to attention logits w.r.t. relative positions

1758 + (1, n_head, time1, time2)

1759 + """

1760 + n_batch = query.size(0)

1761 +

1762 + q = self.linear_q(query).view(n_batch, -1, self.h,

1763 + self.d_k) # (b, t, d)

1764 + k = self.linear_k(key).view(n_batch, -1, self.h_k,

1765 + self.d_k) # (b, t, d)

1766 + v = self.linear_v(value).view(n_batch, -1, self.h_k, self.d_k)

1767 + q = (q.transpose(1, 2) if self.use_pt_scaled_dot_product_attention

1768 + and not torch.jit.is_scripting() else q.transpose(1, 2) *

1769 + self.inv_sqrt_d_k)

1770 + k = k.transpose(1, 2) # (batch, head_k, time2, d_k)

1771 + v = v.transpose(1, 2) # (batch, head_k, time2, d_k)

1772 +

1773 + if (self.use_pt_scaled_dot_product_attention

1774 + and not torch.jit.is_scripting()):

1775 + attn_mask = None

1776 + if mask is not None:

1777 + mask = mask.unsqueeze(1)

1778 + if relative_attention_bias is not None:

1779 + attn_mask = mask + relative_attention_bias

1780 + else:

1781 + attn_mask = mask

1782 + if mask.dtype != q.dtype:

1783 + attn_mask = attn_mask.to(q.dtype)

1784 +

1785 + with torch.backends.cuda.sdp_kernel(enable_flash=True,

1786 + enable_math=True,

1787 + enable_mem_efficient=True):

1788 + x = torch.nn.functional.scaled_dot_product_attention(

1789 + q,

1790 + k,

1791 + v,

1792 + attn_mask=attn_mask,

1793 + dropout_p=self.dropout_rate,

1794 + )

1795 + else:

1796 + if self.h != self.h_k:

1797 + q = q.reshape(n_batch, self.g, self.h_k, -1, self.d_k)

1798 + A = torch.einsum("b g h t d, b h s d -> b h t s", q, k)

1799 + else:

1800 + A = torch.matmul(q, k.transpose(-2, -1))

1801 + if pos_k is not None:

1802 + if self.h != self.h_k:

1803 + B = torch.einsum("b g h t d, t s d -> b h t s", q, pos_k)

1804 + else:

1805 + reshape_q = (q.contiguous().view(n_batch * self.h, -1,

1806 + self.d_k).transpose(0, 1)

1807 + ) # (t1,nh,dk)

1808 + B = torch.matmul(reshape_q,

1809 + pos_k.transpose(-2,

1810 + -1)) # pos_k: (t1,dk,t2)

1811 + B = B.transpose(0, 1).view(n_batch, self.h, pos_k.size(0),

1812 + pos_k.size(1))

1813 + scores = A + B

1814 + else:

1815 + scores = A

1816 +

1817 + if relative_attention_bias is not None:

1818 + scores = scores + relative_attention_bias

1819 +

1820 + attn = masked_softmax(scores, mask) # (batch, head, time1, time2)

1821 +

1822 + self.attn = attn

1823 +

1824 + p_attn = self.dropout(attn)

1825 + x = torch.matmul(p_attn.to(v.dtype),

1826 + v) # (batch, head, time1, d_k)

1827 + if pos_v is not None:

1828 + reshape_attn = (p_attn.contiguous().view(

1829 + n_batch * self.h, pos_v.size(0),

1830 + pos_v.size(1)).transpose(0, 1)) # (t1, bh, t2)

1831 +

1832 + attn_v = (torch.matmul(reshape_attn, pos_v).transpose(

1833 + 0, 1).contiguous().view(n_batch, self.h, pos_v.size(0),

1834 + self.d_k))

1835 + x = x + attn_v

1836 + x = (x.transpose(1, 2).contiguous().view(n_batch, -1,

1837 + self.h_k * self.d_k)

1838 + ) # (batch, time1, d_model)

1839 +

1840 + return self.linear_out(x) # (batch, time1, d_model)

1841 +

1842 +

1843 +def validate_checkpointing_config(activation_checkpointing):

1844 + """validate activation checkpointing configuration"""

1845 + if isinstance(activation_checkpointing, str):

1846 + assert activation_checkpointing in (

1847 + "",

1848 + "checkpoint",

1849 + "offload",

1850 + ), "activation_checkpointing has to be a dict or a str in "\

1851 + "('', 'checkpoint', 'offload')."

1852 + elif isinstance(activation_checkpointing, dict):

1853 + assert activation_checkpointing.get("module", "transformer") in (

1854 + "transformer",

1855 + "attention",

1856 + ), "module in activation_checkpointing has to be in "\

1857 + "('transformer', 'attention')."

1858 + else:

1859 + raise ValueError("activation_checkpointing has to be a str"\

1860 + " or dict.")

1861 +

1862 +

1863 +def embedding_checkpoint_wrapper(

1864 + activation_checkpointing: Union[str, Dict], ) -> Callable:

1865 + """return encoder embedding activation checkpoint wrapper"""

1866 + validate_checkpointing_config(activation_checkpointing)

1867 +

1868 + if isinstance(activation_checkpointing, str):

1869 + if activation_checkpointing:

1870 + if activation_checkpointing == "offload":

1871 + return offload_wrapper

1872 + return partial(checkpoint_wrapper)

1873 + return lambda x: x

1874 +

1875 + if isinstance(activation_checkpointing, dict):

1876 + enabled = activation_checkpointing.get("embed", False)

1877 + if enabled:

1878 + offloading = activation_checkpointing.get("offload", False)

1879 + if offloading:

1880 + return offload_wrapper

1881 + impl = (CheckpointImpl.REENTRANT if activation_checkpointing.get(

1882 + "reentrant", False) else CheckpointImpl.NO_REENTRANT)

1883 + return partial(checkpoint_wrapper, checkpoint_impl=impl)

1884 + return lambda x: x

1885 + raise ValueError("Invalid activation_checkpointing config")

1886 +

1887 +

1888 +def attn_checkpointing(activation_checkpointing: Union[str, Dict],

1889 + i) -> Union[str, Dict]:

1890 + """return activation checkpointing config for attention layer"""

1891 + if isinstance(activation_checkpointing, str):

1892 + return ""

1893 +

1894 + if isinstance(activation_checkpointing, dict):

1895 + target_layer_cls = activation_checkpointing.get(

1896 + "module", "transformer")

1897 + checkpointing_interval = activation_checkpointing.get("interval", 1)

1898 + if target_layer_cls == "attention" and i % checkpointing_interval == 0:

1899 + return activation_checkpointing

1900 + return ""

1901 +

1902 + raise ValueError("Invalid activation_checkpointing config")

1903 +

1904 +

1905 +class MultiSequential(torch.nn.Sequential):

1906 + """Multi-input multi-output torch.nn.Sequential"""

1907 +

1908 + @torch.jit.ignore

1909 + def forward(self, *args):

1910 + """Forward method implementation."""

1911 + for m in self:

1912 + args = m(*args)

1913 + return args

1914 +

1915 +

1916 +def repeat(repeat_num, module_gen_fn):

1917 + """repeat module N times

1918 +

1919 + :param int repeat_num: repeat time

1920 + :param function module_gen_fn: function to generate module

1921 + :return: repeated modules

1922 + :rtype: MultiSequential

1923 + """

1924 + return MultiSequential(*[module_gen_fn(i) for i in range(repeat_num)])

1925 +

1926 +

1927 +def get_offset(input_layer: str, time_reduction: int):

1928 + """Get an offset. We will use the offset for determining #frames of a

1929 + subsampled feature.

1930 +

1931 + Args:

1932 + input_layer (str): Type of an input layer

1933 + time_reduction (int): time reduction factor for downsampling a feature

1934 + Returns:

1935 + int: offset

1936 + """

1937 + if input_layer in ("conv2d", "nemo_conv") and time_reduction == 4:

1938 + return 3

1939 + if input_layer in ("conv2d", ) and time_reduction == 6:

1940 + return 1

1941 + if input_layer in ("conv2d", "nemo_conv") and time_reduction == 8:

1942 + return 7

1943 + return 0

1944 +

1945 +

1946 +def unfold_tensor(xs_pad, max_seq_len):

1947 + """

1948 + For a given tensor with shape of (N, T, D), if sequence length T is

1949 + longer than max_seq_len, this function unfold it to a

1950 + (NT', max_seq_len, D) where T' is T // max_seq_len.

1951 + Args:

1952 + xs_pad: N, T, D

1953 + """

1954 + _, _, D = xs_pad.shape

1955 + xs_pad = xs_pad.transpose(-1, -2) # convert to N, D, T

1956 + # N x D x 1 x T => N x (D x max_seq_len) x T'

1957 + xs_pad = F.unfold(

1958 + xs_pad[..., None, :],

1959 + kernel_size=(1, max_seq_len),

1960 + stride=(1, max_seq_len),

1961 + )

1962 + new_bsz, _, slen = xs_pad.shape

1963 + # N x D x max_seq_len x T'

1964 + xs_pad = xs_pad.view(new_bsz, -1, max_seq_len, slen)

1965 + # N x T' x max_seq_len x D

1966 + xs_pad = xs_pad.permute(0, 3, 2, 1).contiguous()

1967 + # NT' x max_seq_len x D

1968 + xs_pad = xs_pad.view(-1, max_seq_len, D)

1969 + return xs_pad

vllm/model_executor/models/registry.py

@@ -182,6 +182,7 @@ _MULTIMODAL_MODELS = {

182 "Qwen2_5_VLForConditionalGeneration": ("qwen2_5_vl", "Qwen2_5_VLForConditionalGeneration"), # noqa: E501

183 "Qwen2AudioForConditionalGeneration": ("qwen2_audio", "Qwen2AudioForConditionalGeneration"), # noqa: E501

184 "UltravoxModel": ("ultravox", "UltravoxModel"),

185 + "Phi4MMForCausalLM": ("phi4mm", "Phi4MMForCausalLM"),

186 # [Encoder-decoder]

187 "Florence2ForConditionalGeneration": ("florence2", "Florence2ForConditionalGeneration"), # noqa: E501

188 "MllamaForConditionalGeneration": ("mllama", "MllamaForConditionalGeneration"), # noqa: E501

vllm/model_executor/models/vision_siglip_navit.py new

+1966

@@ -0,0 +1,1966 @@

1 +# SPDX-License-Identifier: Apache-2.0

3 +#

4 +# Licensed under the Apache License, Version 2.0 (the "License");

5 +# you may not use this file except in compliance with the License.

6 +# You may obtain a copy of the License at

7 +#

8 +# http://www.apache.org/licenses/LICENSE-2.0

9 +#

10 +# Unless required by applicable law or agreed to in writing, software

11 +# distributed under the License is distributed on an "AS IS" BASIS,

12 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

13 +# See the License for the specific language governing permissions and

14 +# limitations under the License.

15 +""" Siglip model configuration"""

16 +

17 +import math

18 +import os

19 +import warnings

20 +from dataclasses import dataclass

21 +from typing import Any, Optional, Tuple, Union

22 +

23 +import numpy as np

24 +import torch

25 +import torch.nn.functional as F

26 +import torch.utils.checkpoint

27 +from torch import nn

28 +from torch.nn.init import _calculate_fan_in_and_fan_out

29 +from transformers.activations import ACT2FN

30 +from transformers.configuration_utils import PretrainedConfig

31 +from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask

32 +from transformers.modeling_outputs import (BaseModelOutput,

33 + BaseModelOutputWithPooling)

34 +from transformers.modeling_utils import PreTrainedModel

35 +from transformers.utils import (ModelOutput, add_start_docstrings,

36 + add_start_docstrings_to_model_forward, logging,

37 + replace_return_docstrings)

38 +

39 +from vllm.platforms import _Backend

40 +

41 +from .vision import get_vit_attn_backend

42 +

43 +logger = logging.get_logger(__name__)

44 +

45 +SIGLIP_PRETRAINED_CONFIG_ARCHIVE_MAP = {

46 + "google/siglip-base-patch16-224":

47 + "https://huggingface.co/google/siglip-base-patch16-224/"\

48 + "resolve/main/config.json",

49 +}

50 +

51 +

52 +class SiglipTextConfig(PretrainedConfig):

53 + r"""

54 + This is the configuration class to store the configuration of a

55 + [`SiglipTextModel`]. It is used to instantiate a Siglip text encoder

56 + according to the specified arguments, defining the model architecture.

57 + Instantiating a configuration with the defaults will yield a similar

58 + configuration to that of the text encoder of the Siglip [google/

59 + siglip-base-patch16-224](https://huggingface.co/google/siglip-base

60 + -patch16-224) architecture.

61 + Configuration objects inherit from [`PretrainedConfig`] and can be used

62 + to control the model outputs. Read the documentation from

63 + [`PretrainedConfig`] for more information.

64 + Args:

65 + vocab_size (`int`, *optional*, defaults to 32000):

66 + Vocabulary size of the Siglip text model. Defines the number of

67 + different tokens that can be represented by the `inputs_ids`

68 + passed when calling [`SiglipModel`].

69 + hidden_size (`int`, *optional*, defaults to 768):

70 + Dimensionality of the encoder layers and the pooler layer.

71 + intermediate_size (`int`, *optional*, defaults to 3072):

72 + Dimensionality of the "intermediate" (i.e., feed-forward) layer

73 + in the Transformer encoder.

74 + num_hidden_layers (`int`, *optional*, defaults to 12):

75 + Number of hidden layers in the Transformer encoder.

76 + num_attention_heads (`int`, *optional*, defaults to 12):

77 + Number of attention heads for each attention layer in the

78 + Transformer encoder.

79 + max_position_embeddings (`int`, *optional*, defaults to 64):

80 + The maximum sequence length that this model might ever be used

81 + with. Typically set this to something large

82 + just in case (e.g., 512 or 1024 or 2048).

83 + hidden_act (`str` or `function`, *optional*, defaults to

84 + `"gelu_pytorch_tanh"`):

85 + The non-linear activation function (function or string) in the

86 + encoder and pooler. If string, `"gelu"`,

87 + `"relu"`, `"selu"` and `"gelu_new"` `"quick_gelu"` are supported.

88 + layer_norm_eps (`float`, *optional*, defaults to 1e-06):

89 + The epsilon used by the layer normalization layers.

90 + attention_dropout (`float`, *optional*, defaults to 0.0):

91 + The dropout ratio for the attention probabilities.

92 + pad_token_id (`int`, *optional*, defaults to 1):

93 + The id of the padding token in the vocabulary.

94 + bos_token_id (`int`, *optional*, defaults to 49406):

95 + The id of the beginning-of-sequence token in the vocabulary.

96 + eos_token_id (`int`, *optional*, defaults to 49407):

97 + The id of the end-of-sequence token in the vocabulary.

98 + Example:

99 + ```python

100 + >>> from transformers import SiglipTextConfig, SiglipTextModel

101 + >>> # Initializing a SiglipTextConfig with google/siglip-base-patch16-224

102 + style configuration

103 + >>> configuration = SiglipTextConfig()

104 + >>> # Initializing a SiglipTextModel (with random weights) from the

105 + google/siglip-base-patch16-224 style configuration

106 + >>> model = SiglipTextModel(configuration)

107 + >>> # Accessing the model configuration

108 + >>> configuration = model.config

109 + ```"""

110 +

111 + model_type = "siglip_text_model"

112 +

113 + def __init__(

114 + self,

115 + vocab_size=32000,

116 + hidden_size=768,

117 + intermediate_size=3072,

118 + num_hidden_layers=12,

119 + num_attention_heads=12,

120 + max_position_embeddings=64,

121 + hidden_act="gelu_pytorch_tanh",

122 + layer_norm_eps=1e-6,

123 + attention_dropout=0.0,

124 + # This differs from `CLIPTokenizer`'s default and from openai/siglip

125 + # See https://github.com/huggingface/transformers/pull/24773#

126 + # issuecomment-1632287538

127 + pad_token_id=1,

128 + bos_token_id=49406,

129 + eos_token_id=49407,

130 + _flash_attn_2_enabled=True,

131 + **kwargs,

132 + ):

133 + super().__init__(pad_token_id=pad_token_id,

134 + bos_token_id=bos_token_id,

135 + eos_token_id=eos_token_id,

136 + **kwargs)

137 +

138 + self.vocab_size = vocab_size

139 + self.hidden_size = hidden_size

140 + self.intermediate_size = intermediate_size

141 + self.num_hidden_layers = num_hidden_layers

142 + self.num_attention_heads = num_attention_heads

143 + self.max_position_embeddings = max_position_embeddings

144 + self.layer_norm_eps = layer_norm_eps

145 + self.hidden_act = hidden_act

146 + self.attention_dropout = attention_dropout

147 + self._flash_attn_2_enabled = _flash_attn_2_enabled

148 +

149 + @classmethod

150 + def from_pretrained(cls, pretrained_model_name_or_path: Union[str,

151 + os.PathLike],

152 + **kwargs) -> "PretrainedConfig":

153 + cls._set_token_in_kwargs(kwargs)

154 +

155 + config_dict, kwargs = cls.get_config_dict(

156 + pretrained_model_name_or_path, **kwargs)

157 +

158 + # get the text config dict if we are loading from SiglipConfig

159 + if config_dict.get("model_type") == "siglip":

160 + config_dict = config_dict["text_config"]

161 +

162 + if "model_type" in config_dict and hasattr(

163 + cls,

164 + "model_type") and config_dict["model_type"] != cls.model_type:

165 + logger.warning(

166 + "You are using a model of type %s to instantiate a model of "

167 + "type %s. This is not supported for all configurations of "

168 + "models and can yield errors.", config_dict['model_type'],

169 + cls.model_type)

170 +

171 + return cls.from_dict(config_dict, **kwargs)

172 +

173 +

174 +class SiglipVisionConfig(PretrainedConfig):

175 + r"""

176 + This is the configuration class to store the configuration of a

177 + [`SiglipVisionModel`]. It is used to instantiate a

178 + Siglip vision encoder according to the specified arguments, defining the

179 + model architecture. Instantiating a configuration with the defaults will

180 + yield a similar configuration to that of the vision encoder of the Siglip

181 + [google/siglip-base-patch16-224](https://huggingface.co/google/

182 + siglip-base-patch16-224) architecture.

183 + Configuration objects inherit from [`PretrainedConfig`] and can be used

184 + to control the model outputs. Read the

185 + documentation from [`PretrainedConfig`] for more information.

186 + Args:

187 + hidden_size (`int`, *optional*, defaults to 768):

188 + Dimensionality of the encoder layers and the pooler layer.

189 + intermediate_size (`int`, *optional*, defaults to 3072):

190 + Dimensionality of the "intermediate" (i.e., feed-forward) layer

191 + in the Transformer encoder.

192 + num_hidden_layers (`int`, *optional*, defaults to 12):

193 + Number of hidden layers in the Transformer encoder.

194 + num_attention_heads (`int`, *optional*, defaults to 12):

195 + Number of attention heads for each attention layer in the

196 + Transformer encoder.

197 + num_channels (`int`, *optional*, defaults to 3):

198 + Number of channels in the input images.

199 + image_size (`int`, *optional*, defaults to 224):

200 + The size (resolution) of each image.

201 + patch_size (`int`, *optional*, defaults to 16):

202 + The size (resolution) of each patch.

203 + hidden_act (`str` or `function`, *optional*, defaults to

204 + `"gelu_pytorch_tanh"`):

205 + The non-linear activation function (function or string) in the

206 + encoder and pooler. If string, `"gelu"`, `"relu"`, `"selu"` and

207 + `"gelu_new"` ``"quick_gelu"` are supported.

208 + layer_norm_eps (`float`, *optional*, defaults to 1e-06):

209 + The epsilon used by the layer normalization layers.

210 + attention_dropout (`float`, *optional*, defaults to 0.0):

211 + The dropout ratio for the attention probabilities.

212 + Example:

213 + ```python

214 + >>> from transformers import SiglipVisionConfig, SiglipVisionModel

215 + >>> # Initializing a SiglipVisionConfig with google/siglip-base-patch16-224

216 + style configuration

217 + >>> configuration = SiglipVisionConfig()

218 + >>> # Initializing a SiglipVisionModel (with random weights) from the

219 + google/siglip-base-patch16-224 style configuration

220 + >>> model = SiglipVisionModel(configuration)

221 + >>> # Accessing the model configuration

222 + >>> configuration = model.config

223 + ```"""

224 +

225 + model_type = "siglip_vision_model"

226 +

227 + def __init__(

228 + self,

229 + hidden_size=768,

230 + intermediate_size=3072,

231 + num_hidden_layers=12,

232 + num_attention_heads=12,

233 + num_channels=3,

234 + image_size=224,

235 + patch_size=16,

236 + hidden_act="gelu_pytorch_tanh",

237 + layer_norm_eps=1e-6,

238 + attention_dropout=0.0,

239 + _flash_attn_2_enabled=True,

240 + **kwargs,

241 + ):

242 + super().__init__(**kwargs)

243 +

244 + self.hidden_size = hidden_size

245 + self.intermediate_size = intermediate_size

246 + self.num_hidden_layers = num_hidden_layers

247 + self.num_attention_heads = num_attention_heads

248 + self.num_channels = num_channels

249 + self.patch_size = patch_size

250 + self.image_size = image_size

251 + self.attention_dropout = attention_dropout

252 + self.layer_norm_eps = layer_norm_eps

253 + self.hidden_act = hidden_act

254 + self._flash_attn_2_enabled = _flash_attn_2_enabled

255 +

256 + @classmethod

257 + def from_pretrained(cls, pretrained_model_name_or_path: Union[str,

258 + os.PathLike],

259 + **kwargs) -> "PretrainedConfig":

260 + cls._set_token_in_kwargs(kwargs)

261 +

262 + config_dict, kwargs = cls.get_config_dict(

263 + pretrained_model_name_or_path, **kwargs)

264 +

265 + # get the vision config dict if we are loading from SiglipConfig

266 + if config_dict.get("model_type") == "siglip":

267 + config_dict = config_dict["vision_config"]

268 +

269 + if "model_type" in config_dict and hasattr(

270 + cls,

271 + "model_type") and config_dict["model_type"] != cls.model_type:

272 + logger.warning(

273 + "You are using a model of type %s to "

274 + "instantiate a model of type %s. This is not"

275 + " supported for all configurations of models and can yield"

276 + " errors.", config_dict['model_type'], cls.model_type)

277 +

278 + return cls.from_dict(config_dict, **kwargs)

279 +

280 +

281 +class SiglipConfig(PretrainedConfig):

282 + r"""

283 + [`SiglipConfig`] is the configuration class to store the configuration of a

284 + [`SiglipModel`]. It is used to instantiate a Siglip model according to the

285 + specified arguments, defining the text model and vision model configs.

286 + Instantiating a configuration with the defaults will yield a similar

287 + configuration to that of the Siglip [google/siglip-base-patch16-224](

288 + https://huggingface.co/google/siglip-base-patch16-224) architecture.

289 + Configuration objects inherit from [`PretrainedConfig`] and can be used to

290 + control the model outputs. Read the documentation from

291 + [`PretrainedConfig`] for more information.

292 + Args:

293 + text_config (`dict`, *optional*):

294 + Dictionary of configuration options used to initialize

295 + [`SiglipTextConfig`].

296 + vision_config (`dict`, *optional*):

297 + Dictionary of configuration options used to initialize

298 + [`SiglipVisionConfig`].

299 + kwargs (*optional*):

300 + Dictionary of keyword arguments.

301 + Example:

302 + ```python

303 + >>> from transformers import SiglipConfig, SiglipModel

304 + >>> # Initializing a SiglipConfig with google/siglip-base-patch16-224

305 + style configuration

306 + >>> configuration = SiglipConfig()

307 + >>> # Initializing a SiglipModel (with random weights) from the

308 + google/siglip-base-patch16-224 style configuration

309 + >>> model = SiglipModel(configuration)

310 + >>> # Accessing the model configuration

311 + >>> configuration = model.config

312 + >>> # We can also initialize a SiglipConfig from a SiglipTextConfig

313 + and a SiglipVisionConfig

314 + >>> from transformers import SiglipTextConfig, SiglipVisionConfig

315 + >>> # Initializing a SiglipText and SiglipVision configuration

316 + >>> config_text = SiglipTextConfig()

317 + >>> config_vision = SiglipVisionConfig()

318 + >>> config = SiglipConfig.from_text_vision_configs(config_text,

319 + config_vision)

320 + ```"""

321 +

322 + model_type = "siglip"

323 +

324 + def __init__(self, text_config=None, vision_config=None, **kwargs):

325 + super().__init__(**kwargs)

326 +

327 + if text_config is None:

328 + text_config = {}

329 + logger.info(

330 + "`text_config` is `None`. Initializing the `SiglipTextConfig`"

331 + " with default values.")

332 +

333 + if vision_config is None:

334 + vision_config = {}

335 + logger.info("`vision_config` is `None`. initializing the "

336 + "`SiglipVisionConfig` with default values.")

337 +

338 + self.text_config = SiglipTextConfig(**text_config)

339 + self.vision_config = SiglipVisionConfig(**vision_config)

340 +

341 + self.initializer_factor = 1.0

342 +

343 + @classmethod

344 + def from_text_vision_configs(cls, text_config: SiglipTextConfig,

345 + vision_config: SiglipVisionConfig, **kwargs):

346 + r"""

347 + Instantiate a [`SiglipConfig`] (or a derived class) from siglip text

348 + model configuration and siglip vision

349 + model configuration.

350 + Returns:

351 + [`SiglipConfig`]: An instance of a configuration object

352 + """

353 +

354 + return cls(text_config=text_config.to_dict(),

355 + vision_config=vision_config.to_dict(),

356 + **kwargs)

357 +

358 +

359 +# coding=utf-8

361 +#

362 +# Licensed under the Apache License, Version 2.0 (the "License");

363 +# you may not use this file except in compliance with the License.

364 +# You may obtain a copy of the License at

365 +#

366 +# http://www.apache.org/licenses/LICENSE-2.0

367 +#

368 +# Unless required by applicable law or agreed to in writing, software

369 +# distributed under the License is distributed on an "AS IS" BASIS,

370 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

371 +# See the License for the specific language governing permissions and

372 +# limitations under the License.

373 +""" PyTorch Siglip model."""

374 +

375 +_CHECKPOINT_FOR_DOC = "google/siglip-base-patch16-224"

376 +

377 +SIGLIP_PRETRAINED_MODEL_ARCHIVE_LIST = [

378 + "google/siglip-base-patch16-224",

379 + # See all SigLIP models at https://huggingface.co/models?filter=siglip

380 +]

381 +

382 +

383 +# Copied from transformers.models.llama.modeling_llama._get_unpad_data

384 +def _get_unpad_data(attention_mask):

385 + seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)

386 + indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()

387 + max_seqlen_in_batch = seqlens_in_batch.max().item()

388 + cu_seqlens = F.pad(

389 + torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))

390 + return (

391 + indices,

392 + cu_seqlens,

393 + max_seqlen_in_batch,

394 + )

395 +

396 +

397 +def _trunc_normal_(tensor, mean, std, a, b):

398 + # Cut & paste from PyTorch official master until it's in a few official

399 + # releases - RW

400 + # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/

401 + # truncated_normal.pdf

402 + def norm_cdf(x):

403 + # Computes standard normal cumulative distribution function

404 + return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0

405 +

406 + if (mean < a - 2 * std) or (mean > b + 2 * std):

407 + warnings.warn(

408 + "mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "

409 + "The distribution of values may be incorrect.",

410 + stacklevel=2,

411 + )

412 +

413 + # Values are generated by using a truncated uniform distribution and

414 + # then using the inverse CDF for the normal distribution.

415 + # Get upper and lower cdf values

416 + l = norm_cdf((a - mean) / std) # noqa

417 + u = norm_cdf((b - mean) / std) # noqa

418 +

419 + # Uniformly fill tensor with values from [l, u], then translate to

420 + # [2l-1, 2u-1].

421 + tensor.uniform_(2 * l - 1, 2 * u - 1)

422 +

423 + # Use inverse cdf transform for normal distribution to get truncated

424 + # standard normal

425 + if tensor.dtype in [torch.float16, torch.bfloat16]:

426 + # The `erfinv_` op is not (yet?) defined in float16+cpu, bfloat16+gpu

427 + og_dtype = tensor.dtype

428 + tensor = tensor.to(torch.float32)

429 + tensor.erfinv_()

430 + tensor = tensor.to(og_dtype)

431 + else:

432 + tensor.erfinv_()

433 +

434 + # Transform to proper mean, std

435 + tensor.mul_(std * math.sqrt(2.0))

436 + tensor.add_(mean)

437 +

438 + # Clamp to ensure it's in the proper range

439 + if tensor.dtype == torch.float16:

440 + # The `clamp_` op is not (yet?) defined in float16+cpu

441 + tensor = tensor.to(torch.float32)

442 + tensor.clamp_(min=a, max=b)

443 + tensor = tensor.to(torch.float16)

444 + else:

445 + tensor.clamp_(min=a, max=b)

446 +

447 +

448 +def trunc_normal_tf_(tensor: torch.Tensor,

449 + mean: float = 0.0,

450 + std: float = 1.0,

451 + a: float = -2.0,

452 + b: float = 2.0) -> torch.Tensor:

453 + """Fills the input Tensor with values drawn from a truncated

454 + normal distribution. The values are effectively drawn from the

455 + normal distribution :math:`\\mathcal{N}(\text{mean}, \text{std}^2)`

456 + with values outside :math:`[a, b]` redrawn until they are within

457 + the bounds. The method used for generating the random values works

458 + best when :math:`a \\leq \text{mean} \\leq b`.

459 + NOTE: this 'tf' variant behaves closer to Tensorflow / JAX impl where

460 + the bounds [a, b] are applied when sampling the normal distribution with

461 + mean=0, std=1.0 and the result is subsequently scaled and shifted by the

462 + mean and std args.

463 + Args:

464 + tensor: an n-dimensional `torch.Tensor`

465 + mean: the mean of the normal distribution

466 + std: the standard deviation of the normal distribution

467 + a: the minimum cutoff value

468 + b: the maximum cutoff value

469 + """

470 + with torch.no_grad():

471 + _trunc_normal_(tensor, 0, 1.0, a, b)

472 + tensor.mul_(std).add_(mean)

473 +

474 +

475 +def variance_scaling_(tensor, scale=1.0, mode="fan_in", distribution="normal"):

476 + fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)

477 + if mode == "fan_in":

478 + denom = fan_in

479 + elif mode == "fan_out":

480 + denom = fan_out

481 + elif mode == "fan_avg":

482 + denom = (fan_in + fan_out) / 2

483 +

484 + variance = scale / denom

485 +

486 + if distribution == "truncated_normal":

487 + # constant is stddev of standard normal truncated to (-2, 2)

488 + trunc_normal_tf_(tensor, std=math.sqrt(variance) / 0.87962566103423978)

489 + elif distribution == "normal":

490 + with torch.no_grad():

491 + tensor.normal_(std=math.sqrt(variance))

492 + elif distribution == "uniform":

493 + bound = math.sqrt(3 * variance)

494 + with torch.no_grad():

495 + tensor.uniform_(-bound, bound)

496 + else:

497 + raise ValueError(f"invalid distribution {distribution}")

498 +

499 +

500 +def lecun_normal_(tensor):

501 + variance_scaling_(tensor, mode="fan_in", distribution="truncated_normal")

502 +

503 +

504 +def default_flax_embed_init(tensor):

505 + variance_scaling_(tensor, mode="fan_in", distribution="normal")

506 +

507 +

508 +@dataclass

509 +# Copied from transformers.models.clip.modeling_clip.CLIPVisionModelOutput with

510 +# CLIP->Siglip

511 +class SiglipVisionModelOutput(ModelOutput):

512 + """

513 + Base class for vision model's outputs that also contains image embeddings

514 + of the pooling of the last hidden states.

515 + Args:

516 + image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)`

517 + *optional* returned when model is initialized with

518 + `with_projection=True`):

519 + The image embeddings obtained by applying the projection layer to

520 + the pooler_output.

521 + last_hidden_state (`torch.FloatTensor` of shape `(batch_size,

522 + sequence_length, hidden_size)`):

523 + Sequence of hidden-states at the output of the last layer of the

524 + model.

525 + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when

526 + `output_hidden_states=True` is passed or when

527 + `config.output_hidden_states=True`):

528 + Tuple of `torch.FloatTensor` (one for the output of the embeddings,

529 + if the model has an embedding layer, + one for the output of each

530 + layer) of shape `(batch_size, sequence_length, hidden_size)`.

531 + Hidden-states of the model at the output of each layer plus the

532 + optional initial embedding outputs.

533 + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when

534 + `output_attentions=True` is passed or when

535 + `config.output_attentions=True`):

536 + Tuple of `torch.FloatTensor` (one for each layer) of shape

537 + `(batch_size, num_heads, sequence_length, sequence_length)`.

538 + Attentions weights after the attention softmax, used to compute the

539 + weighted average in the self-attention heads.

540 + """

541 +

542 + image_embeds: Optional[torch.FloatTensor] = None

543 + last_hidden_state: torch.FloatTensor = None

544 + hidden_states: Optional[Tuple[torch.FloatTensor]] = None

545 + attentions: Optional[Tuple[torch.FloatTensor]] = None

546 +

547 +

548 +@dataclass

549 +# Copied from transformers.models.clip.modeling_clip.CLIPTextModelOutput with

550 +# CLIP->Siglip

551 +class SiglipTextModelOutput(ModelOutput):

552 + """

553 + Base class for text model's outputs that also contains a pooling of the

554 + last hidden states.

555 + Args:

556 + text_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)`

557 + *optional* returned when model is initialized with

558 + `with_projection=True`):

559 + The text embeddings obtained by applying the projection layer to

560 + model.

561 + the pooler_output.

562 + last_hidden_state (`torch.FloatTensor` of shape `(batch_size,

563 + sequence_length, hidden_size)`):

564 + Sequence of hidden-states at the output of the last layer of the

565 + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when

566 + `output_hidden_states=True` is passed or when

567 + `config.output_hidden_states=True`):

568 + Tuple of `torch.FloatTensor` (one for the output of the

569 + embeddings, if the model has an embedding layer, + one for the

570 + output of each layer) of shape `(batch_size, sequence_length,

571 + hidden_size)`.

572 + Hidden-states of the model at the output of each layer plus the

573 + optional initial embedding outputs.

574 + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when

575 + `output_attentions=True` is passed or when

576 + `config.output_attentions=True`):

577 + Tuple of `torch.FloatTensor` (one for each layer) of shape

578 + `(batch_size, num_heads, sequence_length, sequence_length)`.

579 + Attentions weights after the attention softmax, used to compute

580 + the weighted average in the self-attention heads.

581 + """

582 +

583 + text_embeds: Optional[torch.FloatTensor] = None

584 + last_hidden_state: torch.FloatTensor = None

585 + hidden_states: Optional[Tuple[torch.FloatTensor]] = None

586 + attentions: Optional[Tuple[torch.FloatTensor]] = None

587 +

588 +

589 +@dataclass

590 +# Copied from transformers.models.clip.modeling_clip.CLIPOutput with

591 +# CLIP->Siglip

592 +class SiglipOutput(ModelOutput):

593 + """

594 + Args:

595 + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when

596 + `return_loss` is `True`):

597 + Contrastive loss for image-text similarity.

598 + logits_per_image:(`torch.FloatTensor` of shape `(image_batch_size,

599 + text_batch_size)`):

600 + The scaled dot product scores between `image_embeds` and

601 + `text_embeds`. This represents the image-text similarity scores.

602 + logits_per_text:(`torch.FloatTensor` of shape `(text_batch_size,

603 + image_batch_size)`):

604 + The scaled dot product scores between `text_embeds` and

605 + `image_embeds`. This represents the text-image similarity scores.

606 + text_embeds(`torch.FloatTensor` of shape `(batch_size, output_dim`):

607 + The text embeddings obtained by applying the projection layer to

608 + the pooled output of [`SiglipTextModel`].

609 + image_embeds(`torch.FloatTensor` of shape `(batch_size, output_dim`):

610 + The image embeddings obtained by applying the projection layer to

611 + the pooled output of [`SiglipVisionModel`].

612 + text_model_output(`BaseModelOutputWithPooling`):

613 + The output of the [`SiglipTextModel`].

614 + vision_model_output(`BaseModelOutputWithPooling`):

615 + The output of the [`SiglipVisionModel`].

616 + """

617 +

618 + loss: Optional[torch.FloatTensor] = None

619 + logits_per_image: torch.FloatTensor = None

620 + logits_per_text: torch.FloatTensor = None

621 + text_embeds: torch.FloatTensor = None

622 + image_embeds: torch.FloatTensor = None

623 + text_model_output: BaseModelOutputWithPooling = None

624 + vision_model_output: BaseModelOutputWithPooling = None

625 +

626 + def to_tuple(self) -> Tuple[Any]:

627 + return tuple(

628 + self[k] if k not in ["text_model_output", "vision_model_output"

629 + ] else getattr(self, k).to_tuple()

630 + for k in self.keys())

631 +

632 +

633 +class SiglipVisionEmbeddings(nn.Module):

634 +

635 + def __init__(self, config: SiglipVisionConfig):

636 + super().__init__()

637 + self.config = config

638 + self.embed_dim = config.hidden_size

639 + self.image_size = config.image_size

640 + self.patch_size = config.patch_size

641 +

642 + self.patch_embedding = nn.Conv2d(

643 + in_channels=config.num_channels,

644 + out_channels=self.embed_dim,

645 + kernel_size=self.patch_size,

646 + stride=self.patch_size,

647 + padding="valid",

648 + )

649 +

650 + self.num_patches_per_side = self.image_size // self.patch_size

651 + self.num_patches = self.num_patches_per_side**2

652 + self.num_positions = self.num_patches

653 + self.position_embedding = nn.Embedding(self.num_positions,

654 + self.embed_dim)

655 +

656 + def forward(self, pixel_values: torch.FloatTensor,

657 + patch_attention_mask: torch.BoolTensor) -> torch.Tensor:

658 + batch_size = pixel_values.size(0)

659 +

660 + patch_embeds = self.patch_embedding(pixel_values)

661 + embeddings = patch_embeds.flatten(2).transpose(1, 2)

662 +

663 + max_im_h, max_im_w = pixel_values.size(2), pixel_values.size(3)

664 + max_nb_patches_h, max_nb_patches_w = max_im_h // self.patch_size, \

665 + max_im_w // self.patch_size

666 + boundaries = torch.arange(1 / self.num_patches_per_side, 1.0,

667 + 1 / self.num_patches_per_side)

668 + position_ids = torch.full(

669 + size=(

670 + batch_size,

671 + max_nb_patches_h * max_nb_patches_w,

672 + ),

673 + fill_value=0,

674 + )

675 +

676 + for batch_idx, p_attn_mask in enumerate(patch_attention_mask):

677 + nb_patches_h = p_attn_mask[:, 0].sum()

678 + nb_patches_w = p_attn_mask[0].sum()

679 +

680 + fractional_coords_h = torch.linspace(0, 1 - 1 / nb_patches_h,

681 + nb_patches_h)

682 + fractional_coords_w = torch.linspace(0, 1 - 1 / nb_patches_w,

683 + nb_patches_w)

684 +

685 + bucket_coords_h = torch.bucketize(fractional_coords_h,

686 + boundaries,

687 + right=True)

688 + bucket_coords_w = torch.bucketize(fractional_coords_w,

689 + boundaries,

690 + right=True)

691 +

692 + pos_ids = (bucket_coords_h[:, None] * self.num_patches_per_side +

693 + bucket_coords_w).flatten()

694 + position_ids[batch_idx][p_attn_mask.view(-1).cpu()] = pos_ids

695 +

696 + position_ids = position_ids.to(self.position_embedding.weight.device)

697 +

698 + embeddings = embeddings + self.position_embedding(position_ids)

699 + return embeddings

700 +

701 +

702 +# Copied from transformers.models.clip.modeling_clip.CLIPTextEmbeddings with

703 +# CLIP->Siglip

704 +class SiglipTextEmbeddings(nn.Module):

705 +

706 + def __init__(self, config: SiglipTextConfig):

707 + super().__init__()

708 + embed_dim = config.hidden_size

709 +

710 + self.token_embedding = nn.Embedding(config.vocab_size, embed_dim)

711 + self.position_embedding = nn.Embedding(config.max_position_embeddings,

712 + embed_dim)

713 +

714 + # position_ids (1, len position emb) is contiguous in memory and

715 + # exported when serialized

716 + self.register_buffer(

717 + "position_ids",

718 + torch.arange(config.max_position_embeddings).expand((1, -1)),

719 + persistent=False)

720 +

721 + def forward(

722 + self,

723 + input_ids: Optional[torch.LongTensor] = None,

724 + position_ids: Optional[torch.LongTensor] = None,

725 + inputs_embeds: Optional[torch.FloatTensor] = None,

726 + ) -> torch.Tensor:

727 + seq_length = input_ids.shape[

728 + -1] if input_ids is not None else inputs_embeds.shape[-2]

729 +

730 + if position_ids is None:

731 + position_ids = self.position_ids[:, :seq_length]

732 +

733 + if inputs_embeds is None:

734 + inputs_embeds = self.token_embedding(input_ids)

735 +

736 + position_embeddings = self.position_embedding(position_ids)

737 + embeddings = inputs_embeds + position_embeddings

738 +

739 + return embeddings

740 +

741 +

742 +class SiglipAttention(nn.Module):

743 + """Multi-headed attention from 'Attention Is All You Need' paper"""

744 +

745 + # Copied from transformers.models.clip.modeling_clip.CLIPAttention.__init__

746 + def __init__(self, config):

747 + super().__init__()

748 + self.config = config

749 + self.embed_dim = config.hidden_size

750 + self.num_heads = config.num_attention_heads

751 + self.head_dim = self.embed_dim // self.num_heads

752 + if self.head_dim * self.num_heads != self.embed_dim:

753 + raise ValueError(

754 + f"embed_dim must be divisible by num_heads (got `embed_dim`:"

755 + f" {self.embed_dim} and `num_heads`: {self.num_heads}).")

756 + self.scale = self.head_dim**-0.5

757 + self.dropout = config.attention_dropout

758 +

759 + self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)

760 + self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)

761 + self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)

762 + self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)

763 +

764 + def forward(

765 + self,

766 + hidden_states: torch.Tensor,

767 + attention_mask: Optional[torch.Tensor] = None,

768 + output_attentions: Optional[bool] = False,

769 + ) -> Tuple[torch.Tensor, Optional[torch.Tensor],

770 + Optional[Tuple[torch.Tensor]]]:

771 + """Input shape: Batch x Time x Channel"""

772 +

773 + batch_size, q_len, _ = hidden_states.size()

774 +

775 + query_states = self.q_proj(hidden_states)

776 + key_states = self.k_proj(hidden_states)

777 + value_states = self.v_proj(hidden_states)

778 +

779 + query_states = query_states.view(batch_size, q_len, self.num_heads,

780 + self.head_dim).transpose(1, 2)

781 + key_states = key_states.view(batch_size, q_len, self.num_heads,

782 + self.head_dim).transpose(1, 2)

783 + value_states = value_states.view(batch_size, q_len, self.num_heads,

784 + self.head_dim).transpose(1, 2)

785 +

786 + k_v_seq_len = key_states.shape[-2]

787 + attn_weights = torch.matmul(query_states, key_states.transpose(

788 + 2, 3)) * self.scale

789 +

790 + if attn_weights.size() != (batch_size, self.num_heads, q_len,

791 + k_v_seq_len):

792 + raise ValueError(

793 + f"Attention weights should be of size "

794 + f"{(batch_size, self.num_heads, q_len, k_v_seq_len)}, but is"

795 + f" {attn_weights.size()}")

796 +

797 + if attention_mask is not None:

798 + if attention_mask.size() != (batch_size, 1, q_len, k_v_seq_len):

799 + raise ValueError(f"Attention mask should be of size "

800 + f"{(batch_size, 1, q_len, k_v_seq_len)}, "

801 + f"but is {attention_mask.size()}")

802 + attn_weights = attn_weights + attention_mask

803 +

804 + # upcast attention to fp32

805 + attn_weights = nn.functional.softmax(attn_weights,

806 + dim=-1,

807 + dtype=torch.float32).to(

808 + query_states.dtype)

809 + attn_weights = nn.functional.dropout(attn_weights,

810 + p=self.dropout,

811 + training=self.training)

812 + attn_output = torch.matmul(attn_weights, value_states)

813 +

814 + if attn_output.size() != (batch_size, self.num_heads, q_len,

815 + self.head_dim):

816 + raise ValueError(

817 + f"`attn_output` should be of size "

818 + f"{(batch_size, self.num_heads, q_len, self.head_dim)}, but is"

819 + f" {attn_output.size()}")

820 +

821 + attn_output = attn_output.transpose(1, 2).contiguous()

822 + attn_output = attn_output.reshape(batch_size, q_len, self.embed_dim)

823 +

824 + attn_output = self.out_proj(attn_output)

825 +

826 + return attn_output, attn_weights

827 +

828 +

829 +class SiglipFlashAttention2(SiglipAttention):

830 + """

831 + Llama flash attention module. This module inherits from `LlamaAttention` as

832 + the weights of the module stays untouched. The only required change would

833 + be on the forward pass where it needs to correctly call the public API of

834 + flash attention and deal with padding tokens in case the input contains any

835 + of them.

836 + """

837 +

838 + def __init__(self, *args, **kwargs):

839 + super().__init__(*args, **kwargs)

840 + self.is_causal = False # Hack to make sure we don't use a causal mask

841 +

842 + def forward(

843 + self,

844 + hidden_states: torch.Tensor,

845 + attention_mask: Optional[torch.LongTensor] = None,

846 + position_ids: Optional[torch.LongTensor] = None,

847 + past_key_value: Optional[Tuple[torch.Tensor]] = None,

848 + output_attentions: bool = False,

849 + use_cache: bool = False,

850 + **kwargs,

851 + ) -> Tuple[torch.Tensor, Optional[torch.Tensor],

852 + Optional[Tuple[torch.Tensor]]]:

853 + output_attentions = False

854 +

855 + bsz, q_len, _ = hidden_states.size()

856 +

857 + query_states = self.q_proj(hidden_states)

858 + key_states = self.k_proj(hidden_states)

859 + value_states = self.v_proj(hidden_states)

860 +

861 + # Flash attention requires the input to have the shape

862 + # batch_size x seq_length x head_dim x hidden_dim

863 + # therefore we just need to keep the original shape

864 + query_states = query_states.view(bsz, q_len, self.num_heads,

865 + self.head_dim).transpose(1, 2)

866 + key_states = key_states.view(bsz, q_len, self.num_heads,

867 + self.head_dim).transpose(1, 2)

868 + value_states = value_states.view(bsz, q_len, self.num_heads,

869 + self.head_dim).transpose(1, 2)

870 +

871 + kv_seq_len = key_states.shape[-2]

872 + if past_key_value is not None:

873 + kv_seq_len += past_key_value.get_usable_length(

874 + kv_seq_len, self.layer_idx)

875 +

876 + # TODO: These transpose are quite inefficient but Flash Attention

877 + # requires the layout [batch_size, sequence_length, num_heads,

878 + # head_dim]. We would need to refactor the KV cache

879 + # to be able to avoid many of these transpose/reshape/view.

880 + query_states = query_states.transpose(1, 2)

881 + key_states = key_states.transpose(1, 2)

882 + value_states = value_states.transpose(1, 2)

883 +

884 + dropout_rate = self.dropout if self.training else 0.0

885 +

886 + # In PEFT, usually we cast the layer norms in float32 for training

887 + # stability reasons therefore the input hidden states gets silently

888 + # casted in float32. Hence, we need cast them back in the correct

889 + # dtype just to be sure everything works as expected.

890 + # This might slowdown training & inference so it is recommended to

891 + # not cast the LayerNorms in fp32. (LlamaRMSNorm handles it correctly)

892 +

893 + input_dtype = query_states.dtype

894 + if input_dtype == torch.float32:

895 + if torch.is_autocast_enabled():

896 + target_dtype = torch.get_autocast_gpu_dtype()

897 + # Handle the case where the model is quantized

898 + elif hasattr(self.config, "_pre_quantization_dtype"):

899 + target_dtype = self.config._pre_quantization_dtype

900 + else:

901 + target_dtype = self.q_proj.weight.dtype

902 +

903 + logger.warning_once(

904 + "The input hidden states seems to be silently casted in "

905 + "float32, this might be related to the fact you have upcasted "

906 + "embedding or layer norm layers in float32. We will cast "

907 + f"back the input in {target_dtype}.")

908 +

909 + query_states = query_states.to(target_dtype)

910 + key_states = key_states.to(target_dtype)

911 + value_states = value_states.to(target_dtype)

912 +

913 + attn_output = self._flash_attention_forward(query_states,

914 + key_states,

915 + value_states,

916 + attention_mask,

917 + q_len,

918 + dropout=dropout_rate)

919 +

920 + attn_output = attn_output.reshape(bsz, q_len,

921 + self.embed_dim).contiguous()

922 + attn_output = self.out_proj(attn_output)

923 +

924 + if not output_attentions:

925 + attn_weights = None

926 +

927 + return attn_output, attn_weights

928 +

929 + def _flash_attention_forward(self,

930 + query_states,

931 + key_states,

932 + value_states,

933 + attention_mask,

934 + query_length,

935 + dropout=0.0,

936 + softmax_scale=None):

937 + """

938 + Calls the forward method of Flash Attention - if the input hidden

939 + states contain at least one padding token first unpad the input,

940 + then computes the attention scores and pad the final attention

941 + scores.

942 + Args:

943 + query_states (`torch.Tensor`):

944 + Input query states to be passed to Flash Attention API

945 + key_states (`torch.Tensor`):

946 + Input key states to be passed to Flash Attention API

947 + value_states (`torch.Tensor`):

948 + Input value states to be passed to Flash Attention API

949 + attention_mask (`torch.Tensor`):

950 + The padding mask - corresponds to a tensor of size

951 + `(batch_size, seq_len)` where 0 stands for the position

952 + of padding tokens and 1 for the position of non-padding

953 + tokens.

954 + dropout (`int`, *optional*):

955 + Attention dropout

956 + softmax_scale (`float`, *optional*):

957 + The scaling of QK^T before applying softmax. Default to 1 /

958 + sqrt(head_dim)

959 + """

960 + from flash_attn import flash_attn_func, flash_attn_varlen_func

961 + from flash_attn.bert_padding import pad_input # noqa

962 +

963 + # TODO: Remove the `query_length != 1` check once Flash Attention for

964 + # RoCm is bumped to 2.1. For details, please see the comment in

965 + # LlamaFlashAttention2 __init__.

966 + causal = self.is_causal and query_length != 1

967 +

968 + # Contains at least one padding token in the sequence

969 + if attention_mask is not None:

970 + batch_size = query_states.shape[0]

971 + query_states, key_states, value_states, indices_q, cu_seq_lens, \

972 + max_seq_lens = self._upad_input(

973 + query_states, key_states, value_states, attention_mask,

974 + query_length)

975 +

976 + cu_seqlens_q, cu_seqlens_k = cu_seq_lens

977 + max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens

978 +

979 + attn_output_unpad = flash_attn_varlen_func(

980 + query_states,

981 + key_states,

982 + value_states,

983 + cu_seqlens_q=cu_seqlens_q,

984 + cu_seqlens_k=cu_seqlens_k,

985 + max_seqlen_q=max_seqlen_in_batch_q,

986 + max_seqlen_k=max_seqlen_in_batch_k,

987 + dropout_p=dropout,

988 + softmax_scale=softmax_scale,

989 + causal=causal,

990 + )

991 +

992 + attn_output = pad_input(attn_output_unpad, indices_q, batch_size,

993 + query_length)

994 + else:

995 + attn_output = flash_attn_func(query_states,

996 + key_states,

997 + value_states,

998 + dropout,

999 + softmax_scale=softmax_scale,

1000 + causal=causal)

1001 +

1002 + return attn_output

1003 +

1004 + def _upad_input(self, query_layer, key_layer, value_layer, attention_mask,

1005 + query_length):

1006 + from flash_attn.bert_padding import index_first_axis, unpad_input

1007 + indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(

1008 + attention_mask)

1009 + batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape

1010 +

1011 + key_layer = index_first_axis(

1012 + key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads,

1013 + head_dim), indices_k)

1014 + value_layer = index_first_axis(

1015 + value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads,

1016 + head_dim), indices_k)

1017 + if query_length == kv_seq_len:

1018 + query_layer = index_first_axis(

1019 + query_layer.reshape(batch_size * kv_seq_len, self.num_heads,

1020 + head_dim), indices_k)

1021 + cu_seqlens_q = cu_seqlens_k

1022 + max_seqlen_in_batch_q = max_seqlen_in_batch_k

1023 + indices_q = indices_k

1024 + elif query_length == 1:

1025 + max_seqlen_in_batch_q = 1

1026 + cu_seqlens_q = torch.arange(

1027 + batch_size + 1, dtype=torch.int32, device=query_layer.device

1028 + ) # There is a memcpy here, that is very bad.

1029 + indices_q = cu_seqlens_q[:-1]

1030 + query_layer = query_layer.squeeze(1)

1031 + else:

1032 + # The -q_len: slice assumes left padding.

1033 + attention_mask = attention_mask[:, -query_length:]

1034 + query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = \

1035 + unpad_input(query_layer, attention_mask)

1036 +

1037 + return (

1038 + query_layer,

1039 + key_layer,

1040 + value_layer,

1041 + indices_q,

1042 + (cu_seqlens_q, cu_seqlens_k),

1043 + (max_seqlen_in_batch_q, max_seqlen_in_batch_k),

1044 + )

1045 +

1046 +

1047 +# Copied from transformers.models.clip.modeling_clip.CLIPMLP with CLIP->Siglip

1048 +class SiglipMLP(nn.Module):

1049 +

1050 + def __init__(self, config):

1051 + super().__init__()

1052 + self.config = config

1053 + self.activation_fn = ACT2FN[config.hidden_act]

1054 + self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)

1055 + self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)

1056 +

1057 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

1058 + hidden_states = self.fc1(hidden_states)

1059 + hidden_states = self.activation_fn(hidden_states)

1060 + hidden_states = self.fc2(hidden_states)

1061 + return hidden_states

1062 +

1063 +

1064 +# Copied from transformers.models.clip.modeling_clip.CLIPEncoderLayer with

1065 +# CLIP->Siglip

1066 +class SiglipEncoderLayer(nn.Module):

1067 +

1068 + def __init__(self, config: SiglipConfig):

1069 + super().__init__()

1070 + self.embed_dim = config.hidden_size

1071 + self.self_attn = (SiglipAttention(config) if

1072 + not getattr(config, "_flash_attn_2_enabled", False)

1073 + else SiglipFlashAttention2(config))

1074 + self.layer_norm1 = nn.LayerNorm(self.embed_dim,

1075 + eps=config.layer_norm_eps)

1076 + self.mlp = SiglipMLP(config)

1077 + self.layer_norm2 = nn.LayerNorm(self.embed_dim,

1078 + eps=config.layer_norm_eps)

1079 +

1080 + def forward(

1081 + self,

1082 + hidden_states: torch.Tensor,

1083 + attention_mask: torch.Tensor,

1084 + output_attentions: Optional[bool] = False,

1085 + ) -> Tuple[torch.FloatTensor]:

1086 + """

1087 + Args:

1088 + hidden_states (`torch.FloatTensor`):

1089 + Input to the layer of shape `(batch, seq_len, embed_dim)`.

1090 + attention_mask (`torch.FloatTensor`):

1091 + Attention mask of shape `(batch, 1, q_len, k_v_seq_len)` where

1092 + padding elements are indicated by very large negative values.

1093 + output_attentions (`bool`, *optional*, defaults to `False`):

1094 + Whether or not to return the attentions tensors of all

1095 + attention layers. See `attentions` under returned tensors for

1096 + more detail.

1097 + """

1098 + residual = hidden_states

1099 +

1100 + hidden_states = self.layer_norm1(hidden_states)

1101 + hidden_states, attn_weights = self.self_attn(

1102 + hidden_states=hidden_states,

1103 + attention_mask=attention_mask,

1104 + output_attentions=output_attentions,

1105 + )

1106 + hidden_states = residual + hidden_states

1107 +

1108 + residual = hidden_states

1109 + hidden_states = self.layer_norm2(hidden_states)

1110 + hidden_states = self.mlp(hidden_states)

1111 + hidden_states = residual + hidden_states

1112 +

1113 + outputs = (hidden_states, )

1114 +

1115 + if output_attentions:

1116 + outputs += (attn_weights, )

1117 +

1118 + return outputs

1119 +

1120 +

1121 +class SiglipPreTrainedModel(PreTrainedModel):

1122 + """

1123 + An abstract class to handle weights initialization and a simple interface

1124 + for downloading and loading pretrained models.

1125 + """

1126 +

1127 + config_class = SiglipConfig

1128 + base_model_prefix = "siglip"

1129 + supports_gradient_checkpointing = True

1130 +

1131 + def _init_weights(self, module):

1132 + """Initialize the weights"""

1133 +

1134 + if isinstance(module, SiglipVisionEmbeddings):

1135 + width = (self.config.vision_config.hidden_size if isinstance(

1136 + self.config, SiglipConfig) else self.config.hidden_size)

1137 + nn.init.normal_(module.position_embedding.weight,

1138 + std=1 / np.sqrt(width))

1139 + elif isinstance(module, nn.Embedding):

1140 + default_flax_embed_init(module.weight)

1141 + elif isinstance(module, SiglipAttention):

1142 + nn.init.normal_(module.q_proj.weight)

1143 + nn.init.normal_(module.k_proj.weight)

1144 + nn.init.normal_(module.v_proj.weight)

1145 + nn.init.normal_(module.out_proj.weight)

1146 + nn.init.zeros_(module.q_proj.bias)

1147 + nn.init.zeros_(module.k_proj.bias)

1148 + nn.init.zeros_(module.v_proj.bias)

1149 + nn.init.zeros_(module.out_proj.bias)

1150 + elif isinstance(module, SiglipMLP):

1151 + nn.init.normal_(module.fc1.weight)

1152 + nn.init.normal_(module.fc2.weight)

1153 + nn.init.normal_(module.fc1.bias, std=1e-6)

1154 + nn.init.normal_(module.fc2.bias, std=1e-6)

1155 + elif isinstance(module, SiglipMultiheadAttentionPoolingHead):

1156 + nn.init.normal_(module.probe.data)

1157 + nn.init.normal_(module.attention.in_proj_weight.data)

1158 + nn.init.zeros_(module.attention.in_proj_bias.data)

1159 + elif isinstance(module, SiglipModel):

1160 + logit_scale_init = torch.tensor(0.0)

1161 + module.logit_scale.data.fill_(logit_scale_init)

1162 + module.logit_bias.data.zero_()

1163 + elif isinstance(module, (nn.Linear, nn.Conv2d)):

1164 + lecun_normal_(module.weight)

1165 + if module.bias is not None:

1166 + nn.init.zeros_(module.bias)

1167 + elif isinstance(module, nn.LayerNorm):

1168 + module.bias.data.zero_()

1169 + module.weight.data.fill_(1.0)

1170 +

1171 +

1172 +SIGLIP_START_DOCSTRING = r"""

1173 + This model inherits from [`PreTrainedModel`]. Check the superclass

1174 + documentation for the generic methods the library implements for all

1175 + its model (such as downloading or saving, resizing the input embeddings,

1176 + pruning heads etc.)

1177 + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/

1178 + stable/nn.html#torch.nn.Module) subclass.

1179 + Use it as a regular PyTorch Module and refer to the PyTorch documentation

1180 + for all matter related to general usage and behavior.

1181 + Parameters:

1182 + config ([`SiglipConfig`]): Model configuration class with all the

1183 + parameters of the model.

1184 + Initializing with a config file does not load the weights

1185 + associated with the model, only the configuration. Check out

1186 + the [`~PreTrainedModel.from_pretrained`] method to load the

1187 + model weights.

1188 +"""

1189 +

1190 +SIGLIP_TEXT_INPUTS_DOCSTRING = r"""

1191 + Args:

1192 + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)

1193 + `):

1194 + Indices of input sequence tokens in the vocabulary. Padding will

1195 + be ignored by default should you provide it.

1196 + Indices can be obtained using [`AutoTokenizer`]. See

1197 + [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`]

1198 + for details. [What are input IDs?](../glossary#input-ids)

1199 + attention_mask (`torch.Tensor` of shape `(batch_size,

1200 + sequence_length)`, *optional*):

1201 + Mask to avoid performing attention on padding token indices. Mask

1202 + values selected in `[0, 1]`:

1203 + - 1 for tokens that are **not masked**,

1204 + - 0 for tokens that are **masked**.

1205 + [What are attention masks?](../glossary#attention-mask)

1206 + position_ids (`torch.LongTensor` of shape `(batch_size,

1207 + sequence_length)`, *optional*):

1208 + Indices of positions of each input sequence tokens in the position

1209 + embeddings. Selected in the range `[0,

1210 + config.max_position_embeddings - 1]`.

1211 + [What are position IDs?](../glossary#position-ids)

1212 + output_attentions (`bool`, *optional*):

1213 + Whether or not to return the attentions tensors of all attention

1214 + layers. See `attentions` under returned tensors for more detail.

1215 + output_hidden_states (`bool`, *optional*):

1216 + Whether or not to return the hidden states of all layers. See

1217 + `hidden_states` under returned tensors for more detail.

1218 + return_dict (`bool`, *optional*):

1219 + Whether or not to return a [`~utils.ModelOutput`] instead of a

1220 + plain tuple.

1221 +"""

1222 +

1223 +SIGLIP_VISION_INPUTS_DOCSTRING = r"""

1224 + Args:

1225 + pixel_values (`torch.FloatTensor` of shape `(batch_size,

1226 + num_channels, height, width)`):

1227 + Pixel values. Padding will be ignored by default should you

1228 + provide it. Pixel values can be obtained using

1229 + [`AutoImageProcessor`]. See [`CLIPImageProcessor.__call__`]

1230 + for details.

1231 + output_attentions (`bool`, *optional*):

1232 + Whether or not to return the attentions tensors of all attention

1233 + layers. See `attentions` under returned

1234 + tensors for more detail.

1235 + output_hidden_states (`bool`, *optional*):

1236 + Whether or not to return the hidden states of all layers. See

1237 + `hidden_states` under returned tensors for more detail.

1238 + return_dict (`bool`, *optional*):

1239 + Whether or not to return a [`~utils.ModelOutput`] instead of a

1240 + plain tuple.

1241 +"""

1242 +

1243 +SIGLIP_INPUTS_DOCSTRING = r"""

1244 + Args:

1245 + input_ids (`torch.LongTensor` of shape `(batch_size,

1246 + sequence_length)`):

1247 + Indices of input sequence tokens in the vocabulary. Padding

1248 + will be ignored by default should you provide it.

1249 + Indices can be obtained using [`AutoTokenizer`]. See

1250 + [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`]

1251 + for details. [What are input IDs?](../glossary#input-ids)

1252 + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`

1253 + , *optional*):

1254 + Mask to avoid performing attention on padding token indices. Mask

1255 + values selected in `[0, 1]`:

1256 + - 1 for tokens that are **not masked**,

1257 + - 0 for tokens that are **masked**.

1258 + [What are attention masks?](../glossary#attention-mask)

1259 + position_ids (`torch.LongTensor` of shape `(batch_size,

1260 + sequence_length)`, *optional*):

1261 + Indices of positions of each input sequence tokens in the position

1262 + embeddings. Selected in the range `[0,

1263 + config.max_position_embeddings - 1]`.

1264 + [What are position IDs?](../glossary#position-ids)

1265 + pixel_values (`torch.FloatTensor` of shape `(batch_size,

1266 + num_channels, height, width)`):

1267 + Pixel values. Padding will be ignored by default should you

1268 + provide it. Pixel values can be obtained using

1269 + [`AutoImageProcessor`]. See [`CLIPImageProcessor.__call__`]

1270 + for details.

1271 + return_loss (`bool`, *optional*):

1272 + Whether or not to return the contrastive loss.

1273 + output_attentions (`bool`, *optional*):

1274 + Whether or not to return the attentions tensors of all attention

1275 + layers. See `attentions` under returned tensors for more detail.

1276 + output_hidden_states (`bool`, *optional*):

1277 + Whether or not to return the hidden states of all layers. See

1278 + `hidden_states` under returned tensors for

1279 + more detail.

1280 + return_dict (`bool`, *optional*):

1281 + Whether or not to return a [`~utils.ModelOutput`] instead of a

1282 + plain tuple.

1283 +"""

1284 +

1285 +

1286 +# Copied from transformers.models.clip.modeling_clip.CLIPEncoder with

1287 +# CLIP->Siglip

1288 +class SiglipEncoder(nn.Module):

1289 + """

1290 + Transformer encoder consisting of `config.num_hidden_layers`

1291 + self attention layers. Each layer is a [`SiglipEncoderLayer`].

1292 + Args:

1293 + config: SiglipConfig

1294 + """

1295 +

1296 + def __init__(self, config: SiglipConfig):

1297 + super().__init__()

1298 + self.config = config

1299 + self.layers = nn.ModuleList([

1300 + SiglipEncoderLayer(config) for _ in range(config.num_hidden_layers)

1301 + ])

1302 + self.gradient_checkpointing = False

1303 +

1304 + # Ignore copy

1305 + def forward(

1306 + self,

1307 + inputs_embeds,

1308 + attention_mask: Optional[torch.Tensor] = None,

1309 + output_attentions: Optional[bool] = None,

1310 + output_hidden_states: Optional[bool] = None,

1311 + return_dict: Optional[bool] = None,

1312 + ) -> Union[Tuple, BaseModelOutput]:

1313 + r"""

1314 + Args:

1315 + inputs_embeds (`torch.FloatTensor` of shape `(batch_size,

1316 + sequence_length, hidden_size)`):

1317 + Optionally, instead of passing `input_ids` you can choose to

1318 + directly pass an embedded representation.

1319 + This is useful if you want more control over how to convert

1320 + `input_ids` indices into associated vectors

1321 + than the model's internal embedding lookup matrix.

1322 + attention_mask (`torch.Tensor` of shape `(batch_size,

1323 + sequence_length)`, *optional*):

1324 + Mask to avoid performing attention on padding token indices.

1325 + Mask values selected in `[0, 1]`:

1326 + - 1 for tokens that are **not masked**,

1327 + - 0 for tokens that are **masked**.

1328 + [What are attention masks?](../glossary#attention-mask)

1329 + output_attentions (`bool`, *optional*):

1330 + Whether or not to return the attentions tensors of all

1331 + attention layers. See `attentions` under returned tensors for

1332 + more detail.

1333 + output_hidden_states (`bool`, *optional*):

1334 + Whether or not to return the hidden states of all layers. See

1335 + `hidden_states` under returned tensors for more detail.

1336 + return_dict (`bool`, *optional*):

1337 + Whether or not to return a [`~utils.ModelOutput`] instead of a

1338 + plain tuple.

1339 + """

1340 + output_attentions = output_attentions if output_attentions \

1341 + is not None else self.config.output_attentions

1342 + output_hidden_states = (output_hidden_states

1343 + if output_hidden_states is not None else

1344 + self.config.output_hidden_states)

1345 + return_dict = return_dict if return_dict is not None else \

1346 + self.config.use_return_dict

1347 +

1348 + encoder_states = () if output_hidden_states else None

1349 + all_attentions = () if output_attentions else None

1350 +

1351 + hidden_states = inputs_embeds

1352 + for encoder_layer in self.layers:

1353 + if output_hidden_states:

1354 + encoder_states = encoder_states + (hidden_states, )

1355 + if self.gradient_checkpointing and self.training:

1356 + layer_outputs = self._gradient_checkpointing_func(

1357 + encoder_layer.__call__,

1358 + hidden_states,

1359 + attention_mask,

1360 + output_attentions,

1361 + )

1362 + else:

1363 + layer_outputs = encoder_layer(

1364 + hidden_states,

1365 + attention_mask,

1366 + output_attentions=output_attentions,

1367 + )

1368 +

1369 + hidden_states = layer_outputs[0]

1370 +

1371 + if output_attentions:

1372 + all_attentions = all_attentions + (layer_outputs[1], )

1373 +

1374 + if output_hidden_states:

1375 + encoder_states = encoder_states + (hidden_states, )

1376 +

1377 + if not return_dict:

1378 + return tuple(

1379 + v for v in [hidden_states, encoder_states, all_attentions]

1380 + if v is not None)

1381 + return BaseModelOutput(last_hidden_state=hidden_states,

1382 + hidden_states=encoder_states,

1383 + attentions=all_attentions)

1384 +

1385 +

1386 +class SiglipTextTransformer(nn.Module):

1387 +

1388 + def __init__(self, config: SiglipTextConfig):

1389 + super().__init__()

1390 + self.config = config

1391 + embed_dim = config.hidden_size

1392 + self.embeddings = SiglipTextEmbeddings(config)

1393 + self.encoder = SiglipEncoder(config)

1394 + self.final_layer_norm = nn.LayerNorm(embed_dim,

1395 + eps=config.layer_norm_eps)

1396 +

1397 + self.head = nn.Linear(embed_dim, embed_dim)

1398 +

1399 + @add_start_docstrings_to_model_forward(SIGLIP_TEXT_INPUTS_DOCSTRING)

1400 + @replace_return_docstrings(output_type=BaseModelOutputWithPooling,

1401 + config_class=SiglipTextConfig)

1402 + def forward(

1403 + self,

1404 + input_ids: Optional[torch.Tensor] = None,

1405 + attention_mask: Optional[torch.Tensor] = None,

1406 + position_ids: Optional[torch.Tensor] = None,

1407 + output_attentions: Optional[bool] = None,

1408 + output_hidden_states: Optional[bool] = None,

1409 + return_dict: Optional[bool] = None,

1410 + ) -> Union[Tuple, BaseModelOutputWithPooling]:

1411 + r"""

1412 + Returns:

1413 + """

1414 + output_attentions = output_attentions if output_attentions \

1415 + is not None else self.config.output_attentions

1416 + output_hidden_states = (output_hidden_states

1417 + if output_hidden_states \

1418 + is not None else

1419 + self.config.output_hidden_states)

1420 + return_dict = return_dict if return_dict is not None else \

1421 + self.config.use_return_dict

1422 +

1423 + if input_ids is None:

1424 + raise ValueError("You have to specify input_ids")

1425 +

1426 + input_shape = input_ids.size()

1427 + input_ids = input_ids.view(-1, input_shape[-1])

1428 +

1429 + hidden_states = self.embeddings(input_ids=input_ids,

1430 + position_ids=position_ids)

1431 +

1432 + # note: SigLIP's text model does not use a causal mask, unlike the

1433 + # original CLIP model.

1434 + # expand attention_mask

1435 + if attention_mask is not None:

1436 + # [batch_size, seq_len] ->

1437 + # [batch_size, 1, tgt_seq_len, src_seq_len]

1438 + attention_mask = _prepare_4d_attention_mask(

1439 + attention_mask, hidden_states.dtype)

1440 +

1441 + encoder_outputs = self.encoder(

1442 + inputs_embeds=hidden_states,

1443 + attention_mask=attention_mask,

1444 + output_attentions=output_attentions,

1445 + output_hidden_states=output_hidden_states,

1446 + return_dict=return_dict,

1447 + )

1448 +

1449 + last_hidden_state = encoder_outputs[0]

1450 + last_hidden_state = self.final_layer_norm(last_hidden_state)

1451 +

1452 + # Assuming "sticky" EOS tokenization, last token is always EOS.

1453 + pooled_output = last_hidden_state[:, -1, :]

1454 + pooled_output = self.head(pooled_output)

1455 +

1456 + if not return_dict:

1457 + return (last_hidden_state, pooled_output) + encoder_outputs[1:]

1458 +

1459 + return BaseModelOutputWithPooling(

1460 + last_hidden_state=last_hidden_state,

1461 + pooler_output=pooled_output,

1462 + hidden_states=encoder_outputs.hidden_states,

1463 + attentions=encoder_outputs.attentions,

1464 + )

1465 +

1466 +

1467 +@add_start_docstrings(

1468 + """The text model from SigLIP without any head or projection on top.""",

1469 + SIGLIP_START_DOCSTRING,

1470 +)

1471 +class SiglipTextModel(SiglipPreTrainedModel):

1472 + config_class = SiglipTextConfig

1473 +

1474 + _no_split_modules = ["SiglipTextEmbeddings", "SiglipEncoderLayer"]

1475 +

1476 + def __init__(self, config: SiglipTextConfig):

1477 + super().__init__(config)

1478 + self.text_model = SiglipTextTransformer(config)

1479 + # Initialize weights and apply final processing

1480 + self.post_init()

1481 +

1482 + def get_input_embeddings(self) -> nn.Module:

1483 + return self.text_model.embeddings.token_embedding

1484 +

1485 + def set_input_embeddings(self, value):

1486 + self.text_model.embeddings.token_embedding = value

1487 +

1488 + @add_start_docstrings_to_model_forward(SIGLIP_TEXT_INPUTS_DOCSTRING)

1489 + @replace_return_docstrings(output_type=BaseModelOutputWithPooling,

1490 + config_class=SiglipTextConfig)

1491 + def forward(

1492 + self,

1493 + input_ids: Optional[torch.Tensor] = None,

1494 + attention_mask: Optional[torch.Tensor] = None,

1495 + position_ids: Optional[torch.Tensor] = None,

1496 + output_attentions: Optional[bool] = None,

1497 + output_hidden_states: Optional[bool] = None,

1498 + return_dict: Optional[bool] = None,

1499 + ) -> Union[Tuple, BaseModelOutputWithPooling]:

1500 + r"""

1501 + Returns:

1502 + Examples:

1503 + ```python

1504 + >>> from transformers import AutoTokenizer, SiglipTextModel

1505 + >>> model = SiglipTextModel.

1506 + from_pretrained("google/siglip-base-patch16-224")

1507 + >>> tokenizer = AutoTokenizer.

1508 + from_pretrained("google/siglip-base-patch16-224")

1509 + >>> # important: make sure to set padding="max_length"

1510 + as that's how the model was trained

1511 + >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"],

1512 + padding="max_length", return_tensors="pt")

1513 + >>> outputs = model(**inputs)

1514 + >>> last_hidden_state = outputs.last_hidden_state

1515 + >>> pooled_output = outputs.pooler_output # pooled (EOS token)

1516 + states

1517 + ```"""

1518 + return_dict = return_dict if return_dict is not None else \

1519 + self.config.use_return_dict

1520 +

1521 + return self.text_model(

1522 + input_ids=input_ids,

1523 + attention_mask=attention_mask,

1524 + position_ids=position_ids,

1525 + output_attentions=output_attentions,

1526 + output_hidden_states=output_hidden_states,

1527 + return_dict=return_dict,

1528 + )

1529 +

1530 +

1531 +class SiglipVisionTransformer(nn.Module):

1532 +

1533 + def __init__(self, config: SiglipVisionConfig):

1534 + super().__init__()

1535 + self.config = config

1536 + embed_dim = config.hidden_size

1537 +

1538 + self.embeddings = SiglipVisionEmbeddings(config)

1539 + self.encoder = SiglipEncoder(config)

1540 + self.post_layernorm = nn.LayerNorm(embed_dim,

1541 + eps=config.layer_norm_eps)

1542 + self.head = SiglipMultiheadAttentionPoolingHead(config)

1543 +

1544 + @add_start_docstrings_to_model_forward(SIGLIP_VISION_INPUTS_DOCSTRING)

1545 + @replace_return_docstrings(output_type=BaseModelOutputWithPooling,

1546 + config_class=SiglipVisionConfig)

1547 + def forward(

1548 + self,

1549 + pixel_values,

1550 + patch_attention_mask: Optional[torch.BoolTensor] = None,

1551 + output_attentions: Optional[bool] = None,

1552 + output_hidden_states: Optional[bool] = None,

1553 + return_dict: Optional[bool] = None,

1554 + ) -> Union[Tuple, BaseModelOutputWithPooling]:

1555 + r"""

1556 + Returns:

1557 + """

1558 + output_attentions = output_attentions if output_attentions is not None\

1559 + else self.config.output_attentions

1560 + output_hidden_states = (output_hidden_states

1561 + if output_hidden_states is not None else

1562 + self.config.output_hidden_states)

1563 + return_dict = return_dict if return_dict is not None \

1564 + else self.config.use_return_dict

1565 +

1566 + batch_size = pixel_values.size(0)

1567 + if patch_attention_mask is None:

1568 + patch_attention_mask = torch.ones(

1569 + size=(

1570 + batch_size,

1571 + pixel_values.size(2) // self.config.patch_size,

1572 + pixel_values.size(3) // self.config.patch_size,

1573 + ),

1574 + dtype=torch.bool,

1575 + device=pixel_values.device,

1576 + )

1577 +

1578 + hidden_states = self.embeddings(

1579 + pixel_values=pixel_values,

1580 + patch_attention_mask=patch_attention_mask)

1581 +

1582 + patch_attention_mask = patch_attention_mask.view(batch_size, -1)

1583 + # The call to `_upad_input` in `_flash_attention_forward` is expensive

1584 + # So when the `patch_attention_mask` is full of 1s (i.e. attending

1585 + # to the whole sequence), avoiding passing the attention_mask, which

1586 + # is equivalent to attending to the full sequence

1587 + if not torch.any(~patch_attention_mask):

1588 + attention_mask = None

1589 + else:

1590 + attention_mask = (_prepare_4d_attention_mask(

1591 + patch_attention_mask, hidden_states.dtype)

1592 + if not self.config._flash_attn_2_enabled else

1593 + patch_attention_mask)

1594 +

1595 + encoder_outputs = self.encoder(

1596 + inputs_embeds=hidden_states,

1597 + attention_mask=attention_mask,

1598 + output_attentions=output_attentions,

1599 + output_hidden_states=output_hidden_states,

1600 + return_dict=return_dict,

1601 + )

1602 +

1603 + last_hidden_state = encoder_outputs[0]

1604 + last_hidden_state = self.post_layernorm(last_hidden_state)

1605 +

1606 + pooled_output = self.head(

1607 + hidden_state=last_hidden_state,

1608 + attention_mask=patch_attention_mask,

1609 + )

1610 +

1611 + if not return_dict:

1612 + return (last_hidden_state, pooled_output) + encoder_outputs[1:]

1613 +

1614 + return BaseModelOutputWithPooling(

1615 + last_hidden_state=last_hidden_state,

1616 + pooler_output=pooled_output,

1617 + hidden_states=encoder_outputs.hidden_states,

1618 + attentions=encoder_outputs.attentions,

1619 + )

1620 +

1621 +

1622 +class SiglipMultiheadAttentionPoolingHead(nn.Module):

1623 + """Multihead Attention Pooling."""

1624 +

1625 + def __init__(self, config: SiglipVisionConfig):

1626 + super().__init__()

1627 +

1628 + self.probe = nn.Parameter(torch.randn(1, 1, config.hidden_size))

1629 + self.attention = torch.nn.MultiheadAttention(

1630 + config.hidden_size, config.num_attention_heads, batch_first=True)

1631 + self.layernorm = nn.LayerNorm(config.hidden_size,

1632 + eps=config.layer_norm_eps)

1633 + self.mlp = SiglipMLP(config)

1634 +

1635 + def forward(self, hidden_state, attention_mask):

1636 + batch_size = hidden_state.shape[0]

1637 + probe = self.probe.repeat(batch_size, 1, 1)

1638 +

1639 + hidden_state = self.attention(query=probe,

1640 + key=hidden_state,

1641 + value=hidden_state,

1642 + key_padding_mask=~attention_mask)[0]

1643 +

1644 + residual = hidden_state

1645 + hidden_state = self.layernorm(hidden_state)

1646 + hidden_state = residual + self.mlp(hidden_state)

1647 +

1648 + return hidden_state[:, 0]

1649 +

1650 +

1651 +@add_start_docstrings(

1652 + """The vision model from SigLIP without any head or projection on top.""",

1653 + SIGLIP_START_DOCSTRING,

1654 +)

1655 +class SiglipVisionModel(SiglipPreTrainedModel):

1656 + config_class = SiglipVisionConfig

1657 + main_input_name = "pixel_values"

1658 +

1659 + def __init__(self, config: SiglipVisionConfig):

1660 + super().__init__(config)

1661 +

1662 + self.vision_model = SiglipVisionTransformer(config)

1663 +

1664 + # Initialize weights and apply final processing

1665 + self.post_init()

1666 +

1667 + def get_input_embeddings(self) -> nn.Module:

1668 + return self.vision_model.embeddings.patch_embedding

1669 +

1670 + @add_start_docstrings_to_model_forward(SIGLIP_VISION_INPUTS_DOCSTRING)

1671 + @replace_return_docstrings(output_type=BaseModelOutputWithPooling,

1672 + config_class=SiglipVisionConfig)

1673 + def forward(

1674 + self,

1675 + pixel_values,

1676 + patch_attention_mask: Optional[torch.BoolTensor] = None,

1677 + output_attentions: Optional[bool] = None,

1678 + output_hidden_states: Optional[bool] = None,

1679 + return_dict: Optional[bool] = None,

1680 + ) -> Union[Tuple, BaseModelOutputWithPooling]:

1681 + r"""

1682 + Returns:

1683 + Examples:

1684 + ```python

1685 + >>> from PIL import Image

1686 + >>> import requests

1687 + >>> from transformers import AutoProcessor, SiglipVisionModel

1688 + >>> model = SiglipVisionModel.from_pretrained(

1689 + "google/siglip-base-patch16-224")

1690 + >>> processor = AutoProcessor.from_pretrained(

1691 + "google/siglip-base-patch16-224")

1692 + >>> url =

1693 + "http://images.cocodataset.org/val2017/000000039769.jpg"

1694 + >>> image = Image.open(requests.get(url, stream=True).raw)

1695 + >>> inputs = processor(images=image, return_tensors="pt")

1696 + >>> outputs = model(**inputs)

1697 + >>> last_hidden_state = outputs.last_hidden_state

1698 + >>> pooled_output = outputs.pooler_output # pooled features

1699 + ```"""

1700 + return_dict = return_dict if return_dict is not None \

1701 + else self.config.use_return_dict

1702 +

1703 + return self.vision_model(

1704 + pixel_values=pixel_values,

1705 + patch_attention_mask=patch_attention_mask,

1706 + output_attentions=output_attentions,

1707 + output_hidden_states=output_hidden_states,

1708 + return_dict=return_dict,

1709 + )

1710 +

1711 +

1712 +@add_start_docstrings(SIGLIP_START_DOCSTRING)

1713 +class SiglipModel(SiglipPreTrainedModel):

1714 + config_class = SiglipConfig

1715 +

1716 + def __init__(self, config: SiglipConfig):

1717 + super().__init__(config)

1718 +

1719 + if not isinstance(config.text_config, SiglipTextConfig):

1720 + raise ValueError("config.text_config is expected to be of type "

1721 + f"SiglipTextConfig but is of type"

1722 + f" {type(config.text_config)}.")

1723 +

1724 + if not isinstance(config.vision_config, SiglipVisionConfig):

1725 + raise ValueError("config.vision_config is expected to be of type "

1726 + "SiglipVisionConfig but is of type"

1727 + f" {type(config.vision_config)}.")

1728 +

1729 + text_config = config.text_config

1730 + vision_config = config.vision_config

1731 +

1732 + self.text_model = SiglipTextTransformer(text_config)

1733 + self.vision_model = SiglipVisionTransformer(vision_config)

1734 +

1735 + self.logit_scale = nn.Parameter(torch.randn(1))

1736 + self.logit_bias = nn.Parameter(torch.randn(1))

1737 +

1738 + # Initialize weights and apply final processing

1739 + self.post_init()

1740 +

1741 + @add_start_docstrings_to_model_forward(SIGLIP_TEXT_INPUTS_DOCSTRING)

1742 + def get_text_features(

1743 + self,

1744 + input_ids: Optional[torch.Tensor] = None,

1745 + attention_mask: Optional[torch.Tensor] = None,

1746 + position_ids: Optional[torch.Tensor] = None,

1747 + output_attentions: Optional[bool] = None,

1748 + output_hidden_states: Optional[bool] = None,

1749 + return_dict: Optional[bool] = None,

1750 + ) -> torch.FloatTensor:

1751 + r"""

1752 + Returns:

1753 + text_features (`torch.FloatTensor` of shape `(batch_size,

1754 + output_dim`): The text embeddings obtained by

1755 + applying the projection layer to the pooled output

1756 + of [`SiglipTextModel`].

1757 + Examples:

1758 + ```python

1759 + >>> from transformers import AutoTokenizer, AutoModel

1760 + >>> import torch

1761 + >>> model = AutoModel.from_pretrained(

1762 + "google/siglip-base-patch16-224")

1763 + >>> tokenizer = AutoTokenizer.from_pretrained(

1764 + "google/siglip-base-patch16-224")

1765 + >>> # important: make sure to set padding="max_length" as that's

1766 + how the model was trained

1767 + >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"],

1768 + padding="max_length", return_tensors="pt")

1769 + >>> with torch.no_grad():

1770 + ... text_features = model.get_text_features(**inputs)

1771 + ```"""

1772 + # Use SigLIP model's config for some fields (if specified) instead

1773 + # of those of vision & text components.

1774 + output_attentions = output_attentions if output_attentions is not None\

1775 + else self.config.output_attentions

1776 + output_hidden_states = (output_hidden_states

1777 + if output_hidden_states is not None else

1778 + self.config.output_hidden_states)

1779 + return_dict = return_dict if return_dict is not None \

1780 + else self.config.use_return_dict

1781 +

1782 + text_outputs = self.text_model(

1783 + input_ids=input_ids,

1784 + attention_mask=attention_mask,

1785 + position_ids=position_ids,

1786 + output_attentions=output_attentions,

1787 + output_hidden_states=output_hidden_states,

1788 + return_dict=return_dict,

1789 + )

1790 +

1791 + pooled_output = text_outputs[1]

1792 +

1793 + return pooled_output

1794 +

1795 + @add_start_docstrings_to_model_forward(SIGLIP_VISION_INPUTS_DOCSTRING)

1796 + def get_image_features(

1797 + self,

1798 + pixel_values: Optional[torch.FloatTensor] = None,

1799 + output_attentions: Optional[bool] = None,

1800 + output_hidden_states: Optional[bool] = None,

1801 + return_dict: Optional[bool] = None,

1802 + ) -> torch.FloatTensor:

1803 + r"""

1804 + Returns:

1805 + image_features (`torch.FloatTensor` of shape `(batch_size,

1806 + output_dim`): The image embeddings obtained by applying the

1807 + projection layer to the pooled output of [`SiglipVisionModel`].

1808 + Examples:

1809 + ```python

1810 + >>> from PIL import Image

1811 + >>> import requests

1812 + >>> from transformers import AutoProcessor, AutoModel

1813 + >>> import torch

1814 + >>> model = AutoModel.from_pretrained("google/siglip-base-patch16-224")

1815 + >>> processor = AutoProcessor.from_pretrained(

1816 + "google/siglip-base-patch16-224")

1817 + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"

1818 + >>> image = Image.open(requests.get(url, stream=True).raw)

1819 + >>> inputs = processor(images=image, return_tensors="pt")

1820 + >>> with torch.no_grad():

1821 + ... image_features = model.get_image_features(**inputs)

1822 + ```"""

1823 + # Use SiglipModel's config for some fields (if specified) instead

1824 + # of those of vision & text components.

1825 + output_attentions = output_attentions if output_attentions \

1826 + is not None else self.config.output_attentions

1827 + output_hidden_states = (output_hidden_states

1828 + if output_hidden_states is not None else

1829 + self.config.output_hidden_states)

1830 + return_dict = return_dict if return_dict is not None else \

1831 + self.config.use_return_dict

1832 +

1833 + vision_outputs = self.vision_model(

1834 + pixel_values=pixel_values,

1835 + output_attentions=output_attentions,

1836 + output_hidden_states=output_hidden_states,

1837 + return_dict=return_dict,

1838 + )

1839 +

1840 + pooled_output = vision_outputs[1]

1841 +

1842 + return pooled_output

1843 +

1844 + @add_start_docstrings_to_model_forward(SIGLIP_INPUTS_DOCSTRING)

1845 + @replace_return_docstrings(output_type=SiglipOutput,

1846 + config_class=SiglipConfig)

1847 + def forward(

1848 + self,

1849 + input_ids: Optional[torch.LongTensor] = None,

1850 + pixel_values: Optional[torch.FloatTensor] = None,

1851 + attention_mask: Optional[torch.Tensor] = None,

1852 + position_ids: Optional[torch.LongTensor] = None,

1853 + return_loss: Optional[bool] = None,

1854 + output_attentions: Optional[bool] = None,

1855 + output_hidden_states: Optional[bool] = None,

1856 + return_dict: Optional[bool] = None,

1857 + ) -> Union[Tuple, SiglipOutput]:

1858 + r"""

1859 + Returns:

1860 + Examples:

1861 + ```python

1862 + >>> from PIL import Image

1863 + >>> import requests

1864 + >>> from transformers import AutoProcessor, AutoModel

1865 + >>> import torch

1866 + >>> model = AutoModel.from_pretrained("google/siglip-base-patch16-224")

1867 + >>> processor = AutoProcessor.from_pretrained(

1868 + "google/siglip-base-patch16-224")

1869 + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"

1870 + >>> image = Image.open(requests.get(url, stream=True).raw)

1871 + >>> texts = ["a photo of 2 cats", "a photo of 2 dogs"]

1872 + >>> # important: we pass `padding=max_length` since the model was

1873 + trained with this

1874 + >>> inputs = processor(text=texts, images=image,

1875 + padding="max_length", return_tensors="pt")

1876 + >>> with torch.no_grad():

1877 + ... outputs = model(**inputs)

1878 + >>> logits_per_image = outputs.logits_per_image

1879 + >>> probs = torch.sigmoid(logits_per_image) # these are the

1880 + probabilities

1881 + >>> print(f"{probs[0][0]:.1%} that image 0 is '{texts[0]}'")

1882 + 31.9% that image 0 is 'a photo of 2 cats'

1883 + ```"""

1884 + # Use SigLIP model's config for some fields (if specified) instead of

1885 + # those of vision & text components.

1886 + output_attentions = output_attentions if output_attentions \

1887 + is not None else self.config.output_attentions

1888 + output_hidden_states = (output_hidden_states

1889 + if output_hidden_states is not None else

1890 + self.config.output_hidden_states)

1891 + return_dict = return_dict if return_dict is not None else \

1892 + self.config.use_return_dict

1893 +

1894 + vision_outputs = self.vision_model(

1895 + pixel_values=pixel_values,

1896 + output_attentions=output_attentions,

1897 + output_hidden_states=output_hidden_states,

1898 + return_dict=return_dict,

1899 + )

1900 +

1901 + text_outputs = self.text_model(

1902 + input_ids=input_ids,

1903 + attention_mask=attention_mask,

1904 + position_ids=position_ids,

1905 + output_attentions=output_attentions,

1906 + output_hidden_states=output_hidden_states,

1907 + return_dict=return_dict,

1908 + )

1909 +

1910 + image_embeds = vision_outputs[1]

1911 + text_embeds = text_outputs[1]

1912 +

1913 + # normalized features

1914 + image_embeds = image_embeds / image_embeds.norm(

1915 + p=2, dim=-1, keepdim=True)

1916 + text_embeds = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True)

1917 +

1918 + # cosine similarity as logits

1919 + logits_per_text = torch.matmul(text_embeds, image_embeds.t(

1920 + )) * self.logit_scale.exp() + self.logit_bias

1921 + logits_per_image = logits_per_text.t()

1922 +

1923 + loss = None

1924 + if return_loss:

1925 + raise NotImplementedError("SigLIP loss to be implemented")

1926 +

1927 + if not return_dict:

1928 + output = (logits_per_image, logits_per_text, text_embeds,

1929 + image_embeds, text_outputs, vision_outputs)

1930 + return ((loss, ) + output) if loss is not None else output

1931 +

1932 + return SiglipOutput(

1933 + loss=loss,

1934 + logits_per_image=logits_per_image,

1935 + logits_per_text=logits_per_text,

1936 + text_embeds=text_embeds,

1937 + image_embeds=image_embeds,

1938 + text_model_output=text_outputs,

1939 + vision_model_output=vision_outputs,

1940 + )

1941 +

1942 +

1943 +def get_siglip_vision_model(_flash_attn_2_enabled=True, **kwargs):

1944 + siglip_vision_config = {

1945 + "hidden_size": 1152,

1946 + "image_size": 448,

1947 + "intermediate_size": 4304,

1948 + "model_type": "siglip_vision_model",

1949 + "num_attention_heads": 16,

1950 + "num_hidden_layers": 27,

1951 + "patch_size": 14,

1952 + }

1953 +

1954 + # Detect attention implementation.

1955 + attn_backend: _Backend = get_vit_attn_backend(support_fa=True)

1956 + if attn_backend != _Backend.FLASH_ATTN:

1957 + _flash_attn_2_enabled = False

1958 +

1959 + model_config = SiglipVisionConfig(

1960 + **siglip_vision_config,

1961 + _flash_attn_2_enabled=_flash_attn_2_enabled,

1962 + **kwargs)

1963 +

1964 + vision_model = SiglipVisionModel(model_config).vision_model

1965 +

1966 + return vision_model