add bytecover

app.py

@@ -1,4 +1,4 @@
-import spaces
+# import spaces
 from pathlib import Path
 import yaml
 import time
@@ -15,6 +15,191 @@ import gradio as gr
 from vampnet.interface import Interface, signal_concat
 from vampnet import mask as pmask
 
+from pyharp import *
+from bytecover.models.train_module import TrainModule
+from bytecover.utils import initialize_logging, load_config
+import pinecone
+import laion_clap
+from tqdm import tqdm
+
+import os
+
+### INIT BYTECOVER
+print(f"Is CUDA available: {torch.cuda.is_available()}")
+# True
+print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
+
+index_clap = pinecone.Index(os.environ["PC_API_KEY"], host=os.environ["CLAP_INDEX"])  #host='https://clap-nathan-500-index-af8053a.svc.us-west1-gcp.pinecone.io')
+index_bytecover = pinecone.Index(os.environ["PC_API_KEY"], host=os.environ["BC_INDEX"]) #host='https://bytecover-nathan-500-index-af8053a.svc.us-west1-gcp.pinecone.io')
+
+print("Loading ByteCover model")
+
+if torch.cuda.is_available():
+    bytecover_config = load_config(config_path="bytecover/config_gpu.yaml")
+else:
+    bytecover_config = load_config(config_path="bytecover/config.yaml")
+bytecover_module = TrainModule(bytecover_config)
+bytecover_model = bytecover_module.model
+if bytecover_module.best_model_path is not None:
+    bytecover_model.load_state_dict(torch.load(bytecover_module.best_model_path), strict=False)
+    print(f"Best model loaded from checkpoint: {bytecover_module.best_model_path}")
+elif bytecover_module.config["test"]["model_ckpt"] is not None:
+    bytecover_model.load_state_dict(torch.load(bytecover_module.config["test"]["model_ckpt"], map_location='cpu'), strict=False)
+    print(f'Model loaded from checkpoint: {bytecover_module.config["test"]["model_ckpt"]}')
+elif bytecover_module.state == "initializing":
+    print("Warning: Running with random weights")
+
+bytecover_model.eval()
+
+print("Loading CLAP model")
+
+if torch.cuda.is_available():
+    clap_model = laion_clap.CLAP_Module(enable_fusion=False, device="cuda:0")
+else:
+    clap_model = laion_clap.CLAP_Module(enable_fusion=False)
+clap_model.load_ckpt() # download the default pretrained checkpoint.
+
+print("Models loaded!")
+def convert_to_npfloat64(original_array):
+    #return np.array(flat_df["flat_vector_embed"][0],dtype=np.float64)
+    return np.array(original_array,dtype=np.float64)
+
+def convert_to_npfloat64_to_list(vector_embed_64):
+    # list(flat_df["flat_vector_embed_64"][0])
+    return list(vector_embed_64)
+
+def flatten_vector_embed(vector_embed):
+    return list(vector_embed.flatten())
+
+def format_time(num_seconds):
+    return f"{num_seconds // 60}:{num_seconds % 60:02d}"
+
+def bytecover(sig, chunk_size=3.0, bytecover_match_ct=3, clap_match_ct=3):
+    """
+    This function defines the audio processing steps
+    Args:
+        input_audio_path (str): the audio filepath to be processed.
+        <YOUR_KWARGS>: additional keyword arguments necessary for processing.
+            NOTE: These should correspond to and match order of UI elements defined below.
+    Returns:
+        output_audio_path (str): the filepath of the processed audio.
+        output_labels (LabelList): any labels to display.
+    """
+
+    """
+    <YOUR AUDIO LOADING CODE HERE>
+    """
+
+
+    """
+    <YOUR AUDIO PROCESSING CODE HERE>
+    """
+    sig_mono = sig.copy().to_mono().audio_data.squeeze(1)
+
+    # Chunk audio to desired length
+
+    chunk_samples = int(chunk_size * sig.sample_rate)
+    print(f"Chunk samples: {chunk_samples}")
+    print(f"Shape of audio: {sig_mono.shape}")
+    chunks = torch.tensor_split(sig_mono, [i for i in range(chunk_samples, sig_mono.shape[1], chunk_samples)], dim=1)
+    if chunks[-1].shape[1] < chunk_samples:
+        print("Cutting last chunk due to length")
+        chunks = tuple(list(chunks)[:-1])
+    print(f"Number of chunks: {len(chunks)}")
+
+    print("Getting Bytecover embeddings")
+    bytecover_embeddings = []
+    for chunk in tqdm(chunks):
+        result = bytecover_model.forward(chunk.to(bytecover_module.config["device"]))['f_t'].detach()
+        bytecover_embeddings.append(result)
+
+    clean_bytecover_embeddings = [convert_to_npfloat64_to_list(convert_to_npfloat64(flatten_vector_embed(embedding.cpu()))) for embedding in bytecover_embeddings]
+
+    print("Getting CLAP embeddings")
+    clap_embeddings = []
+    for chunk in tqdm(chunks):
+        result = clap_model.get_audio_embedding_from_data(chunk.numpy())
+        clap_embeddings.append(result)
+
+    clean_clap_embeddings = [convert_to_npfloat64_to_list(convert_to_npfloat64(flatten_vector_embed(embedding))) for embedding in clap_embeddings]
+
+    clap_matches = []
+    bytecover_matches = []
+    match_metadatas = {}
+
+    output_labels = LabelList()
+
+    times = {}
+
+    for clean_embeddings, pinecone_index, match_list, embedding_num, num_matches in zip([clean_bytecover_embeddings, clean_clap_embeddings], [index_bytecover, index_clap], [bytecover_matches, clap_matches], range(2), [bytecover_match_ct, clap_match_ct]):
+
+        for i, embedding in enumerate(clean_embeddings):
+
+            print(f"Getting match {i + 1} of {len(clean_embeddings)}")
+            matches = pinecone_index.query(
+                vector=embedding,
+                top_k=10,
+                #include_values=False,
+                    include_metadata=True
+                )['matches']
+
+            # Store matches as [score, time, id]
+
+            for match in matches:
+                id = match['id']
+                if id not in match_metadatas:
+                    match_metadatas[id] = match['metadata']
+                match_list.append([match['score'], i * chunk_size, id])
+
+
+        print("Matches obtained!")
+
+        top_matches = sorted(match_list, key=lambda item: item[0], reverse=True)
+    
+        for i, match in enumerate(top_matches[:int(num_matches)]):
+            metadata = match_metadatas[match[2]]
+            song_artists = metadata['artists']
+            if type(song_artists) is list:
+                artists = ' and '.join(artists)
+
+
+            song_title = metadata['song']
+
+            song_link = f"https://open.spotify.com/track/{metadata['spotify_id'].split(':')[2]}"
+
+            embed_name = ['ByteCover', 'CLAP'][embedding_num]
+
+            match_time = match[1]
+            times[match_time] = times.get(match_time, 0) + 1
+
+            label = AudioLabel(
+                t=match_time, 
+                label=f'{song_title}', 
+                duration=chunk_size, 
+                link=song_link, 
+                description=f'Embedding: {embed_name}\n{song_title} by {song_artists}\nClick the tag to view on Spotify!', 
+                amplitude=1.0 - 0.5 * (times[match_time] - 1), 
+                color=AudioLabel.rgb_color_to_int(200, 170, 3, 10) if embedding_num == 1 else 0
+            )
+
+            # if embedding_num == 1:
+            #     label.rgb_color_to_int(200, 170, 3, 240)
+            # else:
+            #     pass
+            #     #label.set_color(204, 52, 235, 240)
+        
+            output_labels.append(label)
+
+    """
+    <YOUR AUDIO SAVING CODE HERE>
+    # Save processed audio and obtain default path
+    output_audio_path = save_audio(signal, None)
+    """
+
+    return output_labels
+
+### END BYTECOVER
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 interface = Interface.default()
@@ -57,7 +242,6 @@ def shift_pitch(signal, interval: int):
     return signal
 
 
