modeling.py

from typing import Optional, Union

import torch
import torch.nn as nn
from transformers import (
    M2M100Config,)

from transformers.models.m2m_100.modeling_m2m_100 import (
    M2M100Encoder,
    M2M100ScaledWordEmbedding,
    M2M100ForConditionalGeneration,
    M2M100Model,
    shift_tokens_right,
    logger)

from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPastAndCrossAttentions, Seq2SeqLMOutput, Seq2SeqModelOutput

from transformers.utils import auto_docstring

from torch.nn import CrossEntropyLoss

class Pooling(nn.Module):
    """
    Pooling layer for sequence representations.
    
    Supports multiple pooling strategies: mean, cls, last, max, and none.
    """
    
    def __init__(self, pooling_type: str = "mean"):
        super().__init__()
        valid_types = {"mean", "cls", "last", "max", "none"}
        if pooling_type not in valid_types:
            raise ValueError(f"pooling_type must be one of {valid_types}, got {pooling_type}")
        self.pooling_type = pooling_type
    
    def forward(
        self, 
        hidden_states: torch.Tensor, 
        attention_mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """
        Apply pooling to hidden states.
        
        Args:
            hidden_states: Tensor of shape (batch_size, seq_len, hidden_size)
            attention_mask: Tensor of shape (batch_size, seq_len), values in {0, 1}
        
        Returns:
            Pooled tensor of shape (batch_size, 1, hidden_size) or (batch_size, seq_len, hidden_size) for none
        """
        if self.pooling_type == "none":
            return hidden_states
        
        if self.pooling_type == "cls":
            return hidden_states[:, 0, :].unsqueeze(1)
        
        elif self.pooling_type == "last":
            return hidden_states[:, -1, :].unsqueeze(1)
        
        elif self.pooling_type in ["mean", "max"]:
            if attention_mask is None:
                raise ValueError(f"attention_mask is required for {self.pooling_type} pooling")
            
            # Expand attention mask to match hidden_states dimensions
            mask = attention_mask.unsqueeze(-1)
            
            if self.pooling_type == "mean":
                # Apply mask and compute mean over valid tokens
                masked_hidden = hidden_states * mask
                sum_hidden = masked_hidden.sum(dim=1, keepdim=True)
                sum_mask = mask.sum(dim=1, keepdim=True)
                
                return sum_hidden / sum_mask
            
            elif self.pooling_type == "max":
                # Apply mask (set masked positions to large negative value)
                masked_hidden = hidden_states.masked_fill(mask == 0, float('-inf'))
                return masked_hidden.max(dim=1, keepdim=True)[0]

class SONARTextEncoder(M2M100Encoder):
    """
    Transformer encoder with pooling capabilities.
    
    Inherits from M2M100Encoder and adds configurable pooling functionality.
    
    Args:
        config: M2M100Config with optional pooling_type attribute
        embed_tokens: Optional embedding layer
    """

    def __init__(self, config: M2M100Config, embed_tokens: Optional[nn.Embedding] = None):
        super().__init__(config, embed_tokens)
        
        # Initialize pooling layer
        pooling_type = getattr(config, 'pooling_type', 'mean')
        self.pooling = Pooling(pooling_type)
        
    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ):
        """
        Forward pass with optional pooling.
        
        Args:
            input_ids: Input token ids
            attention_mask: Attention mask for padding tokens
            head_mask: Mask for attention heads
            inputs_embeds: Pre-computed embeddings
            output_attentions: Whether to return attention weights
            output_hidden_states: Whether to return all hidden states
            return_dict: Whether to return ModelOutput object
            pool: Pooling strategy override (if None, uses config default)
        
        Returns:
            Model output with pooled representations
        """
        # Get encoder output
        encoder_output = super().forward(
            input_ids=input_ids,
            attention_mask=attention_mask,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        
        # Extract hidden states
        if return_dict:
            hidden_states = encoder_output.last_hidden_state

        else:
            hidden_states = encoder_output[0]
        
        pooled_output = self.pooling(hidden_states, attention_mask)

        if return_dict:
            encoder_output.last_hidden_state = pooled_output
        else:
            encoder_output = (pooled_output,) + encoder_output[1:]

        return encoder_output

class SONARModel(M2M100Model):
    """SONAR model based on M2M100."""

    def __init__(self, config: M2M100Config):
        super().__init__(config)

        self.encoder = SONARTextEncoder(config, self.shared)

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        decoder_input_ids: Optional[torch.LongTensor] = None,
        decoder_attention_mask: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        decoder_head_mask: Optional[torch.Tensor] = None,
        cross_attn_head_mask: Optional[torch.Tensor] = None,
        encoder_outputs: Optional[tuple[tuple[torch.FloatTensor]]] = None,
        past_key_values: Optional[tuple[tuple[torch.FloatTensor]]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        decoder_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.Tensor] = None,
        return_logits: Optional[bool] = False,
    ) -> Union[tuple[torch.Tensor], Seq2SeqModelOutput]:
        r"""
        decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
            Indices of decoder input sequence tokens in the vocabulary.

            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
            [`PreTrainedTokenizer.__call__`] for details.

            [What are decoder input IDs?](../glossary#decoder-input-ids)

            M2M100 uses the `eos_token_id` as the starting token for `decoder_input_ids` generation. If
            `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
            `past_key_values`).
        decoder_attention_mask (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):    
            Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
            be used by default.
        cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):
            Mask to nullify selected heads of the cross-attention modules in the decoder. Mask values selected in `[0,
            1]`:

            - 1 indicates the head is **not masked**,
            - 0 indicates the head is **masked**.
        """     
        
        outputs = super().forward(
            input_ids,
            attention_mask=attention_mask,
            decoder_input_ids=decoder_input_ids,
            decoder_attention_mask=decoder_attention_mask,
            head_mask=head_mask,
            decoder_head_mask=decoder_head_mask,
            cross_attn_head_mask=cross_attn_head_mask,
            encoder_outputs=encoder_outputs,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            decoder_inputs_embeds=decoder_inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            cache_position=cache_position,
        )
        
        if return_logits:
            lm_logits = self.decoder.lm_head(outputs[0])
            if not return_dict:
                outputs = (lm_logits,) + outputs[1:]
            else:
                outputs.last_hidden_state = lm_logits
                
        return outputs
        
        


    @auto_docstring
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        decoder_input_ids: Optional[torch.LongTensor] = None,
        decoder_attention_mask: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        decoder_head_mask: Optional[torch.Tensor] = None,
        cross_attn_head_mask: Optional[torch.Tensor] = None,
        encoder_outputs: Optional[tuple[tuple[torch.FloatTensor]]] = None,
        past_key_values: Optional[tuple[tuple[torch.FloatTensor]]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        decoder_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.Tensor] = None,
    ) -> Union[tuple[torch.Tensor], Seq2SeqModelOutput]:
        r"""
        decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
            Indices of decoder input sequence tokens in the vocabulary.

            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
            [`PreTrainedTokenizer.__call__`] for details.

            [What are decoder input IDs?](../glossary#decoder-input-ids)

            M2M100 uses the `eos_token_id` as the starting token for `decoder_input_ids` generation. If
            `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
            `past_key_values`).
        decoder_attention_mask (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
            Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
            be used by default.
        cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):
            Mask to nullify selected heads of the cross-attention modules in the decoder. Mask values selected in `[0,
            1]`:

            - 1 indicates the head is **not masked**,
            - 0 indicates the head is **masked**.
        """
        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 encoder_outputs is None:
            encoder_outputs = self.encoder(
                input_ids=input_ids,
                attention_mask=attention_mask,
                head_mask=head_mask,
                inputs_embeds=inputs_embeds,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                return_dict=return_dict,
            )
            

        # If the user passed a tuple for encoder_outputs, we wrap it in a BaseModelOutput when return_dict=True
        elif return_dict and not isinstance(encoder_outputs, BaseModelOutput):
            encoder_outputs = BaseModelOutput(
                last_hidden_state=encoder_outputs[0],
                hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
                attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
            )
        if attention_mask is not None:
            if (encoder_outputs[0].size(1) != 1 and encoder_outputs[0].dim() == 3):
                logger.warning_once(
                    f"Encoder is not pooled"
                )
                encoder_attention_mask = attention_mask
            else:
                encoder_attention_mask = attention_mask[:, :1]
        
        # decoder outputs consists of (dec_features, past_key_value, dec_hidden, dec_attn)
        decoder_outputs = self.decoder(
            input_ids=decoder_input_ids,
            attention_mask=decoder_attention_mask,
            encoder_hidden_states=encoder_outputs[0],
            encoder_attention_mask=encoder_attention_mask,
            head_mask=decoder_head_mask,
            cross_attn_head_mask=cross_attn_head_mask,
            past_key_values=past_key_values,
            inputs_embeds=decoder_inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            cache_position=cache_position,
        )

        if not return_dict:
            return decoder_outputs + encoder_outputs

        return Seq2SeqModelOutput(
            last_hidden_state=decoder_outputs.last_hidden_state,
            past_key_values=decoder_outputs.past_key_values,
            decoder_hidden_states=decoder_outputs.hidden_states,
            decoder_attentions=decoder_outputs.attentions,
            cross_attentions=decoder_outputs.cross_attentions,
            encoder_last_hidden_state=encoder_outputs.last_hidden_state,
            encoder_hidden_states=encoder_outputs.hidden_states,
            encoder_attentions=encoder_outputs.attentions,
        )

