models/CLAP/open_clip/model.py

""" CLAP Model

Adapted from CLIP: https://github.com/openai/CLIP. Originally MIT License, Copyright (c) 2021 OpenAI.
Adapted to the Audio Task.
"""

from dataclasses import dataclass

import numpy as np
import torch
import torch.nn.functional as F
from torch import nn

import logging
from .htsat import create_htsat_model
from transformers import RobertaModel, RobertaConfig
# removed: OrderedDict, email.mime, typing, freeze_batch_norm_2d, BertModel, BartModel, BatchEncoding — dead after removing visual/transformer/bert/bart branches


class MLPLayers(nn.Module):
    def __init__(self, units=[512, 512, 512], nonlin=nn.ReLU(), dropout=0.1):
        super(MLPLayers, self).__init__()
        self.nonlin = nonlin
        self.dropout = dropout

        sequence = []
        for u0, u1 in zip(units[:-1], units[1:]):
            sequence.append(nn.Linear(u0, u1))
            sequence.append(self.nonlin)
            sequence.append(nn.Dropout(self.dropout))
        sequence = sequence[:-2]

        self.sequential = nn.Sequential(*sequence)

    def forward(self, X):
        X = self.sequential(X)
        return X


# removed: Bottleneck, AttentionPool2d, ModifiedResNet, LayerNorm, QuickGELU, ResidualAttentionBlock, Transformer, VisualTransformer
# — all CLIP visual encoder / transformer text encoder classes, not used with HTSAT+roberta inference




# Audio Config Class
@dataclass
class CLAPAudioCfp:
    model_type: str = "PANN"
    model_name: str = "Cnn14"
    sample_rate: int = 48000
    # Param
    audio_length: int = 1024
    window_size: int = 1024
    hop_size: int = 1024
    fmin: int = 50
    fmax: int = 14000
    class_num: int = 527
    mel_bins: int = 64
    clip_samples: int = 480000


@dataclass
class CLAPTextCfg:
    context_length: int
    vocab_size: int
    width: int
    heads: int
    layers: int
    model_type: str


class CLAP(nn.Module):
    def __init__(
        self,
        embed_dim: int,
        audio_cfg: CLAPAudioCfp,
        text_cfg: CLAPTextCfg,
        quick_gelu: bool = False,
        enable_fusion: bool = False,
        fusion_type: str = "None",
        joint_embed_shape: int = 512,
        mlp_act: str = "relu",
    ):
        super().__init__()
        if isinstance(audio_cfg, dict):
            audio_cfg = CLAPAudioCfp(**audio_cfg)
        if isinstance(text_cfg, dict):
            text_cfg = CLAPTextCfg(**text_cfg)

        self.audio_cfg = audio_cfg
        self.text_cfg = text_cfg
        self.enable_fusion = enable_fusion
        self.fusion_type = fusion_type
        self.joint_embed_shape = joint_embed_shape
        self.mlp_act = mlp_act

        self.context_length = text_cfg.context_length

        # removed: act_layer/QuickGELU — only used by transformer text branch
        if mlp_act == "relu":
            mlp_act_layer = nn.ReLU()
        elif mlp_act == "gelu":
            mlp_act_layer = nn.GELU()
        else:
            raise NotImplementedError

        # audio branch — removed: PANN branch, only HTSAT used
        self.audio_branch = create_htsat_model(audio_cfg, enable_fusion, fusion_type)

        # text branch — removed: transformer/bert/bart branches, only roberta used
        if text_cfg.model_type == "roberta":
            self.text_branch = RobertaModel.from_pretrained("roberta-base")

            self.text_transform = MLPLayers(
                units=[
                    self.joint_embed_shape,
                    self.joint_embed_shape,
                    self.joint_embed_shape,
                ],
                dropout=0.1,
            )
            self.text_projection = nn.Sequential(
                nn.Linear(768, self.joint_embed_shape),
                mlp_act_layer,
                nn.Linear(self.joint_embed_shape, self.joint_embed_shape),
            )
        # removed: bart branch — only roberta used
        else:
            logging.error(f"Model config for {text_cfg.model_type} not found")
            raise RuntimeError(f"Model config for {text_cfg.model_type} not found.")
        self.text_branch_type = text_cfg.model_type
        # text branch parameters

        # audio branch parameters
        self.audio_transform = MLPLayers(
            units=[
                self.joint_embed_shape,
                self.joint_embed_shape,
                self.joint_embed_shape,
            ],
            dropout=0.1,
        )

        # below here is text branch parameters

        # ============================================================================================================
        self.audio_projection = nn.Sequential(
            nn.Linear(embed_dim, self.joint_embed_shape),
            mlp_act_layer,
            nn.Linear(self.joint_embed_shape, self.joint_embed_shape),
        )

        self.logit_scale_a = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
        self.logit_scale_t = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
        # removed: attn_mask buffer — only used by transformer text branch

        self.init_text_branch_parameters()

    def init_text_branch_parameters(self):
        # removed: transformer/bert/bart init branches
        nn.init.constant_(self.logit_scale_a, np.log(1 / 0.07))
        nn.init.constant_(self.logit_scale_t, np.log(1 / 0.07))

    # removed: build_attention_mask — only used by transformer text branch

    # removed: encode_audio — audio modality not used in inference

    def encode_text(self, text, device):
        # removed: transformer/bert/bart branches — only roberta used
        if self.text_branch_type == "roberta":
            x = self.text_branch(
                input_ids=text["input_ids"].to(device=device, non_blocking=True),
                attention_mask=text["attention_mask"].to(
                    device=device, non_blocking=True
                ),
            )["pooler_output"]
            x = self.text_projection(x)
        else:
            raise RuntimeError(f"Model type {self.text_branch_type} not found.")
        return x

    # removed: forward, get_logit_scale — training/contrastive loss helpers, not used in inference

    def get_text_embedding(self, data):
        """Get the text embedding from the model

        Parameters
        ----------
        data: torch.Tensor
            a tensor of text embedding

        Returns
        ----------
        text_embed: torch.Tensor
            a tensor of text_embeds (N, D)

        """
        device = next(self.parameters()).device
        for k in data:
            data[k] = data[k].to(device)
        text_embeds = self.encode_text(data, device=device)
        text_embeds = F.normalize(text_embeds, dim=-1)

        return text_embeds

    # removed: get_audio_embedding, audio_infer — audio modality not used in inference