-@spaces.GPU
 def _vamp(
         seed, input_audio, model_choice, 
         pitch_shift_amt, periodic_p, 
@@ -168,7 +352,7 @@ def api_vamp(data):
 
 OUT_DIR = Path("gradio-outputs")
 OUT_DIR.mkdir(exist_ok=True)
-def harp_vamp(input_audio_file, periodic_p, n_mask_codebooks):
+def harp_vamp(input_audio_file, periodic_p, n_mask_codebooks, chunk_size=3.0, bytecover_match_ct=3, clap_match_ct=3):
     sig = at.AudioSignal(input_audio_file)
     sr, samples = sig.sample_rate, sig.samples[0][0].detach().cpu().numpy()
     # convert to int32
@@ -192,7 +376,11 @@ def harp_vamp(input_audio_file, periodic_p, n_mask_codebooks):
         stretch_factor=1,
     )
     
-    sig = at.AudioSignal(samples, sr)
+    sig = at.AudioSignal(samples, sr).cpu()
+
+    # run bytecover
+    labels = bytecover(sig, chunk_size, bytecover_match_ct, clap_match_ct)
+
     # write to file
     # clear the outdir
     for p in OUT_DIR.glob("*"):
@@ -200,10 +388,7 @@ def harp_vamp(input_audio_file, periodic_p, n_mask_codebooks):
     OUT_DIR.mkdir(exist_ok=True)
     outpath = OUT_DIR / f"{uuid.uuid4()}.wav"
     sig.write(outpath)
-    from pyharp import AudioLabel, LabelList
-    output_labels = LabelList()
-    output_labels.append(AudioLabel(label='~', t=0.0, amplitude=0.5, description='generated audio'))
-    return outpath, output_labels
+    return outpath, labels
     
 
 with gr.Blocks() as demo:
@@ -425,19 +610,47 @@ with gr.Blocks() as demo:
 
     from pyharp import ModelCard, build_endpoint
     card = ModelCard(
-        name="vampnet", 
+        name="vampnet + aitribution", 
         description="vampnet! is a model for generating audio from audio",
         author="hugo flores garcía", 
         tags=["music generation"], 
         midi_in=False, 
         midi_out=False
     )
+
+    # BYTECOVER
+    # Define Gradio Components
+    components = [
+        # <YOUR UI ELEMENTS HERE>
+        gr.Slider(
+            minimum=1.0,
+            maximum=10.0,
+            step=0.5,
+            value=3.0,
+            label="Sample size (s)"
+        ),
+        gr.Slider(
+            minimum=0,
+            maximum=5,
+            step=1,
+            value=3,
+            label="Bytecover matches to generate"
+        ),
+        gr.Slider(
+            minimum=0,
+            maximum=5,
+            step=1,
+            value=3,
+            label="CLAP matches to generate"
+        )
+    ]
         
     # Build a HARP-compatible endpoint
     app = build_endpoint(model_card=card,
                          components=[
                             periodic_p, 
                             n_mask_codebooks,
+                            *components
                         ],
                          process_fn=harp_vamp)
 

bytecover/__main__.py

@@ -0,0 +1,17 @@
+import wandb
+
+from bytecover.models.train_module import TrainModule
+from bytecover.utils import initialize_logging, load_config
+
+config = load_config(config_path="config/config.yaml")
+initialize_logging(config_path="config/logging_config.yaml", debug=False)
+if config["wandb"]:
+    wandb.init(
+        # set the wandb project where this run will be logged
+        project="ByteCover",
+        # track hyperparameters and run metadata
+        config=config["train"],
+    )
+trainer = TrainModule(config)
+trainer.pipeline()
+trainer.test()

bytecover/config.yaml

@@ -0,0 +1,51 @@
+data_path: bytecover_data
+file_extension: mp3
+dataset_path: /bytecover_data/shs100k
+save_model_path: models
+debug: True
+progress_bar: True
+device: cpu
+num_channels: 1
+wandb: False
+
+train:
+  mixed_precision: True
+  target_sr: 22050
+  compress_ratio: 20
+  max_seq_len: [100, 150, 200]
+  num_classes: 10000
+  triplet_margin: 0.3
+  smooth_factor: 0.1
+  model_ckpt: null
+  batch_size: 8
+  num_workers: 0
+  shuffle: True
+  drop_last: True
+  epochs: 1
+  learning_rate: 0.0001
+  patience: 4
+  tempo_factors: [0.7, 1.3]
+  log_steps: 20
+
+val:
+  target_sr: 22050
+  compress_ratio: 20
+  save_val_outputs: True
+  max_seq_len: -1
+  batch_size: 1
+  num_workers: 8
+  shuffle: False
+  drop_last: False
+  output_dir: outputs_val
+
+test:
+  target_sr: 22050
+  compress_ratio: 20
+  save_test_outputs: True
+  model_ckpt: orfium-bytecover.pt
+  max_seq_len: -1
+  batch_size: 1
+  num_workers: 8
+  shuffle: False
+  drop_last: False
+  output_dir: outputs_test

bytecover/config_gpu.yaml

@@ -0,0 +1,51 @@
+data_path: bytecover_data
+file_extension: mp3
+dataset_path: /bytecover_data/shs100k
+save_model_path: models
+debug: True
+progress_bar: True
+device: cuda:0
+num_channels: 1
+wandb: False
+
+train:
+  mixed_precision: True
+  target_sr: 22050
+  compress_ratio: 20
+  max_seq_len: [100, 150, 200]
+  num_classes: 10000
+  triplet_margin: 0.3
+  smooth_factor: 0.1
+  model_ckpt: null
+  batch_size: 8
+  num_workers: 0
+  shuffle: True
+  drop_last: True
+  epochs: 1
+  learning_rate: 0.0001
+  patience: 4
+  tempo_factors: [0.7, 1.3]
+  log_steps: 20
+
+val:
+  target_sr: 22050
+  compress_ratio: 20
+  save_val_outputs: True
+  max_seq_len: -1
+  batch_size: 1
+  num_workers: 8
+  shuffle: False
+  drop_last: False
+  output_dir: outputs_val
+
+test:
+  target_sr: 22050
+  compress_ratio: 20
+  save_test_outputs: True
+  model_ckpt: orfium-bytecover.pt
+  max_seq_len: -1
+  batch_size: 1
+  num_workers: 8
+  shuffle: False
+  drop_last: False
+  output_dir: outputs_test

bytecover/models/data_loader.py

