Add CM3P model

configuration_cm3p.py

@@ -13,7 +13,7 @@ class CM3PMetadataConfig(PretrainedConfig):
 
     def __init__(
         self,
-        cls_embed=False,
+        cls_embed=True,
 
         projection_dim=512,
         initializer_factor=1.0,
@@ -177,6 +177,7 @@ class CM3PAudioConfig(PretrainedConfig):
 
 class CM3PBeatmapConfig(PretrainedConfig):
     model_type = "CM3PBeatmap"
+    is_composition = True
     base_config_key = "beatmap_config"
     sub_configs = {"audio_config": CM3PAudioConfig}
 
@@ -186,7 +187,7 @@ class CM3PBeatmapConfig(PretrainedConfig):
         audio_sos_token_id=3164,
         audio_eos_token_id=3165,
         audio_token_id=3166,
-        cls_embed=False,
+        cls_embed=True,
 
         projection_dim=512,
         initializer_factor=1.0,
@@ -222,12 +223,15 @@ class CM3PBeatmapConfig(PretrainedConfig):
         sparse_pred_ignore_index=-100,
         reference_compile=None,
         repad_logits_with_grad=False,
+
+        attn_implementation: str = None,
         **kwargs,
     ):
         super().__init__(
             pad_token_id=pad_token_id,
             bos_token_id=bos_token_id,
             eos_token_id=eos_token_id,
+            attn_implementation=attn_implementation,
             **kwargs,
         )
 
@@ -235,7 +239,11 @@ class CM3PBeatmapConfig(PretrainedConfig):
             audio_config = {}
             logger.info("`audio_config` is `None`. Initializing the `CM3PAudioConfig` with default values.")
 
-        self.audio_config = CM3PAudioConfig(**audio_config)
+        self.audio_config = CM3PAudioConfig(
+            attn_implementation=attn_implementation,
+            **audio_config
+        )
+
         self.audio_sos_token_id = audio_sos_token_id
         self.audio_eos_token_id = audio_eos_token_id
         self.audio_token_id = audio_token_id
@@ -280,6 +288,7 @@ class CM3PBeatmapConfig(PretrainedConfig):
 
 class CM3PConfig(PretrainedConfig):
     model_type = "CM3P"
+    is_composition = True
     sub_configs = {"metadata_config": CM3PMetadataConfig, "beatmap_config": CM3PBeatmapConfig}
 
     def __init__(
@@ -291,9 +300,15 @@ class CM3PConfig(PretrainedConfig):
         initializer_factor=1.0,
         initializer_range=0.02,
         loss_type=None,
+        has_decoder_head=False,
+
+        attn_implementation: str = None,
         **kwargs
     ):
-        super().__init__(**kwargs)
+        super().__init__(
+            attn_implementation=attn_implementation,
+            **kwargs
+        )
 
         if metadata_config is None:
             metadata_config = {}
@@ -303,14 +318,21 @@ class CM3PConfig(PretrainedConfig):
             beatmap_config = {}
             logger.debug("`beatmap_config` is `None`. initializing the `CM3PBeatmapConfig` with default values.")
 
-        self.metadata_config = CM3PMetadataConfig(**metadata_config)
-        self.beatmap_config = CM3PBeatmapConfig(**beatmap_config)
+        self.metadata_config = CM3PMetadataConfig(
+            attn_implementation=attn_implementation,
+            **metadata_config
+        )
+        self.beatmap_config = CM3PBeatmapConfig(
+            attn_implementation=attn_implementation,
+            **beatmap_config
+        )
 
         self.projection_dim = projection_dim
         self.logit_scale_init_value = logit_scale_init_value
         self.initializer_factor = initializer_factor
         self.initializer_range = initializer_range
         self.loss_type = loss_type
+        self.has_decoder_head = has_decoder_head
 
 
 AutoConfig.register("CM3PMetadata", CM3PMetadataConfig)

modeling_cm3p.py

