| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| """ PyTorch KPhi-3 model.""" |
|
|
| import inspect |
| import math |
| import warnings |
| from typing import List, Optional, Tuple, Union |
|
|
| import torch |
| import torch.nn.functional as F |
| import torch.utils.checkpoint |
| from torch import nn |
| from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss |
|
|
| from transformers.activations import ACT2FN |
| from transformers.generation import GenerationMixin |
| from transformers.cache_utils import Cache, DynamicCache, SlidingWindowCache, StaticCache |
| from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask |
| from transformers.modeling_outputs import ( |
| BaseModelOutputWithPast, |
| CausalLMOutputWithPast, |
| SequenceClassifierOutputWithPast, |
| TokenClassifierOutput, |
| ) |
| from transformers.modeling_utils import PreTrainedModel |
| from transformers.utils import ( |
| add_code_sample_docstrings, |
| add_start_docstrings, |
| add_start_docstrings_to_model_forward, |
| is_flash_attn_2_available, |
| is_flash_attn_greater_or_equal_2_10, |
| logging, |
| replace_return_docstrings, |
| ) |
| from .configuration_kphi3 import KPhi3Config |
|
|
| if is_flash_attn_2_available(): |
| from transformers.modeling_flash_attention_utils import _flash_attention_forward |
|
|
| def get_max_acceptable_common_divisor(a, b, max_acceptable=1000000): |
| """ |
| This is an inefficient max acceptable common divisor implementation to be improved. |
| # Arguments |
| a: is an integer. |
| b: is an integer. |
| max_acceptable: maximum acceptable common divisor. |
| """ |
| divisor = max(1, min(a, b, max_acceptable)) |
| while divisor > 0: |
| if a % divisor == 0 and b % divisor == 0: |
| return divisor |
| break |
| divisor -= 1 |
|
|
| class InterleaveChannels(nn.Module): |
| """ |
| This layer interleaves channels stepping according to the number passed as parameter. |
| This layer assumes "channel last". |
| """ |
| def __init__(self, step_size=2, last_dim=3): |
| super().__init__() |
| self.step_size = step_size if step_size >= 2 else 1 |
| self.last_dim = last_dim |
|
|
| def forward(self, x): |
| if (self.last_dim==3): |
| return torch.cat([x[:, :, :, shift_pos::self.step_size] for shift_pos in range(self.step_size)], dim=3) |
| else: |
| if (self.last_dim==2): |
| return torch.cat([x[:, :, shift_pos::self.step_size] for shift_pos in range(self.step_size)], dim=2) |
|
|
| def SignedSquareRoot1(x): |
| """ |
| Custom activation function that implements: |
| f(x) = sqrt(x) for x > 1 |
| f(x) = sqrt(-x) for x < -1 |
| f(x) = x for -1 ≤ x ≤ 1 |
| """ |
| return torch.where(x > 1, |
| torch.sqrt(x), |
| torch.where(x < -1, torch.sqrt(-x), x) |
| ) |
|
|
| |
| class InterleaveChannelsFast(nn.Module): |
| """ |
| This layer interleaves channels stepping according to the number passed as parameter. |
| This layer assumes "channel last". |
| """ |
| def __init__(self, step_size=2, last_dim=3): |
| super().__init__() |
| self.step_size = max(step_size, 1) |
| self.last_dim = last_dim |
|
|
| def forward(self, x): |
| if self.last_dim == 3: |
| N, H, W, C = x.shape |
| if C % self.step_size != 0: |
| raise ValueError("Number of channels must be divisible by step_size") |
| |
| x = x.view(N, H, W, self.step_size, C // self.step_size) |
| |
| x = x.permute(0, 1, 2, 4, 3) |
| |
| x = x.reshape(N, H, W, C) |
| return x |
| elif self.last_dim == 2: |
| N, H, W = x.shape |
| if W % self.step_size != 0: |
| raise ValueError("Width must be divisible by step_size") |
| x = x.view(N, H, self.step_size, W // self.step_size) |
| x = x.permute(0, 1, 3, 2) |
| x = x.reshape(N, H, W) |
| return x |
| else: |
| raise ValueError("last_dim must be 2 or 3") |
|
|
| class GroupedLinear(nn.Module): |
| """ |
| Similarly to a grouped pointwise convolution, this layer is a grouped linear layer. |
| This layer assumes "channel last". |
| """ |
| def __init__(self, in_features, out_features, num_groups=1, bias=True): |
| super().__init__() |
| self.in_features = in_features |
| self.out_features = out_features |
| self.num_groups = num_groups |
| self.bias = bias |
|
|
| |
| if in_features % num_groups != 0: |
| raise ValueError("Input features must be divisible by num_groups.") |
| if out_features % num_groups != 0: |
| raise ValueError("Output features must be divisible by num_groups.") |
|
|
| self.in_features_per_group = in_features // num_groups |
| self.out_features_per_group = out_features // num_groups |
|
|
| |
| self.group_layers = nn.ModuleList([ |
| nn.Linear(self.in_features_per_group, self.out_features_per_group, bias=bias) |
| for _ in range(num_groups) |
| ]) |
|
|
| def forward(self, x): |
| |
| |
| |
| |
| |
| |
| if self.in_features != x.shape[-1]: |
| raise ValueError( |
| "GroupedLinear error: "+ |
| "expected in_feautures "+str(self.in_features)+ |
| " but got "+str(x.shape[-1]) |
| ) |
| |
| x_groups = x.chunk(self.num_groups, dim=-1) |
|
|
| |
| |
|
|
| |
| out_groups = [layer(group) for layer, group in zip(self.group_layers, x_groups)] |
|
|
| |
| out = torch.cat(out_groups, dim=-1) |
| if self.out_features != out.shape[-1]: |
| raise ValueError( |
| "GroupedLinear error: "+ |
| "expected out_feautures "+str(self.out_features)+ |
| " but got "+str(out.shape[-1]) |
| ) |
| |
| return out |
|
|
| class GroupedLinearFast(nn.Module): |
| """ |
| Optimized grouped linear layer. |
| This layer assumes "channel last". |
| """ |
| def __init__(self, in_features, out_features, num_groups=1, bias=True): |
| super().__init__() |
| self.in_features = in_features |
| self.out_features = out_features |
| self.num_groups = num_groups |
|
|
| |
| if in_features % num_groups != 0: |
| raise ValueError("Input features must be divisible by num_groups.") |
| if out_features % num_groups != 0: |
| raise ValueError("Output features must be divisible by num_groups.") |
|
|
| self.in_features_per_group = in_features // num_groups |
| self.out_features_per_group = out_features // num_groups |
|
|
| |
| self.weight = nn.Parameter( |
| torch.Tensor(num_groups, self.in_features_per_group, self.out_features_per_group) |
| ) |
| if bias: |
| self.bias = nn.Parameter(torch.Tensor(num_groups, self.out_features_per_group)) |
| else: |
| self.register_parameter('bias', None) |
|
|
| self.reset_parameters() |
|
|
| def reset_parameters(self): |
| |
| nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) |
| if self.bias is not None: |
| |
| fan_in = self.in_features_per_group |
| bound = 1 / math.sqrt(fan_in) |
| nn.init.uniform_(self.bias, -bound, bound) |
|
|
| def forward(self, x): |
| in_shape = x.shape |
| if x.shape[-1] != self.in_features: |
| raise ValueError( |
| f"GroupedLinear error: expected in_features {self.in_features}, but got {x.shape[-1]}" |
| ) |
| |
| x = x.view(*x.shape[:-1], self.num_groups, self.in_features_per_group) |
|
|
| |
| |
| |
|
|
| |
| |
| |
| |
| out = torch.einsum('...ni,niq->...nq', x, self.weight) |
| |
| |
|
|
| |
| if self.bias is not None: |
| out += self.bias |
|
|
| |
| |
| out = out.contiguous().view(*out.shape[:-2], self.out_features) |
| return out |
|
|
| class GroupedPointwiseConvolutionBlock(nn.Module): |
| """ |
| This layer is composed by a grouped pointwise convolution followed by interleaving and another grouped pointwise comvolution with skip connection. This basic architecture can |
| vary according to the input tensor and its parameters. This is the basic building block for the papers: |
| https://www.researchgate.net/publication/360226228_Grouped_Pointwise_Convolutions_Reduce_Parameters_in_Convolutional_Neural_Networks |
| https://www.researchgate.net/publication/355214501_Grouped_Pointwise_Convolutions_Significantly_Reduces_Parameters_in_EfficientNet |
| This layer assumes "channel last". |
| """ |
| def __init__(self, in_features, out_features, min_channels_per_group=32, last_dim=2, use_bias=False, activation=None, has_batch_norm=False, has_batch_scale=False): |
| super().__init__() |
| self.in_features = in_features |
| self.out_features = out_features |
| self.min_channels_per_group = min_channels_per_group |
| self.last_dim = last_dim |
| self.activation = activation |
| self.has_batch_norm = has_batch_norm |
| self.has_batch_scale = has_batch_scale |
| self.has_interleaving = False |
| self.use_bias = use_bias |
| self.grouped = False |
| self.second_conv = False |
| self.first_pointwise_conv = None |
| self.second_pointwise_conv = None |
| self.interleave_layer = None |
| |
| self.weight = torch.Tensor(1, 1, 1) |
| self.bias = torch.Tensor(1, 1, 1) |
|
|
| prev_layer_channel_count = in_features |
| output_channel_count = out_features |
| max_acceptable_divisor = (prev_layer_channel_count//min_channels_per_group) |
| group_count = get_max_acceptable_common_divisor(prev_layer_channel_count, output_channel_count, max_acceptable = max_acceptable_divisor) |
| if group_count is None: group_count=1 |
| self.output_group_size = output_channel_count // group_count |
| |
|
|
| if (group_count>1): |
| self.grouped = True |
| |
| |
| |
| self.first_pointwise_conv = GroupedLinearFast(in_features=in_features, out_features=out_features, num_groups=group_count, bias=use_bias) |
| if self.output_group_size > 1: |
| self.has_interleaving = True |
| self.interleave_layer = InterleaveChannelsFast(self.output_group_size, last_dim=last_dim) |
| if (prev_layer_channel_count >= output_channel_count): |
| |
| self.second_conv = True |
| |
| |
| self.second_pointwise_conv = GroupedLinearFast(in_features=out_features, out_features=out_features, num_groups=group_count, bias=use_bias) |
| else: |
| |
| |
| self.first_pointwise_conv = GroupedLinear(in_features=in_features, out_features=out_features, num_groups=group_count, bias=use_bias) |
|
|
| def forward(self, x): |
| if (self.grouped): |
| output_tensor = self.first_pointwise_conv(x) |
| if self.activation is not None: |
| output_tensor = self.activation(output_tensor) |
| compression_tensor = output_tensor |
| if self.has_interleaving: |
| output_tensor = self.interleave_layer(output_tensor) |
| if self.second_conv: |
| output_tensor = self.second_pointwise_conv(output_tensor) |
| if self.activation is not None: |
| output_tensor = self.activation(output_tensor) |
| output_tensor = output_tensor + compression_tensor |
| else: |
| output_tensor = self.first_pointwise_conv(x) |
| if self.activation is not None: |
| output_tensor = self.activation(output_tensor) |
| return output_tensor |
|
|
| class GroupedPointwiseConvolutionBlockIO(nn.Module): |
| """ |
| This layer is composed by a grouped pointwise convolution followed by interleaving and another grouped pointwise comvolution with skip connection. This basic architecture can |
| vary according to the input tensor and its parameters. This is the basic building block for the papers: |
| https://www.researchgate.net/publication/360226228_Grouped_Pointwise_Convolutions_Reduce_Parameters_in_Convolutional_Neural_Networks |
| https://www.researchgate.net/publication/355214501_Grouped_Pointwise_Convolutions_Significantly_Reduces_Parameters_in_EfficientNet |
| This layer assumes "channel last". |
| """ |
| def __init__(self, in_features, out_features, min_channels_per_group=32, last_dim=2, use_bias=False, activation=None, has_batch_norm=False, has_batch_scale=False): |
| super().__init__() |
| self.in_features = in_features |
| self.out_features = out_features |
| self.min_channels_per_group = min_channels_per_group |
| self.last_dim = last_dim |
| self.activation = activation |
| self.has_batch_norm = has_batch_norm |
| self.has_batch_scale = has_batch_scale |
| self.has_interleaving = False |
| self.use_bias = use_bias |
| self.grouped = False |
| self.second_conv = False |
| self.first_pointwise_conv = None |
| self.second_pointwise_conv = None |
| self.interleave_layer = None |
| |
| self.weight = torch.Tensor(1, 1, 1) |
| self.bias = torch.Tensor(1, 1, 1) |
|
|
| prev_layer_channel_count = in_features |
| output_channel_count = out_features |
| max_acceptable_divisor = (prev_layer_channel_count//min_channels_per_group) |
| group_count = get_max_acceptable_common_divisor(prev_layer_channel_count, output_channel_count, max_acceptable = max_acceptable_divisor) |
| if group_count is None: group_count=1 |
| self.output_group_size = output_channel_count // group_count |
| self.input_group_size = prev_layer_channel_count // group_count |
|
|
| if (group_count>1): |
| self.grouped = True |
| self.has_interleaving = True |
| self.second_conv = True |
| |
| |
| |
| self.first_pointwise_conv = GroupedLinearFast(in_features=in_features, out_features=out_features, num_groups=group_count, bias=use_bias) |
| self.interleave_layer = InterleaveChannelsFast(self.input_group_size, last_dim=last_dim) |
| |
| |
| |
| self.second_pointwise_conv = GroupedLinearFast(in_features=in_features, out_features=out_features, num_groups=group_count, bias=use_bias) |
| else: |
| |
| |
| self.first_pointwise_conv = GroupedLinear(in_features=in_features, out_features=out_features, num_groups=group_count, bias=use_bias) |
|
|
| def forward(self, x): |
| if (self.grouped): |
| if self.activation is not None: |
| return self.activation(self.first_pointwise_conv(x) + self.second_pointwise_conv(self.interleave_layer(x))) |
| else: |
| return self.first_pointwise_conv(x) + self.second_pointwise_conv(self.interleave_layer(x)) |
| else: |
| if self.activation is not None: |
| return self.activation(self.first_pointwise_conv(x)) |
| else: |
| return self.first_pointwise_conv(x) |
|
|
| logger = logging.get_logger(__name__) |
|
|
| |
| |
| _flash_supports_window_size = False |
| try: |
| from flash_attn import flash_attn_func, flash_attn_varlen_func |
| from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input |
|
|
| _flash_supports_window_size = "window_size" in list(inspect.signature(flash_attn_func).parameters) |
| except ImportError as error: |
| logger.warning( |
| f"`flash-attention` package not found, consider installing for better performance: {error}." |
| ) |
| if not _flash_supports_window_size: |
| logger.warning( |
| "Current `flash-attention` does not support `window_size`. Either upgrade or use `attn_implementation='eager'`." |
| ) |
|
|
| _CHECKPOINT_FOR_DOC = "microsoft/Phi-3-mini-4k-instruct" |
| _CONFIG_FOR_DOC = "KPhi3Config" |
|
|
| |
| def _prepare_4d_causal_attention_mask_with_cache_position( |
| attention_mask: torch.Tensor, |
| sequence_length: int, |
| target_length: int, |
| dtype: torch.dtype, |
| device: torch.device, |
| min_dtype: float, |
| cache_position: torch.Tensor, |
| batch_size: int, |
| ): |
| """ |
| Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape |
| `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. |
| |
| Args: |
| attention_mask (`torch.Tensor`): |
| A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`. |
| sequence_length (`int`): |
| The sequence length being processed. |
| target_length (`int`): |
| The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet. |
| dtype (`torch.dtype`): |
| The dtype to use for the 4D attention mask. |
| device (`torch.device`): |
| The device to plcae the 4D attention mask on. |
| min_dtype (`float`): |
| The minimum value representable with the dtype `dtype`. |
| cache_position (`torch.Tensor`): |
| Indices depicting the position of the input sequence tokens in the sequence. |
| batch_size (`torch.Tensor`): |
| Batch size. |
| """ |
| if attention_mask is not None and attention_mask.dim() == 4: |
| |
| causal_mask = attention_mask |
| else: |
| causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device) |
| if sequence_length != 1: |
| causal_mask = torch.triu(causal_mask, diagonal=1) |
| causal_mask *= torch.arange(target_length, device=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, :] |
| padding_mask = padding_mask == 0 |
| causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( |
| padding_mask, min_dtype |
| ) |
|
|
| return causal_mask |
|
|
| |
| class Phi3RMSNorm(nn.Module): |
| def __init__(self, hidden_size, eps=1e-6): |
| """ |
| Phi3RMSNorm is equivalent to T5LayerNorm |
| """ |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(hidden_size)) |
| self.variance_epsilon = eps |
|
|
| def forward(self, hidden_states): |
| input_dtype = hidden_states.dtype |
| hidden_states = hidden_states.to(torch.float32) |
| variance = hidden_states.pow(2).mean(-1, keepdim=True) |
| hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) |
| return self.weight * hidden_states.to(input_dtype) |
|
|
| def extra_repr(self): |
| return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" |
|
|
|
|
| |
| class Phi3RotaryEmbedding(nn.Module): |
| def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None): |
| super().__init__() |
|
|
| self.dim = dim |
| self.max_position_embeddings = max_position_embeddings |
| self.base = base |
|
|
| inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float() / self.dim)) |
| self.register_buffer("inv_freq", tensor=inv_freq, persistent=False) |
|
|
| @torch.no_grad() |
| def forward(self, x, position_ids, seq_len=None): |
| |
| self.inv_freq.to(x.device) |
| inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) |
| position_ids_expanded = position_ids[:, None, :].float() |
| |
| |
| device_type = x.device.type |
| device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu" |
| with torch.autocast(device_type=device_type, enabled=False): |
| freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) |
| emb = torch.cat((freqs, freqs), dim=-1) |
| cos = emb.cos() |
| sin = emb.sin() |
| return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) |
|
|
|
|
| class Phi3SuScaledRotaryEmbedding(Phi3RotaryEmbedding): |
| def __init__(self, dim, config, device=None): |
| warnings.warn( |
| "The class Phi3SuScaledRotaryEmbedding is deprecated and will be removed in version 5 of Transformers. Please" |
| " use Phi3LongRoPEScaledRotaryEmbedding instead.", |
| FutureWarning, |
| ) |
| super().__init__(dim, config.max_position_embeddings, config.rope_theta, device) |
|
|
| self.short_factor = config.rope_scaling["short_factor"] |
| self.long_factor = config.rope_scaling["long_factor"] |
| self.original_max_position_embeddings = config.original_max_position_embeddings |
|
|
| @torch.no_grad() |
| def forward(self, x, position_ids, seq_len=None): |
| seq_len = torch.max(position_ids) + 1 |
| if seq_len > self.original_max_position_embeddings: |
| ext_factors = torch.tensor(self.long_factor, dtype=torch.float32, device=x.device) |
| else: |
| ext_factors = torch.tensor(self.short_factor, dtype=torch.float32, device=x.device) |
| inv_freq_shape = torch.arange(0, self.dim, 2, dtype=torch.int64, device=x.device).float() / self.dim |
| self.inv_freq = 1.0 / (ext_factors * self.base**inv_freq_shape) |
| inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) |
| position_ids_expanded = position_ids[:, None, :].float() |
| |
| |
| device_type = x.device.type |
| device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu" |
| with torch.autocast(device_type=device_type, enabled=False): |
| freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) |
| emb = torch.cat((freqs, freqs), dim=-1) |
| scale = self.max_position_embeddings / self.original_max_position_embeddings |
| if scale <= 1.0: |
| scaling_factor = 1.0 |
| else: |
| scaling_factor = math.sqrt(1 + math.log(scale) / math.log(self.original_max_position_embeddings)) |
| cos = emb.cos() * scaling_factor |
| sin = emb.sin() * scaling_factor |
| return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) |
|
|
|
|
| class Phi3YarnScaledRotaryEmbedding(Phi3RotaryEmbedding): |
| def __init__(self, dim, config, device=None): |
| warnings.warn( |
| "The class Phi3YarnScaledRotaryEmbedding is deprecated and will be removed in version 5 of Transformers", |
| FutureWarning, |
| ) |
| super().__init__(dim, config.max_position_embeddings, config.rope_theta, device) |
|
|
| self.short_factor = config.rope_scaling["short_factor"] |
| self.long_factor = config.rope_scaling["long_factor"] |
| self.original_max_position_embeddings = config.original_max_position_embeddings |
|
|
| @torch.no_grad() |
| def forward(self, x, position_ids, seq_len=None): |
| seq_len = torch.max(position_ids) + 1 |
| if seq_len > self.original_max_position_embeddings: |
| ext_factors = torch.tensor(self.long_factor, dtype=torch.float32, device=x.device) |
| else: |
| ext_factors = torch.tensor(self.short_factor, dtype=torch.float32, device=x.device) |
|
|
| inv_freq_shape = torch.arange(0, self.dim, 2, dtype=torch.int64, device=x.device).float() / self.dim |
| self.inv_freq = 1.0 / (ext_factors * self.base**inv_freq_shape) |
|
|
| inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) |
| position_ids_expanded = position_ids[:, None, :].float() |
|
|
| |
| |
| device_type = x.device.type |
| device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu" |
| with torch.autocast(device_type=device_type, enabled=False): |
| freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) |
| emb = torch.cat((freqs, freqs), dim=-1) |
|
|
| scale = self.max_position_embeddings / self.original_max_position_embeddings |
| if scale <= 1.0: |
| scaling_factor = 1.0 |
| else: |
| scaling_factor = 0.1 * math.log(scale) + 1.0 |
|
|
| cos = emb.cos() * scaling_factor |
| sin = emb.sin() * scaling_factor |
| return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) |
|
|
|
|
| class Phi3LongRoPEScaledRotaryEmbedding(Phi3RotaryEmbedding): |
| def __init__(self, dim, config, device=None): |
| super().__init__(dim, config.max_position_embeddings, config.rope_theta, device) |
|
|
| self.short_factor = config.rope_scaling["short_factor"] |
| self.long_factor = config.rope_scaling["long_factor"] |
| self.original_max_position_embeddings = config.original_max_position_embeddings |
|
|
| @torch.no_grad() |
| def forward(self, x, position_ids, seq_len=None): |
| seq_len = torch.max(position_ids) + 1 |
| if seq_len > self.original_max_position_embeddings: |
| ext_factors = torch.tensor(self.long_factor, dtype=torch.float32, device=x.device) |
| else: |
| ext_factors = torch.tensor(self.short_factor, dtype=torch.float32, device=x.device) |
|
|
| inv_freq_shape = torch.arange(0, self.dim, 2, dtype=torch.int64, device=x.device).float() / self.dim |
| self.inv_freq = 1.0 / (ext_factors * self.base**inv_freq_shape) |
|
|
| inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) |
| position_ids_expanded = position_ids[:, None, :].float() |
|
|
| |
| |
| device_type = x.device.type |
| device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu" |
| with torch.autocast(device_type=device_type, enabled=False): |
| freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) |
| emb = torch.cat((freqs, freqs), dim=-1) |
|
|
| scale = self.max_position_embeddings / self.original_max_position_embeddings |
| if scale <= 1.0: |
| scaling_factor = 1.0 |
| else: |
| scaling_factor = math.sqrt(1 + math.log(scale) / math.log(self.original_max_position_embeddings)) |
|
|
| cos = emb.cos() * scaling_factor |
| sin = emb.sin() * scaling_factor |
| return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) |
|
|
|
|
| |
| def rotate_half(x): |
| """Rotates half the hidden dims of the input.""" |
| x1 = x[..., : x.shape[-1] // 2] |
| x2 = x[..., x.shape[-1] // 2 :] |
| return torch.cat((-x2, x1), dim=-1) |
|
|
|
|
| |
| def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): |
| """Applies Rotary Position Embedding to the query and key tensors. |
| |
| Args: |
| q (`torch.Tensor`): The query tensor. |
| k (`torch.Tensor`): The key tensor. |
| cos (`torch.Tensor`): The cosine part of the rotary embedding. |
| sin (`torch.Tensor`): The sine part of the rotary embedding. |
| position_ids (`torch.Tensor`, *optional*): |
| Deprecated and unused. |
| unsqueeze_dim (`int`, *optional*, defaults to 1): |
| The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and |
| sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note |
| that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and |
| k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes |
| cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have |
| the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. |
| Returns: |
| `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. |
| """ |
| cos = cos.unsqueeze(unsqueeze_dim) |
| sin = sin.unsqueeze(unsqueeze_dim) |
| q_embed = (q * cos) + (rotate_half(q) * sin) |
| k_embed = (k * cos) + (rotate_half(k) * sin) |
| return q_embed, k_embed |
|
|
|
|
| class KPhi3MLP(nn.Module): |
| def __init__(self, config): |
| super().__init__() |
| self.config = config |
| self.activation_fn = ACT2FN[config.hidden_act] |
| if self.config.min_channels_per_group >= 0: |
| self.gate_up_proj = GroupedPointwiseConvolutionBlockIO(in_features=config.hidden_size, out_features=(2*config.intermediate_size), min_channels_per_group=self.config.min_channels_per_group , last_dim=2, use_bias=False) |
| self.down_proj = GroupedPointwiseConvolutionBlockIO(in_features=config.intermediate_size, out_features=config.hidden_size, min_channels_per_group=self.config.min_channels_per_group, last_dim=2, use_bias=False) |
| else: |
| self.gate_up_proj = nn.Linear(config.hidden_size, 2 * config.intermediate_size, bias=False) |
| self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) |
|
|
| def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor: |
| up_states = self.gate_up_proj(hidden_states) |
|
|
| gate, up_states = up_states.chunk(2, dim=-1) |
| up_states = up_states * self.activation_fn(gate) |
|
|
| return self.down_proj(up_states) |
|
|
|
|
| |
| def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: |
| """ |
| This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, |
| num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) |
| """ |
| batch, num_key_value_heads, slen, head_dim = hidden_states.shape |
| if n_rep == 1: |
| return hidden_states |
| hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) |
| return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) |
|
|
|
|
| class KPhi3Attention(nn.Module): |
| """Multi-headed attention from 'Attention Is All You Need' paper""" |
|
|
| def __init__(self, config: KPhi3Config, layer_idx: Optional[int] = None): |
| super().__init__() |
| self.config = config |
| self.layer_idx = layer_idx |
| if layer_idx is None: |
| logger.warning_once( |
| f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will " |
| "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` " |
| "when creating this class." |
| ) |
|
|
| self.attention_dropout = config.attention_dropout |
| self.hidden_size = config.hidden_size |
| self.num_heads = config.num_attention_heads |
| self.head_dim = self.hidden_size // self.num_heads |
| self.num_key_value_heads = config.num_key_value_heads |
| self.num_key_value_groups = self.num_heads // self.num_key_value_heads |
| self.max_position_embeddings = config.max_position_embeddings |
| self.original_max_position_embeddings = config.original_max_position_embeddings |
| self.rope_theta = config.rope_theta |
| self.rope_scaling = config.rope_scaling |
| self.is_causal = True |
|
|
| if (self.head_dim * self.num_heads) != self.hidden_size: |
| raise ValueError( |
| f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}" |
| f" and `num_heads`: {self.num_heads})." |
| ) |
|
|
| op_size = self.num_heads * self.head_dim + 2 * (self.num_key_value_heads * self.head_dim) |
| if self.config.min_channels_per_group >= 0: |
| self.o_proj = GroupedPointwiseConvolutionBlockIO(in_features=self.num_heads * self.head_dim, out_features=self.hidden_size, min_channels_per_group=self.config.min_channels_per_group, last_dim=2, use_bias=False) |
| self.qkv_proj = GroupedPointwiseConvolutionBlockIO(in_features=self.hidden_size, out_features=op_size, min_channels_per_group=self.config.min_channels_per_group, last_dim=2, use_bias=False) |
| else: |
| self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) |
| self.qkv_proj = nn.Linear(self.hidden_size, op_size, bias=False) |
| self._init_rope() |
|
|
| def _init_rope(self): |
| if self.rope_scaling is None: |
| self.rotary_emb = Phi3RotaryEmbedding( |
| self.head_dim, |
| max_position_embeddings=self.max_position_embeddings, |
| base=self.rope_theta, |
| ) |
| else: |
| scaling_type = self.config.rope_scaling["type"] |
| if scaling_type == "longrope": |
| self.rotary_emb = Phi3LongRoPEScaledRotaryEmbedding(self.head_dim, self.config) |
| else: |
| raise ValueError(f"Unknown RoPE scaling type {scaling_type}") |
|
|
| def forward( |
| self, |
| hidden_states: torch.Tensor, |
| attention_mask: Optional[torch.Tensor] = None, |
| position_ids: Optional[torch.LongTensor] = None, |
| past_key_value: Optional[Cache] = None, |
| output_attentions: bool = False, |
| use_cache: bool = False, |
| cache_position: Optional[torch.LongTensor] = None, |
| ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: |
| logger.warning_once("You are not running the flash-attention implementation, expect numerical differences.") |
|
|
| bsz, q_len, _ = hidden_states.size() |
|
|
| qkv = self.qkv_proj(hidden_states) |
| query_pos = self.num_heads * self.head_dim |
| query_states = qkv[..., :query_pos] |
| key_states = qkv[..., query_pos : query_pos + self.num_key_value_heads * self.head_dim] |
| value_states = qkv[..., query_pos + self.num_key_value_heads * self.head_dim :] |
|
|
| query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) |
| key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) |
| value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) |
|
|
| kv_seq_len = key_states.shape[-2] |
| if past_key_value is not None: |
| if self.layer_idx is None: |
| raise ValueError( |
| f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " |
| "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " |
| "with a layer index." |
| ) |
| kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) |
| cos, sin = self.rotary_emb(value_states, position_ids, seq_len=kv_seq_len) |
|
|
| query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids) |
|
|
| if past_key_value is not None: |
| cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} |
| key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) |
|
|
| |
| key_states = repeat_kv(key_states, self.num_key_value_groups) |
| value_states = repeat_kv(value_states, self.num_key_value_groups) |
|
|
| attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) |
|
|
| if attention_mask is not None: |
| causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] |
| attn_weights += causal_mask |
|
|
| |
| attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(value_states.dtype) |
| |
| |
| attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training) |
|
|
| attn_output = torch.matmul(attn_weights, value_states) |
| |
|
|
| if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): |
| raise ValueError( |
| f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" |
| f" {attn_output.size()}" |
| ) |
|
|
| attn_output = attn_output.transpose(1, 2).contiguous() |
| attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) |
|
|
| attn_output = self.o_proj(attn_output) |
| |
|
|
| if not output_attentions: |
| attn_weights = None |
|
|
| return attn_output, attn_weights, past_key_value |
|
|
|
|
| class KPhi3FlashAttention2(KPhi3Attention): |
| """ |
| KPhi-3 flash attention module. This module inherits from `KPhi3Attention` as the weights of the module stays |
| untouched. The only required change would be on the forward pass where it needs to correctly call the public API of |
| flash attention and deal with padding tokens in case the input contains any of them. |
| """ |
|
|
| |
| def __init__(self, *args, **kwargs): |
| super().__init__(*args, **kwargs) |
|
|
| |
| |
| |
| self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10() |
|
|
| def forward( |
| self, |
| hidden_states: torch.Tensor, |
| attention_mask: Optional[torch.LongTensor] = None, |
| position_ids: Optional[torch.LongTensor] = None, |
| past_key_value: Optional[Cache] = None, |
| output_attentions: bool = False, |
| use_cache: bool = False, |
| cache_position: Optional[torch.LongTensor] = None, |
| ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: |
| |
|
|
| output_attentions = False |
|
|
| bsz, q_len, _ = hidden_states.size() |
|
|
| qkv = self.qkv_proj(hidden_states) |
| query_pos = self.num_heads * self.head_dim |
| query_states = qkv[..., :query_pos] |
| key_states = qkv[..., query_pos : query_pos + self.num_key_value_heads * self.head_dim] |
| value_states = qkv[..., query_pos + self.num_key_value_heads * self.head_dim :] |
|
|
| |
| |
| |
| query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) |
| key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) |
| value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) |
|
|
| kv_seq_len = key_states.shape[-2] |
| if past_key_value is not None: |
| if self.layer_idx is None: |
| raise ValueError( |
| f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " |
| "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " |
| "with a layer index." |
| ) |
| kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) |
|
|
| |
| rotary_seq_len = ( |
| max(kv_seq_len, position_ids[:, -1].max().item() + 1) if position_ids is not None else kv_seq_len |
| ) |
|
|
| cos, sin = self.rotary_emb(value_states, seq_len=rotary_seq_len, position_ids=position_ids) |
|
|
| query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids) |
|
|
| if past_key_value is not None: |
| |
| cache_has_contents = past_key_value.get_seq_length(self.layer_idx) > 0 |
| if ( |
| getattr(self.config, "sliding_window", None) is not None |
| and kv_seq_len > self.config.sliding_window |
| and cache_has_contents |
| ): |
| slicing_tokens = 1 - self.config.sliding_window |
|
|
| past_key = past_key_value[self.layer_idx][0] |
| past_value = past_key_value[self.layer_idx][1] |
|
|
| past_key = past_key[:, :, slicing_tokens:, :].contiguous() |
| past_value = past_value[:, :, slicing_tokens:, :].contiguous() |
|
|
| if past_key.shape[-2] != self.config.sliding_window - 1: |
| raise ValueError( |
| f"past key must have a shape of (`batch_size, num_heads, self.config.sliding_window-1, head_dim`), got" |
| f" {past_key.shape}" |
| ) |
|
|
| if attention_mask is not None: |
| attention_mask = attention_mask[:, slicing_tokens:] |
| attention_mask = torch.cat([attention_mask, torch.ones_like(attention_mask[:, -1:])], dim=-1) |
|
|
| cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} |
| key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) |
|
|
| |
| key_states = repeat_kv(key_states, self.num_key_value_groups) |
| value_states = repeat_kv(value_states, self.num_key_value_groups) |
|
|
| attn_dropout = self.attention_dropout if self.training else 0.0 |
|
|
| |
| |
| |
| |
| |
|
|
| if query_states.dtype == torch.float32: |
| if torch.is_autocast_enabled(): |
| target_dtype = torch.get_autocast_gpu_dtype() |
| |
| elif hasattr(self.config, "_pre_quantization_dtype"): |
| target_dtype = self.config._pre_quantization_dtype |
| else: |
| target_dtype = self.qkv_proj.weight.dtype |
|
|
| logger.warning_once( |
| f"The input hidden states seems to be silently casted in float32, this might be related to" |
| f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" |
| f" {target_dtype}." |
| ) |
|
|
| query_states = query_states.to(target_dtype) |
| key_states = key_states.to(target_dtype) |
| value_states = value_states.to(target_dtype) |
|
|
| |
| query_states = query_states.transpose(1, 2) |
| key_states = key_states.transpose(1, 2) |
| value_states = value_states.transpose(1, 2) |
|
|
| attn_output = _flash_attention_forward( |
| query_states, |
| key_states, |
| value_states, |
| attention_mask, |
| q_len, |
| position_ids=position_ids, |
| dropout=attn_dropout, |
| sliding_window=getattr(self.config, "sliding_window", None), |
| use_top_left_mask=self._flash_attn_uses_top_left_mask, |
| is_causal=self.is_causal, |
| ) |
|
|
| attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous() |
| attn_output = self.o_proj(attn_output) |
|
|
| if not output_attentions: |
| attn_weights = None |
|
|
| return attn_output, attn_weights, past_key_value |
|
|
|
|
| |
| |
| class KPhi3SdpaAttention(KPhi3Attention): |
| """ |
| KPhi3 attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from |
| `KPhi3Attention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to |
| SDPA API. |
| """ |
|
|
| |
| def forward( |
| self, |
| hidden_states: torch.Tensor, |
| attention_mask: Optional[torch.Tensor] = None, |
| position_ids: Optional[torch.LongTensor] = None, |
| past_key_value: Optional[Cache] = None, |
| output_attentions: bool = False, |
| use_cache: bool = False, |
| cache_position: Optional[torch.LongTensor] = None, |
| ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: |
| if output_attentions: |
| |
| logger.warning_once( |
| "Phi3Model is using Phi3SdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, " |
| 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' |
| ) |
| return super().forward( |
| hidden_states=hidden_states, |
| attention_mask=attention_mask, |
| position_ids=position_ids, |
| past_key_value=past_key_value, |
| output_attentions=output_attentions, |
| use_cache=use_cache, |
| ) |
|
|
| bsz, q_len, _ = hidden_states.size() |
|
|
| qkv = self.qkv_proj(hidden_states) |
| query_pos = self.num_heads * self.head_dim |
| query_states = qkv[..., :query_pos] |
| key_states = qkv[..., query_pos : query_pos + self.num_key_value_heads * self.head_dim] |
| value_states = qkv[..., query_pos + self.num_key_value_heads * self.head_dim :] |
|
|
| query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) |
| key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) |
| value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) |
|
|
| kv_seq_len = key_states.shape[-2] |
| if past_key_value is not None: |
| kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) |
| cos, sin = self.rotary_emb(value_states, position_ids, seq_len=kv_seq_len) |
|
|
| query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids) |
|
|
| if past_key_value is not None: |
| cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} |
| key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) |
|
|
| key_states = repeat_kv(key_states, self.num_key_value_groups) |
| value_states = repeat_kv(value_states, self.num_key_value_groups) |
|
|
| causal_mask = attention_mask |
| if attention_mask is not None: |
| causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] |
|
|
| |
| |
| if query_states.device.type == "cuda" and attention_mask is not None: |
| query_states = query_states.contiguous() |
| key_states = key_states.contiguous() |
| value_states = value_states.contiguous() |
|
|
| |
| |
| |
| is_causal = True if causal_mask is None and q_len > 1 else False |
|
|
| attn_output = torch.nn.functional.scaled_dot_product_attention( |
| query_states, |
| key_states, |
| value_states, |
| attn_mask=causal_mask, |
| dropout_p=self.attention_dropout if self.training else 0.0, |
| is_causal=is_causal, |
| ) |
|
|
| attn_output = attn_output.transpose(1, 2).contiguous() |
| attn_output = attn_output.view(bsz, q_len, self.hidden_size) |
|
|
| attn_output = self.o_proj(attn_output) |
|
|
| return attn_output, None, past_key_value |
|
|
|
|
| KPHI3_ATTENTION_CLASSES = { |
| "eager": KPhi3Attention, |
| "flash_attention_2": KPhi3FlashAttention2, |
| "sdpa": KPhi3SdpaAttention, |
| } |
|
|
|
|
| class KPhi3DecoderLayer(nn.Module): |
| def __init__(self, config: KPhi3Config, layer_idx: int): |
| super().__init__() |
|
|
| self.config = config |
| self.self_attn = KPHI3_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx=layer_idx) |
|
|
| self.mlp = KPhi3MLP(config) |
| self.input_layernorm = Phi3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) |
|
|
| self.resid_attn_dropout = nn.Dropout(config.resid_pdrop) |
| self.resid_mlp_dropout = nn.Dropout(config.resid_pdrop) |
| self.post_attention_layernorm = Phi3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) |
|
|
| def forward( |
| self, |
| hidden_states: torch.Tensor, |
| attention_mask: Optional[torch.Tensor] = None, |
| position_ids: Optional[torch.LongTensor] = None, |
| past_key_value: Optional[Tuple[torch.Tensor]] = None, |
| output_attentions: Optional[bool] = False, |
| use_cache: Optional[bool] = False, |
| cache_position: Optional[torch.LongTensor] = None, |
| **kwargs, |
| ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: |
| """ |
| Args: |
| hidden_states (`torch.FloatTensor`): |
| input to the layer of shape `(batch, seq_len, embed_dim)` |
| attention_mask (`torch.FloatTensor`, *optional*): attention mask of size |
| `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. |
| position_ids (`torch.LongTensor` of shape `({0})`, *optional*): |
| Indices of positions of each input sequence tokens in the position embeddings. Selected in the range |
| `[0, config.n_positions - 1]`. [What are position IDs?](../glossary#position-ids) |
| output_attentions (`bool`, *optional*): |
| Whether or not to return the attentions tensors of all attention layers. See `attentions` under |
| returned tensors for more detail. |
| use_cache (`bool`, *optional*): |
| If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding |
| (see `past_key_values`). |
| past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states |
| cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): |
| Indices depicting the position of the input sequence tokens in the sequence |
| kwargs (`dict`, *optional*): |
| Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code |
| into the model |
| """ |
|
|
| residual = hidden_states |
|
|
| hidden_states = self.input_layernorm(hidden_states) |
|
|
| |
| attn_outputs, self_attn_weights, present_key_value = self.self_attn( |
| hidden_states=hidden_states, |
| attention_mask=attention_mask, |
| position_ids=position_ids, |
| past_key_value=past_key_value, |
| output_attentions=output_attentions, |
| use_cache=use_cache, |
| cache_position=cache_position, |
| ) |
|
|
| hidden_states = residual + self.resid_attn_dropout(attn_outputs) |
|
|
| residual = hidden_states |
| hidden_states = self.post_attention_layernorm(hidden_states) |
| hidden_states = self.mlp(hidden_states) |
| hidden_states = residual + self.resid_mlp_dropout(hidden_states) |
|
|
| outputs = (hidden_states,) |
|
|
| if output_attentions: |
| outputs += (self_attn_weights,) |
|
|
| if use_cache: |
| outputs += (present_key_value,) |
|
|
| return outputs |
|
|
|
|
| KPHI3_START_DOCSTRING = r""" |
| This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the |
| library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads |
| etc.) |
| |
| This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. |
| Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage |
| and behavior. |
| |
| Parameters: |
| config ([`KPhi3Config`]): |
| Model configuration class with all the parameters of the model. Initializing with a config file does not |
| load the weights associated with the model, only the configuration. Check out the |
| [`~PreTrainedModel.from_pretrained`] method to load the model weights. |
| """ |
|
|
|
|
| @add_start_docstrings( |
| "The bare Phi-3 model outputting raw hidden-states without any specific head on top.", |
| KPHI3_START_DOCSTRING, |
| ) |
| class KPhi3PreTrainedModel(PreTrainedModel, GenerationMixin): |
| config_class = KPhi3Config |
| base_model_prefix = "model" |
| supports_gradient_checkpointing = True |
| _no_split_modules = ["KPhi3DecoderLayer"] |
| _skip_keys_device_placement = "past_key_values" |
| _supports_flash_attn_2 = True |
| _supports_sdpa = False |
| _supports_cache_class = True |
|
|
| _version = "1.0.0" |
|
|
| def _init_weights(self, module): |
| std = self.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, nn.Embedding): |
| module.weight.data.normal_(mean=0.0, std=std) |
| if module.padding_idx is not None: |
| module.weight.data[module.padding_idx].zero_() |
|
|
|
|
| KPHI3_INPUTS_DOCSTRING = r""" |
| Args: |
| input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): |
| Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide |
| it. |
| |
| Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and |
| [`PreTrainedTokenizer.__call__`] for details. |
| |
| [What are input IDs?](../glossary#input-ids) |
| attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): |
| Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: |
| |
| - 1 for tokens that are **not masked**, |
| - 0 for tokens that are **masked**. |
| |
| [What are attention masks?](../glossary#attention-mask) |
| |
| Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and |
| [`PreTrainedTokenizer.__call__`] for details. |
| |
| If `past_key_values` is used, optionally only the last `input_ids` have to be input (see |
| `past_key_values`). |
| |
| If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`] |
| and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more |
| information on the default strategy. |
| |
| - 1 indicates the head is **not masked**, |
| - 0 indicates the head is **masked**. |
| position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): |
| Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, |
| config.n_positions - 1]`. |
| |
| [What are position IDs?](../glossary#position-ids) |
| past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*): |
| Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention |
| blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values` |
| returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`. |
| |
| Two formats are allowed: |
| - a [`~cache_utils.Cache`] instance; |
| - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of |
| shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy |
| cache format. |
| |
| The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the |
| legacy cache format will be returned. |
| |
| If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't |
| have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids` |
| of shape `(batch_size, sequence_length)`. |
| inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): |
| Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This |
| is useful if you want more control over how to convert `input_ids` indices into associated vectors than the |
| model's internal embedding lookup matrix. |
| use_cache (`bool`, *optional*): |
| If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see |
| `past_key_values`). |
| output_attentions (`bool`, *optional*): |
| Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned |
| tensors for more detail. |
| output_hidden_states (`bool`, *optional*): |
| Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for |
| more detail. |
| return_dict (`bool`, *optional*): |
| Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. |
| cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): |
| Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`, |
| this tensor is not affected by padding. It is used to update the cache in the correct position and to infer |
| the complete sequence length. |
| """ |
|
|
|
|
| @add_start_docstrings( |
| "The bare KPhi-3 model outputting raw hidden-states without any specific head on top.", |
| KPHI3_START_DOCSTRING, |
| ) |
| class KPhi3Model(KPhi3PreTrainedModel): |
| """ |
| Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`KPhi3DecoderLayer`] |
| |
| Args: |
| config: KPhi3Config |
| """ |
|
|
| def __init__(self, config: KPhi3Config): |
| super().__init__(config) |
| self.padding_idx = config.pad_token_id |
| self.vocab_size = config.vocab_size |
| self.activation_fn = ACT2FN[config.hidden_act] |
|
|
| self.embed_tokens = nn.Embedding(config.vocab_size, config.embed_size, self.padding_idx) |
| if config.embed_size != config.hidden_size: |
| self.embed_to_hidden = GroupedPointwiseConvolutionBlockIO(config.embed_size, config.hidden_size, config.min_channels_per_group) |
| self.embed_dropout = nn.Dropout(config.embd_pdrop) |
| self.layers = nn.ModuleList( |
| [KPhi3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] |
| ) |
| self._attn_implementation = config._attn_implementation |
| self.norm = Phi3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) |
|
|
| self.gradient_checkpointing = False |
| |
| self.post_init() |
|
|
| def get_input_embeddings(self): |
| return self.embed_tokens |
|
|
| def set_input_embeddings(self, value): |
| self.embed_tokens = value |
|
|
| @add_start_docstrings_to_model_forward(KPHI3_INPUTS_DOCSTRING) |
| def forward( |
| self, |
| input_ids: torch.LongTensor = None, |
| attention_mask: Optional[torch.Tensor] = None, |
| position_ids: Optional[torch.LongTensor] = None, |
| past_key_values: Optional[List[torch.FloatTensor]] = None, |
| inputs_embeds: Optional[torch.FloatTensor] = None, |
| use_cache: Optional[bool] = None, |
| output_attentions: Optional[bool] = None, |
| output_hidden_states: Optional[bool] = None, |
| return_dict: Optional[bool] = None, |
| cache_position: Optional[torch.LongTensor] = None, |
| ) -> Union[Tuple, BaseModelOutputWithPast]: |
| output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions |
| output_hidden_states = ( |
| output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states |
| ) |
| use_cache = use_cache if use_cache is not None else self.config.use_cache |
|
|
| return_dict = return_dict if return_dict is not None else self.config.use_return_dict |
|
|
| if (input_ids is None) ^ (inputs_embeds is not None): |
| raise ValueError( |
| "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" |
| ) |
|
|
| if self.gradient_checkpointing and self.training: |
| if use_cache: |
| logger.warning_once( |
| "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." |
| ) |
| use_cache = False |
|
|
| use_legacy_cache = False |
| if use_cache and not isinstance(past_key_values, Cache) and not self.training: |
| use_legacy_cache = True |
| past_key_values = DynamicCache.from_legacy_cache(past_key_values) |
| logger.warning_once( |
| "We detected that you are passing `past_key_values` as a tuple and this is deprecated and will be removed in v4.43. " |
| "Please use an appropriate `Cache` class (https://huggingface.co/docs/transformers/v4.41.3/en/internal/generation_utils#transformers.Cache)" |
| ) |
|
|
| if inputs_embeds is None: |
| inputs_embeds = self.embed_tokens(input_ids) |
|
|
| if cache_position is None: |
| past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 |
| cache_position = torch.arange( |
| past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device |
| ) |
| if position_ids is None: |
| position_ids = cache_position.unsqueeze(0) |
|
|
| causal_mask = self._update_causal_mask( |
| attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions |
| ) |
| inputs_embeds = self.embed_dropout(inputs_embeds) |
|
|
| if self.config.embed_size != self.config.hidden_size: |
| hidden_states = self.embed_to_hidden(inputs_embeds) |
| |
| else: |
| hidden_states = inputs_embeds |
|
|
| |
| all_hidden_states = () if output_hidden_states else None |
| all_self_attns = () if output_attentions else None |
| next_decoder_cache = None |
|
|
| for decoder_layer in self.layers: |
| if output_hidden_states: |
| all_hidden_states += (hidden_states,) |
|
|
| if self.gradient_checkpointing and self.training: |
| layer_outputs = self._gradient_checkpointing_func( |
| decoder_layer.__call__, |
| hidden_states, |
| causal_mask, |
| position_ids, |
| past_key_values, |
| output_attentions, |
| use_cache, |
| cache_position, |
| ) |
| else: |
| layer_outputs = decoder_layer( |
| hidden_states, |
| attention_mask=causal_mask, |
| position_ids=position_ids, |
| past_key_value=past_key_values, |
| output_attentions=output_attentions, |
| use_cache=use_cache, |
| cache_position=cache_position, |
| ) |
|
|
| hidden_states = layer_outputs[0] |
|
|
| if use_cache: |
| next_decoder_cache = layer_outputs[2 if output_attentions else 1] |
|
|
| if output_attentions: |
| all_self_attns += (layer_outputs[1],) |
|
|
| hidden_states = self.norm(hidden_states) |
|
|
| |
| if output_hidden_states: |
| all_hidden_states += (hidden_states,) |
|
|
| next_cache = None |
| if use_cache: |
| next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache |
| if not return_dict: |
| return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) |
| return BaseModelOutputWithPast( |
| last_hidden_state=hidden_states, |
| past_key_values=next_cache, |
| hidden_states=all_hidden_states, |
| attentions=all_self_attns, |
| ) |
|
|
| |
| def _update_causal_mask( |
| self, |
| attention_mask: torch.Tensor, |
| input_tensor: torch.Tensor, |
| cache_position: torch.Tensor, |
| past_key_values: Cache, |
| output_attentions: bool, |
| ): |
| |
| |
| |
| |
|
|
| if self.config._attn_implementation == "flash_attention_2": |
| if attention_mask is not None and 0.0 in attention_mask: |
| return attention_mask |
| return None |
|
|
| |
| |
| |
| past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 |
| using_static_cache = isinstance(past_key_values, StaticCache) |
|
|
| |
| if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions: |
| if AttentionMaskConverter._ignore_causal_mask_sdpa( |
| attention_mask, |
| inputs_embeds=input_tensor, |
| past_key_values_length=past_seen_tokens, |
| is_training=self.training, |
| ): |
| return None |
|
|
| dtype, device = input_tensor.dtype, input_tensor.device |
| min_dtype = torch.finfo(dtype).min |
| sequence_length = input_tensor.shape[1] |
| if using_static_cache: |
| target_length = past_key_values.get_max_length() |
| else: |
| target_length = ( |
| attention_mask.shape[-1] |
| if isinstance(attention_mask, torch.Tensor) |
| else past_seen_tokens + sequence_length + 1 |
| ) |
|
|
| |
| causal_mask = _prepare_4d_causal_attention_mask_with_cache_position( |
| attention_mask, |
| sequence_length=sequence_length, |
| target_length=target_length, |
| dtype=dtype, |
| device=device, |
| min_dtype=min_dtype, |
| cache_position=cache_position, |
| batch_size=input_tensor.shape[0], |
| ) |
|
|
| if ( |
| self.config._attn_implementation == "sdpa" |
| and attention_mask is not None |
| and attention_mask.device.type == "cuda" |
| and not output_attentions |
| ): |
| |
| |
| |
| causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) |
|
|
| return causal_mask |
|
|
|
|
| class KPhi3ForCausalLM(KPhi3PreTrainedModel): |
| _tied_weights_keys = ["lm_head.weight"] |
|
|
| |
| def __init__(self, config): |
| super().__init__(config) |
| self.model = KPhi3Model(config) |
| self.vocab_size = config.vocab_size |
| if config.embed_size != config.hidden_size: |
| self.hidden_to_embed = GroupedPointwiseConvolutionBlockIO(config.hidden_size, config.embed_size, config.min_channels_per_group) |
| self.lm_head = nn.Linear(config.embed_size, config.vocab_size, bias=False) |
|
|
| |
| self.post_init() |
|
|
| |
| def get_input_embeddings(self): |
| return self.model.embed_tokens |
|
|
| |
| def set_input_embeddings(self, value): |
| self.model.embed_tokens = value |
|
|
| |
| def get_output_embeddings(self): |
| 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 |
|
|
| |
| @add_start_docstrings_to_model_forward(KPHI3_INPUTS_DOCSTRING) |
| @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC) |
| def forward( |
| self, |
| input_ids: torch.LongTensor = None, |
| attention_mask: Optional[torch.Tensor] = None, |
| position_ids: Optional[torch.LongTensor] = None, |
| past_key_values: Optional[List[torch.FloatTensor]] = None, |
| inputs_embeds: Optional[torch.FloatTensor] = None, |
| labels: Optional[torch.LongTensor] = None, |
| use_cache: Optional[bool] = None, |
| output_attentions: Optional[bool] = None, |
| output_hidden_states: Optional[bool] = None, |
| return_dict: Optional[bool] = None, |
| cache_position: Optional[torch.LongTensor] = None, |
| ) -> Union[Tuple, CausalLMOutputWithPast]: |
| r""" |
| Args: |
| labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): |
| Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., |
| config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored |
| (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. |
| |
| Returns: |
| |
| Example: |
| |
| ```python |
| >>> from transformers import AutoTokenizer, Phi3ForCausalLM |
| |
| >>> model = KPhi3ForCausalLM.from_pretrained("microsoft/phi-3-mini-4k-instruct") |
| >>> tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-3-mini-4k-instruct") |
| |
| >>> prompt = "This is an example script ." |
| >>> inputs = tokenizer(prompt, return_tensors="pt") |
| |
| >>> # Generate |
| >>> generate_ids = model.generate(inputs.input_ids, max_length=30) |
| >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] |
| 'This is an example script .\n Certainly! Below is a sample script that demonstrates a simple task, such as calculating the sum' |
| ```""" |
|
|
| 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 |
|
|
| |
| outputs = self.model( |
| input_ids=input_ids, |
| attention_mask=attention_mask, |
| position_ids=position_ids, |
| past_key_values=past_key_values, |
| inputs_embeds=inputs_embeds, |
| use_cache=use_cache, |
| output_attentions=output_attentions, |
| output_hidden_states=output_hidden_states, |
| return_dict=return_dict, |
| ) |
|
|
| hidden_states = outputs[0] |
| if self.config.embed_size != self.config.hidden_size: |
| hidden_states = self.hidden_to_embed(hidden_states) |
| logits = self.lm_head(hidden_states) |
| logits = logits.float() |
|
|
| loss = None |
| if labels is not None: |
| |
| shift_logits = logits[..., :-1, :].contiguous() |
| shift_labels = labels[..., 1:].contiguous() |
| |
| loss_fct = CrossEntropyLoss() |
| shift_logits = shift_logits.view(-1, self.config.vocab_size) |
| shift_labels = shift_labels.view(-1) |
| |
| shift_labels = shift_labels.to(shift_logits.device) |
| loss = loss_fct(shift_logits, shift_labels) |
|
|
| if not return_dict: |
| output = (logits,) + outputs[1:] |
| return (loss,) + output if loss is not None else output |
|
|
| return CausalLMOutputWithPast( |
| loss=loss, |
| logits=logits, |
| past_key_values=outputs.past_key_values, |
| hidden_states=outputs.hidden_states, |
| attentions=outputs.attentions, |
| ) |
|
|
| |
| def prepare_inputs_for_generation( |
| self, |
| input_ids, |
| past_key_values=None, |
| attention_mask=None, |
| inputs_embeds=None, |
| cache_position=None, |
| position_ids=None, |
| use_cache=True, |
| **kwargs, |
| ): |
| |
| |
| |
| if past_key_values is not None: |
| if inputs_embeds is not None: |
| input_ids = input_ids[:, -cache_position.shape[0] :] |
| elif input_ids.shape[1] != cache_position.shape[0]: |
| input_ids = input_ids[:, cache_position] |
|
|
| if attention_mask is not None and position_ids is None: |
| |
| position_ids = attention_mask.long().cumsum(-1) - 1 |
| position_ids.masked_fill_(attention_mask == 0, 1) |
| if past_key_values: |
| position_ids = position_ids[:, -input_ids.shape[1] :] |
|
|
| |
| position_ids = position_ids.clone(memory_format=torch.contiguous_format) |
|
|
| |
| if inputs_embeds is not None and cache_position[0] == 0: |
| model_inputs = {"inputs_embeds": inputs_embeds, "input_ids": None} |
| else: |
| |
| model_inputs = {"input_ids": input_ids.clone(memory_format=torch.contiguous_format), "inputs_embeds": None} |
|
|
| if isinstance(past_key_values, StaticCache) and attention_mask.ndim == 2: |
| if model_inputs["inputs_embeds"] is not None: |
| batch_size, sequence_length, _ = model_inputs["inputs_embeds"].shape |
| device = model_inputs["inputs_embeds"].device |
| else: |
| batch_size, sequence_length = model_inputs["input_ids"].shape |
| device = model_inputs["input_ids"].device |
|
|
| dtype = self.lm_head.weight.dtype |
| min_dtype = torch.finfo(dtype).min |
|
|
| attention_mask = _prepare_4d_causal_attention_mask_with_cache_position( |
| attention_mask, |
| sequence_length=sequence_length, |
| target_length=past_key_values.get_max_length(), |
| dtype=dtype, |
| device=device, |
| min_dtype=min_dtype, |
| cache_position=cache_position, |
| batch_size=batch_size, |
| ) |
|
|
| model_inputs.update( |
| { |
| "position_ids": position_ids, |
| "cache_position": cache_position, |
| "past_key_values": past_key_values, |
| "use_cache": use_cache, |
| "attention_mask": attention_mask, |
| } |
| ) |
| return model_inputs |
|
|
|
|
| @add_start_docstrings( |
| """ |
| The [`KPhi3Model`] with a sequence classification head on top (linear layer). |
| |
| [`KPhi3ForSequenceClassification`] uses the last token in order to do the classification, as other causal models |
| (e.g. GPT-2) do. |
| |
| Since it does classification on the last token, it requires to know the position of the last token. If a |
| `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If |
| no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the |
| padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in |
| each row of the batch). |
| """, |
| KPHI3_START_DOCSTRING, |
| ) |
| |
| class KPhi3ForSequenceClassification(KPhi3PreTrainedModel): |
| def __init__(self, config): |
| super().__init__(config) |
| self.num_labels = config.num_labels |
| self.model = KPhi3Model(config) |
| self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False) |
|
|
| |
| self.post_init() |
|
|
| def get_input_embeddings(self): |
| return self.model.embed_tokens |
|
|
| def set_input_embeddings(self, value): |
| self.model.embed_tokens = value |
|
|
| @add_start_docstrings_to_model_forward(KPHI3_INPUTS_DOCSTRING) |
| def forward( |
| self, |
| input_ids: Optional[torch.LongTensor] = None, |
| attention_mask: Optional[torch.Tensor] = None, |
| position_ids: Optional[torch.LongTensor] = None, |
| past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, |
| inputs_embeds: Optional[torch.FloatTensor] = None, |
| labels: Optional[torch.LongTensor] = None, |
| use_cache: Optional[bool] = None, |
| output_attentions: Optional[bool] = None, |
| output_hidden_states: Optional[bool] = None, |
| return_dict: Optional[bool] = None, |
| ) -> Union[Tuple, SequenceClassifierOutputWithPast]: |
| r""" |
| labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): |
| Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., |
| config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If |
| `config.num_labels > 1` a classification loss is computed (Cross-Entropy). |
| """ |
| return_dict = return_dict if return_dict is not None else self.config.use_return_dict |
|
|
| model_outputs = self.model( |
| input_ids, |
| attention_mask=attention_mask, |
| position_ids=position_ids, |
| past_key_values=past_key_values, |
| inputs_embeds=inputs_embeds, |
| use_cache=use_cache, |
| output_attentions=output_attentions, |
| output_hidden_states=output_hidden_states, |
| return_dict=return_dict, |
| ) |
| hidden_states = model_outputs[0] |
| logits = self.score(hidden_states) |
|
|
| if input_ids is not None: |
| batch_size = input_ids.shape[0] |
| else: |
| batch_size = inputs_embeds.shape[0] |
|
|
| if self.config.pad_token_id is None and batch_size != 1: |
| raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") |
| if self.config.pad_token_id is None: |
| sequence_lengths = -1 |
| else: |
| if input_ids is not None: |
| |
| sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1 |
| sequence_lengths = sequence_lengths % input_ids.shape[-1] |
| sequence_lengths = sequence_lengths.to(logits.device) |
| else: |
| sequence_lengths = -1 |
|
|
| pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths] |
|
|
| loss = None |
| if labels is not None: |
| labels = labels.to(logits.device) |
| if self.config.problem_type is None: |
| if self.num_labels == 1: |
| self.config.problem_type = "regression" |
| elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): |
| self.config.problem_type = "single_label_classification" |
| else: |
| self.config.problem_type = "multi_label_classification" |
|
|
| if self.config.problem_type == "regression": |
| loss_fct = MSELoss() |
| if self.num_labels == 1: |
| loss = loss_fct(pooled_logits.squeeze(), labels.squeeze()) |
| else: |
| loss = loss_fct(pooled_logits, labels) |
| elif self.config.problem_type == "single_label_classification": |
| loss_fct = CrossEntropyLoss() |
| loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1)) |
| elif self.config.problem_type == "multi_label_classification": |
| loss_fct = BCEWithLogitsLoss() |
| loss = loss_fct(pooled_logits, labels) |
| if not return_dict: |
| output = (pooled_logits,) + model_outputs[1:] |
| return ((loss,) + output) if loss is not None else output |
|
|
| return SequenceClassifierOutputWithPast( |
| loss=loss, |
| logits=pooled_logits, |
| past_key_values=model_outputs.past_key_values, |
| hidden_states=model_outputs.hidden_states, |
| attentions=model_outputs.attentions, |
| ) |
|
|
|
|
| @add_start_docstrings( |
| """ |
| [`KPhi3Model`] with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for |
| Named-Entity-Recognition (NER) tasks. |
| """, |
| KPHI3_START_DOCSTRING, |
| ) |
| |
| class KPhi3ForTokenClassification(KPhi3PreTrainedModel): |
| def __init__(self, config: KPhi3Config): |
| super().__init__(config) |
| self.num_labels = config.num_labels |
|
|
| self.model = KPhi3Model(config) |
| if hasattr(config, "classifier_dropout") and config.classifier_dropout is not None: |
| classifier_dropout = config.classifier_dropout |
| elif hasattr(config, "hidden_dropout") and config.hidden_dropout is not None: |
| classifier_dropout = config.hidden_dropout |
| else: |
| classifier_dropout = 0.1 |
| self.dropout = nn.Dropout(classifier_dropout) |
| self.classifier = nn.Linear(config.hidden_size, config.num_labels) |
|
|
| |
| self.post_init() |
|
|
| @add_start_docstrings_to_model_forward(KPHI3_INPUTS_DOCSTRING) |
| @add_code_sample_docstrings( |
| checkpoint=_CHECKPOINT_FOR_DOC, |
| output_type=TokenClassifierOutput, |
| config_class=_CONFIG_FOR_DOC, |
| ) |
| def forward( |
| self, |
| input_ids: Optional[torch.LongTensor] = None, |
| past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None, |
| attention_mask: Optional[torch.Tensor] = None, |
| inputs_embeds: Optional[torch.Tensor] = None, |
| labels: Optional[torch.Tensor] = None, |
| use_cache: Optional[bool] = None, |
| output_attentions: Optional[bool] = None, |
| output_hidden_states: Optional[bool] = None, |
| return_dict: Optional[bool] = None, |
| **deprecated_arguments, |
| ) -> Union[Tuple[torch.Tensor], TokenClassifierOutput]: |
| r""" |
| labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): |
| Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., |
| config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If |
| `config.num_labels > 1` a classification loss is computed (Cross-Entropy). |
| """ |
| return_dict = return_dict if return_dict is not None else self.config.use_return_dict |
|
|
| model_outputs = self.model( |
| input_ids, |
| past_key_values=past_key_values, |
| attention_mask=attention_mask, |
| inputs_embeds=inputs_embeds, |
| use_cache=use_cache, |
| output_attentions=output_attentions, |
| output_hidden_states=output_hidden_states, |
| return_dict=return_dict, |
| ) |
|
|
| hidden_states = model_outputs[0] |
| hidden_states = self.dropout(hidden_states) |
| logits = self.classifier(hidden_states) |
|
|
| loss = None |
| if labels is not None: |
| |
| labels = labels.to(logits.device) |
| batch_size, seq_length = labels.shape |
| loss_fct = CrossEntropyLoss() |
| loss = loss_fct( |
| logits.view(batch_size * seq_length, self.num_labels), labels.view(batch_size * seq_length) |
| ) |
|
|
| if not return_dict: |
| output = (logits,) + model_outputs[2:] |
| return ((loss,) + output) if loss is not None else output |
|
|
| return TokenClassifierOutput( |
| loss=loss, |
| logits=logits, |
| hidden_states=model_outputs.hidden_states, |
| attentions=model_outputs.attentions, |
| ) |
