| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| import math |
| from dataclasses import dataclass |
| from typing import Optional, Union |
|
|
| import numpy as np |
| import torch |
| from torch import nn |
|
|
| from transformers.activations import GELUActivation |
|
|
| from transformers.generation import GenerationMixin |
| from transformers.image_processing_utils import select_best_resolution |
| from transformers.modeling_flash_attention_utils import FlashAttentionKwargs |
| from transformers.modeling_outputs import BaseModelOutputWithPast, ModelOutput |
| from transformers.modeling_utils import PreTrainedModel |
| from transformers.models.auto import AutoModel |
| from transformers.processing_utils import Unpack |
| from transformers.utils import ( |
| can_return_tuple, |
| is_torchdynamo_compiling, |
| logging, |
| ) |
| from .configuration_bee import BeeConfig |
|
|
| logger = logging.get_logger(__name__) |
|
|
|
|
| @dataclass |
| class BeeModelOutputWithPast(BaseModelOutputWithPast): |
|
|
| image_hidden_states: Optional[torch.FloatTensor] = None |
|
|
|
|
| @dataclass |
| class BeeCausalLMOutputWithPast(ModelOutput): |
|
|
| loss: Optional[torch.FloatTensor] = None |
| logits: Optional[torch.FloatTensor] = None |
| past_key_values: Optional[list[torch.FloatTensor]] = None |
| hidden_states: Optional[tuple[torch.FloatTensor]] = None |
| attentions: Optional[tuple[torch.FloatTensor]] = None |
| image_hidden_states: Optional[torch.FloatTensor] = None |
|
|
|
|
| def get_anyres_image_grid_shape(image_size, grid_pinpoints, patch_size): |
| if not isinstance(grid_pinpoints, list): |
| raise TypeError("grid_pinpoints should be a list of tuples or lists") |
|
|
| |
| if not isinstance(image_size, (list, tuple)): |
| if not isinstance(image_size, (torch.Tensor, np.ndarray)): |
| raise TypeError( |
| f"image_size invalid type: {type(image_size)} not valid, should be either list, tuple, np.ndarray or tensor" |
| ) |
| image_size = image_size.tolist() |
|
|
| height, width = select_best_resolution(image_size, grid_pinpoints) |
| return height // patch_size, width // patch_size |
|
|
|
|
| def image_size_to_num_patches(image_size, grid_pinpoints, patch_size: int): |
| if not isinstance(grid_pinpoints, list): |
| raise TypeError("grid_pinpoints should be a list of tuples or lists") |
|
|
| |
| if not isinstance(image_size, (list, tuple)): |
| if not isinstance(image_size, (torch.Tensor, np.ndarray)): |
| raise TypeError( |
| f"image_size invalid type {type(image_size)} with value {image_size}" |
| ) |
| image_size = image_size.tolist() |
|
|
| best_resolution = select_best_resolution(image_size, grid_pinpoints) |
| height, width = best_resolution |
| num_patches = 0 |
| |
| for i in range(0, height, patch_size): |
| for j in range(0, width, patch_size): |
| num_patches += 1 |
| |
| num_patches += 1 |
| return num_patches |
|
|
|
|
| def unpad_image(tensor, original_size): |
| if not isinstance(original_size, (list, tuple)): |
| if not isinstance(original_size, (torch.Tensor, np.ndarray)): |
| raise TypeError( |
| f"image_size invalid type: {type(original_size)} not valid, should be either list, tuple, np.ndarray or tensor" |
| ) |
| original_size = original_size.tolist() |
| original_height, original_width = original_size |
| current_height, current_width = tensor.shape[1:] |
|
|
| original_aspect_ratio = original_width / original_height |
| current_aspect_ratio = current_width / current_height |
|
|
| if original_aspect_ratio > current_aspect_ratio: |
| scale_factor = current_width / original_width |
| new_height = int(round(original_height * scale_factor, 7)) |
| padding = (current_height - new_height) // 2 |
| unpadded_tensor = tensor[:, padding:current_height - padding, :] |
| else: |
| scale_factor = current_height / original_height |
| new_width = int(round(original_width * scale_factor, 7)) |
| padding = (current_width - new_width) // 2 |
| unpadded_tensor = tensor[:, :, padding:current_width - padding] |
|
|
| return unpadded_tensor |
|
|
|
|
| class BeePreTrainedModel(PreTrainedModel): |
| config_class = BeeConfig |
| base_model_prefix = "" |
| supports_gradient_checkpointing = True |
| |
| _no_split_modules = [ |
| "SiglipEncoderLayer", |
| "Qwen3DecoderLayer", |
| ] |
| _skip_keys_device_placement = "past_key_values" |
| _supports_cache_class = True |
| _supports_flash_attn_2 = True |
| _supports_sdpa = True |
| _supports_quantized_cache = True |
| _supports_static_cache = True |
| _supports_flex_attn = True |
| _supports_attention_backend = True |
|
|
| def _init_weights(self, module): |
| std = getattr(self.config, "initializer_range", |
| self.config.get_text_config().initializer_range) |
|
|
| if isinstance(module, nn.Linear): |
| module.weight.data.normal_(mean=0.0, std=std) |
| if module.bias is not None: |
| module.bias.data.zero_() |
| elif isinstance(module, BeeModel): |
| embed_std = 1 / math.sqrt(self.config.text_config.hidden_size) |
| module.image_newline.data.normal_(mean=0.0, std=embed_std) |
|
|
|
|
| class BeeMultiModalProjector(nn.Module): |
|
|
| def __init__(self, config): |
| super().__init__() |
|
|
| self.pre_norm = torch.nn.LayerNorm(config.vision_config.hidden_size, |
| eps=1e-06) |
| self.linear_1 = nn.Linear(config.vision_config.hidden_size, |
| config.text_config.hidden_size * 4, |
| bias=True) |
| self.act = GELUActivation() |
| self.linear_2 = nn.Linear(config.text_config.hidden_size * 4, |
| config.text_config.hidden_size, |
| bias=True) |
|
|
| def forward(self, image_feature: torch.Tensor) -> torch.Tensor: |
| image_feature = self.pre_norm(image_feature) |
| hidden_states = self.linear_1(image_feature) |
| hidden_states = self.act(hidden_states) |
| hidden_states = self.linear_2(hidden_states) |
|
|
| return hidden_states |
|
|
|
|
| class BeeModel(BeePreTrainedModel): |
| _checkpoint_conversion_mapping = {"language_model.model": "language_model"} |
|
|
| def __init__(self, config): |
| super().__init__(config) |
| self.vision_tower = AutoModel.from_config(config.vision_config) |
| self.multi_modal_projector = BeeMultiModalProjector(config) |
| embed_std = 1 / math.sqrt(config.text_config.hidden_size) |
| self.image_newline = nn.Parameter( |
| torch.randn(config.text_config.hidden_size, dtype=self.dtype) * |
| embed_std) |
|
|
| self.vocab_size = config.text_config.vocab_size |
| self.language_model = AutoModel.from_config(config.text_config) |
| self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1 |
| self.post_init() |
|
|
| def get_input_embeddings(self): |
| return self.language_model.get_input_embeddings() |
|
|
| def set_input_embeddings(self, value): |
| self.language_model.set_input_embeddings(value) |
|
|
| def pack_image_features(self, |
| image_features, |
| image_sizes, |
| image_newline=None, |
| vision_aspect_ratio="anyres"): |
| new_image_features = [] |
| feature_lens = [] |
| for image_idx, image_feature in enumerate(image_features): |
| if image_feature.shape[0] > 1: |
| base_image_feature = image_feature[0] |
| image_feature = image_feature[1:] |
| height = width = self.config.vision_config.image_size // self.config.vision_config.patch_size |
| if height * width != base_image_feature.shape[0]: |
| raise ValueError( |
| "The number of patches is not consistent with the image size." |
| ) |
| num_patch_height, num_patch_width = get_anyres_image_grid_shape( |
| image_sizes[image_idx], |
| self.config.image_grid_pinpoints, |
| self.config.vision_config.image_size, |
| ) |
| image_feature = image_feature.view(num_patch_height, |
| num_patch_width, height, |
| width, -1) |
| image_feature = image_feature.permute(4, 0, 2, 1, |
| 3).contiguous() |
| image_feature = image_feature.flatten(1, 2).flatten(2, 3) |
| image_feature = unpad_image(image_feature, |
| image_sizes[image_idx]) |
| try: |
| max_num_patches = int( |
| vision_aspect_ratio.strip("anyres_max_")) |
| channels, curr_height, curr_width = image_feature.shape |
| ratio = math.sqrt(curr_height * curr_width / |
| (max_num_patches * height**2)) |
| if ratio > 1.1: |
| image_feature = image_feature[None] |
| image_feature = nn.functional.interpolate( |
| image_feature, [ |
| int(curr_height // ratio), |
| int(curr_width // ratio) |
| ], |
| mode="bilinear")[0] |
| except: |
| pass |
| if image_newline is not None: |
| image_feature = torch.cat( |
| ( |
| image_feature, |
| image_newline[:, None, None].expand( |
| *image_feature.shape[:-1], 1).to( |
| image_feature.device, image_feature.dtype), |
| ), |
| dim=-1, |
| ) |
| image_feature = image_feature.flatten(1, 2).transpose(0, 1) |
| image_feature = torch.cat((base_image_feature, image_feature), |
| dim=0) |
| else: |
| image_feature = image_feature[0] |
| if image_newline is not None: |
| image_feature = torch.cat( |
| (image_feature, image_newline[None].to(image_feature)), |
| dim=0) |
| image_feature = image_feature.flatten(0, 1) |
| new_image_features.append(image_feature) |
| feature_lens.append(image_feature.size(0)) |
| feature_lens = torch.tensor(feature_lens, |
| dtype=torch.long, |
| device=image_features[0].device) |
| return new_image_features, feature_lens |
|
|
| def get_image_features( |
| self, |
| pixel_values: torch.FloatTensor, |
| image_sizes: torch.Tensor, |
| vision_feature_layer: Optional[Union[int, list[int]]] = None, |
| vision_feature_select_strategy: Optional[str] = None, |
| vision_aspect_ratio: Optional[str] = None, |
| batch_num_images: Optional[torch.LongTensor] = None, |
| ): |
| vision_feature_layer = (vision_feature_layer |
| if vision_feature_layer is not None else |
| self.config.vision_feature_layer) |
| vision_feature_select_strategy = ( |
| vision_feature_select_strategy if vision_feature_select_strategy |
| is not None else self.config.vision_feature_select_strategy) |
| vision_aspect_ratio = (vision_aspect_ratio |
| if vision_aspect_ratio is not None else |
| self.config.vision_aspect_ratio) |
|
|
| if batch_num_images is None: |
| |
| need_patching = [True] * len(image_sizes) |
| else: |
| need_patching = [ |
| n == 1 for n in batch_num_images for _ in range(n) |
| ] |
| image_num_patches = [ |
| image_size_to_num_patches( |
| image_size=imsize, |
| grid_pinpoints=self.config.image_grid_pinpoints, |
| patch_size=self.config.vision_config.image_size, |
| ) if should_patch else 1 |
| for imsize, should_patch in zip(image_sizes, need_patching) |
| ] |
|
|
| if isinstance(pixel_values, torch.Tensor): |
| if pixel_values.dim() == 5: |
| |
| _pixel_values_list = [ |
| pix_val[:num_patch] for pix_val, num_patch in zip( |
| pixel_values, image_num_patches) |
| ] |
| pixel_values = torch.cat(_pixel_values_list, dim=0) |
| elif pixel_values.dim() != 4: |
| |
| raise ValueError( |
| f"pixel_values of shape {pixel_values.shape}, expect to be of 4 or 5 dimensions" |
| ) |
| elif isinstance(pixel_values, list): |
| |
| assert len(pixel_values) == len(image_num_patches), ( |
| f"pixel_values is a list of {len(pixel_values)} tensors, but image_num_patches is of length {len(image_num_patches)}" |
| ) |
| _pixel_values_list = [ |
| pix_val.squeeze(0)[:num_patch] |
| for pix_val, num_patch in zip(pixel_values, image_num_patches) |
| ] |
|
|
| pixel_values = torch.cat(_pixel_values_list, dim=0) |
|
|
| image_features = self.vision_tower(pixel_values, |
| output_hidden_states=True) |
| |
| |
| if isinstance(vision_feature_layer, int): |
| selected_image_feature = image_features.hidden_states[ |
| vision_feature_layer] |
| else: |
| hs_pool = [ |
| image_features.hidden_states[layer_idx] |
| for layer_idx in vision_feature_layer |
| ] |
| selected_image_feature = torch.cat(hs_pool, dim=-1) |
|
|
| if vision_feature_select_strategy == "default": |
| selected_image_feature = selected_image_feature[:, 1:] |
| elif vision_feature_select_strategy == "full": |
| selected_image_feature = selected_image_feature |
| image_features = self.multi_modal_projector(selected_image_feature) |
|
|
| image_features = torch.split(image_features, image_num_patches, dim=0) |
|
|
| image_features, feature_lens = self.pack_image_features( |
| image_features, |
| image_sizes, |
| image_newline=self.image_newline, |
| vision_aspect_ratio=vision_aspect_ratio, |
| ) |
|
|
| return image_features |
|
|
| @can_return_tuple |
| def forward( |
| self, |
| input_ids: torch.LongTensor = None, |
| pixel_values: torch.FloatTensor = None, |
| image_sizes: Optional[torch.LongTensor] = None, |
| attention_mask: Optional[torch.Tensor] = None, |
| position_ids: Optional[torch.LongTensor] = None, |
| past_key_values: Optional[list[torch.FloatTensor]] = None, |
| inputs_embeds: Optional[torch.FloatTensor] = None, |
| vision_feature_layer: Optional[Union[int, list[int]]] = None, |
| vision_feature_select_strategy: Optional[str] = None, |
| vision_aspect_ratio: Optional[str] = None, |
| batch_num_images: Optional[torch.LongTensor] = None, |
| use_cache: Optional[bool] = None, |
| output_attentions: Optional[bool] = None, |
| output_hidden_states: Optional[bool] = None, |
| return_dict: Optional[bool] = None, |
| cache_position: Optional[torch.LongTensor] = None, |
| **kwargs: Unpack[FlashAttentionKwargs], |
| ) -> Union[tuple, BeeModelOutputWithPast]: |
|
|
| output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions |
| output_hidden_states = (output_hidden_states |
| if output_hidden_states is not None else |
| self.config.output_hidden_states) |
| return_dict = return_dict if return_dict is not None else self.config.use_return_dict |
| vision_feature_layer = (vision_feature_layer |
| if vision_feature_layer is not None else |
| self.config.vision_feature_layer) |
| vision_feature_select_strategy = ( |
| vision_feature_select_strategy if vision_feature_select_strategy |
| is not None else self.config.vision_feature_select_strategy) |
| vision_aspect_ratio = (vision_aspect_ratio |
| if vision_aspect_ratio is not None else |
| self.config.vision_aspect_ratio) |
|
|
| if (input_ids is None) ^ (inputs_embeds is not None): |
| raise ValueError( |
| "You must specify exactly one of input_ids or inputs_embeds") |
|
|
| if pixel_values is not None and inputs_embeds is not None: |
| raise ValueError( |
| "You cannot specify both `pixel_values` and `inputs_embeds` at the same time, " |
| "and must specify either one") |
| if inputs_embeds is None: |
| inputs_embeds = self.get_input_embeddings()(input_ids) |
|
|
| |
|
|
| if pixel_values is not None: |
| image_features = self.get_image_features( |
| pixel_values, |
| image_sizes, |
| vision_feature_layer=vision_feature_layer, |
| vision_feature_select_strategy=vision_feature_select_strategy, |
| batch_num_images=batch_num_images, |
| ) |
| image_features = torch.cat(image_features, dim=0) |
|
|
| special_image_mask = ( |
| input_ids == self.config.image_token_id).unsqueeze(-1) |
| special_image_mask = special_image_mask.expand_as( |
| inputs_embeds).to(inputs_embeds.device) |
| if not is_torchdynamo_compiling() and inputs_embeds[ |
| special_image_mask].numel() != image_features.numel(): |
| n_image_tokens = ( |
| input_ids == self.config.image_token_id).sum() |
| n_image_features = image_features.shape[0] |
| raise ValueError( |
| f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}" |
| ) |
| image_features = image_features.to(inputs_embeds.device, |
| inputs_embeds.dtype) |
| inputs_embeds = inputs_embeds.masked_scatter( |
| special_image_mask, image_features) |
|
|
| outputs = self.language_model( |
| attention_mask=attention_mask, |
| position_ids=position_ids, |
| past_key_values=past_key_values, |
| inputs_embeds=inputs_embeds, |
| use_cache=use_cache, |
| output_attentions=output_attentions, |
| output_hidden_states=output_hidden_states, |
| return_dict=True, |
| cache_position=cache_position, |
| **kwargs, |
| ) |
|
|
| return BeeModelOutputWithPast( |
| last_hidden_state=outputs.last_hidden_state, |
| past_key_values=outputs.past_key_values, |
| hidden_states=outputs.hidden_states, |
| attentions=outputs.attentions, |
| image_hidden_states=image_features |
| if pixel_values is not None else None, |
| ) |
|
|
| def apply_pooling(self, image_features): |
| height = width = self.config.vision_config.image_size // self.config.vision_config.patch_size |
| batch_frames, seq_len, dim = image_features.shape |
| image_features = image_features.view(batch_frames, height, width, -1) |
| image_features = image_features.permute(0, 3, 1, 2).contiguous() |
|
|
| height, width = image_features.shape[2:] |
| scaled_shape = [math.ceil(height / 2), math.ceil(width / 2)] |
| image_features = nn.functional.interpolate(image_features, |
| size=scaled_shape, |
| mode="bilinear") |
|
|
| image_features = image_features.permute(0, 2, 3, 1) |
| image_features = image_features.view(batch_frames, -1, dim) |
| return image_features |
|
|
|
|
| class BeeForConditionalGeneration(BeePreTrainedModel, GenerationMixin): |
| _checkpoint_conversion_mapping = { |
| "^language_model.model": "model.language_model", |
| "^vision_tower": "model.vision_tower", |
| "^multi_modal_projector": "model.multi_modal_projector", |
| "^image_newline": "model.image_newline", |
| "^language_model.lm_head": "lm_head", |
| } |
| _tied_weights_keys = ["lm_head.weight"] |
|
|
| def __init__(self, config: BeeConfig): |
| super().__init__(config) |
| self.model = BeeModel(config) |
| self.lm_head = nn.Linear(config.text_config.hidden_size, |
| config.text_config.vocab_size, |
| bias=False) |
| self.post_init() |
|
|
| def get_input_embeddings(self): |
| return self.model.get_input_embeddings() |
|
|
| def set_input_embeddings(self, value): |
| self.model.set_input_embeddings(value) |
|
|
| def get_output_embeddings(self) -> nn.Module: |
| return self.lm_head |
|
|
| def set_output_embeddings(self, new_embeddings): |
| self.lm_head = new_embeddings |
|
|
| def set_decoder(self, decoder): |
| self.model = decoder |
|
|
| def get_decoder(self): |
| return self.model |
|
|
| def pack_image_features(self, |
| image_features, |
| image_sizes, |
| vision_feature_select_strategy, |
| image_newline=None): |
| return self.model.pack_image_features( |
| image_features=image_features, |
| image_sizes=image_sizes, |
| vision_feature_select_strategy=vision_feature_select_strategy, |
| image_newline=image_newline, |
| ) |
|
|
| def get_image_features( |
| self, |
| pixel_values: torch.FloatTensor, |
| image_sizes: torch.Tensor, |
| vision_feature_layer: Optional[Union[int, list[int]]] = None, |
| vision_feature_select_strategy: Optional[str] = None, |
| ): |
| return self.model.get_image_features( |
| pixel_values=pixel_values, |
| image_sizes=image_sizes, |
| vision_feature_layer=vision_feature_layer, |
| vision_feature_select_strategy=vision_feature_select_strategy, |
| ) |
|
|
| |
| @property |
| def language_model(self): |
| return self.model.language_model |
|
|
| @property |
| def vision_tower(self): |
| return self.model.vision_tower |
|
|
| @property |
| def multi_modal_projector(self): |
| return self.model.multi_modal_projector |
|
|
| @can_return_tuple |
| def forward( |
| self, |
| input_ids: torch.LongTensor = None, |
| pixel_values: torch.FloatTensor = None, |
| image_sizes: Optional[torch.LongTensor] = None, |
| attention_mask: Optional[torch.Tensor] = None, |
| position_ids: Optional[torch.LongTensor] = None, |
| past_key_values: Optional[list[torch.FloatTensor]] = None, |
| inputs_embeds: Optional[torch.FloatTensor] = None, |
| vision_feature_layer: Optional[Union[int, list[int]]] = None, |
| vision_feature_select_strategy: Optional[str] = None, |
| vision_aspect_ratio: Optional[str] = None, |
| batch_num_images: Optional[torch.LongTensor] = None, |
| labels: Optional[torch.LongTensor] = None, |
| use_cache: Optional[bool] = None, |
| output_attentions: Optional[bool] = None, |
| output_hidden_states: Optional[bool] = None, |
| return_dict: Optional[bool] = None, |
| cache_position: Optional[torch.LongTensor] = None, |
| logits_to_keep: Union[int, torch.Tensor] = 0, |
| **kwargs, |
| ) -> Union[tuple, BeeCausalLMOutputWithPast]: |
|
|
| output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions |
| output_hidden_states = (output_hidden_states |
| if output_hidden_states is not None else |
| self.config.output_hidden_states) |
| return_dict = return_dict if return_dict is not None else self.config.use_return_dict |
| vision_feature_layer = (vision_feature_layer |
| if vision_feature_layer is not None else |
| self.config.vision_feature_layer) |
| vision_feature_select_strategy = ( |
| vision_feature_select_strategy if vision_feature_select_strategy |
| is not None else self.config.vision_feature_select_strategy) |
| vision_aspect_ratio = (vision_aspect_ratio |
| if vision_aspect_ratio is not None else |
| self.config.vision_aspect_ratio) |
|
|
| outputs = self.model( |
| input_ids=input_ids, |
| pixel_values=pixel_values, |
| image_sizes=image_sizes, |
| vision_aspect_ratio=vision_aspect_ratio, |
| vision_feature_layer=vision_feature_layer, |
| vision_feature_select_strategy=vision_feature_select_strategy, |
| batch_num_images=batch_num_images, |
| attention_mask=attention_mask, |
| position_ids=position_ids, |
| past_key_values=past_key_values, |
| inputs_embeds=inputs_embeds, |
| use_cache=use_cache, |
| output_attentions=output_attentions, |
| output_hidden_states=output_hidden_states, |
| return_dict=True, |
| cache_position=cache_position, |
| logits_to_keep=logits_to_keep, |
| **kwargs, |
| ) |
|
|
| hidden_states = outputs[0] |
| |
| slice_indices = slice(-logits_to_keep, None) if isinstance( |
| logits_to_keep, int) else logits_to_keep |
| logits = self.lm_head(hidden_states[:, slice_indices, :]) |
|
|
| loss = None |
| if labels is not None: |
| loss = self.loss_function( |
| logits=logits, |
| labels=labels, |
| vocab_size=self.config.text_config.vocab_size, |
| **kwargs) |
|
|
| return BeeCausalLMOutputWithPast( |
| loss=loss, |
| logits=logits, |
| past_key_values=outputs.past_key_values, |
| hidden_states=outputs.hidden_states, |
| attentions=outputs.attentions, |
| image_hidden_states=outputs.image_hidden_states, |
| ) |
|
|
| def prepare_inputs_for_generation( |
| self, |
| input_ids, |
| past_key_values=None, |
| inputs_embeds=None, |
| pixel_values=None, |
| image_sizes=None, |
| attention_mask=None, |
| cache_position=None, |
| logits_to_keep=None, |
| **kwargs, |
| ): |
| |
|
|
| model_inputs = super().prepare_inputs_for_generation( |
| input_ids, |
| past_key_values=past_key_values, |
| inputs_embeds=inputs_embeds, |
| attention_mask=attention_mask, |
| cache_position=cache_position, |
| logits_to_keep=logits_to_keep, |
| **kwargs, |
| ) |
|
|
| if cache_position[0] == 0: |
| |
| |
| model_inputs["pixel_values"] = pixel_values |
| model_inputs["image_sizes"] = image_sizes |
|
|
| return model_inputs |
|
|
| @staticmethod |
| def _prepare_4d_causal_attention_mask_with_cache_position( |
| attention_mask: torch.Tensor, |
| sequence_length: int, |
| target_length: int, |
| dtype: torch.dtype, |
| cache_position: torch.Tensor, |
| batch_size: int, |
| **kwargs, |
| ): |
|
|
| if attention_mask is not None and attention_mask.dim() == 4: |
| |
| causal_mask = attention_mask |
| else: |
| min_dtype = torch.finfo(dtype).min |
| causal_mask = torch.full((sequence_length, target_length), |
| fill_value=min_dtype, |
| dtype=dtype, |
| device=cache_position.device) |
| if sequence_length != 1: |
| causal_mask = torch.triu(causal_mask, diagonal=1) |
| causal_mask *= torch.arange( |
| target_length, |
| device=cache_position.device) > cache_position.reshape(-1, 1) |
| causal_mask = causal_mask[None, None, :, :].expand( |
| batch_size, 1, -1, -1) |
| if attention_mask is not None: |
| causal_mask = causal_mask.clone( |
| ) |
| mask_length = attention_mask.shape[-1] |
| padding_mask = causal_mask[:, :, :, : |
| mask_length] + attention_mask[:, |
| None, |
| None, :].to( |
| causal_mask |
| . |
| device |
| ) |
| padding_mask = padding_mask == 0 |
| causal_mask[:, :, :, : |
| mask_length] = causal_mask[:, :, :, : |
| mask_length].masked_fill( |
| padding_mask, |
| min_dtype) |
|
|
| return causal_mask |
|
|
|
|
| __all__ = ["BeeModel", "BeeForConditionalGeneration", "BeePreTrainedModel"] |
|
|