@@ -24,7 +24,7 @@ logger = logging.get_logger(__name__)
 
 # contrastive loss function, adapted from
 # https://sachinruk.github.io/blog/2021-03-07-clip.html
-def contrastive_loss(logits: torch.Tensor, target: torch.LongTensor = None) -> torch.Tensor:
+def contrastive_loss(logits: torch.Tensor, target: torch.Tensor = None) -> torch.Tensor:
     target = target if target is not None else torch.arange(len(logits), device=logits.device)
     return nn.functional.cross_entropy(logits, target)
 
@@ -192,7 +192,7 @@ class CM3PBeatmapModelOutput(BaseModelOutputWithPooling):
     """
 
     beatmap_embeds: Optional[torch.FloatTensor] = None
-    audio_model_output: CM3PAudioModelOutput = None
+    audio_model_output: Optional[CM3PAudioModelOutput] = None
 
 
 @dataclass
@@ -235,8 +235,8 @@ class CM3POutput(ModelOutput):
     """
 
     loss: Optional[torch.FloatTensor] = None
-    logits_per_beatmap: Optional[torch.FloatTensor] = None
-    logits_per_metadata: Optional[torch.FloatTensor] = None
+    logits_per_beatmap: Optional[torch.Tensor] = None
+    logits_per_metadata: Optional[torch.Tensor] = None
     metadata_embeds: Optional[torch.FloatTensor] = None
     beatmap_embeds: Optional[torch.FloatTensor] = None
     logits: Optional[torch.FloatTensor] = None
@@ -301,6 +301,7 @@ class CM3PMetadataTransformer(nn.Module):
     def __init__(self, config: CM3PMetadataConfig):
         super().__init__()
         self.config = config
+        # noinspection PyTypeChecker
         self.encoder = ModernBertModel(config)
 
     def get_input_embeddings(self):
@@ -486,6 +487,7 @@ class CM3PAudioEncoder(nn.Module):
         self.config = config
         self.conv1 = nn.Conv1d(config.n_mels, config.hidden_size, kernel_size=3, padding=1)
         self.conv2 = nn.Conv1d(config.hidden_size, config.hidden_size, kernel_size=3, stride=2, padding=1)
+        # noinspection PyTypeChecker
         self.encoder = ModernBertModel(config)
         self.multi_modal_projector = CM3PMultiModalProjector(config)
 
@@ -531,6 +533,7 @@ class CM3PBeatmapTransformer(nn.Module):
         super().__init__()
         self.config = config
         self.audio_encoder = CM3PAudioEncoder(config.audio_config)
+        # noinspection PyTypeChecker
         self.encoder = ModernBertModel(config)
 
     def get_input_embeddings(self):
@@ -590,7 +593,7 @@ class CM3PBeatmapTransformer(nn.Module):
 
         audio_model_outputs = None
         if input_features is not None:
-            audio_model_outputs: CM3PAudioModelOutput = self.audio_encoder(
+            audio_model_outputs = self.audio_encoder(
                 input_features=input_features,
                 output_attentions=output_attentions,
                 output_hidden_states=output_hidden_states,
@@ -744,9 +747,9 @@ class CM3PModel(CM3PPreTrainedModel):
         metadata_config = config.metadata_config
         beatmap_config = config.beatmap_config
 
-        self.projection_dim = config.projection_dim
-        self.metadata_embed_dim = metadata_config.hidden_size
-        self.beatmap_embed_dim = beatmap_config.hidden_size
+        self.projection_dim: int = config.projection_dim
+        self.metadata_embed_dim: int = metadata_config.hidden_size
+        self.beatmap_embed_dim: int = beatmap_config.hidden_size
         self.loss_type = config.loss_type
 
         metadata_model = CM3PMetadataModel._from_config(metadata_config)
@@ -759,8 +762,9 @@ class CM3PModel(CM3PPreTrainedModel):
         self.metadata_projection = nn.Linear(self.metadata_embed_dim, self.projection_dim, bias=False)
         self.logit_scale = nn.Parameter(torch.tensor(self.config.logit_scale_init_value))
 
-        self.head = CM3PPredictionHead(beatmap_config)
-        self.decoder = nn.Linear(beatmap_config.hidden_size, beatmap_config.vocab_size, bias=beatmap_config.decoder_bias)
+        if config.has_decoder_head:
+            self.head = CM3PPredictionHead(beatmap_config)
+            self.decoder = nn.Linear(beatmap_config.hidden_size, beatmap_config.vocab_size, bias=beatmap_config.decoder_bias)
 
         # Initialize weights and apply final processing
         self.post_init()
@@ -861,6 +865,7 @@ class CM3PModel(CM3PPreTrainedModel):
         return_loss: Optional[bool] = True,
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
+        output_logits: Optional[bool] = None,
         **kwargs,
     ) -> CM3POutput:
         r"""
@@ -886,16 +891,22 @@ class CM3PModel(CM3PPreTrainedModel):
             Sequence length of the input sequences including padding tokens. Used to pad the output tensors.
         return_loss (`bool`, *optional*):
             Whether to return the contrastive loss.
+        output_logits (`bool`, *optional*):
+            Whether to return the logits from the decoder head.
         """
         # Use CM3P model's config for some fields (if specified) instead of those of beatmap & metadata components.
         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
         )
+        output_logits = output_logits if output_logits is not None else self.config.has_decoder_head
 
-        if metadata_ids.dim() == 3 and return_loss and metadata_variation_classes is None:
+        if metadata_ids is not None and metadata_ids.dim() == 3 and return_loss and metadata_variation_classes is None:
             raise ValueError("When providing multiple metadata variations, metadata_variation_classes must be provided in order to compute loss correctly.")
 
+        if output_logits and not self.config.has_decoder_head:
+            raise ValueError("Cannot return logits when the model is not configured with a decoder head.")
+
         # noinspection PyProtectedMember
         if self.config._attn_implementation == "flash_attention_2":
             if indices is None and cu_seqlens is None and max_seqlen is None:
@@ -919,65 +930,75 @@ class CM3PModel(CM3PPreTrainedModel):
                         inputs=inputs_embeds, attention_mask=attention_mask, position_ids=position_ids, labels=labels
                     )
 
-        beatmap_outputs: BaseModelOutputWithPooling = self.beatmap_model(
-            input_ids=input_ids,
-            input_features=input_features,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            inputs_embeds=inputs_embeds,
-            indices=indices,
-            cu_seqlens=cu_seqlens,
-            max_seqlen=max_seqlen,
-            batch_size=batch_size,
-            seq_len=seq_len,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-        )
-
-        metadata_outputs: BaseModelOutputWithPooling = self.metadata_model(
-            input_ids=metadata_ids,
-            attention_mask=metadata_attention_mask,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-        )
+        beatmap_embeds = None
+        beatmap_outputs = None
+        metadata_embeds = None
+        metadata_outputs = None
+        logits_per_beatmap = None
+        logits_per_metadata = None
+        loss = 0 if return_loss else None
+        logits = None
+
+        if input_ids is not None:
+            beatmap_outputs: BaseModelOutputWithPooling = self.beatmap_model(
+                input_ids=input_ids,
+                input_features=input_features,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                inputs_embeds=inputs_embeds,
+                indices=indices,
+                cu_seqlens=cu_seqlens,
+                max_seqlen=max_seqlen,
+                batch_size=batch_size,
+                seq_len=seq_len,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+            )
 
-        beatmap_embeds = beatmap_outputs.pooler_output
-        beatmap_embeds = self.beatmap_projection(beatmap_embeds)
+            beatmap_embeds = beatmap_outputs.pooler_output
+            beatmap_embeds = self.beatmap_projection(beatmap_embeds)
+            beatmap_embeds = beatmap_embeds / _get_vector_norm(beatmap_embeds)
 