def convert_weights_to_fp16(model: nn.Module):
    """Convert applicable model parameters to fp16"""

    def _convert_weights_to_fp16(l):
        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
            l.weight.data = l.weight.data.half()
            if l.bias is not None:
                l.bias.data = l.bias.data.half()

        if isinstance(l, nn.MultiheadAttention):
            for attr in [
                *[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]],
                "in_proj_bias",
                "bias_k",
                "bias_v",
            ]:
                tensor = getattr(l, attr)
                if tensor is not None:
                    tensor.data = tensor.data.half()

        for name in ["text_projection", "proj"]:
            if hasattr(l, name):
                attr = getattr(l, name)
                if attr is not None:
                    attr.data = attr.data.half()

    model.apply(_convert_weights_to_fp16)


# Ignore the state dict of the vision part
def build_model_from_openai_state_dict(
    state_dict: dict, model_cfg, enable_fusion: bool = False, fusion_type: str = "None"
):

    embed_dim = model_cfg["embed_dim"]
    audio_cfg = model_cfg["audio_cfg"]
    text_cfg = model_cfg["text_cfg"]
    context_length = state_dict["positional_embedding"].shape[0]
    vocab_size = state_dict["token_embedding.weight"].shape[0]
    transformer_width = state_dict["ln_final.weight"].shape[0]
    transformer_heads = transformer_width // 64
    transformer_layers = len(
        set(
            k.split(".")[2]
            for k in state_dict
            if k.startswith(f"transformer.resblocks")
        )
    )

    audio_cfg = CLAPAudioCfp(**audio_cfg)
    text_cfg = CLAPTextCfg(**text_cfg)

    model = CLAP(
        embed_dim,
        audio_cfg=audio_cfg,
        text_cfg=text_cfg,
        quick_gelu=True,  # OpenAI models were trained with QuickGELU
        enable_fusion=enable_fusion,
        fusion_type=fusion_type,
    )
    state_dict["logit_scale_a"] = state_dict["logit_scale"]
    state_dict["logit_scale_t"] = state_dict["logit_scale"]
    pop_keys = list(state_dict.keys())[::]
    # pop the visual branch saved weights
    for key in pop_keys:
        if key.startswith("visual."):
            state_dict.pop(key, None)

    for key in ["logit_scale", "input_resolution", "context_length", "vocab_size"]:
        state_dict.pop(key, None)

    # not use fp16
    # convert_weights_to_fp16(model)
    model.load_state_dict(state_dict, strict=False)
    return model.eval()


# removed: trace_model — JIT tracing utility, not used in inference