{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ratio_mask(tokens, mask_ratio, pad_token=4000):\n",
    "    if not 0 <= mask_ratio <= 1:\n",
    "        raise ValueError(\"mask_ratio must be between 0 and 1.\")\n",
    "\n",
    "    if not isinstance(tokens, torch.Tensor):\n",
    "        tokens = torch.tensor(tokens)\n",
    "\n",
    "    if mask_ratio == 0:\n",
    "        return tokens.clone()\n",
    "\n",
    "    if len(tokens) == 0:\n",
    "        return tokens\n",
    "\n",
    "    step = 1 / mask_ratio\n",
    "    indices = torch.arange(0, len(tokens), step).long()\n",
    "    indices = torch.unique(indices)\n",
    "\n",
    "    masked = tokens.clone()\n",
    "    masked[indices] = pad_token\n",
    "    return masked"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "semantic_codes = torch.ones(25) * 1111#torch.randint(0, 4000, (25,))\n",
    "print(semantic_codes)\n",
    "\n",
    "masked_codes = ratio_mask(semantic_codes, 0.1)\n",
    "print(masked_codes)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class DiffusionGenerationConfig:\n",
    "    semantic_mask_ratio: float = 0.0\n",
    "\n",
    "\n",
    "def ratio_mask_semantic_codes(\n",
    "    semantic_codes_chunk,\n",
    "    start_index: int,\n",
    "    generation_config: DiffusionGenerationConfig,\n",
    "):\n",
    "    if not 0 <= generation_config.semantic_mask_ratio <= 1:\n",
    "        raise ValueError(\"semantic_mask_ratio must be 0 and 1.\")\n",
    "\n",
    "    if not isinstance(semantic_codes_chunk, torch.Tensor):\n",
    "        semantic_codes_chunk = torch.tensor(semantic_codes_chunk)\n",
    "\n",
    "    if generation_config.semantic_mask_ratio == 0:\n",
    "        return semantic_codes_chunk.clone()\n",
    "\n",
    "    if semantic_codes_chunk.shape[1] == 0:\n",
    "        return semantic_codes_chunk\n",
    "\n",
    "    # Use start_index to offset the masking pattern for consistent masking across chunks\n",
    "    step = 1 / generation_config.semantic_mask_ratio\n",
    "    print(\"step\", step)\n",
    "    indices = torch.arange(start_index, start_index + semantic_codes_chunk.shape[1], step).long()\n",
    "    indices = indices[indices < semantic_codes_chunk.shape[1]]  # Ensure indices are within bounds\n",
    "    indices = torch.unique(indices)\n",
    "    print(\"indices\", indices)\n",
    "    masked = semantic_codes_chunk.clone()\n",
    "    masked[indices] = 4000\n",
    "    return masked"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "semantic_codes = torch.ones(25) * 1111#torch.randint(0, 4000, (25,))\n",
    "print(semantic_codes)\n",
    "\n",
    "cfg = DiffusionGenerationConfig()\n",
    "cfg.semantic_mask_ratio = 0.5\n",
    "\n",
    "masked_codes = ratio_mask_semantic_codes(semantic_codes, 12, cfg)\n",
    "print(masked_codes)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a json file for evaluation \n",
    "\n",
    "# get all files in the reference-audio-wav directory\n",
    "import glob\n",
    "import json\n",
    "from suno_utils.utils.text import write_jsonl\n",
    "filepaths = glob.glob(\"/home/christian/audio/reference-audio-wav/*.wav\")\n",
    "print(filepaths)\n",
    "\n",
    "s3_basename = \"s3://suno-data/christian/reference-audio-wav/\"\n",
    "\n",
    "# look for the corresponding txt file in the reference-audio-lyrics directory\n",
    "import os\n",
    "txt_filepaths = glob.glob(\"/home/christian/audio/reference-audio-lyrics/*.txt\")\n",
    "print(txt_filepaths)\n",
    "# create a dict that maps from the filename to the txt file\n",
    "filename_to_txt = {os.path.basename(filepath): filepath for filepath in txt_filepaths}\n",
    "print(filename_to_txt)\n",
    "\n",
    "examples = []\n",
    "# create a json file for evaluation \n",
    "for filepath in filepaths:\n",
    "    print(filepath)\n",
    "    # find corresponding txt file\n",
    "    txt_filename = os.path.basename(filepath).replace(\".wav\", \".txt\")\n",
    "    print(txt_filename)\n",
    "    txt_filepath = filename_to_txt.get(txt_filename, None)\n",
    "    if txt_filepath is None:\n",
    "        print(f\"No txt file found for {filepath}\")\n",
    "        txt = \"[Instrumental]\"\n",
    "    else:\n",
    "        # read the txt file\n",
    "        with open(txt_filepath, \"r\") as f:\n",
    "            txt = f.read()\n",
    "\n",
    "    # lets use the filename to get a safe id\n",
    "    # the id needs to be safe for a url\n",
    "    # so we replace spaces with - and remove any non-alphanumeric characters\n",
    "    import re\n",
    "    clip_id = re.sub(r'[^\\w\\-]', '-', os.path.basename(filepath).replace(\".wav\", \"\")).strip('-')\n",
    "    # Remove consecutive dashes and ensure no leading/trailing dashes\n",
    "    clip_id = re.sub(r'-+', '-', clip_id).strip('-')\n",
    "    print(clip_id)\n",
    "\n",
    "    example = {\n",
    "        \"id\": clip_id,\n",
    "        \"s3_filepath\": s3_basename + os.path.basename(filepath),\n",
    "        \"tags\": [],\n",
    "        \"lyrics\": txt,\n",
    "    }\n",
    "    examples.append(example)\n",
    "\n",
    "# save the examples to a json file\n",
    "write_jsonl(examples, \"/home/christian/audio/reference-audio.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": 2
}