cryceleb_titanet.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OpY2Tm59VWmD"
      },
      "outputs": [],
      "source": [
        "!pip install speechbrain\n",
        "!pip install online_triplet_loss\n",
        "!pip install nemo_toolkit[\"all\"]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SvgXw0UiVf0R"
      },
      "outputs": [],
      "source": [
        "from huggingface_hub import notebook_login\n",
        "\n",
        "notebook_login()"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import random\n",
        "import zipfile\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import seaborn as sns\n",
        "import torch\n",
        "from huggingface_hub import hf_hub_download\n",
        "from nemo.collections import asr as nemo_asr\n",
        "from online_triplet_loss.losses import *\n",
        "from speechbrain.dataio.dataio import read_audio\n",
        "from speechbrain.nnet.losses import AdditiveAngularMargin, LogSoftmaxWrapper\n",
        "from speechbrain.utils.metric_stats import EER\n",
        "from torch import nn, optim\n",
        "from torch.nn import functional as F\n",
        "from torch.utils.data import DataLoader, Dataset\n",
        "from tqdm import tqdm"
      ],
      "metadata": {
        "id": "XyYO2MEnB52J"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0jkK1YasVZcm"
      },
      "outputs": [],
      "source": [
        "class StreamingDataset(Dataset):\n",
        "    def __init__(self, samples, targets, cut_interval=None, num_copies=1):\n",
        "        self.samples = [torch.as_tensor(x) for x in samples]\n",
        "        self.targets = targets\n",
        "        self.cut_interval = cut_interval\n",
        "        self.num_copies = num_copies\n",
        "        self.targets = [torch.as_tensor(x) for x in targets]\n",
        "        self.samples, self.targets = zip(\n",
        "            *sorted(zip(self.samples, self.targets), key=lambda pair: len(pair[0]))\n",
        "        )\n",
        "        if cut_interval is not None:\n",
        "            tmp_samples = []\n",
        "            tmp_targets = []\n",
        "            for sample, target in zip(self.samples, self.targets):\n",
        "                for i in range(num_copies):\n",
        "                    tmp_samples.append(sample)\n",
        "                    tmp_targets.append(target)\n",
        "            self.samples = tmp_samples\n",
        "            self.targets = tmp_targets\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "        sample = self.samples[index]\n",
        "        target = self.targets[index]\n",
        "        if self.cut_interval is not None:\n",
        "            duration = round(random.uniform(*self.cut_interval) * 16000)\n",
        "            start_idx = random.randint(0, max(0, len(sample) - duration))\n",
        "            end_idx = start_idx + duration\n",
        "            sample = sample[start_idx:end_idx]\n",
        "        return sample, target\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.samples)\n",
        "\n",
        "\n",
        "def download_data(dest=\"data\"):\n",
        "    if os.path.exists(os.path.join(dest, \"audio\", \"train\")):\n",
        "        print(\n",
        "            f\"It appears that data is already downloaded. \\n\"\n",
        "            f\"If you think it should be re-downloaded, remove {dest} directory and re-run\"\n",
        "        )\n",
        "        return\n",
        "\n",
        "    # Download data from Huggingface\n",
        "    for file_name in [\n",
        "        \"metadata.csv\",\n",
        "        \"audio.zip\",\n",
        "        \"dev_pairs.csv\",\n",
        "        \"test_pairs.csv\",\n",
        "        \"sample_submission.csv\",\n",
        "    ]:\n",
        "        hf_hub_download(\n",
        "            repo_id=\"Ubenwa/CryCeleb2023\",\n",
        "            filename=file_name,\n",
        "            local_dir=dest,\n",
        "            repo_type=\"dataset\",\n",
        "        )\n",
        "\n",
        "    with zipfile.ZipFile(os.path.join(dest, \"audio.zip\"), \"r\") as zip_ref:\n",
        "        zip_ref.extractall(dest)\n",
        "\n",
        "    print(f\"Data downloaded to {dest}/ directory\")\n",
        "\n",
        "\n",
        "def get_baby_ids_with_both_periods(manifest_df):\n",
        "    count_of_periods_per_baby = manifest_df.groupby(\"baby_id\")[\"period\"].count()\n",
        "    baby_ids_with_both_periods = count_of_periods_per_baby[\n",
        "        count_of_periods_per_baby == 2\n",
        "    ].index\n",
        "    return baby_ids_with_both_periods\n",
        "\n",
        "\n",
        "def collate_with_padding(batch):\n",
        "    batch = (batch,) if not isinstance(batch[0], tuple) else tuple(zip(*batch))\n",
        "    samples = batch[0]\n",
        "    lengths = torch.as_tensor([x.shape[0] for x in samples])\n",
        "    max_length = max(lengths)\n",
        "    lengths = lengths / max_length\n",
        "    samples = [F.pad(x, [0, max_length - x.shape[0]]) for x in samples]\n",
        "    return torch.stack(samples), lengths, *(torch.stack(x) for x in batch[1:])\n",
        "\n",
        "\n",
        "def compute_cosine_similarity_score(row, cry_dict):\n",
        "    cos = torch.nn.CosineSimilarity(dim=-1)\n",
        "    similarity_score = cos(\n",
        "        cry_dict[(row[\"baby_id_B\"], \"B\")][\"cry_encoded\"],\n",
        "        cry_dict[(row[\"baby_id_D\"], \"D\")][\"cry_encoded\"],\n",
        "    )\n",
        "    return similarity_score.item()\n",
        "\n",
        "\n",
        "def compute_eer_and_plot_verification_scores(pairs_df):\n",
        "    # pairs_df must have `score` and `label` columns\n",
        "    positive_scores = pairs_df.loc[pairs_df[\"label\"] == 1][\"score\"].values\n",
        "    negative_scores = pairs_df.loc[pairs_df[\"label\"] == 0][\"score\"].values\n",
        "    eer, threshold = EER(torch.tensor(positive_scores), torch.tensor(negative_scores))\n",
        "    ax = sns.histplot(pairs_df, x=\"score\", hue=\"label\", stat=\"percent\", common_norm=False)\n",
        "    ax.set_title(f\"EER={round(eer, 4)} - Thresh={round(threshold, 4)}\")\n",
        "    plt.axvline(x=[threshold], color=\"red\", ls=\"--\")\n",
        "    plt.show()\n",
        "    return eer, threshold"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Download data\n",
        "dataset_path = \"data\"\n",
        "download_data(dataset_path)\n",
        "\n",
        "# Read manifest\n",
        "metadata = pd.read_csv(\n",
        "    f\"{dataset_path}/metadata.csv\", dtype={\"baby_id\": str, \"chronological_index\": str}\n",
        ")\n",
        "\n",
        "# Load train data\n",
        "train_metadata = metadata.loc[metadata[\"split\"] == \"train\"].copy()\n",
        "train_metadata[\"cry\"] = train_metadata.apply(\n",
        "    lambda row: read_audio(f'{dataset_path}/{row[\"file_name\"]}').numpy(), axis=1\n",
        ")\n",
        "# Concatenate all segments for each (baby_id, period) group\n",
        "manifest_df = pd.DataFrame(\n",
        "    train_metadata.groupby([\"baby_id\", \"period\"])[\"cry\"].agg(lambda x: np.concatenate(x.values)),\n",
        "    columns=[\"cry\"],\n",
        ").reset_index()\n",
        "\n",
        "# Load dev data\n",
        "dev_metadata = metadata.loc[metadata[\"split\"] == \"dev\"].copy()\n",
        "dev_pairs = pd.read_csv(\n",
        "    f\"{dataset_path}/dev_pairs.csv\", dtype={\"baby_id_B\": str, \"baby_id_D\": str}\n",
        ")\n",
        "dev_metadata[\"cry\"] = dev_metadata.apply(\n",
        "    lambda row: read_audio(f'{dataset_path}/{row[\"file_name\"]}').numpy(), axis=1\n",
        ")\n",
        "# Concatenate all segments for each (baby_id, period) group\n",
        "cry_dict = pd.DataFrame(\n",
        "    dev_metadata.groupby([\"baby_id\", \"period\"])[\"cry\"].agg(lambda x: np.concatenate(x.values)),\n",
        "    columns=[\"cry\"],\n",
        ").to_dict(orient=\"index\")\n",
        "\n",
        "# Load test data\n",
        "test_metadata = metadata.loc[metadata[\"split\"] == \"test\"].copy()\n",
        "test_pairs = pd.read_csv(f\"{dataset_path}/test_pairs.csv\")\n",
        "test_metadata[\"cry\"] = test_metadata.apply(\n",
        "    lambda row: read_audio(f'{dataset_path}/{row[\"file_name\"]}').numpy(), axis=1\n",
        ")\n",
        "# Concatenate all segments for each (baby_id, period) group\n",
        "cry_dict_test = pd.DataFrame(\n",
        "    test_metadata.groupby([\"baby_id\", \"period\"])[\"cry\"].agg(lambda x: np.concatenate(x.values)),\n",
        "    columns=[\"cry\"],\n",
        ").to_dict(orient=\"index\")"
      ],
      "metadata": {
        "id": "onzErscJNQy3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "CUT_INTERVAL = (3, 5)\n",
        "TRAIN_BATCH_SIZE = 32  #160\n",
        "NUM_EPOCHS = 200\n",
        "LR = 1e-4\n",
        "MARGIN = 0.2  # 0.5\n",
        "SCALE = 30\n",
        "EVAL_FREQ = 1\n",
        "\n",
        "\n",
        "# Define device\n",
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "\n",
        "# Define dataloaders\n",
        "baby_ids_with_both_periods = get_baby_ids_with_both_periods(manifest_df)\n",
        "manifest_df = manifest_df.loc[manifest_df[\"baby_id\"].isin(baby_ids_with_both_periods)]\n",
        "all_targets = list([x for x in manifest_df[\"baby_id\"]])\n",
        "id2label = {v: k for k, v in enumerate(np.unique(all_targets))}\n",
        "train_samples = []\n",
        "train_targets = []\n",
        "for cry, baby_id, period in zip(manifest_df[\"cry\"], manifest_df[\"baby_id\"], manifest_df[\"period\"]):\n",
        "    if period == \"B\":\n",
        "        train_samples.append(cry)\n",
        "        train_targets.append(baby_id)\n",
        "train_targets = [torch.as_tensor(id2label[x]) for x in train_targets]\n",
        "train_data = StreamingDataset(train_samples, train_targets, CUT_INTERVAL)\n",
        "train_dataloader = DataLoader(\n",
        "    train_data, TRAIN_BATCH_SIZE, collate_fn=collate_with_padding, drop_last=True, shuffle=False,\n",
        ")\n",
        "\n",
        "# Define model\n",
        "encoder = nemo_asr.models.EncDecSpeakerLabelModel.from_pretrained(\n",
        "    \"nvidia/speakerverification_en_titanet_large\"\n",
        ")\n",
        "encoder.requires_grad_(True)\n",
        "classifier = nn.Linear(192, len(id2label))\n",
        "model = nn.ModuleList([encoder, classifier]).to(device)\n",
        "\n",
        "# Load checkpoint\n",
        "checkpoint_path = None #os.path.join(\"checkpoints\", \"titanet_checkpoint_epoch=00_eer=0.41987180709838867.pt\")\n",
        "if checkpoint_path is not None:\n",
        "    checkpoint = torch.load(checkpoint_path)\n",
        "    model.load_state_dict(checkpoint)\n",
        "\n",
        "# Define optimizer\n",
        "optimizer = optim.AdamW(model.parameters(), lr=LR)\n",
        "\n",
        "# Margin loss\n",
        "criterion = LogSoftmaxWrapper(AdditiveAngularMargin(margin=MARGIN, scale=SCALE))"
      ],
      "metadata": {
        "id": "-zRbWbrlsHjw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_losses = []\n",
        "for i in range(NUM_EPOCHS):\n",
        "    model.train()\n",
        "    total_loss = 0.\n",
        "    for samples, lengths, targets in train_dataloader:\n",
        "        samples, lengths, targets = samples.to(device), lengths.to(device), targets.to(device)\n",
        "        _, embeddings = encoder(input_signal=samples, input_signal_length=lengths * samples.shape[1])\n",
        "        logits = classifier(embeddings)\n",
        "        loss = criterion(logits[:, None, :], targets[:, None], lengths)\n",
        "        #margin = 0.2\n",
        "        #loss = batch_all_triplet_loss(targets, embeddings, margin=margin)[0]\n",
        "        #loss = batch_hard_triplet_loss(targets, embeddings, margin=margin)\n",
        "        optimizer.zero_grad(set_to_none=True)\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "        total_loss += loss.item()\n",
        "    train_losses.append(total_loss / len(train_dataloader))\n",
        "\n",
        "    if i % EVAL_FREQ == 0:\n",
        "        ##### VALIDATION #####\n",
        "        # Compute embeddings\n",
        "        model.eval()\n",
        "        with torch.no_grad():\n",
        "            for (baby_id, period), d in tqdm(cry_dict.items()):\n",
        "                samples = torch.as_tensor(d[\"cry\"][None], device=device)\n",
        "                lengths = torch.ones(1, device=device)\n",
        "                _, embeddings = encoder(input_signal=samples, input_signal_length=lengths * samples.shape[1])\n",
        "                embedding = embeddings[0].to(\"cpu\")\n",
        "                d[\"cry_encoded\"] = embedding\n",
        "\n",
        "        # Compute scores\n",
        "        dev_pairs[\"score\"] = dev_pairs.apply(\n",
        "            lambda row: compute_cosine_similarity_score(row=row, cry_dict=cry_dict), axis=1\n",
        "        )\n",
        "\n",
        "        eer, threshold = compute_eer_and_plot_verification_scores(pairs_df=dev_pairs)\n",
        "\n",
        "        print(eer, threshold)\n",
        "        ##########\n",
        "\n",
        "        ##### TEST #####\n",
        "        # Compute embeddings\n",
        "        model.eval()\n",
        "        with torch.no_grad():\n",
        "            for (baby_id, period), d in tqdm(cry_dict_test.items()):\n",
        "                samples = torch.as_tensor(d[\"cry\"][None], device=device)\n",
        "                lengths = torch.ones(1, device=device)\n",
        "                _, embeddings = encoder(input_signal=samples, input_signal_length=lengths * samples.shape[1])\n",
        "                embedding = embeddings[0].to(\"cpu\")\n",
        "                d[\"cry_encoded\"] = embedding\n",
        "\n",
        "        # Compute scores\n",
        "        test_pairs[\"score\"] = test_pairs.apply(\n",
        "            lambda row: compute_cosine_similarity_score(row=row, cry_dict=cry_dict_test), axis=1\n",
        "        )\n",
        "\n",
        "        # Write submission file\n",
        "        submission = test_pairs[[\"id\", \"score\"]]\n",
        "        os.makedirs(\"submissions\", exist_ok=True)\n",
        "        submission.to_csv(\n",
        "            os.path.join(\"submissions\", f\"titanet_submission_epoch={i:02d}_eer={eer}.csv\"),\n",
        "            index=False,\n",
        "        )\n",
        "        ##########\n",
        "\n",
        "        # Save model\n",
        "        os.makedirs(\"checkpoints\", exist_ok=True)\n",
        "        torch.save(\n",
        "            model.state_dict(),\n",
        "            os.path.join(\"checkpoints\", f\"titanet_checkpoint_epoch={i:02d}_eer={eer}.pt\")\n",
        "        )\n",
        "\n",
        "# Plot loss\n",
        "plt.plot(train_losses)\n",
        "plt.xlabel(\"Epoch\")\n",
        "plt.ylabel(\"Training loss\")\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "LHDn_Fs7EH7u"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}