First pass at captioning functionality through web app

.gitignore

@@ -0,0 +1 @@
+*.pyc

app.py

@@ -1,9 +1,10 @@
 import gradio as gr
+from infer import infer
 
 def greet(name):
     return f'Hello {name}!!'
 
-demo = gr.Interface(fn=greet,
+demo = gr.Interface(fn=infer,
                     inputs=gr.Audio(sources='upload', type='filepath'),
                     outputs='text')
 demo.launch()

audiocaptioner.py

@@ -0,0 +1,68 @@
+from lib import *
+
+import contextlib
+import io
+import laion_clap
+import torch
+
+class AudioCaptioner(torch.nn.Module):
+
+    def get_dummy_token(self, batch_size: int) -> torch.Tensor:
+        return torch.zeros(batch_size, self.prefix_length, dtype=torch.int64)
+
+    def embed_waveform(self, waveform):
+        # compute the prefix
+        input_dict = {
+            'waveform': waveform # you can add more key-values
+        }
+        audio_embeds = self.clap_model.model.encode_audio(
+            input_dict,
+            device=waveform.device
+        )
+
+        # get BxD-dim embedding (last layer) D = 1024 -> 512 after audio projection
+        audio_embedding = torch.nn.functional.normalize(self.clap_model.model.audio_projection(audio_embeds['embedding']), dim=-1)
+        return audio_embedding
+
+    def create_prefix(self, waveform, batch_size):
+        if waveform is not None:
+            audio_embedding = self.embed_waveform(waveform)
+        else:
+            audio_embedding = torch.zeros(batch_size, self.prefix_size).cuda()
+        # project the prefix through map net and append it
+        prefix_projections = self.clip_project(audio_embedding).view(-1, self.prefix_length, self.gpt_embedding_size)
+        return prefix_projections
+
+    def forward(self, tokens: torch.Tensor, waveform: torch.Tensor, mask: Optional[torch.Tensor] = None,
+                labels: Optional[torch.Tensor] = None, freeze_gpt = False):
+        # embed the text
+        embedding_text = self.gpt.transformer.wte(tokens)
+        prefix_projections = self.create_prefix(waveform, tokens.shape[0])
+        embedding_text = torch.cat((prefix_projections, embedding_text), dim=1)
+        # offset labels
+        if labels is not None:
+            dummy_token = self.get_dummy_token(tokens.shape[0], tokens.device)
+            labels = torch.cat((dummy_token, tokens), dim=1)
+        # push through GPT
+        if freeze_gpt:
+            with torch.no_grad():
+                out = self.gpt(inputs_embeds=embedding_text, labels=labels, attention_mask=mask)
+        else:
+            out = self.gpt(inputs_embeds=embedding_text, labels=labels, attention_mask=mask)
+        return out
+
+    def __init__(self, prefix_length: int, clip_length: Optional[int] = None, prefix_size: int = 512,
+                 num_layers: int = 8):
+        super(AudioCaptioner, self).__init__()
+        self.prefix_size = prefix_size
+        self.prefix_length = prefix_length
+        self.gpt = GPT2LMHeadModel.from_pretrained('gpt2')
+        self.gpt_embedding_size = self.gpt.transformer.wte.weight.shape[1]
+        self.clip_project = MLP((prefix_size, (self.gpt_embedding_size * prefix_length) // 2,
+                                 self.gpt_embedding_size * prefix_length))
+        self.clap_model = laion_clap.CLAP_Module(
+            enable_fusion=False,
+            amodel = 'HTSAT-base'
+        )
+        with contextlib.redirect_stdout(io.StringIO()):
+            self.clap_model.load_ckpt(ckpt = '/graft1/datasets/kechen/clap_ckpt/music_audioset_epoch_15_esc_90.14.pt')

clipcap.py

@@ -0,0 +1,405 @@
+#####################################################################
+### Credit: Ron Mokady / rmokady                                  ###
+### Original Repo: https://github.com/rmokady/CLIP_prefix_caption ###
+#####################################################################
+
+from enum import Enum
+from collections import defaultdict
+import os
+from torch import nn
+import numpy as np
+import torch
+import torch.nn.functional as nnf
+import sys
+from typing import Tuple, List, Union, Optional
+from transformers import (
+    GPT2Tokenizer,
+    GPT2LMHeadModel,
+    AdamW,
+    get_linear_schedule_with_warmup,
+)
+
+# import torch
+
+N = type(None)
+V = np.array
+ARRAY = np.ndarray
+ARRAYS = Union[Tuple[ARRAY, ...], List[ARRAY]]
+VS = Union[Tuple[V, ...], List[V]]
+VN = Union[V, N]
+VNS = Union[VS, N]
+T = torch.Tensor
+TS = Union[Tuple[T, ...], List[T]]
+TN = Optional[T]
+TNS = Union[Tuple[TN, ...], List[TN]]
+TSN = Optional[TS]
+TA = Union[T, ARRAY]
+
+WEIGHTS_PATHS = {
+    "coco": "coco_weights.pt",
+    "conceptual-captions": "conceptual_weights.pt",
+}
+
+class MappingType(Enum):
+    MLP = 'mlp'
+    Transformer = 'transformer'
+
+class MLP(nn.Module):
+    def forward(self, x: T) -> T:
+        return self.model(x)
+
+    def __init__(self, sizes: Tuple[int, ...], bias=True, act=nn.Tanh):
+        super(MLP, self).__init__()
+        layers = []
+        for i in range(len(sizes) - 1):
+            layers.append(nn.Linear(sizes[i], sizes[i + 1], bias=bias))
+            if i < len(sizes) - 2:
+                layers.append(act())
+        self.model = nn.Sequential(*layers)
+
+class MlpTransformer(nn.Module):
+    def __init__(self, in_dim, h_dim, out_d: Optional[int] = None, act=nnf.relu, dropout=0.):
+        super().__init__()
+        out_d = out_d if out_d is not None else in_dim
+        self.fc1 = nn.Linear(in_dim, h_dim)
+        self.act = act
+        self.fc2 = nn.Linear(h_dim, out_d)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.dropout(x)
+        x = self.fc2(x)
+        x = self.dropout(x)
+        return x
+
+class MultiHeadAttention(nn.Module):
+
+    def __init__(self, dim_self, dim_ref, num_heads, bias=True, dropout=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim_self // num_heads
+        self.scale = head_dim ** -0.5
+        self.to_queries = nn.Linear(dim_self, dim_self, bias=bias)
+        self.to_keys_values = nn.Linear(dim_ref, dim_self * 2, bias=bias)
+        self.project = nn.Linear(dim_self, dim_self)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, y=None, mask=None):
+        y = y if y is not None else x
+        b, n, c = x.shape
+        _, m, d = y.shape
+        # b n h dh
+        queries = self.to_queries(x).reshape(b, n, self.num_heads, c // self.num_heads)
+        # b m 2 h dh
+        keys_values = self.to_keys_values(y).reshape(b, m, 2, self.num_heads, c // self.num_heads)
+        keys, values = keys_values[:, :, 0], keys_values[:, :, 1]
+        attention = torch.einsum('bnhd,bmhd->bnmh', queries, keys) * self.scale
+        if mask is not None:
+            if mask.dim() == 2:
+                mask = mask.unsqueeze(1)
+            attention = attention.masked_fill(mask.unsqueeze(3), float("-inf"))
+        attention = attention.softmax(dim=2)
+        out = torch.einsum('bnmh,bmhd->bnhd', attention, values).reshape(b, n, c)
+        out = self.project(out)
+        return out, attention
+
+
+class TransformerLayer(nn.Module):
+
+    def forward_with_attention(self, x, y=None, mask=None):
+        x_, attention = self.attn(self.norm1(x), y, mask)
+        x = x + x_
+        x = x + self.mlp(self.norm2(x))
+        return x, attention
+
+    def forward(self, x, y=None, mask=None):
+        x = x + self.attn(self.norm1(x), y, mask)[0]
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+    def __init__(self, dim_self, dim_ref, num_heads, mlp_ratio=4., bias=False, dropout=0., act=nnf.relu,
+                 norm_layer: nn.Module = nn.LayerNorm):
+        super().__init__()
+        self.norm1 = norm_layer(dim_self)
+        self.attn = MultiHeadAttention(dim_self, dim_ref, num_heads, bias=bias, dropout=dropout)
+        self.norm2 = norm_layer(dim_self)
+        self.mlp = MlpTransformer(dim_self, int(dim_self * mlp_ratio), act=act, dropout=dropout)
+
+
+class Transformer(nn.Module):
+
+    def forward_with_attention(self, x, y=None, mask=None):
+        attentions = []
+        for layer in self.layers:
+            x, att = layer.forward_with_attention(x, y, mask)
+            attentions.append(att)
+        return x, attentions
+
+    def forward(self, x, y=None, mask=None):
+        for i, layer in enumerate(self.layers):
+            if i % 2 == 0 and self.enc_dec: # cross
+                x = layer(x, y)
+            elif self.enc_dec:  # self
+                x = layer(x, x, mask)
+            else:  # self or cross
+                x = layer(x, y, mask)
+        return x
+
+    def __init__(self, dim_self: int, num_heads: int, num_layers: int, dim_ref: Optional[int] = None,
+                 mlp_ratio: float = 2., act=nnf.relu, norm_layer: nn.Module = nn.LayerNorm, enc_dec: bool = False):
+        super(Transformer, self).__init__()
+        dim_ref = dim_ref if dim_ref is not None else dim_self
+        self.enc_dec = enc_dec
+        if enc_dec:
+            num_layers = num_layers * 2
+        layers = []
+        for i in range(num_layers):
+            if i % 2 == 0 and enc_dec:  # cross
+                layers.append(TransformerLayer(dim_self, dim_ref, num_heads, mlp_ratio, act=act, norm_layer=norm_layer))
+            elif enc_dec:  # self
+                layers.append(TransformerLayer(dim_self, dim_self, num_heads, mlp_ratio, act=act, norm_layer=norm_layer))
+            else:  # self or cross
+                layers.append(TransformerLayer(dim_self, dim_ref, num_heads, mlp_ratio, act=act, norm_layer=norm_layer))
+        self.layers = nn.ModuleList(layers)
+
+
+class TransformerMapper(nn.Module):
+
+    def forward(self, x):
+        x = self.linear(x).view(x.shape[0], self.clip_length, -1)
+        prefix = self.prefix_const.unsqueeze(0).expand(x.shape[0], *self.prefix_const.shape)
+        prefix = torch.cat((x, prefix), dim=1)
+        out = self.transformer(prefix)[:, self.clip_length:]
+        return out
+
+    def __init__(self, dim_clip: int, dim_embedding: int, prefix_length: int, clip_length: int, num_layers: int = 8):
+        super(TransformerMapper, self).__init__()
+        self.clip_length = clip_length
+        self.transformer = Transformer(dim_embedding, 8, num_layers)
+        self.linear = nn.Linear(dim_clip, clip_length * dim_embedding)
+        self.prefix_const = nn.Parameter(torch.randn(prefix_length, dim_embedding), requires_grad=True)
+
+
+class ClipCaptionModel(nn.Module):
+
+    def get_dummy_token(self, batch_size: int, device: torch.device) -> torch.Tensor:
+        return torch.zeros(batch_size, self.prefix_length, dtype=torch.int64, device=device)
+
+    def forward(self, tokens: torch.Tensor, prefix: torch.Tensor, mask: Optional[torch.Tensor] = None,
+                labels: Optional[torch.Tensor] = None):
+        embedding_text = self.gpt.transformer.wte(tokens)
+        if prefix is not None:
+            prefix_projections = self.clip_project(prefix).view(-1, self.prefix_length, self.gpt_embedding_size)
+            embedding_text = torch.cat((prefix_projections, embedding_text), dim=1)
+        if labels is not None:
+            dummy_token = self.get_dummy_token(tokens.shape[0], tokens.device)
+            labels = torch.cat((dummy_token, tokens), dim=1)
+        out = self.gpt(inputs_embeds=embedding_text, labels=labels, attention_mask=mask)
+        return out
+
+    def __init__(self, prefix_length: int, clip_length: Optional[int] = None, prefix_size: int = 512,
+                 num_layers: int = 8, mapping_type: MappingType = MappingType.MLP):
+        super(ClipCaptionModel, self).__init__()
+        self.prefix_size = prefix_size
+        self.prefix_length = prefix_length
+        self.gpt = GPT2LMHeadModel.from_pretrained('gpt2')
+        self.gpt_embedding_size = self.gpt.transformer.wte.weight.shape[1]
+        if mapping_type == MappingType.MLP:
+            self.clip_project = MLP((prefix_size, (self.gpt_embedding_size * prefix_length) // 2,
+                                     self.gpt_embedding_size * prefix_length))
+        else:
+            self.clip_project = TransformerMapper(prefix_size, self.gpt_embedding_size, prefix_length,
+                                                                     clip_length, num_layers)
+class ClipCaptionPrefix(ClipCaptionModel):
+    def parameters(self, recurse: bool = True):
+        return self.clip_project.parameters()
+
+    def train(self, mode: bool = True):
+        super(ClipCaptionPrefix, self).train(mode)
+        self.gpt.eval()
+        return self
+
+
+def generate_beam(
+    model,
+    tokenizer,
+    beam_size: int = 5,
+    prompt=None,
+    embed=None,
+    #entry_length=67,
+    entry_length=150,
+    #temperature=1.0,
+    temperature=0.7,
+    stop_token: str = ".",
+    no_repeat_ngram = 3,
+    #no_repeat_ngram = None,
+):
+
+    model.eval()
+    stop_token_index = tokenizer.encode(stop_token)[0]
+    tokens = None
+    scores = None
+    device = next(model.parameters()).device
+    seq_lengths = torch.ones(beam_size, device=device)
+    is_stopped = torch.zeros(beam_size, device=device, dtype=torch.bool)
+    filter_value = -float("Inf")
+    with torch.no_grad():
+        if embed is not None:
+            generated = embed
+        else:
+            if tokens is None:
+                tokens = torch.tensor(tokenizer.encode(prompt))
+                tokens = tokens.unsqueeze(0).to(device)
+                generated = model.gpt.transformer.wte(tokens)
+
+        stop_seq = tokenizer.encode('<STOP>')
+
+        for i in range(entry_length):
+            outputs = model.gpt(inputs_embeds=generated)
+            logits = outputs.logits
+            logits = logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
+            logits = logits.softmax(-1).log()
+            # prevent repeated ngrams
+            if no_repeat_ngram is not None:
+                if tokens is not None:
+                    for b in range(beam_size):
+                        tokens_list = tokens[b].tolist()
+                        for idx in range(len(tokens_list) - no_repeat_ngram):
+                            subseq = tokens_list[idx:idx+no_repeat_ngram]
+                            if tokens_list[-no_repeat_ngram+1:] == subseq[:-1] and subseq[-1] not in stop_seq:
+                                logits[b, subseq[-1]] = filter_value
+            if scores is None:
+                scores, next_tokens = logits.topk(beam_size, -1)
+                generated = generated.expand(beam_size, *generated.shape[1:])
+                next_tokens, scores = next_tokens.permute(1, 0), scores.squeeze(0)
+                if tokens is None:
+                    tokens = next_tokens
+                else:
+                    tokens = tokens.expand(beam_size, *tokens.shape[1:])
+                    tokens = torch.cat((tokens, next_tokens), dim=1)
+            else:
+                logits[is_stopped] = -float(np.inf)
+                logits[is_stopped, 0] = 0
+                scores_sum = scores[:, None] + logits
+                seq_lengths[~is_stopped] += 1
+                scores_sum_average = scores_sum / seq_lengths[:, None]
+                scores_sum_average, next_tokens = scores_sum_average.view(-1).topk(
+                    beam_size, -1
+                )
+                next_tokens_source = next_tokens // scores_sum.shape[1]
+                seq_lengths = seq_lengths[next_tokens_source]
+                next_tokens = next_tokens % scores_sum.shape[1]
+                next_tokens = next_tokens.unsqueeze(1)
+                tokens = tokens[next_tokens_source]
+                tokens = torch.cat((tokens, next_tokens), dim=1)
+                generated = generated[next_tokens_source]
+                scores = scores_sum_average * seq_lengths
+                is_stopped = is_stopped[next_tokens_source]
+            next_token_embed = model.gpt.transformer.wte(next_tokens.squeeze()).view(
+                generated.shape[0], 1, -1
+            )
+            generated = torch.cat((generated, next_token_embed), dim=1)
+            is_stopped = is_stopped + next_tokens.eq(stop_token_index).squeeze()
+            if is_stopped.all():
+                break
+    scores = scores / seq_lengths
+    output_list = tokens.cpu().numpy()
+    output_texts = [
+        tokenizer.decode(output[: int(length)])
+        for output, length in zip(output_list, seq_lengths)
+    ]
+    order = scores.argsort(descending=True)
+    output_texts = [output_texts[i] for i in order]
+    return output_texts
+
+
+def generate2(
+    model,
+    tokenizer,
+    tokens=None,
+    prompt=None,
+    embed=None,
+    entry_count=1,
+    #entry_length=67,  # maximum number of words
+    entry_length=150,  # maximum number of words
+    top_p=0.8,
+    nucleus=False,
+    #temperature=1.0,
+    temperature=0.7,
+    stop_token: str = ".",
+    no_repeat_ngram = 3,
+):
+    model.eval()
+    generated_num = 0
+    generated_list = []
+    stop_token_index = tokenizer.encode(stop_token)[0]
+    filter_value = -1e10
+    device = next(model.parameters()).device
+
+    with torch.no_grad():
+
+        for entry_idx in range(entry_count):
+            if embed is not None:
+                generated = embed
+            else:
+                if tokens is None:
+                    tokens = torch.tensor(tokenizer.encode(prompt))
+                    tokens = tokens.unsqueeze(0).to(device)
+
+                generated = model.gpt.transformer.wte(tokens)
+
+            ngrams = defaultdict(lambda: set())
+            stop_seq = tokenizer.encode('<STOP>')
+
+            for i in range(entry_length):
+
+                outputs = model.gpt(inputs_embeds=generated)
+                logits = outputs.logits
+                logits = logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
+                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+                cumulative_probs = torch.cumsum(
+                    nnf.softmax(sorted_logits, dim=-1), dim=-1
+                )
+                sorted_indices_to_remove = cumulative_probs > top_p
+                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[
+                    ..., :-1
+                ].clone()
+                sorted_indices_to_remove[..., 0] = 0
+
+                indices_to_remove = sorted_indices[sorted_indices_to_remove]
+                logits[:, indices_to_remove] = filter_value
+                # remove any potential ngram repeats, unless part of <STOP>
+                if no_repeat_ngram is not None:
+                    if tokens is not None:
+                        for token in ngrams[tuple(tokens[0][-no_repeat_ngram+1:].tolist())]:
+                            if token not in stop_seq:
+                                logits[:, token] = filter_value
+                # either sample or argmax
+                if nucleus:
+                    distr = torch.distributions.categorical.Categorical(logits=logits.squeeze())
+                    next_token = distr.sample().unsqueeze(0).unsqueeze(0)
+                else:
+                    next_token = torch.argmax(logits, -1).unsqueeze(0)
+                next_token_embed = model.gpt.transformer.wte(next_token)
+                if logits[:, next_token].item() == filter_value:
+                    break
+                # add to our set of ngrams
+                if no_repeat_ngram is not None:
+                    if tokens is not None and len(tokens[0]) >= no_repeat_ngram - 1:
+                        ngrams[tuple(tokens[0][-no_repeat_ngram+1:].tolist())].add(next_token.item())
+                if tokens is None:
+                    tokens = next_token
+                else:
+                    tokens = torch.cat((tokens, next_token), dim=1)
+                generated = torch.cat((generated, next_token_embed), dim=1)
+                if stop_token_index == next_token.item():
+                    break
+
+
+            output_list = tokens.cpu().tolist()[0]
+            output_text = tokenizer.decode(output_list)
+            generated_list.append(output_text)
+
+    return generated_list[0]

data_module.py

@@ -0,0 +1,382 @@
+'''
+Ke Chen | [email protected] & Nikita Srivatsan | [email protected]
+Load the mp3 format data from audiostock-full dataset
+'''
+import json
+import numpy as np
+import os
+import pandas as pd
+from pathlib import PurePosixPath
+import random
+import torch
+import torchaudio
+from torch.utils.data import Dataset
+import sys
+
+from lib import *
+from utils import *
+
+import torch.utils.data
+
+def int16_to_float32(x):
+    return (x / 32767.0).type(torch.float)
+
+
+def float32_to_int16(x):
+    x = torch.clip(x, min=-1., max=1.)
+    return (x * 32767.).type(torch.int16)
+
+def my_collate(batch):
+    batch = [x for x in batch if x is not None]
+    if len(batch) == 0:
+        return batch
+    else:
+        return torch.utils.data.dataloader.default_collate(batch)
+
+class AudiostockDataset(Dataset):
+    '''
+    Args:
+        dataset_path (str): the dataset folder path
+        train (bool): if True, we randomly return a 10-sec chunk from each audio file; if False, we return the middle 10-sec chunk (fixed)
+        split (str): a txt file to assign the idx in this dataset (for trainng, validation and testing)
+        factor (float): how many time we need to loop the whole dataset, this is to increase the number of training data batches in each epoch
+        whole_track (bool): if True, the dataset will return the full length of the audio file. However, this means the batch_size = 1, and it is usually in the test/validation case
+    '''
+    def __init__(self, dataset_path, tweet_prefix=True, prefix_length=10, normalize=False, dupefile='dupes.pkl', train = True, split = None, factor = 1.0, whole_track = False, verbose=True, dedup=True, file_list=[]):
+        super().__init__()
+        # set up parameters
+        self.max_seq_len = 150
+        self.tweet_prefix = tweet_prefix
+        if self.tweet_prefix:
+            self.max_seq_len *= 2
+        self.tokenizer = GPT2Tokenizer.from_pretrained('gpt2', local_files_only=True)
+        self.prefix_length = prefix_length
+        self.normalize = normalize
+        self.id2neighbor = defaultdict(lambda: '')
+
+        if dedup:
+            if dupefile is not None and os.path.exists(dupefile):
+                with open(dupefile, 'rb') as dupefile:
+                    self.is_rep = pickle.load(dupefile).is_rep
+            elif dupefile == 'both':
+                with open('dupes.pkl', 'rb') as dupefile:
+                    dupes1 = pickle.load(dupefile)
+                with open('dupes_audio.pkl', 'rb') as dupefile:
+                    dupes2 = pickle.load(dupefile)
+                self.is_rep = defaultdict(lambda: True)
+                for k,v in dupes1.is_rep.items():
+                    self.is_rep[k] = v
+                for k,v in dupes2.is_rep.items():
+                    self.is_rep[k] = v
+            else:
+                sys.exit('Could not find duplicate file')
+
+        subfolders = [f'audiostock-part-{i}' for i in range(1,9)]
+        self.label_path = os.path.join(dataset_path, 'audiostock-full-label')
+        self.whole_track = whole_track
+        self.file_list = file_list
+        
+        # select out the elements for this split
+        if self.file_list == []:
+            temp_file_list = []
+            for subfolder in subfolders:
+                temp_file_list += [os.path.join(dataset_path, subfolder, f) for f in os.listdir(os.path.join(dataset_path, subfolder)) if not dedup or self.is_rep[os.path.basename(f).split('.')[0]]]
+            if split is not None:
+                split = set(np.loadtxt(split, dtype = str))
+                self.file_list = [f for f in temp_file_list if os.path.basename(f).split('.')[0] in split]
+            else:
+                self.file_list = temp_file_list
+        
+        self.train = train
+        self.total_len = int(len(self.file_list) * factor)
+        if verbose:
+            print(f'Dataset Loaded | File Num.: {len(self.file_list)} | Batches per epoch: {self.total_len}')
+
+    def precompute_rand(self, candidate_set=None):
+        self.id2neighbor = defaultdict(lambda: '')
+        # if train
+        if candidate_set is None:
+            my_ids = []
+            candidate_caps = []
+            temp_loader = DataLoader(self, batch_size=32, shuffle=False, num_workers=32, drop_last=False, collate_fn=my_collate)
+            for batch in temp_loader:
+                my_ids += batch['id']
+                candidate_caps += batch['short_text']
+            for idx in my_ids:
+                self.id2neighbor[idx] = random.choice(candidate_caps)
+        # if test
+        else:
+            temp_loader = DataLoader(candidate_set, batch_size=32, shuffle=False, num_workers=32, drop_last=False, collate_fn=my_collate)
+            candidate_caps = []
+            for batch in temp_loader:
+                candidate_caps += batch['short_text']
+            temp_loader = DataLoader(self, batch_size=32, shuffle=False, num_workers=32, drop_last=False, collate_fn=my_collate)
+            my_ids = []
+            for batch in temp_loader:
+                my_ids += batch['id']
+            for idx in my_ids:
+                self.id2neighbor[idx] = random.choice(candidate_caps)
+
+    def precompute_gold(self):
+        self.id2neighbor = defaultdict(lambda: '')
+        temp_loader = DataLoader(self, batch_size=32, shuffle=False, num_workers=32, drop_last=False, collate_fn=my_collate)
+        for batch in temp_loader:
+            for idx,short_text in zip(batch['id'], batch['short_text']):
+                self.id2neighbor[idx] = short_text
+
+    def precompute_blank(self):
+        self.id2neighbor = defaultdict(lambda: '\n')
+
+    def precompute_neighbors(self, model, candidate_set=None):
+        print('Precomputing neighbors')
+        self.id2neighbor = defaultdict(lambda: '')
+        # if train and model given
+        if candidate_set is None:
+            # compute waveform embeddings for each song
+            cand_features = None
+            cand_ids = []
+            cand_caps = []
+            temp_loader = DataLoader(self, batch_size=32, shuffle=False, num_workers=32, drop_last=False, collate_fn=my_collate)
+            progress = tqdm(total=len(temp_loader), dynamic_ncols=True)
+            for batch in temp_loader:
+                with torch.no_grad():
+                    batch_features = model.embed_waveform(batch['waveform'].cuda())
+                    if cand_features is not None:
+                        cand_features = torch.cat([cand_features, batch_features])
+                    else:
+                        cand_features = batch_features
+                    cand_ids += batch['id']
+                    cand_caps += batch['short_text']
+                progress.update()
+            progress.close()
+            my_features = cand_features
+            my_ids = cand_ids
+        # if test and model given
+        else:
+            # check if we already precomputed the embeddings
+            pickle_filename = 'nn_features.pkl'
+            if os.path.isfile(pickle_filename):
+                with open(pickle_filename, 'rb') as f:
+                    (cand_features, cand_ids, cand_caps) = pickle.load(f)
+            else:
+                # build the features from the provided set instead of self
+                cand_features = None
+                cand_ids = []
+                cand_caps = []
+                temp_loader = DataLoader(candidate_set, batch_size=32, shuffle=False, num_workers=32, drop_last=False, collate_fn=my_collate)
+                progress = tqdm(total=len(temp_loader), dynamic_ncols=True)
+                for batch in temp_loader:
+                    with torch.no_grad():
+                        batch_features = model.embed_waveform(batch['waveform'].cuda())
+                        if cand_features is not None:
+                            cand_features = torch.cat([cand_features, batch_features])
+                        else:
+                            cand_features = batch_features
+                        cand_ids += batch['id']
+                        #cand_caps += [' '.join(x.split()[:10]) for x in batch['short_text']]
+                        cand_caps += batch['short_text']
+                    progress.update()
+                progress.close()
+                # dump to pickle so we don't have to redo this each time
+                with open(pickle_filename, 'wb') as f:
+                    pickle.dump((cand_features, cand_ids, cand_caps), f)
+            # load up my own ids and features
+            my_features = None
+            my_ids = []
+            temp_loader = DataLoader(self, batch_size=32, shuffle=False, num_workers=32, drop_last=False, collate_fn=my_collate)
+            progress = tqdm(total=len(temp_loader), dynamic_ncols=True)
+            for batch in temp_loader:
+                with torch.no_grad():
+                    batch_features = model.embed_waveform(batch['waveform'].cuda())
+                    if my_features is not None:
+                        my_features = torch.cat([my_features, batch_features])
+                    else:
+                        my_features = batch_features
+                    my_ids += batch['id']
+                progress.update()
+            progress.close()
+        is_self_sim = my_ids == cand_ids
+        for idx,audio_id in tqdm(enumerate(my_ids), total=len(my_ids), dynamic_ncols=True):
+            features = my_features[idx]
+            similarities = features @ cand_features.T
+            # remove identical matches
+            if is_self_sim:
+                similarities[idx] = float('-inf')
+            best_idx = torch.argmax(similarities)
+            most_similar_caption = cand_caps[best_idx]
+            self.id2neighbor[my_ids[idx]] = most_similar_caption
+
+    def pad_tokens(self, tokens, tokens_tweet):
+        tweet_text_len = 0
+        if self.tweet_prefix:
+            tweet_text_len = tokens_tweet[:self.max_seq_len // 2].shape[0]
+            tokens = torch.cat((tokens_tweet[:tweet_text_len], tokens))
+        padding = self.max_seq_len - tokens.shape[0]
+        if padding > 0:
+            tokens = torch.cat((tokens, torch.zeros(padding, dtype=torch.int64) - 1))
+        elif padding < 0:
+            tokens = tokens[:self.max_seq_len]
+        mask = tokens.ge(0)  # mask is zero where we out of sequence
+        tokens[~mask] = 0
+        mask = mask.float()
+        mask = torch.cat((torch.ones(self.prefix_length), mask), dim=0)  # adding prefix mask
+        return tokens, mask, tweet_text_len
+
+    def read_wav(self, filename):
+        stem = PurePosixPath(filename).stem
+        picklefile = f'wt-{self.whole_track}-t-{self.train}-{stem}.pt'
+        picklepath = f'/trunk/datasets/nsrivats/audiostock_proc/{picklefile}'
+        if os.path.exists(picklepath):
+            y = torch.load(picklepath)
+        else:
+            # chunk
+            try:
+                num_frames = torchaudio.info(filename).num_frames
+            except:
+                return None
+            # make sure it wasn't empty, if so die
+            if num_frames == 0:
+                return None
+            sta = 0
+            if not self.whole_track:
+                if self.train: 
+                    sta = random.randint(0, num_frames - 441001)
+                else:
+                    sta = (num_frames - 441001) // 2
+                num_frames = 441000
+
+            y, sr = torchaudio.load(filename, frame_offset=sta, num_frames=num_frames)
+            # resample
+            y = torchaudio.functional.resample(y, sr, 48000)
+            y = y[:, :441000]
+            # mono
+            y = y.mean(dim=0)
+            # normalize
+            y = int16_to_float32(float32_to_int16(y))
+            # save
+            torch.save(y, picklepath)
+        return y
+
+    def __getitem__(self, index):
+        idx = index % len(self.file_list)
+        data_dict = {}
+        f = self.file_list[idx]
+        lf = os.path.join(self.label_path, os.path.basename(f).split('.')[0] + '.json')
+        data_dict['waveform'] = self.read_wav(f)
+        if os.path.isfile(lf):
+            with open(lf,'r') as label_file:
+                label_data = json.load(label_file)
+                data_dict['id'] = label_data['id']
+                data_dict['short_text'] = label_data['short_text']
+                if self.normalize:
+                    data_dict['short_text'] = ' '.join(muscaps_tokenize(data_dict['short_text']))
+                if 'long_text' in label_data and label_data['long_text'] is not None:
+                    data_dict['long_text'] = label_data['long_text']
+                else:
+                    data_dict['long_text'] = ''
+                '''
+                data_dict['tag'] = label_data['tag']
+                data_dict['impression'] = label_data['impression']
+                data_dict['purpose'] = label_data['purpose']
+                '''
+        else:
+            data_dict['id'] = os.path.basename(f).split('.')[0]
+            data_dict['short_text'] = ''
+            data_dict['long_text'] = ''
+
+        # tokenize the caption
+        caption_proc = preproc(data_dict['short_text'], self.tokenizer)
+        tokens = torch.tensor(caption_proc, dtype=torch.int64)
+        tweet_text = self.id2neighbor[data_dict['id']] if self.tweet_prefix else ''
+        tweet_proc = preproc(tweet_text, self.tokenizer, stop=False)
+        tokens_tweet = torch.tensor(tweet_proc, dtype=torch.int64)
+        tokens, mask, tweet_text_len = self.pad_tokens(tokens, tokens_tweet)
+        data_dict['tokens'] = tokens
+        data_dict['mask'] = mask
+        data_dict['tweet_text_len'] = tweet_text_len
+        data_dict['tweet_text'] = tweet_text
+
+        if (data_dict['id'] is None or
+            data_dict['short_text'] is None or
+            data_dict['long_text'] is None or
+            data_dict['tokens'] is None or
+            data_dict['mask'] is None or
+            data_dict['tweet_text_len'] is None or
+            data_dict['tweet_text'] is None or
+            data_dict['waveform'] is None
+        ):
+            return None
+        else:
+            return data_dict
+
+    def __len__(self): 
+        return self.total_len
+
+class MusicCapsDataset(AudiostockDataset):
+    def __init__(self, dataset_path, args, train = True, split = None, factor = 1.0, whole_track = False, verbose=True, dedup=True):
+        super(AudiostockDataset, self).__init__()
+        # set up parameters
+        self.max_seq_len = 150
+        self.tweet_prefix = args.tweet_prefix
+        if self.tweet_prefix:
+            self.max_seq_len *= 2
+        self.tokenizer = GPT2Tokenizer.from_pretrained('gpt2', local_files_only=True)
+        self.prefix_length = args.prefix_length
+        self.normalize = args.normalize
+        self.whole_track = whole_track
+
+        self.label_path = os.path.join(dataset_path, 'audio')
+        self.file_list = []
+        self.label_data = []
+        label_reader = pd.read_csv(f'{dataset_path}/musiccaps-resplit.csv')
+        for idx,row in label_reader.iterrows():
+            if (row['is_audioset_eval'] == 1 and split == 'musiccaps_eval') \
+              or (row['is_audioset_eval'] == 0 and split == 'musiccaps_train') \
+              or (row['is_audioset_eval'] == 2 and split == 'musiccaps_dev'):
+                data_dict = {}
+                data_dict['id'] = row['ytid']
+                self.file_list.append(f"{dataset_path}/audio/{data_dict['id']}.wav")
+                data_dict['short_text'] = row['caption']
+                if self.normalize:
+                    data_dict['short_text'] = ' '.join(muscaps_tokenize(data_dict['short_text']))
+                data_dict['long_text'] = ''
+                data_dict['tag'] = row['aspect_list']
+                self.label_data.append(data_dict)
+
+        self.train = train
+        self.total_len = int(len(self.file_list) * factor)
+        if verbose:
+            print(f'Dataset Loaded | File Num.: {len(self.file_list)} | Batches per epoch: {self.total_len}')
+
+    def __getitem__(self, index):
+        idx = index % len(self.file_list)
+        data_dict = {}
+        f = self.file_list[idx]
+        data_dict['waveform'] = self.read_wav(f)
+        for k,v in self.label_data[idx].items():
+            data_dict[k] = v
+
+        # tokenize the caption
+        caption_proc = preproc(data_dict['short_text'], self.tokenizer)
+        tokens = torch.tensor(caption_proc, dtype=torch.int64)
+        tweet_text = self.id2neighbor[data_dict['id']] if self.tweet_prefix else ''
+        tweet_proc = preproc(tweet_text, self.tokenizer, stop=False)
+        tokens_tweet = torch.tensor(tweet_proc, dtype=torch.int64)
+        tokens, mask, tweet_text_len = self.pad_tokens(tokens, tokens_tweet)
+        data_dict['tokens'] = tokens
+        data_dict['mask'] = mask
+        data_dict['tweet_text_len'] = tweet_text_len
+        data_dict['tweet_text'] = tweet_text
+
+        if (data_dict['id'] is None or
+            data_dict['short_text'] is None or
+            data_dict['long_text'] is None or
+            data_dict['tokens'] is None or
+            data_dict['mask'] is None or
+            data_dict['tweet_text_len'] is None or
+            data_dict['tweet_text'] is None or
+            data_dict['waveform'] is None
+        ):
+            return None
+        else:
+            return data_dict

dupes.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e83b71d63cd11dc8840b44bcea625d1c618c8b421e4c6ec6c65580af5109c7bd
+size 1807022

infer.py

@@ -0,0 +1,55 @@
+from audiocaptioner import AudioCaptioner
+from data_module import AudiostockDataset
+from utils import *
+
+def infer(input_filename):
+    device = get_device(0)
+    # connect to GCS
+    gcs = CheckpointManager()
+    # create and/or load model
+    tokenizer = GPT2Tokenizer.from_pretrained('gpt2', local_files_only=False)
+    prefix_dim = 512
+    prefix_length = 10
+    prefix_length_clip = 10
+    num_layers = 8
+    checkpoint = 'checkpoints/ZRIUE-BEST.pt'
+    model = AudioCaptioner(prefix_length, clip_length=prefix_length_clip, prefix_size=prefix_dim, num_layers=num_layers).to(device)
+    model.load_state_dict(gcs.get_checkpoint(checkpoint))
+    print(f'Loaded from {checkpoint}')
+    model.eval()
+    # read in the wav file and precompute neighbors
+    #dataset_path = '/graft1/datasets/kechen/audiostock-full'
+    dataset_path = ''
+    train_dataset = AudiostockDataset(
+            dataset_path=dataset_path,
+            train=False,
+            split='audiostock-train-240k.txt',
+            factor=1.0,
+            verbose=False,
+            file_list=open('audiostock-train-240k.txt', 'r').read().split()
+    )
+    print('Reading in file', input_filename)
+    dataset = AudiostockDataset(
+        dataset_path=dataset_path,
+        train=False,
+        split=None,
+        factor=1.0,
+        verbose=False,
+        file_list=[input_filename] # manually override file list
+    )
+    dataset.precompute_neighbors(model, candidate_set=train_dataset)
+    waveform = dataset.read_wav(input_filename).unsqueeze(0).to(device, dtype=torch.float32)
+    # predict
+    with torch.no_grad():
+        prefix_embed = model.create_prefix(waveform, 1)
+        tweet_tokens = torch.tensor(preproc(dataset.id2neighbor[os.path.basename(input_filename).split('.')[0]], tokenizer, stop=False), dtype=torch.int64).to(device)[:150]
+        tweet_embed = model.gpt.transformer.wte(tweet_tokens)
+        prefix_embed = torch.cat([prefix_embed, tweet_embed.unsqueeze(0)], dim=1)
+        candidates = generate_beam(model, tokenizer, embed=prefix_embed, beam_size=5)
+        generated_text = candidates[0]
+        generated_text = postproc(generated_text)
+    print('=======================================')
+    print(generated_text)
+
+if __name__ == '__main__':
+    infer('../MusicCaptioning/sample_inputs/sisters.mp3')

lib.py

@@ -0,0 +1,19 @@
+from collections import defaultdict
+import json
+import numpy as np
+import os
+import pandas as pd
+import dill as pickle
+pickle._dill._reverse_typemap['ClassType'] = type
+import random
+import string
+import sys
+import torch
+from torch import nn
+import torch.nn.functional as nnf
+from torch.utils.data import Dataset, DataLoader
+from tqdm import tqdm
+from transformers import GPT2Tokenizer, GPT2LMHeadModel, AdamW, get_linear_schedule_with_warmup
+from typing import Tuple, List, Union, Optional
+
+from clipcap import *

utils.py

@@ -0,0 +1,45 @@
+from lib import *
+from twokenize import tokenizeRawTweetText
+import re
+
+def muscaps_tokenize(raw):
+    raw = raw.lower()
+    for punc in string.punctuation:
+        raw = raw.replace(punc, ' ')
+    tokens = raw.split()
+    return tokens
+
+def get_device(device_id: int) -> torch.device:
+    if not torch.cuda.is_available():
+        return torch.device('cpu')
+    device_id = min(torch.cuda.device_count() - 1, device_id)
+    return torch.device(f'cuda:{device_id}')    
+
+def preproc(caption, tokenizer, stop=True):
+    caption = caption.replace('.', '<STOP>')
+    caption_proc = tokenizer.encode(caption)
+    if stop:
+        caption_proc += tokenizer.encode('.')
+    return caption_proc
+
+def postproc(caption):
+    caption = caption.replace('<STOP>', '.')
+    if caption[-1] == '.':
+        caption = caption[:-1]
+    return caption
+
+class CheckpointManager:
+    def __init__(self):
+        self.checkpoint_dir = '/home/nsrivats/Repositories/MusicCaptioning/checkpoints'
+
+    def get_checkpoint(self, checkpoint):
+        with open(checkpoint, 'rb') as infile:
+            return torch.load(infile)
+
+    def save_checkpoint(self, state_dict, checkpoint):
+        filename = f'{self.checkpoint_dir}/{checkpoint}'
+        with open(filename, 'wb') as outfile:
+            torch.save(state_dict, outfile)
+
+    def save_logs(self, logdir):
+        pass