@@ -0,0 +1,152 @@
+import os
+from typing import Dict, Literal, Tuple
+
+import ffmpeg
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, Dataset
+from torchvision import transforms
+
+from bytecover.models.data_model import BatchDict
+from bytecover.utils import bcolors
+
+
+class ByteCoverDataset(Dataset):
+    def __init__(
+        self,
+        data_path: str,
+        file_ext: str,
+        dataset_path: str,
+        data_split: Literal["TRAIN", "VAL", "TEST"],
+        debug: bool,
+        target_sr: int,
+        max_len: int,
+    ) -> None:
+        super().__init__()
+        self.data_path = data_path
+        self.file_ext = file_ext
+        self.dataset_path = dataset_path
+        self.data_split = data_split
+        self.debug = debug
+        self.target_sr = target_sr
+        self.max_len = max_len
+        self._load_data()
+        self.pipeline = transforms.Compose([self._read_audio, self._pad_or_trim_audio])
+
+    def __len__(self) -> int:
+        return len(self.track_ids)
+
+    def __getitem__(self, index: int) -> BatchDict:
+        track_id = self.track_ids[index]
+        anchor_audio = self.pipeline(track_id)
+
+        clique_id, pos_id, neg_id = self._triplet_sampling(track_id)
+
+        if self.data_split == "TRAIN":
+            positive_audio = self.pipeline(pos_id)
+            negative_audio = self.pipeline(neg_id)
+        else:
+            positive_audio = torch.empty(0)
+            negative_audio = torch.empty(0)
+        return dict(
+            anchor_id=track_id,
+            anchor=anchor_audio,
+            anchor_label=torch.tensor(clique_id, dtype=torch.float),
+            positive_id=pos_id,
+            positive=positive_audio,
+            negative_id=neg_id,
+            negative=negative_audio,
+        )
+
+    def _triplet_sampling(self, track_id: str) -> Tuple[int, str, str]:
+        clique_id = self.labels.loc[track_id, "clique"]
+        versions = self.versions.loc[clique_id, "versions"]
+        np.random.shuffle(versions)
+        pos_list = np.setdiff1d(versions, track_id)
+        pos_id = np.random.choice(pos_list, 1)[0]
+
+        neg_id = self.labels[~self.labels.index.isin(versions)].sample(1).index[0]
+
+        return (clique_id, pos_id, neg_id)
+
+    def _load_data(self) -> None:
+        self.track_ids = np.load(
+            os.path.join(self.data_path, "splits", f"{self.data_split.lower()}_ids.npy"), allow_pickle=True
+        )
+
+        self.labels = pd.read_csv(os.path.join(self.data_path, "interim", "shs100k.csv"), usecols=["clique", "id"])
+        self.labels = self.labels[self.labels["id"].isin(self.track_ids)]
+        self.labels.dropna(inplace=True)
+        self.labels.set_index("id", inplace=True)
+        cliques = self.labels["clique"].unique()
+        mapping = {}
+        for k, clique in enumerate(cliques):
+            mapping[clique] = k
+        self.labels["clique"] = self.labels["clique"].map(lambda x: mapping[x])
+
+        self.versions = pd.read_csv(
+            os.path.join(self.data_path, "interim", "versions.csv"), converters={"versions": eval}
+        )
+        self.versions.dropna(inplace=True)
+        self.versions = self.versions[self.versions["clique"].isin(cliques)]
+        self.versions["clique"] = self.versions["clique"].map(lambda x: mapping[x])
+        self.versions.set_index("clique", inplace=True)
+
+    def _read_audio(self, track_id: str) -> torch.Tensor:
+        if self.debug:
+            seq_len = np.random.randint(10, 200) if self.max_len <= 0 else self.max_len
+            return torch.rand(seq_len * self.target_sr)
+        filename = os.path.join(self.dataset_path, f"{track_id}.{self.file_ext}")
+
+        try:
+            # This launches a subprocess to decode audio while down-mixing and resampling as necessary.
+            # Requires the ffmpeg CLI and `ffmpeg-python` package to be installed.
+            out, _ = (
+                ffmpeg.input(filename, threads=0)
+                .output("-", format="s16le", acodec="pcm_s16le", ac=1, ar=self.target_sr)
+                .run(cmd=["ffmpeg", "-nostdin"], capture_stdout=True, capture_stderr=True)
+            )
+        except ffmpeg.Error as e:
+            raise RuntimeError(
+                f"{bcolors.WARNING}Failed to load audio:{bcolors.FAIL + filename + bcolors.ENDC}\n{e.stderr.decode()}"
+            ) from e
+
+        # int16 ranges between -2^15 and +2^15 (±32768). By convention, floating point audio data is
+        # normalized to the range of [-1.0, 1.0]
+        audio = np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0
+
+        return torch.from_numpy(audio)
+
+    def _pad_or_trim_audio(self, audio: torch.Tensor) -> torch.Tensor:
+        if self.max_len <= 0:
+            return audio
+
+        if (self.data_split == "TRAIN") and (audio.shape[-1] <= self.max_len * self.target_sr):
+            return F.pad(audio, (0, self.max_len * self.target_sr - audio.shape[-1]))
+
+        max_offset = audio.shape[-1] - self.max_len * self.target_sr
+        offset = np.random.randint(max_offset) if max_offset > 0 else 0
+
+        return audio[offset : (offset + self.max_len * self.target_sr)]
+
+
+def bytecover_dataloader(
+    data_path: str,
+    file_ext: str,
+    dataset_path: str,
+    data_split: Literal["TRAIN", "VAL", "TEST"],
+    debug: bool,
+    max_len: int,
+    batch_size: int,
+    target_sr: int,
+    **config: Dict,
+) -> DataLoader:
+    return DataLoader(
+        ByteCoverDataset(data_path, file_ext, dataset_path, data_split, debug, target_sr=target_sr, max_len=max_len),
+        batch_size=batch_size if max_len > 0 else 1,
+        num_workers=config["num_workers"],
+        shuffle=config["shuffle"],
+        drop_last=config["drop_last"],
+    )

bytecover/models/data_model.py

@@ -0,0 +1,39 @@
+from typing import TypedDict
+
+import torch
+
+
+class ValDict(TypedDict):
+    anchor_id: str
+    positive_id: str
+    negative_id: str
+    f_t: torch.Tensor
+    f_c: torch.Tensor
+
+
+class BatchDict(TypedDict):
+    anchor_id: str
+    anchor: torch.Tensor
+    anchor_label: torch.Tensor
+    positive_id: str
+    positive: torch.Tensor
+    negative_id: str
+    negative: torch.Tensor
+
+
+class Postfix(TypedDict):
+    Epoch: int
+    train_loss: float
+    train_loss_step: float
+    train_cls_loss: float
+    train_cls_loss_step: float
+    train_triplet_loss: float
+    train_triplet_loss_step: float
+    val_loss: float
+    mr1: float
+    mAP: float
+
+
+class TestResults(TypedDict):
+    test_mr1: float
+    test_mAP: float

bytecover/models/early_stopper.py

@@ -0,0 +1,19 @@
+import numpy as np
+
+
+class EarlyStopper:
+    def __init__(self, patience: int = 1, min_delta: int = 0):
+        self.patience = patience
+        self.min_delta = min_delta
+        self.counter = 0
+        self.min_validation_loss = np.inf
+
+    def __call__(self, validation_loss) -> bool:
+        if validation_loss < self.min_validation_loss:
+            self.min_validation_loss = validation_loss
+            self.counter = 0
+        elif validation_loss >= (self.min_validation_loss + self.min_delta):
+            self.counter += 1
+            if self.counter >= self.patience:
+                return True
+        return False

bytecover/models/modules.py

