Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,083 Bytes
6efc863 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
import numpy as np
import torch
import torch.nn.functional as F
from torch import nn
from transformers import AutoModel
from .audio import get_audio_encoder
class Projection(nn.Module):
def __init__(self, d_in: int, d_out: int, p: float=0.5) -> None:
super().__init__()
self.linear1 = nn.Linear(d_in, d_out, bias=False)
self.linear2 = nn.Linear(d_out, d_out, bias=False)
self.layer_norm = nn.LayerNorm(d_out)
self.drop = nn.Dropout(p)
def forward(self, x: torch.Tensor) -> torch.Tensor:
embed1 = self.linear1(x)
embed2 = self.drop(self.linear2(F.gelu(embed1)))
embeds = self.layer_norm(embed1 + embed2)
return embeds
class AudioEncoder(nn.Module):
def __init__(self, audioenc_name:str, d_in: int, d_out: int, sample_rate: int, window_size: int,
hop_size: int, mel_bins: int, fmin: int, fmax: int, classes_num: int) -> None:
super().__init__()
audio_encoder = get_audio_encoder(audioenc_name)
self.base = audio_encoder(
sample_rate, window_size,
hop_size, mel_bins, fmin, fmax,
classes_num, d_in)
self.projection = Projection(d_in, d_out)
def forward(self, x):
out_dict = self.base(x)
audio_features, audio_classification_output = out_dict['embedding'], out_dict['clipwise_output']
projected_vec = self.projection(audio_features)
return projected_vec, audio_classification_output
class TextEncoder(nn.Module):
def __init__(self, d_out: int, text_model: str, transformer_embed_dim: int) -> None:
super().__init__()
self.base = AutoModel.from_pretrained(text_model)
self.projection = Projection(transformer_embed_dim, d_out)
def forward(self, x):
out = self.base(**x)[0]
out = out[:, 0, :] # get CLS token output
projected_vec = self.projection(out)
return projected_vec
class CLAP(nn.Module):
def __init__(self,
# audio
audioenc_name: str,
sample_rate: int,
window_size: int,
hop_size: int,
mel_bins: int,
fmin: int,
fmax: int,
classes_num: int,
out_emb: int,
# text
text_model: str,
transformer_embed_dim: int,
# common
d_proj: int,
):
super().__init__()
self.audio_encoder = AudioEncoder(
audioenc_name, out_emb, d_proj,
sample_rate, window_size, hop_size, mel_bins, fmin, fmax, classes_num)
self.caption_encoder = TextEncoder(
d_proj, text_model, transformer_embed_dim
)
self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
def forward(self, audio, text):
audio_embed, _ = self.audio_encoder(audio)
caption_embed = self.caption_encoder(text)
return caption_embed, audio_embed, self.logit_scale.exp() |