class SONARForText2Text(M2M100ForConditionalGeneration):
    """SONAR model for conditional generation tasks."""
    
    # _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight"]

    def __init__(self, config: M2M100Config):
        super().__init__(config)
        
        self.model = SONARModel(config)

        self.cross_entropy_loss = CrossEntropyLoss(
            label_smoothing=0.1,
            ignore_index=-100
        )
        
        self.mse_loss = nn.MSELoss()
        

        self.mse_ratio = getattr(config, 'mse_ratio', 0.2)
        
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        target_ids: Optional[torch.LongTensor] = None,
        target_attention_mask: Optional[torch.Tensor] = None,
        mse_mask: Optional[torch.Tensor] = None,
        decoder_input_ids: Optional[torch.LongTensor] = None,
        decoder_attention_mask: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        decoder_head_mask: Optional[torch.Tensor] = None,
        cross_attn_head_mask: Optional[torch.Tensor] = None,
        encoder_outputs: Optional[tuple[tuple[torch.FloatTensor]]] = None,
        past_key_values: Optional[tuple[tuple[torch.FloatTensor]]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        decoder_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.Tensor] = None,
    ) -> Union[tuple[torch.Tensor], Seq2SeqLMOutput]:
        r"""
        decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
            Indices of decoder input sequence tokens in the vocabulary.

            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
            [`PreTrainedTokenizer.__call__`] for details.

            [What are decoder input IDs?](../glossary#decoder-input-ids)

            M2M100 uses the `eos_token_id` as the starting token for `decoder_input_ids` generation. If
            `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
            `past_key_values`).
        decoder_attention_mask (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
            Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
            be used by default.
        cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):
            Mask to nullify selected heads of the cross-attention modules in the decoder. Mask values selected in `[0,
            1]`:

            - 1 indicates the head is **not masked**,
            - 0 indicates the head is **masked**.
        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]`.

        Example Translation:

        ```python
        >>> from transformers import AutoTokenizer, M2M100ForConditionalGeneration

        >>> model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_418M")
        >>> tokenizer = AutoTokenizer.from_pretrained("facebook/m2m100_418M")

        >>> text_to_translate = "Life is like a box of chocolates"
        >>> model_inputs = tokenizer(text_to_translate, return_tensors="pt")

        >>> # translate to French
        >>> gen_tokens = model.generate(**model_inputs, forced_bos_token_id=tokenizer.get_lang_id("fr"))
        >>> print(tokenizer.batch_decode(gen_tokens, skip_special_tokens=True))
        ```
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if labels is not None:
            if decoder_input_ids is None:
                decoder_input_ids = shift_tokens_right(
                    labels, self.config.pad_token_id, self.config.decoder_start_token_id
                )

        outputs = self.model(
            input_ids,
            attention_mask=attention_mask,
            decoder_input_ids=decoder_input_ids,
            encoder_outputs=encoder_outputs,
            decoder_attention_mask=decoder_attention_mask,
            head_mask=head_mask,
            decoder_head_mask=decoder_head_mask,
            cross_attn_head_mask=cross_attn_head_mask,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            decoder_inputs_embeds=decoder_inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            cache_position=cache_position,
        )
        lm_logits = self.lm_head(outputs[0])
        
        masked_lm_loss = None
        if labels is not None:

            labels = labels.to(lm_logits.device)

            masked_lm_loss = self.cross_entropy_loss(lm_logits.view(-1, self.config.vocab_size), labels.view(-1))
            
            # print(f"Cross Entropy Loss: {masked_lm_loss if masked_lm_loss is not None else 'N/A'}")

            masked_lm_loss = masked_lm_loss.mean()

        if mse_mask is not None and target_ids is None:
            mse_mask = mse_mask.view(-1, 1, 1).to(outputs.encoder_last_hidden_state.device)
            encoder_outputs = outputs.encoder_last_hidden_state.squeeze()
            batch_size = encoder_outputs.size(0)
    
            # Reshape to pair structure: [batch//2, 2, seq_len, hidden]
            paired = encoder_outputs[:batch_size//2*2].view(batch_size//2, 2, *encoder_outputs.shape[1:]) * mse_mask

            mse_loss = self.mse_loss(paired[:, 0], paired[:, 1])
            masked_lm_loss += self.mse_ratio * mse_loss

        if target_ids is not None and labels is not None:

            target_ids = target_ids.to(lm_logits.device)
            target_encoder_outputs = self.model.encoder(
                input_ids=target_ids,
                attention_mask=target_attention_mask,
                return_dict=return_dict,
            )

            mse_loss = self.mse_loss(outputs.encoder_last_hidden_state, target_encoder_outputs.last_hidden_state)

            masked_lm_loss += self.mse_ratio * mse_loss
            
        # print(f"Masked LM Loss: {masked_lm_loss.item() if masked_lm_loss is not None else 'N/A'}")
        # print(f"MSE Loss: {mse_loss.item() if mse_loss is not None else 'N/A'}")

        if not return_dict:
            output = (lm_logits,) + outputs[1:]
            return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output

        return Seq2SeqLMOutput(
            loss=masked_lm_loss,
            logits=lm_logits,
            past_key_values=outputs.past_key_values,
            decoder_hidden_states=outputs.decoder_hidden_states,
            decoder_attentions=outputs.decoder_attentions,
            cross_attentions=outputs.cross_attentions,
            encoder_last_hidden_state=outputs.encoder_last_hidden_state,
            encoder_hidden_states=outputs.encoder_hidden_states,
            encoder_attentions=outputs.encoder_attentions,
        )
        
    @classmethod
    def from_m2m100_pretrained(cls, pretrained_model_name_or_path: str, *model_args, **kwargs):
        model = M2M100ForConditionalGeneration.from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
        
        generation_config = model.generation_config
        
        generation_config.early_stopping = True
        generation_config.num_beams = 5
        generation_config.max_length = 500
        
        config = model.config
        config.pooling_type = getattr(config, 'pooling_type', 'mean')
        
        sonar_model = cls(model.config)
        sonar_model.load_state_dict(model.state_dict(), strict=False)
        sonar_model.generation_config = generation_config
        return sonar_model