from transformers import Starcoder2Model
import sys
from .config import ModularStarEncoderConfig
import os
from dataclasses import dataclass
from typing import Optional, Tuple, Union, List
import sys
import torch
import torch.utils.checkpoint
from torch import nn
from torch.nn import CrossEntropyLoss
from transformers.activations import ACT2FN
from transformers.modeling_utils import PreTrainedModel
from transformers.utils import (
ModelOutput,
logging,
)
logger = logging.get_logger(__name__)
class StarEncoder2PreTrainedModel(PreTrainedModel):
"""
An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
models.
"""
config_class = ModularStarEncoderConfig
base_model_prefix = "ModularStarEncoder"
model_type = "ModularStarEncoder"
supports_gradient_checkpointing = True
_supports_flash_attn_2 = True
_supports_sdpa = True
_supports_cache_class = True
def _init_weights(self, module):
"""Initialize the weights"""
if isinstance(module, nn.Linear):
# Slightly different from the TF version which uses truncated_normal for initialization
# cf https://github.com/pytorch/pytorch/pull/5617
module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
if module.bias is not None:
module.bias.data.zero_()
elif isinstance(module, nn.Embedding):
module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
if module.padding_idx is not None:
module.weight.data[module.padding_idx].zero_()
elif isinstance(module, nn.LayerNorm):
module.bias.data.zero_()
module.weight.data.fill_(1.0)
class StarEncoder2Pooler(nn.Module):
def __init__(self, config):
super().__init__()
self.dense = nn.Linear(config.hidden_size, config.hidden_size)
self.activation = nn.Tanh()
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
# We "pool" the model by simply taking the hidden state corresponding
# to the last token.
last_token_tensor = hidden_states[:, -1]
pooled_output = self.dense(last_token_tensor)
pooled_output = self.activation(pooled_output)
return pooled_output
@dataclass
class ModularStarEncoderOutput(ModelOutput):
"""
Output type of [`BertForPreTraining`].
Args:
loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`):
Total loss as the sum of the masked language modeling loss and the next sequence prediction
(classification) loss.
prediction_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
seq_relationship_logits (`torch.FloatTensor` of shape `(batch_size, 2)`):
Prediction scores of the in context classification (classification) head (scores of True/False continuation
before SoftMax).
hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
shape `(batch_size, sequence_length, hidden_size)`.
Hidden-states of the model at the output of each layer plus the initial embedding outputs.
attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
sequence_length)`.
Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
heads.
"""
projected_pooled_normalized: Optional[List[torch.FloatTensor]] = None
raw_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
attentions: Optional[Tuple[torch.FloatTensor]] = None
def forward(self, sequence_output, pooled_output,idx_layer: Optional[torch.Tensor] = None):
if self.is_matryoshka:
device_sequence = sequence_output.get_device()
if device_sequence<0:
device_sequence = "cpu"
prediction_scores = self.predictions(torch.cat([sequence_output , self.conditional_embeddings(torch.tensor(idx_layer,device=device_sequence).int()).expand(sequence_output.size()[0],sequence_output.size()[1],-1)],dim=-1))
seq_relationship_score = self.seq_relationship(torch.cat([pooled_output , self.conditional_embeddings(torch.tensor(idx_layer,device=device_sequence).int()).expand(pooled_output.size()[0],-1)],dim=-1))
else:
prediction_scores = self.predictions(sequence_output)
seq_relationship_score = self.seq_relationship(pooled_output)
return prediction_scores, seq_relationship_score
def normalize(my_tensor):
embedding_norms = my_tensor.norm(dim=0)
normalizing_factor = torch.where( # Only normalize embeddings with norm > 1.0.
embedding_norms > 1.0, embedding_norms, torch.tensor(1)
)
normalized_tensor = my_tensor / normalizing_factor
return normalized_tensor
def pooling(x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
"""Pools a batch of vector sequences into a batch of vector global representations.
It does so by taking the average representation of the sequence, as indicated by the mask.
Args:
x (torch.Tensor): Batch of vector sequences with shape [B, T, F].
mask (torch.Tensor): Batch of masks with shape [B, T].
Returns:
torch.Tensor: Pooled version of the input batch with shape [B, F].
"""
# Expand the mask to match the feature dimensions for proper masking
mask_expanded = mask.unsqueeze(-1) # Shape [B, T, 1]
# Apply the mask to the input tensor
masked_x = x * mask_expanded # Shape [B, T, F]
# Sum along the time dimension
sum_x = masked_x.sum(dim=1) # Shape [B, F]
# Calculate the length of valid (non-padded) elements
valid_lengths = mask.sum(dim=1).clamp(min=1).unsqueeze(-1) # Shape [B, 1]
# Calculate the average pooling, avoiding division by zero
pooled_x = sum_x / valid_lengths # Shape [B, F]
return pooled_x
def pool_and_normalize(
features_sequence: torch.Tensor,
attention_masks: torch.Tensor,
return_norms: bool = False,
) -> Union[torch.Tensor, List[torch.Tensor]]:
"""Temporal ooling of sequences of vectors and projection onto the unit sphere.
Args:
features_sequence (torch.Tensor): Inpute features with shape [B, T, F].
attention_masks (torch.Tensor): Pooling masks with shape [B, T, F].
return_norms (bool, optional): Whether to additionally return the norms. Defaults to False.
Returns:
Union[torch.Tensor, List[torch.Tensor]]: Pooled and normalized vectors with shape [B, F].
"""
pooled_embeddings = pooling(features_sequence, attention_masks)
embedding_norms = pooled_embeddings.norm(dim=1)
normalizing_factor = torch.where( # Only normalize embeddings with norm > 1.0.
embedding_norms > 1.0, embedding_norms, torch.ones_like(embedding_norms)
)
pooled_normalized_embeddings = pooled_embeddings / normalizing_factor[:, None]
if return_norms:
return pooled_normalized_embeddings, embedding_norms
else:
return pooled_normalized_embeddings
def get_pooling_mask(
input_ids: torch.Tensor, sep_token_id: Union[int, float]
) -> torch.Tensor:
"""Gets pooling masks. For a sequence of input tokens, the mask will be
a sequence of zeros up until the first [SEP] occurrence, and 1 after that.
Args:
input_ids (torch.Tensor): Batch of input ids with shape [B, T].
sep_token_id (Union[int, float]): Id for [SEP] token.
Returns:
torch.Tensor: Batch of pooling masks with shape [B, T]
"""
# idx indicates the first occurrence of sep_token_id per along dim 0 of input_ids
idx = (input_ids == sep_token_id).float().flip(1).argmax(1)
idx = input_ids.size(-1)-idx-1
repeated_idx = idx.unsqueeze(1).repeat(1, input_ids.size(1))
DEVICE = input_ids.get_device()
if DEVICE<0:
DEVICE = "cpu"
ranges = torch.arange(input_ids.size(1)).repeat(input_ids.size(0), 1)
ranges = ranges.to(DEVICE)
pooling_mask = (repeated_idx <= ranges).long()
return pooling_mask
def adapt_model(model,config,till_layer:27):
model = model.starEncoder2
encoder_config = config
layers = encoder_config.matryoshka_layers
feature_dim = encoder_config.hidden_size
model.projection_heads = torch.nn.ModuleList()
if till_layer:
print(f"ATTENTION: till layer is on, you are pruning the model keeping just the first {till_layer} layers")
model.layers = model.layers[:till_layer]
model.projection_heads.append(torch.nn.Sequential(
torch.nn.Linear(feature_dim, feature_dim),
torch.nn.LeakyReLU(),
torch.nn.Linear(feature_dim, feature_dim),
))
else:
for layer in layers:
model.projection_heads.append(torch.nn.Sequential(
torch.nn.Linear(feature_dim, feature_dim),
torch.nn.LeakyReLU(),
torch.nn.Linear(feature_dim, feature_dim),
))
#setting off causal masking
for layer in model.layers:
layer.self_attn.is_causal=False
model.temperature_coef = torch.nn.Parameter(torch.Tensor([10.0]),requires_grad=False)
return model
class ModularStarEncoder(StarEncoder2PreTrainedModel):
_tied_weights_keys = ["predictions.decoder.bias", "cls.predictions.decoder.weight"]
config_class = ModularStarEncoderConfig
def __init__(self, config):
super().__init__(config)
self.model_type = "ModularStarEncoder"
for element in dir(config):
value = getattr(config, element) # Get the attribute value
if (isinstance(value, tuple) or isinstance(value, list)) and len(value)>0:
setattr(config, element, value[0])
self.layer_matryoshka_loss = config.layer_matryoshka_loss
self.matryoshka_layers = config.matryoshka_layers
self.starEncoder2 = Starcoder2Model(config)
#setting off causal masking
for layer in self.starEncoder2.layers:
layer.self_attn.is_causal=False
# Initialize weights and apply final processing
self.post_init()
self.till_layer= 18
self.starEncoder2 = adapt_model(self ,config=config,till_layer=self.till_layer)
def forward(
self,
input_ids: Optional[torch.Tensor] = None,
attention_mask: Optional[torch.Tensor] = None,
#token_type_ids: Optional[torch.Tensor] = None,
position_ids: Optional[torch.Tensor] = None,
head_mask: Optional[torch.Tensor] = None,
inputs_embeds: Optional[torch.Tensor] = None,
labels: Optional[torch.Tensor] = None,
next_sentence_label: Optional[torch.Tensor] = None,
output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
sep_token_id:Optional[int] = 49152,
) -> Union[Tuple[torch.Tensor], ModularStarEncoderOutput]:
r"""
labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
config.vocab_size]` (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]`
next_sentence_label (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
This label is assigned to the in context loss:
- 0 indicates sequence B belongs to the same repository of A,
- 1 indicates sequence B is a random repository.
kwargs (`Dict[str, any]`, optional, defaults to *{}*):
Used to hide legacy arguments that have been deprecated.
"""
source_embedding = self.starEncoder2(
input_ids,
attention_mask=attention_mask,
position_ids=position_ids,
inputs_embeds=inputs_embeds,
output_attentions=output_attentions,
output_hidden_states=True,
return_dict=True,
)
DEVICE = source_embedding.hidden_states[-1].get_device()
if DEVICE<0:
DEVICE = "cpu"
try:
projection_fn = self.starEncoder2.module.projection_heads
temp_coef = self.starEncoder2.module.temperature_coef
except AttributeError:
projection_fn = self.starEncoder2.projection_heads
temp_coef = self.starEncoder2.temperature_coef
for head in projection_fn:
head.to(DEVICE)
temp_coef.to(DEVICE)
pooling_mask_source_targtes = get_pooling_mask(
input_ids, sep_token_id
) # Pooling masks indicate the second [SEP] occurrence, 0 till SEP, then all ones.
if self.till_layer:
self.matryoshka_layers=[self.till_layer]
pooled_and_normalized = []
for idx,matr_layer in enumerate(self.matryoshka_layers):
source_embedding_proj = projection_fn[idx](source_embedding.hidden_states[matr_layer])
normalized_source_embedding, embedding_norms = pool_and_normalize(
source_embedding_proj,
pooling_mask_source_targtes,
return_norms=True,
)
pooled_and_normalized.append(normalized_source_embedding)
if not self.till_layer:
return ModularStarEncoderOutput(
projected_pooled_normalized = pooled_and_normalized,
raw_hidden_states=source_embedding.hidden_states,
attentions=source_embedding.attentions,
)
else:
return ModularStarEncoderOutput(
projected_pooled_normalized = pooled_and_normalized[0],
raw_hidden_states=source_embedding.hidden_states,
attentions=source_embedding.attentions,
)