-        metadata_embeds = metadata_outputs.pooler_output
-        metadata_embeds = self.metadata_projection(metadata_embeds)
+        if metadata_ids is not None:
+            metadata_outputs: BaseModelOutputWithPooling = self.metadata_model(
+                input_ids=metadata_ids,
+                attention_mask=metadata_attention_mask,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+            )
 
-        # normalized features
-        beatmap_embeds = beatmap_embeds / _get_vector_norm(beatmap_embeds)
-        metadata_embeds = metadata_embeds / _get_vector_norm(metadata_embeds)
+            metadata_embeds = metadata_outputs.pooler_output
+            metadata_embeds = self.metadata_projection(metadata_embeds)
+            metadata_embeds = metadata_embeds / _get_vector_norm(metadata_embeds)
 
-        # cosine similarity as logits
-        logits_per_metadata = torch.matmul(metadata_embeds, beatmap_embeds.t().to(metadata_embeds.device))
-        logits_per_metadata = logits_per_metadata * self.logit_scale.exp().to(metadata_embeds.device)
+        if metadata_embeds is not None and beatmap_embeds is not None:
+            # cosine similarity as logits
+            logits_per_metadata = torch.matmul(metadata_embeds, beatmap_embeds.t().to(metadata_embeds.device))
+            logits_per_metadata = logits_per_metadata * self.logit_scale.exp().to(metadata_embeds.device)
 
-        if logits_per_metadata.dim() == 3:
-            logits_per_beatmap = logits_per_metadata.permute(2, 0, 1)
-        else:
-            logits_per_beatmap = logits_per_metadata.t()
+            if logits_per_metadata.dim() == 3:
+                logits_per_beatmap = logits_per_metadata.permute(2, 0, 1)
+            else:
+                logits_per_beatmap = logits_per_metadata.t()
 
-        loss = None
-        if return_loss:
-            loss = cm3p_loss(logits_per_metadata, metadata_variation_classes)
+            if return_loss:
+                loss = cm3p_loss(logits_per_metadata, metadata_variation_classes)
 
-        logits = (
-            self.compiled_head(beatmap_outputs.last_hidden_state)
-            if self.config.beatmap_config.reference_compile
-            else self.decoder(self.head(beatmap_outputs.last_hidden_state))
-        )
+        if output_logits:
+            logits = (
+                self.compiled_head(beatmap_outputs.last_hidden_state)
+                if self.config.beatmap_config.reference_compile
+                else self.decoder(self.head(beatmap_outputs.last_hidden_state))
+            )
 
-        if labels is not None and return_loss:
-            mlm_loss = self.loss_function(logits, labels, vocab_size=self.config.beatmap_config.vocab_size, **kwargs)
-            loss += 0.5 * mlm_loss
+            if labels is not None and return_loss:
+                mlm_loss = self.loss_function(logits, labels, vocab_size=self.config.beatmap_config.vocab_size, **kwargs)
+                loss += 0.5 * mlm_loss
 
-        # noinspection PyProtectedMember
-        if self.config._attn_implementation == "flash_attention_2":
-            with nullcontext() if self.config.beatmap_config.repad_logits_with_grad or labels is None else torch.no_grad():
-                logits = _pad_cm3p_output(inputs=logits, indices=indices, batch=batch_size, seqlen=seq_len)
+            # noinspection PyProtectedMember
+            if self.config._attn_implementation == "flash_attention_2":
+                with nullcontext() if self.config.beatmap_config.repad_logits_with_grad or labels is None else torch.no_grad():
+                    logits = _pad_cm3p_output(inputs=logits, indices=indices, batch=batch_size, seqlen=seq_len)
 
         return CM3POutput(
             loss=loss,
@@ -1372,4 +1393,5 @@ __all__ = [
     "CM3PBeatmapModel",
     "CM3PBeatmapModelWithProjection",
     "CM3PForBeatmapClassification",
+    "CM3PForMaskedLM",
 ]