Upload modeling_delulu.py with huggingface_hub

modeling_delulu.py

@@ -1,27 +1,11 @@
-"""
-DELULU Model
-
-DELULU (Discriminative Embedding Learning Using Latent Units) is a speaker-aware
-self-supervised speech foundational model based on HuBERT architecture.
-
-Paper: https://arxiv.org/abs/2510.17662
-Authors: Massa Baali, Rita Singh, Bhiksha Raj
-
-This implementation wraps torchaudio's wav2vec2_model for compatibility with
-Hugging Face's AutoModel interface.
-"""
-
+"""DELULU Model"""
 import torch
 import torch.nn as nn
 from typing import Optional, Tuple, Union
-from dataclasses import dataclass
-
 from transformers import PreTrainedModel
 from transformers.modeling_outputs import BaseModelOutput
-
 from .configuration_delulu import DELULUConfig
 
-# Try to import torchaudio
 try:
     from torchaudio.models.wav2vec2 import wav2vec2_model
     TORCHAUDIO_AVAILABLE = True
@@ -29,79 +13,25 @@ except ImportError:
     TORCHAUDIO_AVAILABLE = False
 
 
-@dataclass
-class DELULUOutput(BaseModelOutput):
-    """
-    Output class for DELULU model.
-    
-    Args:
-        last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
-            Sequence of hidden-states at the output of the last layer of the model.
-        hidden_states (`tuple(torch.FloatTensor)`, *optional*):
-            Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for each layer)
-            of shape `(batch_size, sequence_length, hidden_size)`.
-        attentions (`tuple(torch.FloatTensor)`, *optional*):
-            Attention weights (not available for torchaudio backend).
-        extract_features (`torch.FloatTensor` of shape `(batch_size, sequence_length, conv_dim[-1])`):
-            Features from the convolutional feature extractor.
-    """
-    last_hidden_state: torch.FloatTensor = None
-    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
-    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
-    extract_features: Optional[torch.FloatTensor] = None
-
-
 class DELULUModel(PreTrainedModel):
-    """
-    DELULU Model for speaker-aware speech representation learning.
-    
-    This model wraps torchaudio's wav2vec2_model with DELULU's custom configuration
-    (modified convolutional strides for 16ms frame shift).
-    
-    Example:
-        ```python
-        from transformers import AutoModel
-        import torch
-        
-        # Load model
-        model = AutoModel.from_pretrained("cmu-mlsp/DELULU", trust_remote_code=True)
-        model.eval()
-        
-        # Process audio (16kHz, mono)
-        waveform = torch.randn(1, 16000)  # 1 second of audio
-        
-        with torch.no_grad():
-            outputs = model(waveform)
-            features = outputs.last_hidden_state  # [1, T, 768]
-            
-            # For speaker verification, use mean pooling
-            speaker_embedding = features.mean(dim=1)  # [1, 768]
-        ```
-    """
-    
     config_class = DELULUConfig
-    base_model_prefix = "delulu"
+    base_model_prefix = "wav2vec2"
     main_input_name = "input_values"
     supports_gradient_checkpointing = False
+    _no_split_modules = []
     
     def __init__(self, config: DELULUConfig):
         super().__init__(config)
-        self.config = config
         
         if not TORCHAUDIO_AVAILABLE:
-            raise ImportError(
-                "torchaudio is required for DELULU model. "
-                "Install with: pip install torchaudio"
-            )
+            raise ImportError("torchaudio required: pip install torchaudio")
         
-        # Build convolutional layer config from DELULU config
         conv_layer_config = list(zip(
             config.conv_dim,
             config.conv_kernel,
             config.conv_stride
         ))
         
-        # Create the underlying torchaudio model
         self.wav2vec2 = wav2vec2_model(
             extractor_mode=config.extractor_mode,
             extractor_conv_layer_config=conv_layer_config,
@@ -120,214 +50,38 @@ class DELULUModel(PreTrainedModel):
             encoder_layer_drop=config.layer_drop,
             aux_num_out=None,
         )
-        
-        # Initialize weights
         self.post_init()
     
+    def _init_weights(self, module):
+        pass
+    
     def forward(
         self,
         input_values: torch.Tensor,
         attention_mask: Optional[torch.Tensor] = None,
         output_hidden_states: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
         return_dict: Optional[bool] = None,
-    ) -> Union[Tuple, DELULUOutput]:
-        """
-        Forward pass of DELULU model.
-        
-        Args:
-            input_values (`torch.FloatTensor` of shape `(batch_size, sequence_length)`):
-                Raw audio waveform at 16kHz sampling rate.
-            attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
-                Mask to avoid performing attention on padding. Not used in current implementation.
-            output_hidden_states (`bool`, *optional*):
-                Whether to return all hidden states.
-            output_attentions (`bool`, *optional*):
-                Whether to return attention weights. Not supported with torchaudio backend.
-            return_dict (`bool`, *optional*):
-                Whether to return a `DELULUOutput` instead of a tuple.
-        
-        Returns:
-            `DELULUOutput` or `tuple`: Model outputs.
-        """
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None 
-            else self.config.output_hidden_states if hasattr(self.config, 'output_hidden_states')
-            else False
-        )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict if hasattr(self.config, 'use_return_dict') else True
+    ) -> Union[Tuple, BaseModelOutput]:
+        return_dict = return_dict if return_dict is not None else True
+        output_hidden_states = output_hidden_states or False
         
-        # Ensure input is 2D: (batch, samples)
         if input_values.dim() == 1:
             input_values = input_values.unsqueeze(0)
         
