{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "703a61c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import torch\n",
    "from suno_utils.utils.text import read_jsonl\n",
    "\n",
    "\n",
    "def apply_specaugment_mask(\n",
    "    latents: np.ndarray,\n",
    "    time_mask_range=(20, 50),\n",
    "    feature_mask_range=(8, 24),\n",
    "    num_time_masks=1,\n",
    "    num_feature_masks=1,\n",
    "    time_mask_prob=0.5,\n",
    "    feature_mask_prob=0.5,\n",
    "    mask_value=0.0,\n",
    "):\n",
    "    \"\"\"\n",
    "    Apply SpecAugment-style masking to VAE latents.\n",
    "\n",
    "    Args:\n",
    "        latents (np.ndarray): Array of shape (T, D).\n",
    "        time_mask_range (tuple): (min, max) width of time masks.\n",
    "        feature_mask_range (tuple): (min, max) width of feature masks.\n",
    "        num_time_masks (int): Number of time masks to attempt.\n",
    "        num_feature_masks (int): Number of feature masks to attempt.\n",
    "        time_mask_prob (float): Probability of applying each time mask.\n",
    "        feature_mask_prob (float): Probability of applying each feature mask.\n",
    "        mask_value (float): Value to use for masking (default: 0.0).\n",
    "    Returns:\n",
    "        np.ndarray: Masked latents.\n",
    "    \"\"\"\n",
    "    T, D = latents.shape\n",
    "    latents = latents.copy()  # avoid modifying original array\n",
    "\n",
    "    # Time masking\n",
    "    for _ in range(num_time_masks):\n",
    "        if np.random.rand() < time_mask_prob:\n",
    "            mask_width = np.random.randint(*time_mask_range)\n",
    "            if T - mask_width > 0:\n",
    "                t = np.random.randint(0, T - mask_width)\n",
    "                latents[t : t + mask_width, :] = mask_value\n",
    "\n",
    "    # Feature masking\n",
    "    for _ in range(num_feature_masks):\n",
    "        if np.random.rand() < feature_mask_prob:\n",
    "            mask_width = np.random.randint(*feature_mask_range)\n",
    "            if D - mask_width > 0:\n",
    "                f = np.random.randint(0, D - mask_width)\n",
    "                latents[:, f : f + mask_width] = mask_value\n",
    "\n",
    "    return latents\n",
    "\n",
    "class RewardModelMemmapDataset(torch.utils.data.Dataset):\n",
    "    def __init__(\n",
    "        self,\n",
    "        dataset_dir,\n",
    "        metas_filename,\n",
    "        vae_memmap_filename,\n",
    "        vae_scale_factor,\n",
    "        mask_prob=0.5,\n",
    "        vae_use_float16=True,\n",
    "        vae_n_tokens=750,\n",
    "        vae_dim=128,\n",
    "    ):\n",
    "        self.metas_filepath = os.path.join(dataset_dir, metas_filename)\n",
    "        self.metas = read_jsonl(self.metas_filepath)\n",
    "        self.vae_scale_factor = vae_scale_factor\n",
    "        self.mask_prob = mask_prob\n",
    "\n",
    "        print(f\"Loaded {len(self.metas)} metas\")\n",
    "\n",
    "        # load the memmap files\n",
    "        vae_data = np.memmap(\n",
    "            os.path.join(dataset_dir, vae_memmap_filename),\n",
    "            dtype=np.float16 if vae_use_float16 else np.float32,\n",
    "            mode=\"r\",\n",
    "        )\n",
    "        vae_data = vae_data.reshape(-1, vae_n_tokens, vae_dim)\n",
    "        self.vae_data = vae_data\n",
    "        print(self.vae_data.shape)\n",
    "\n",
    "        assert len(self.metas) == self.vae_data.shape[0]\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.metas)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "\n",
    "        # only use even indices\n",
    "        # 0, 2, 4, ...\n",
    "        # so we have to convert idx to an even index using modulo\n",
    "        # if idx is odd, we need to subtract 1\n",
    "        if idx % 2 == 1:\n",
    "            idx -= 1\n",
    "\n",
    "        print(idx)\n",
    "        meta = self.metas[idx]\n",
    "        negative_latents = self.vae_data[idx] * self.vae_scale_factor\n",
    "        positive_latents = self.vae_data[idx + 1] * self.vae_scale_factor\n",
    "\n",
    "        if self.mask_prob > 0:\n",
    "            positive_latents = apply_specaugment_mask(\n",
    "                positive_latents,\n",
    "                time_mask_prob=self.mask_prob,\n",
    "                feature_mask_prob=self.mask_prob,\n",
    "            )\n",
    "            negative_latents = apply_specaugment_mask(\n",
    "                negative_latents,\n",
    "                time_mask_prob=self.mask_prob,\n",
    "                feature_mask_prob=self.mask_prob,\n",
    "            )\n",
    "\n",
    "        return torch.from_numpy(positive_latents), torch.from_numpy(negative_latents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45d5c8d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = RewardModelMemmapDataset(\n",
    "    dataset_dir=\"/app2/suno/data/christian/outputs/v3-distill-data-ctx-t1/memmaps/t1_labels_0_6_cut_history_4x\",\n",
    "    metas_filename=\"metas_val.jsonl\",\n",
    "    vae_memmap_filename=\"data_vae_val.bin\",\n",
    "    vae_scale_factor=0.4,\n",
    "    mask_prob=0.5,\n",
    ")\n",
    "\n",
    "pos_vae, neg_vae = dataset[0]\n",
    "\n",
    "print(pos_vae.shape, neg_vae.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9197d5ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "for idx, batch in enumerate(dataset):\n",
    "    print(idx, batch[0].shape, batch[1].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1da053fc",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_diff",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}