""" 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