-        # Handle lengths for torchaudio model
-        lengths = None
-        if attention_mask is not None:
-            lengths = attention_mask.sum(dim=-1)
+        lengths = attention_mask.sum(-1) if attention_mask is not None else None
         
-        # Extract features using torchaudio model
         if output_hidden_states:
-            # Get all layer outputs
-            features, lengths_out = self.wav2vec2.extract_features(
-                input_values, 
-                lengths=lengths
-            )
-            # features is a list of tensors, one per layer
-            hidden_states = tuple(features)
-            last_hidden_state = features[-1]
-        else:
-            # Just get final output
-            outputs, lengths_out = self.wav2vec2(input_values, lengths=lengths)
-            last_hidden_state = outputs
-            hidden_states = None
+            features, _ = self.wav2vec2.extract_features(input_values, lengths=lengths)
+            return BaseModelOutput(last_hidden_state=features[-1], hidden_states=tuple(features))
         
-        # Get convolutional features (before transformer)
-        extract_features = self.wav2vec2.feature_extractor(input_values, lengths)[0]
+        output, _ = self.wav2vec2(input_values, lengths=lengths)
         
         if not return_dict:
-            outputs = (last_hidden_state,)
-            if output_hidden_states:
-                outputs = outputs + (hidden_states,)
-            return outputs
-        
-        return DELULUOutput(
-            last_hidden_state=last_hidden_state,
-            hidden_states=hidden_states,
-            attentions=None,  # torchaudio doesn't expose attention weights
-            extract_features=extract_features,
-        )
+            return (output,)
+        return BaseModelOutput(last_hidden_state=output)
     
-    def extract_features(
-        self, 
-        input_values: torch.Tensor,
-        lengths: Optional[torch.Tensor] = None
-    ) -> Tuple[torch.Tensor, ...]:
-        """
-        Extract features from all layers.
-        
-        Args:
-            input_values: Audio waveform of shape (batch, samples)
-            lengths: Optional lengths for each sample in batch
-            
-        Returns:
-            Tuple of tensors, one per layer (including CNN output)
-        """
+    def extract_features(self, input_values: torch.Tensor):
         if input_values.dim() == 1:
             input_values = input_values.unsqueeze(0)
-        
-        features, _ = self.wav2vec2.extract_features(input_values, lengths=lengths)
+        features, _ = self.wav2vec2.extract_features(input_values)
         return tuple(features)
-    
-    def get_speaker_embedding(
-        self,
-        input_values: torch.Tensor,
-        pooling: str = "mean"
-    ) -> torch.Tensor:
-        """
-        Extract speaker embedding from audio.
-        
-        Args:
-            input_values: Audio waveform of shape (batch, samples)
-            pooling: Pooling method - "mean", "max", or "first"
-            
-        Returns:
-            Speaker embedding of shape (batch, hidden_size)
-        """
-        outputs = self.forward(input_values, return_dict=True)
-        features = outputs.last_hidden_state
-        
-        if pooling == "mean":
-            return features.mean(dim=1)
-        elif pooling == "max":
-            return features.max(dim=1).values
-        elif pooling == "first":
-            return features[:, 0, :]
-        else:
-            raise ValueError(f"Unknown pooling method: {pooling}")
-    
-    def _init_weights(self, module):
-        """Initialize weights - mostly handled by torchaudio."""
-        pass
-
-
-class DELULUForSequenceClassification(PreTrainedModel):
-    """
-    DELULU with a classification head for speaker verification and other tasks.
-    
-    Example:
-        ```python
-        from transformers import AutoModel
-        
-        model = AutoModel.from_pretrained(
-            "cmu-mlsp/DELULU", 
-            trust_remote_code=True,
-            num_labels=1251  # Number of speakers in VoxCeleb2
-        )
-        ```
-    """
-    
-    config_class = DELULUConfig
-    base_model_prefix = "delulu"
-    
-    def __init__(self, config: DELULUConfig):
-        super().__init__(config)
-        
-        self.delulu = DELULUModel(config)
-        self.projector = nn.Linear(config.hidden_size, config.hidden_size)
-        
-        num_labels = getattr(config, 'num_labels', None)
-        if num_labels:
-            self.classifier = nn.Linear(config.hidden_size, num_labels)
-        else:
-            self.classifier = None
-        
-        self.post_init()
-    
-    def forward(
-        self,
-        input_values: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        labels: Optional[torch.Tensor] = None,
-        return_dict: Optional[bool] = None,
-    ):
-        return_dict = return_dict if return_dict is not None else True
-        
-        outputs = self.delulu(
-            input_values,
-            attention_mask=attention_mask,
-            return_dict=True
-        )
-        
-        # Pool features
-        hidden_states = outputs.last_hidden_state
-        pooled = hidden_states.mean(dim=1)
-        
-        # Project
-        embeddings = self.projector(pooled)
-        
-        # Classify if head exists
-        logits = None
-        if self.classifier is not None:
-            logits = self.classifier(embeddings)
-        
-        loss = None
-        if labels is not None and logits is not None:
-            loss_fct = nn.CrossEntropyLoss()
-            loss = loss_fct(logits, labels)
-        
-        if not return_dict:
-            output = (logits, embeddings) + (outputs.last_hidden_state,)
-            return ((loss,) + output) if loss is not None else output
-        
-        return {
-            "loss": loss,
-            "logits": logits,
-            "embeddings": embeddings,
-            "last_hidden_state": outputs.last_hidden_state,
-        }
-
-
-# Register for auto classes
-DELULUConfig.register_for_auto_class()
-DELULUModel.register_for_auto_class("AutoModel")