modularStarEncoder.py

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= 27
        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,
            )