@@ -0,0 +1,200 @@
+from typing import Tuple
+
+import nnAudio.features.cqt as nnAudio
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchaudio.transforms as T
+
+
+class GeM(nn.Module):
+    def __init__(self, p=3, eps=1e-6):
+        super(GeM, self).__init__()
+        self.p = nn.Parameter(torch.ones(1) * p)
+        self.eps = eps
+
+    def forward(self, x):
+        return self.gem(x, p=self.p, eps=self.eps)
+
+    def gem(self, x, p=3, eps=1e-6):
+        return F.avg_pool2d(x.clamp(min=eps).pow(p), (x.size(-2), x.size(-1))).pow(1.0 / p)
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}(p={self.p.data.tolist()[0]:.4f}, eps={str(self.eps)})"
+
+
+class IBN(nn.Module):
+    r"""Instance-Batch Normalization layer from
+    `"Two at Once: Enhancing Learning and Generalization Capacities via IBN-Net"
+    <https://arxiv.org/pdf/1807.09441.pdf>`
+    Args:
+        planes (int): Number of channels for the input tensor
+        ratio (float): Ratio of instance normalization in the IBN layer
+    """
+
+    def __init__(self, planes, ratio):
+        super(IBN, self).__init__()
+        self.half = int(planes * ratio)
+        self.IN = nn.InstanceNorm2d(self.half, affine=True)
+        self.BN = nn.BatchNorm2d(planes - self.half)
+
+    def forward(self, x):
+        split = torch.split(x, self.half, 1)
+        out1 = self.IN(split[0].contiguous())
+        out2 = self.BN(split[1].contiguous())
+        out = torch.cat((out1, out2), 1)
+        return out
+
+
+class Bottleneck(nn.Module):
+
+    expansion: int = 4
+
+    def __init__(
+        self, in_channels: int, out_channels: int, last: bool = False, downsample=None, stride=1, bias: bool = True
+    ):
+        super(Bottleneck, self).__init__()
+
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=bias)
+        if not last:
+            # Apply Instance normalization in first half channels (ratio=0.5)
+            self.ibn = IBN(out_channels, ratio=0.5)
+        else:
+            self.ibn = nn.BatchNorm2d(out_channels)
+
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=bias)
+        self.batch_norm2 = nn.BatchNorm2d(out_channels)
+
+        self.conv3 = nn.Conv2d(
+            out_channels, out_channels * self.expansion, kernel_size=1, stride=1, padding=0, bias=bias
+        )
+        self.batch_norm3 = nn.BatchNorm2d(out_channels * self.expansion)
+
+        self.downsample = downsample
+        self.stride = stride
+        self.relu = nn.ReLU()
+
+    def forward(self, x: torch.Tensor):
+        residual = x.clone()
+
+        x = self.conv1(x)
+        x = self.ibn(x)
+        x = self.relu(x)
+
+        x = self.conv2(x)
+        x = self.batch_norm2(x)
+        x = self.relu(x)
+
+        x = self.conv3(x)
+        x = self.batch_norm3(x)
+        x = self.relu(x)
+
+        if self.downsample is not None:
+            residual = self.downsample(residual)
+
+        out = residual + x
+        out = self.relu(out)
+
+        return out
+
+
+class Resnet50(nn.Module):
+    def __init__(
+        self,
+        ResBlock: Bottleneck,
+        emb_dim: int = 2048,
+        num_channels: int = 1,
+        num_classes: int = 8858,
+        sr: int = 22050,
+        hop_lenght: int = 512,
+        n_bins=84,
+        bins_per_octave=12,
+        window="hann",
+        compress_ratio: int = 20,
+        tempo_factors: Tuple[float, float] = None,
+    ) -> None:
+
+        super(Resnet50, self).__init__()
+        self.in_channels = 64
+
+        self.cqt = nnAudio.CQT2010v2(
+            sr=sr,
+            hop_length=hop_lenght,
+            n_bins=n_bins,
+            bins_per_octave=bins_per_octave,
+            window=window,
+            output_format="Complex",
+            verbose=False,
+        )
+        self.compress = nn.AvgPool2d((1, compress_ratio))
+        self.time_strech = T.TimeStretch(n_freq=n_bins)
+        self.tempo_factors = tempo_factors
+
+        self.conv1 = nn.Conv2d(
+            in_channels=num_channels, out_channels=64, kernel_size=7, stride=2, padding=3, bias=False
+        )
+        self.batch_norm1 = nn.BatchNorm2d(num_features=64)
+        self.relu = nn.ReLU()
+        self.max_pool1 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+
+        self.layer1 = self._make_layer(ResBlock, blocks=3, planes=64, stride=1)
+        self.layer2 = self._make_layer(ResBlock, blocks=4, planes=128, stride=2)
+        self.layer3 = self._make_layer(ResBlock, blocks=6, planes=256, stride=2)
+        self.layer4 = self._make_layer(ResBlock, blocks=3, planes=512, stride=1, last=True)
+
+        self.gem_pool = GeM()
+
+        self.bn_fc = nn.BatchNorm1d(emb_dim)
+        self.fc = nn.Linear(emb_dim, num_classes, bias=False)
+        nn.init.kaiming_normal_(self.fc.weight)
+
+    def _make_layer(self, ResBlock: Bottleneck, blocks: int, planes: int, stride: int = 1, last: bool = False):
+        downsample = None
+        if stride != 1 or self.in_channels != planes * ResBlock.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.in_channels, planes * ResBlock.expansion, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(planes * ResBlock.expansion),
+            )
+        layers = []
+        layers.append(
+            ResBlock(in_channels=self.in_channels, out_channels=planes, stride=stride, downsample=downsample, last=last)
+        )
+        self.in_channels = planes * ResBlock.expansion
+        for _ in range(1, blocks):
+            layers.append(ResBlock(in_channels=self.in_channels, out_channels=planes, last=last))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x: torch.Tensor):
+
+        x = self.cqt(x)
+        # Time-strech requires complex tensors, that's why cqt function returns complex
+        x = torch.view_as_complex(x)
+        if self.tempo_factors is not None:
+            rate = abs(self.tempo_factors[1] - self.tempo_factors[0]) * torch.rand(1).item() + min(self.tempo_factors)
+            strech = (
+                abs(1 - rate) > 7e-2
+            )  # if the strech ratio is too close to 1 (i.e. 0.93 < ratio < 1.07), skip time streching
+            if self.training and strech:
+                x = self.time_strech(x, rate)
+        # Compress the magnitude of the CQT
+        x = self.compress(torch.abs(x))
+
+        # Unsqueeze to simulate 1-channel image
+        x = self.conv1(x.unsqueeze(1))
+        x = self.batch_norm1(x)
+        x = self.relu(x)
+        x = self.max_pool1(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        f_t = self.gem_pool(x)
+        f_t = torch.flatten(f_t, start_dim=1)
+
+        f_c = self.bn_fc(f_t)
+        cls = self.fc(f_c)
+
+        return dict(f_t=f_t, f_c=f_c, cls=cls)

bytecover/models/train_module.py

@@ -0,0 +1,319 @@
+import logging
+import os
+from copy import deepcopy
+from typing import Dict, List
+
+import numpy as np
+import torch
+import torch.nn as nn
+import wandb
+from tqdm import tqdm, trange
+
+from bytecover.models.data_model import BatchDict, Postfix, TestResults, ValDict
+from bytecover.models.early_stopper import EarlyStopper
+from bytecover.models.modules import Bottleneck, Resnet50
+from bytecover.models.utils import (
+    calculate_ranking_metrics,
+    dataloader_factory,
+    dir_checker,
+    save_best_log,
+    save_logs,
+    save_predictions,
+)
+
+logger: logging.Logger = logging.getLogger()  # The logger used to log output
+
+
+class TrainModule:
+    def __init__(self, config: Dict) -> None:
+        self.config = config
+        self.state = "initializing"
+        self.best_model_path: str = None
+        self.num_classes = self.config["train"]["num_classes"]
+        self.max_len = self.config["train"]["max_seq_len"][0]
+
+        self.model = Resnet50(
+            Bottleneck,
+            num_channels=self.config["num_channels"],
+            num_classes=self.num_classes,
+            compress_ratio=self.config["train"]["compress_ratio"],
+            tempo_factors=self.config["train"]["tempo_factors"],
+        )
+        self.model.to(self.config["device"])
+        if self.config["wandb"]:
+            wandb.watch(self.model)
+        self.postfix: Postfix = {}
+
+        self.triplet_loss = nn.TripletMarginLoss(margin=config["train"]["triplet_margin"])
+        self.cls_loss = nn.CrossEntropyLoss(label_smoothing=config["train"]["smooth_factor"])
+
+        self.early_stop = EarlyStopper(patience=self.config["train"]["patience"])
+        self.optimizer = self.configure_optimizers()
+        if self.config["device"] != "cpu":
+            self.scaler = torch.cuda.amp.GradScaler(enabled=self.config["train"]["mixed_precision"])
+
+    def pipeline(self) -> None:
+        self.config["val"]["output_dir"] = dir_checker(self.config["val"]["output_dir"])
+
+        if self.config["train"]["model_ckpt"] is not None:
+            self.model.load_state_dict(torch.load(self.config["train"]["model_ckpt"]), strict=False)
+            logger.info(f'Model loaded from checkpoint: {self.config["train"]["model_ckpt"]}')
+
+        self.t_loaders = dataloader_factory(config=self.config, data_split="TRAIN")
+        self.v_loader = dataloader_factory(config=self.config, data_split="VAL")[0]
+
+        self.state = "running"
+
+        self.pbar = trange(
+            self.config["train"]["epochs"], disable=(not self.config["progress_bar"]), position=0, leave=True
+        )
+        for epoch in self.pbar:
+            if self.state in ["early_stopped", "interrupted", "finished"]:
+                return
+
+            self.postfix["Epoch"] = epoch
+            self.pbar.set_postfix(self.postfix)
+
+            try:
+                self.train_procedure()
+            except KeyboardInterrupt:
+                logger.warning("\nKeyboard Interrupt detected. Attempting gracefull shutdown...")
+                self.state = "interrupted"
+            except Exception as err:
+                raise (err)
+
+            if self.state == "interrupted":
+                self.validation_procedure()
+                self.pbar.set_postfix(
+                    {k: self.postfix[k] for k in self.postfix.keys() & {"train_loss_step", "mr1", "mAP"}}
+                )
+
+        self.state = "finished"
+
+    def validate(self) -> None:
+        self.v_loader = dataloader_factory(config=self.config, data_split="VAL")[0]
+        self.state = "running"
+        self.validation_procedure()
+        self.state = "finished"
+
+    def test(self) -> None:
+        self.test_loader = dataloader_factory(config=self.config, data_split="TEST")[0]
+        self.test_results: TestResults = {}
+
+        if self.best_model_path is not None:
+            self.model.load_state_dict(torch.load(self.best_model_path), strict=False)
+            print(f"Best model loaded from checkpoint: {self.best_model_path}")
+        elif self.config["test"]["model_ckpt"] is not None:
+            self.model.load_state_dict(torch.load(self.config["test"]["model_ckpt"], map_location='cpu'), strict=False)
+            print(f'Model loaded from checkpoint: {self.config["test"]["model_ckpt"]}')
+        elif self.state == "initializing":
+            print("Warning: Testing with random weights")
+
+        self.state = "running"
+        self.test_procedure()
+        self.state = "finished"
+
+    def train_procedure(self) -> None:
+        self.model.train()
+        pbar_loaders = tqdm(self.t_loaders, disable=(not self.config["progress_bar"]), position=1, leave=False)
+        for _, t_loader in enumerate(pbar_loaders):
+            train_loss_list = []
+            train_cls_loss_list = []
+            train_triplet_loss_list = []
+            self.max_len = t_loader.dataset.max_len
+            pbar_loaders.set_postfix_str(f"max_seq_len={self.max_len}")
+            for step, batch in tqdm(
+                enumerate(t_loader),
+                total=len(t_loader),
+                disable=(not self.config["progress_bar"]),
+                position=2,
+                leave=False,
+            ):
+                train_step = self.training_step(batch)
+                self.postfix["train_loss_step"] = float(f"{train_step['train_loss_step']:.3f}")
+                train_loss_list.append(train_step["train_loss_step"])
+                self.postfix["train_cls_loss_step"] = float(f"{train_step['train_cls_loss']:.3f}")
+                train_cls_loss_list.append(train_step["train_cls_loss"])
+                self.postfix["train_triplet_loss_step"] = float(f"{train_step['train_triplet_loss']:.3f}")
+                train_triplet_loss_list.append(train_step["train_triplet_loss"])
+                self.pbar.set_postfix(
+                    {k: self.postfix[k] for k in self.postfix.keys() & {"train_loss_step", "mr1", "mAP"}}
+                )
+                if self.config["wandb"]:
+                    wandb.log(self.postfix)
+                if step % self.config["train"]["log_steps"] == 0:
+                    save_logs(
+                        dict(
+                            epoch=self.postfix["Epoch"],
+                            seq_len=self.max_len,
+                            step=step,
+                            train_loss_step=f"{train_step['train_loss_step']:.3f}",
+                            train_cls_loss_step=f"{train_step['train_cls_loss']:.3f}",
+                            train_triplet_loss_step=f"{train_step['train_triplet_loss']:.3f}",
+                        ),
+                        output_dir=self.config["val"]["output_dir"],
+                        name="log_steps",
+                    )
+            train_loss = torch.tensor(train_loss_list)
+            train_cls_loss = torch.tensor(train_cls_loss_list)
+            train_triplet_loss = torch.tensor(train_triplet_loss_list)
+            self.postfix["train_loss"] = train_loss.mean().item()
+            self.postfix["train_cls_loss"] = train_cls_loss.mean().item()
+            self.postfix["train_triplet_loss"] = train_triplet_loss.mean().item()
+            if self.config["wandb"]:
+                wandb.log(self.postfix)
+            self.validation_procedure()
+            if self.config["wandb"]:
+                wandb.log(self.postfix)
+            self.overfit_check()
+            self.pbar.set_postfix({k: self.postfix[k] for k in self.postfix.keys() & {"train_loss_step", "mr1", "mAP"}})
+
+    def training_step(self, batch: BatchDict) -> Dict[str, float]:
+        with torch.autocast(
+            device_type=self.config["device"].split(":")[0], enabled=self.config["train"]["mixed_precision"]
+        ):
+            anchor = self.model.forward(batch["anchor"].to(self.config["device"]))
+            positive = self.model.forward(batch["positive"].to(self.config["device"]))
+            negative = self.model.forward(batch["negative"].to(self.config["device"]))
+            l1 = self.triplet_loss(anchor["f_t"], positive["f_t"], negative["f_t"])
+            labels = nn.functional.one_hot(batch["anchor_label"].long(), num_classes=self.num_classes)
+            l2 = self.cls_loss(anchor["cls"], labels.float().to(self.config["device"]))
+            loss = l1 + l2
+
+        self.optimizer.zero_grad()
+        if self.config["device"] != "cpu":
+            self.scaler.scale(loss).backward()
+            self.scaler.step(self.optimizer)
+            self.scaler.update()
+        else:
+            loss.backward()
+            self.optimizer.step()
+
+        return {"train_loss_step": loss.item(), "train_triplet_loss": l1.item(), "train_cls_loss": l2.item()}
+
+    def validation_procedure(self) -> None:
+        self.model.eval()
+        embeddings: Dict[str, torch.Tensor] = {}
+        for batch in tqdm(self.v_loader, disable=(not self.config["progress_bar"]), position=1, leave=False):
+            val_dict = self.validation_step(batch)
+            if val_dict["f_t"].ndim == 1:
+                val_dict["f_c"] = val_dict["f_c"].unsqueeze(0)
+                val_dict["f_t"] = val_dict["f_t"].unsqueeze(0)
+            for anchor_id, triplet_embedding, embedding in zip(val_dict["anchor_id"], val_dict["f_t"], val_dict["f_c"]):
+                embeddings[anchor_id] = torch.stack([triplet_embedding, embedding])
+
+        val_outputs = self.validation_epoch_end(embeddings)
+        logger.info(
+            f"\n{' Validation Results ':=^50}\n"
+            + "\n".join([f'"{key}": {value}' for key, value in self.postfix.items()])
+            + f"\n{' End of Validation ':=^50}\n"
+        )
+
+        if self.config["val"]["save_val_outputs"]:
+            val_outputs["val_embeddings"] = torch.stack(list(embeddings.values()))[:, 1].numpy()
+            save_predictions(val_outputs, output_dir=self.config["val"]["output_dir"])
+            save_logs(self.postfix, output_dir=self.config["val"]["output_dir"])
+        self.model.train()
+
+    def validation_epoch_end(self, outputs: Dict[str, torch.Tensor]) -> Dict[str, np.ndarray]:
+        val_loss = torch.zeros(len(outputs))
+        pos_ids = []
+        neg_ids = []
+        clique_ids = []
+        for k, (anchor_id, embeddings) in enumerate(outputs.items()):
+            clique_id, pos_id, neg_id = self.v_loader.dataset._triplet_sampling(anchor_id)
+            val_loss[k] = self.triplet_loss(embeddings[0], outputs[pos_id][0], outputs[neg_id][0]).item()
+            pos_ids.append(pos_id)
+            neg_ids.append(neg_id)
+            clique_ids.append(clique_id)
+        anchor_ids = np.stack(list(outputs.keys()))
+        preds = torch.stack(list(outputs.values()))[:, 1]
+        self.postfix["val_loss"] = val_loss.mean().item()
+        ranks, average_precisions = calculate_ranking_metrics(embeddings=preds.numpy(), cliques=clique_ids)
+        self.postfix["mr1"] = ranks.mean()
+        self.postfix["mAP"] = average_precisions.mean()
+        return {
+            "triplet_ids": np.stack(list(zip(clique_ids, anchor_ids, pos_ids, neg_ids))),
+            "ranks": ranks,
+            "average_precisions": average_precisions,
+        }
+
+    def validation_step(self, batch: BatchDict) -> ValDict:
+        anchor_id = batch["anchor_id"]
+        positive_id = batch["positive_id"]
+        negative_id = batch["negative_id"]
+
+        features = self.model.forward(batch["anchor"].to(self.config["device"]))
+
+        return {
+            "anchor_id": anchor_id,
+            "positive_id": positive_id,
+            "negative_id": negative_id,
+            "f_t": features["f_t"].squeeze(0).detach().cpu(),
+            "f_c": features["f_c"].squeeze(0).detach().cpu(),
+        }
+
+    def test_procedure(self) -> None:
+        self.model.eval()
+        clique_ids = []
+        embeddings: Dict[str, torch.Tensor] = {}
+        for batch in tqdm(self.test_loader, disable=(not self.config["progress_bar"])):
+            clique_ids_batch = self.test_loader.dataset.labels.loc[batch["anchor_id"], "clique"]
+            test_dict = self.validation_step(batch)
+            if test_dict["f_c"].ndim == 1:
+                test_dict["f_c"] = test_dict["f_c"].unsqueeze(0)
+            for anchor_id, clique_id, embedding in zip(test_dict["anchor_id"], clique_ids_batch, test_dict["f_c"]):
+                embeddings[anchor_id] = embedding
+                clique_ids.append(clique_id)
+
+        test_outputs = self.test_epoch_end(embeddings, clique_ids)
+        logger.info(
+            f"\n{' Test Results ':=^50}\n"
+            + "\n".join([f'"{key}": {value}' for key, value in self.test_results.items()])
+            + f"\n{' End of Testing ':=^50}\n"
+        )
+
+        if self.config["test"]["save_test_outputs"]:
+            test_outputs["test_embeddings"] = torch.stack(list(embeddings.values())).numpy()
+            save_predictions(test_outputs, output_dir=self.config["test"]["output_dir"])
+            save_logs(self.test_results, output_dir=self.config["test"]["output_dir"])
+
+    def test_epoch_end(self, outputs: Dict[str, torch.Tensor], clique_ids: List[int]) -> Dict[str, np.ndarray]:
+        anchor_ids = np.stack(list(outputs.keys()))
+        preds = torch.stack(list(outputs.values()))
+        ranks, average_precisions = calculate_ranking_metrics(embeddings=preds.numpy(), cliques=clique_ids)
+        self.test_results["test_mr1"] = ranks.mean()
+        self.test_results["test_mAP"] = average_precisions.mean()
+        return {
+            "anchor_ids": np.stack(list(zip(clique_ids, anchor_ids))),
+            "ranks": ranks,
+            "average_precisions": average_precisions,
+        }
+
+    def overfit_check(self) -> None:
+        if self.early_stop(self.postfix["val_loss"]):
+            logger.info(f"\nValidation not improved for {self.early_stop.patience} consecutive epochs. Stopping...")
+            self.state = "early_stopped"
+
+        if self.early_stop.counter > 0:
+            logger.info("\nValidation loss was not improved")
+        else:
+            logger.info(f"\nMetric improved. New best score: {self.early_stop.min_validation_loss:.3f}")
+            save_best_log(self.postfix, output_dir=self.config["val"]["output_dir"])
+
+            logger.info("Saving model...")
+            epoch = self.postfix["Epoch"]
+            max_secs = self.max_len
+            prev_model = deepcopy(self.best_model_path)
+            self.best_model_path = os.path.join(
+                self.config["val"]["output_dir"], "model", f"best-model-{epoch=}-{max_secs=}.pt"
+            )
+            os.makedirs(os.path.dirname(self.best_model_path), exist_ok=True)
+            torch.save(deepcopy(self.model.state_dict()), self.best_model_path)
+            if prev_model is not None:
+                os.remove(prev_model)
+
+    def configure_optimizers(self) -> torch.optim.Optimizer:
+        optimizer = torch.optim.Adam(self.model.parameters(), lr=self.config["train"]["learning_rate"])
+
+        return optimizer

bytecover/models/utils.py

@@ -0,0 +1,119 @@
+import glob
+import json
+import os
+import re
+from typing import Dict, List, Tuple
+
+import jsonlines
+import numpy as np
+import pandas as pd
+from sklearn.metrics import pairwise_distances
+from torch.utils.data import DataLoader
+
+from bytecover.models.data_loader import bytecover_dataloader
+from bytecover.models.data_model import Postfix
+
+
+def dataloader_factory(config: Dict, data_split: str) -> List[DataLoader]:
+    seq_len_key = "max_seq_len"
+
+    if data_split == "TRAIN":
+        t_loaders = []
+        for L in config["train"][seq_len_key]:
+            t_loaders.append(
+                bytecover_dataloader(
+                    data_path=config["data_path"],
+                    file_ext=config["file_extension"],
+                    dataset_path=config["dataset_path"],
+                    data_split=data_split,
+                    debug=config["debug"],
+                    max_len=L,
+                    **config["train"],
+                )
+            )
+        return t_loaders
+    L = config[data_split.lower()][seq_len_key]
+    return [
+        bytecover_dataloader(
+            data_path=config["data_path"],
+            file_ext=config["file_extension"],
+            dataset_path=config["dataset_path"],
+            data_split=data_split,
+            debug=config["debug"],
+            max_len=L,
+            **config[data_split.lower()],
+        )
+    ]
+
+
+def validation_triplet_sampling(anchor_id: str, val_ids: List[str], df: pd.DataFrame) -> Dict[str, int]:
+    np.random.shuffle(df.loc[anchor_id, "versions"])
+    pos_list = np.setdiff1d(df.loc[anchor_id, "versions"], anchor_id)
+    pos_id = np.random.choice(pos_list, 1)[0]
+    pos_id = val_ids.index(pos_id)
+
+    neg_id = df.loc[~df.index.isin([anchor_id] + list(pos_list))].sample(1).index[0]
+    neg_id = val_ids.index(neg_id)
+
+    return dict(pos_id=pos_id, neg_id=neg_id)
+
+
+def calculate_ranking_metrics(embeddings: np.ndarray, cliques: List[int]) -> Tuple[np.ndarray, np.ndarray]:
+    distances = pairwise_distances(embeddings)
+    s_distances = np.argsort(distances, axis=1)
+    cliques = np.array(cliques)
+    query_cliques = cliques[s_distances[:, 0]]
+    search_cliques = cliques[s_distances[:, 1:]]
+
+    query_cliques = np.tile(query_cliques, (search_cliques.shape[-1], 1)).T
+    mask = np.equal(search_cliques, query_cliques)
+
+    ranks = mask.argmax(axis=1)
+
+    cumsum = np.cumsum(mask, axis=1)
+    mask2 = mask * cumsum
+    mask2 = mask2 / np.arange(1, mask2.shape[-1] + 1)
+    average_precisions = np.sum(mask2, axis=1) / np.sum(mask, axis=1)
+
+    return (ranks, average_precisions)
+
+
+def dir_checker(output_dir: str) -> str:
+    output_dir = re.sub(r"run-[0-9]+/*", "", output_dir)
+    runs = glob.glob(os.path.join(output_dir, "run-*"))
+    if runs != []:
+        max_run = max(map(lambda x: int(x.split("-")[-1]), runs))
+        run = max_run + 1
+    else:
+        run = 0
+    outdir = os.path.join(output_dir, f"run-{run}")
+    return outdir
+
+
+def save_predictions(outputs: Dict[str, np.ndarray], output_dir: str) -> None:
+    os.makedirs(output_dir, exist_ok=True)
+    for key in outputs:
+        if "_ids" in key:
+            with jsonlines.open(os.path.join(output_dir, f"{key}.jsonl"), "w") as f:
+                if len(outputs[key][0]) == 4:
+                    for clique, anchor, pos, neg in outputs[key]:
+                        f.write({"clique_id": clique, "anchor_id": anchor, "positive_id": pos, "negative_id": neg})
+                else:
+                    for clique, anchor in outputs[key]:
+                        f.write({"clique_id": clique, "anchor_id": anchor})
+        else:
+            np.save(os.path.join(output_dir, f"{key}.npy"), outputs[key])
+
+
+def save_logs(outputs: dict, output_dir: str, name: str = "log") -> None:
+    os.makedirs(output_dir, exist_ok=True)
+    log_file = os.path.join(output_dir, f"{name}.jsonl")
+    with jsonlines.open(log_file, "a") as f:
+        f.write(outputs)
+
+
+def save_best_log(outputs: Postfix, output_dir: str) -> None:
+    os.makedirs(output_dir, exist_ok=True)
+    log_file = os.path.join(output_dir, "best-log.json")
+    with open(log_file, "w") as f:
+        json.dump(outputs, f, indent=2)

bytecover/utils.py

@@ -0,0 +1,36 @@
+from dataclasses import dataclass
+from logging import config as logging_config
+from typing import Dict
+
+from yaml import FullLoader, load, safe_load
+
+
+@dataclass
+class bcolors:
+    OKGREEN: str = "\033[92m"
+    WARNING: str = "\033[93m"
+    FAIL: str = "\033[91m"
+    ENDC: str = "\033[0m"
+
+
+def load_config(config_path: str) -> Dict:
+    with open(config_path) as file:
+        config = safe_load(file)
+
+    if config["device"] == "gpu":
+        config["device"] = "cuda:0"
+
+    return config
+
+
+def initialize_logging(config_path: str, debug: bool) -> None:
+    """
+    Setup logging according to the configuration in the given file.
+    :param str config_path: The path to the file containing the logging configuration
+    :return:
+    """
+    with open(config_path) as yaml_fh:
+        config_description = load(yaml_fh, Loader=FullLoader)
+        if debug:
+            config_description["root"]["level"] = "DEBUG"
+        logging_config.dictConfig(config_description)

orfium-bytecover.pt

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

pinecone_generate.py

@@ -0,0 +1,275 @@
+# Generate ByteCover and CLAP Embeddings for a dataset and put to Pinecone
+
+import argparse
+import os
+from typing import Iterator
+import time
+
+import numpy as np
+import pandas as pd
+import torch
+import laion_clap
+from tqdm import tqdm
+from pinecone.grpc import PineconeGRPC as Pinecone
+from pinecone import PodSpec, PineconeApiException
+import ffmpeg
+
+from bytecover.models.train_module import TrainModule
+from bytecover.models.data_loader import ByteCoverDataset
+from bytecover.utils import load_config
+
+class BatchGenerator:
+    #
+    def __init__(self, batch_size: int = 10) -> None:
+        self.batch_size = batch_size
+    #
+    # Makes chunks out of an input DataFrame
+    def to_batches(self, df: pd.DataFrame) -> Iterator[pd.DataFrame]:
+        splits = self.splits_num(df.shape[0])
+        if splits <= 1:
+            yield df
+        else:
+            for chunk in np.array_split(df, splits):
+                yield chunk
+    #
+    # Determines how many chunks DataFrame contains
+    def splits_num(self, elements: int) -> int:
+        return round(elements / self.batch_size)
+    #
+    __call__ = to_batches
+
+# quantization
+def int16_to_float32(x):
+    return (x / 32767.0).astype(np.float32)
+
+def float32_to_int16(x):
+    x = np.clip(x, a_min=-1., a_max=1.)
+    return (x * 32767.).astype(np.int16)
+
+def flatten_vector_embed(vector_embed):
+    return list(vector_embed.flatten())
+
+def grab_song_title(vector_name):
+    return vector_name.split("_")[0]
+
+def convert_to_npfloat64(original_array):
+    #return np.array(flat_df["flat_vector_embed"][0],dtype=np.float64)
+    return np.array(original_array,dtype=np.float64)
+
+def convert_to_npfloat64_to_list(vector_embed_64):
+    # list(flat_df["flat_vector_embed_64"][0])
+    return list(vector_embed_64)
+
+def look_up_metadata(track_id, meta_dataframe, meta_col_interest):
+    # track_id: form = spotify:track:id_##,mp3
+    # meta_datframe: df of all the metavalues
+    # column options = album, artist_names, popularity, release_date, genre
+    df_id = track_id.split("_")[0]
+    meta_row = meta_dataframe[meta_dataframe['uri'] == df_id].reset_index(drop=True)
+    try: 
+        return meta_row[meta_col_interest][0]
+    except:
+        return "unknown"
+    #return meta_row[meta_col_interest][0]
+
+def strip_year_from_date(full_date):
+    if type(full_date) == int:
+        return str(full_date)
+    else:
+        try:
+            return full_date[:4]
+        except:
+            return "CHECK_THIS"
+
+def strip_vector_clip(vector_name):
+    return vector_name.split(".")[0].split("_")[1]
+
+def get_triplet_num(vector_name_str):
+    return str(int(vector_name_str.split("_")[2].split(".")[0]) + 1)
+
+def generate(audio_dir, metadata_dir, index_naming_conv):
+
+    # FILE AND METADATA LOADING
+
+    file_list = [f for f in os.listdir(audio_dir)]
+    print(f"Found {len(file_list)} files")
+
+    meta_list = [f for f in os.listdir(metadata_dir)]
+    meta_list = sorted(meta_list)
+
+    meta_df = pd.read_json(metadata_dir + "/" + meta_list[0])
+
+    for i in range(1, len(meta_list)-1):
+        new_row = pd.read_json(metadata_dir + "/" + meta_list[i])
+        meta_df = pd.concat([meta_df, new_row]).reset_index(drop = True)
+    
+    meta_df["year"] = meta_df.apply(lambda row: strip_year_from_date(row['release_date']),axis=1)
+
+    # BYTECOVER MODEL INITIALIZATION
+            
+    print("Loading ByteCover model")
+
+    bytecover_config = load_config(config_path="bytecover/config.yaml")
+    bytecover_module = TrainModule(bytecover_config)
+    bytecover_model = bytecover_module.model
+    if bytecover_module.best_model_path is not None:
+        bytecover_model.load_state_dict(torch.load(bytecover_module.best_model_path), strict=False)
+        print(f"Best model loaded from checkpoint: {bytecover_module.best_model_path}")
+    elif bytecover_module.config["test"]["model_ckpt"] is not None:
+        bytecover_model.load_state_dict(torch.load(bytecover_module.config["test"]["model_ckpt"], map_location='cpu'), strict=False)
+        print(f'Model loaded from checkpoint: {bytecover_module.config["test"]["model_ckpt"]}')
+    elif bytecover_module.state == "initializing":
+        print("Warning: Running with random weights")
+
+    bytecover_model.eval()
+
+    # BYTECOVER EMBEDDING GENERATION
+
+    audio_dict_bytecover = {}
+    for file in tqdm(file_list, desc="Generating Bytecover Embeddings"):
+        file_path = audio_dir + file
+        # try statement here allows you to skip to items you haven't yet embedded if you stop this step midway (if a key exists, you move on to next key)
+        try:
+            audio_dict_bytecover[file]
+        except:
+            # Load audio
+            try:
+                # This launches a subprocess to decode audio while down-mixing and resampling as necessary.
+                # Requires the ffmpeg CLI and `ffmpeg-python` package to be installed.
+                out, _ = (
+                    ffmpeg.input(file_path, threads=0)
+                    .output("-", format="s16le", acodec="pcm_s16le", ac=1, ar=22050)
+                    .run(cmd=["ffmpeg", "-nostdin"], capture_stdout=True, capture_stderr=True)
+                )
+            except ffmpeg.Error as e:
+                raise RuntimeError(
+                    f"Failed to load audio:{file_path}\n{e.stderr.decode()}"
+                ) from e
+            audio = np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0
+            song_tensor = torch.from_numpy(audio)
+            # this step grabs a ByteCover embedding
+            audio_embed = bytecover_model.forward(song_tensor.to(bytecover_module.config["device"]))['f_t'].detach()
+            audio_dict_bytecover[file] = audio_embed.squeeze()
+
+    # CLAP MODEL INITIALIZATION
+
+    print("Loading CLAP model")
+
+    clap_model = laion_clap.CLAP_Module(enable_fusion=False)
+    clap_model.load_ckpt() # download the default pretrained checkpoint.
+
+    # CLAP EMBEDDING GENERATION
+
+    audio_dict_CLAP = {}
+    for file in tqdm(file_list, desc="Generating CLAP Embeddings"):
+        # try statement here allows you to skip to items you haven't yet embedded if you stop this step midway (if a key exists, you move on to next key)
+        try:
+            audio_dict_CLAP[file]
+        except:
+            # Get audio embeddings from audio data
+            full_path = audio_dir + "/" + file
+            # this step grabs a CLAP embedding from laion_clap library
+            audio_embed = clap_model.get_audio_embedding_from_filelist(x = [full_path], use_tensor=False)
+            audio_dict_CLAP[file] = audio_embed
+
+    # DATAFRAME GENERATION
+            
+    flat_dfs = []
+            
+    for audio_dict in [audio_dict_CLAP, audio_dict_bytecover]:
+
+        flat_df = pd.DataFrame(audio_dict.items(), columns=['vector_name','vector_embed']).reset_index()
+        flat_df.columns=['vector_id','vector_name','vector_embed']
+
+        flat_df["song_title"] = flat_df.apply(lambda row: grab_song_title(row['vector_name']),axis=1)
+        flat_df["flat_vector_embed"] = flat_df.apply(lambda row: flatten_vector_embed(row['vector_embed']),axis=1)
+
+        flat_df["flat_vector_embed_64"] = flat_df.apply(lambda row: convert_to_npfloat64(row['flat_vector_embed']),axis=1)
+        flat_df["flat_vector_embed_64_list"] = flat_df.apply(lambda row: convert_to_npfloat64_to_list(row['flat_vector_embed_64']),axis=1)
+
+        flat_df["genre"] = flat_df.apply(lambda row: look_up_metadata(row['vector_name'], meta_df, 'genre'),axis=1)
+        flat_df["album"] = flat_df.apply(lambda row: look_up_metadata(row['vector_name'], meta_df, 'album'),axis=1)
+        flat_df["name"] = flat_df.apply(lambda row: look_up_metadata(row['vector_name'], meta_df, 'name'),axis=1)
+        flat_df["artist"] = flat_df.apply(lambda row: look_up_metadata(row['vector_name'], meta_df, 'artist_names'),axis=1)
+        flat_df["year"] = flat_df.apply(lambda row: look_up_metadata(row['vector_name'], meta_df, 'year'),axis=1)
+        flat_df["vector_clip_num"] = flat_df.apply(lambda row: strip_vector_clip(row['vector_name']),axis=1)
+        flat_df['embedding_triplet_num'] = flat_df.vector_name.apply(get_triplet_num)
+
+        flat_dfs.append(flat_df)
+
+        print("unique songs:", len(flat_df.song_title.unique()))
+
+    # PINECONE UPLOAD
+
+    api_key = os.environ['PC_API_KEY']
+    pc = Pinecone(api_key=api_key)
+
+    index_name_clap = f'clap-{index_naming_conv}' # free (comes with plan, can have 100k records)
+    index_name_bytecover = f'bytecover-{index_naming_conv}' # free (comes with plan, can have 100k records)
+    index_env = 'us-west1-gcp' # NOT free (take down when not in use)
+    pod_type = 'p1.x1' # NOT free (take down when not in use)
+
+    for index_name, flat_df, index_dim in zip([index_name_clap, index_name_bytecover], flat_dfs, [512, 2048]):
+        
+        try:
+            pc.create_index(
+                name=index_name,
+                dimension=index_dim,
+                metric="cosine",
+                spec=PodSpec(
+                    environment=index_env,
+                    pod_type=pod_type,
+                    pods=1
+                ),
+                deletion_protection="disabled"
+            )
+        except PineconeApiException:
+            print(f"WARNING: INDEX {index_name} ALREADY EXISTS")
+            time.sleep(5)
+
+        index = pc.Index(index_name)
+
+        batch_id = 0
+
+        df_batcher = BatchGenerator(64)
+
+        for batch_df in tqdm(df_batcher(flat_df), desc="Uploading batches"):
+            #print(batch_df)
+            batch_id = batch_id + 1
+            index.upsert(vectors=list(zip(batch_df["vector_name"],batch_df["flat_vector_embed_64_list"])))
+
+        failed_list_update_metadata = []
+
+        for vec_id in tqdm(range(0,len(flat_df)), desc="Adding metadata"):
+            try:
+                row = flat_df.iloc[vec_id]
+                index.update(id=str(row['vector_name']),
+                                set_metadata={"genre": row['genre'],
+                                            "song" : row['name'],
+                                            "album": row['album'],
+                                            "artists": row['artist'],
+                                            "year" : str(row['year']),
+                                            "clip_num" : row['vector_clip_num'],
+                                            "triplet_num": str(row['embedding_triplet_num']),
+                                            "spotify_id" : row['song_title']
+                                        })
+            except:
+                print("failed on:", vec_id)
+                failed_list_update_metadata.append(vec_id)
+
+        pc.create_collection(index_name, index_name)
+        
+
+
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Generate ByteCover and CLAP Embeddings for a dataset and put to Pinecone")
+
+    parser.add_argument('audio_dir')
+    parser.add_argument('metadata_dir')
+    parser.add_argument('index_name')
+
+    args = parser.parse_args()
+
+    generate(args.audio_dir, args.metadata_dir, args.index_name)

requirements.txt

@@ -5,6 +5,24 @@ loralib
 wavebeat @ git+https://github.com/hugofloresgarcia/wavebeat
 lac @ git+https://github.com/hugofloresgarcia/lac.git
 descript-audiotools @ git+https://github.com/hugofloresgarcia/audiotools.git
--e git+https://github.com/audacitorch/pyharp.git@develop#egg=pyharp
 torch_pitch_shift
-gradio
+gradio
+
+--extra-index-url https://download.pytorch.org/whl/cu113
+-e git+https://github.com/TEAMuP-dev/pyharp.git@np/overlaycolor#egg=pyharp
+pinecone
+# towhee
+av
+laion_clap
+# pyshorteners
+nnAudio
+ffmpeg-python
+torchvision
+torch
+jsonlines
+wandb
+tqdm
+# For pinecone_generate only
+google-api-python-client
+protoc-gen-openapiv2-protoc3-19
+transformers==4.30.0