{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 184,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import re\n",
    "import random\n",
    "\n",
    "\n",
    "MAX_TAG_LEN = 256  # characters, not tokens\n",
    "MAX_TOT_TAGS_LEN = 1024  # characters, not tokens\n",
    "MAX_N_TAGS = 10  # to avoid overfitting to an artist\n",
    "\n",
    "CASE_AUGMENT_FUNCS = [\n",
    "    str.upper,\n",
    "    str.lower,\n",
    "    str.capitalize,\n",
    "    str.title,\n",
    "]\n",
    "\n",
    "def _space_repl(m):\n",
    "    s = m.group()\n",
    "    n_newline = s.count(\"\\n\")\n",
    "    if n_newline >= 2:\n",
    "        return \"\\n\\n\"\n",
    "    elif n_newline == 1:\n",
    "        return \"\\n\"\n",
    "    return \" \"\n",
    "\n",
    "def _simplify_whitespace(text, retain_newlines=True):\n",
    "    \"\"\"simplify while respecting up to 2 newlines\"\"\"\n",
    "    if retain_newlines:\n",
    "        text = re.sub(r\"\\s+\", _space_repl, text).strip()\n",
    "    else:\n",
    "        text = re.sub(r\"\\s+\", \" \", text).strip()\n",
    "    return text\n",
    "\n",
    "\n",
    "def _augment_tag(s):\n",
    "    # case augment\n",
    "    if random.random() >= 0.8:\n",
    "        s = random.choice(CASE_AUGMENT_FUNCS)(s)\n",
    "    # other misc formatting\n",
    "    if random.random() >= 0.5:\n",
    "        s = s.replace(\"-\", \" \").strip()\n",
    "    return s\n",
    "\n",
    "def _augment_tags(tags):\n",
    "    random.shuffle(tags)\n",
    "    if random.random() <= 0.5:\n",
    "        tags = tags[: random.randint(0, len(tags))]\n",
    "        tags = [_augment_tag(tag) for tag in tags]\n",
    "    return tags\n",
    "\n",
    "\n",
    "def _clean_tags(tags):\n",
    "    return [\n",
    "        clean_tag[:MAX_TAG_LEN]\n",
    "        for tag in tags\n",
    "        if len(clean_tag := _simplify_whitespace(tag, retain_newlines=False)) > 0\n",
    "    ]\n",
    "\n",
    "\n",
    "def _clean_inline_tags(m):\n",
    "    tags = m.group(2).split(\";\")\n",
    "    tags = _clean_tags(tags)\n",
    "    ts = \";\".join(tags)[:MAX_TOT_TAGS_LEN]\n",
    "    if len(ts) > 0:\n",
    "        return f\"[{m.group(1)}: {ts}]\"\n",
    "    return f\"[{m.group(1)}]\"\n",
    "\n",
    "\n",
    "def _augment_inline_tags(m):\n",
    "    tags = m.group(2).split(\";\")\n",
    "    tags = _augment_tags(tags)\n",
    "    merge_char = random.choice([\", \", \" \", \"; \", \",\", \";\", \". \"])\n",
    "    ts = merge_char.join(tags)\n",
    "    if len(ts) > 0:\n",
    "        return f\"[{m.group(1)}: {ts}]\"\n",
    "    return f\"[{m.group(1)}]\"\n",
    "\n",
    "\n",
    "def _get_control_tags(duration_s, sample_vocal_start_s, do_augment=True):\n",
    "    control_tags = []\n",
    "    control_tags.append(f\"duration:{int(round(duration_s))}\")\n",
    "    for min_duration in [2, 4, 6, 8]:\n",
    "        min_duration = min_duration * 60\n",
    "        if duration_s >= min_duration:\n",
    "            control_tags.append(f\"min_duration:{min_duration}\")\n",
    "    for max_duration in [4, 6, 8]:\n",
    "        max_duration = max_duration * 60\n",
    "        if duration_s <= max_duration:\n",
    "            control_tags.append(f\"max_duration:{max_duration}\")\n",
    "    if sample_vocal_start_s is not None:\n",
    "        if sample_vocal_start_s <= 5:\n",
    "            control_tags.append(\"vocals:fast\")\n",
    "        if sample_vocal_start_s <= 15:\n",
    "            control_tags.append(\"vocals:normal\")\n",
    "        if 10 <= sample_vocal_start_s <= 20:\n",
    "            control_tags.append(\"vocals:intro\")\n",
    "    if do_augment:\n",
    "        if random.random() >= 0.5:\n",
    "            random.shuffle(control_tags)\n",
    "            control_tags = control_tags[: random.randint(0, len(control_tags))]\n",
    "    if len(control_tags) == 0:\n",
    "        return None\n",
    "    return \"{\" + \";\".join(control_tags) + \"}\"\n",
    "\n",
    "\n",
    "def build_text(\n",
    "    tags, text, duration_s, sample_vocal_start_s, inference, suppress_text, enable_control_tags=True\n",
    "):\n",
    "    if suppress_text:\n",
    "        return \"\"\n",
    "    text_elements = []\n",
    "    # collect tags\n",
    "    tags = _clean_tags(tags)\n",
    "    if not inference:\n",
    "        tags = _augment_tags(tags)\n",
    "    merge_char = random.choice([\", \", \" \", \"; \", \",\", \";\", \". \"])\n",
    "    tags_str = f\"{merge_char.join(tags[:MAX_N_TAGS])}\"[:MAX_TOT_TAGS_LEN]\n",
    "    if len(tags_str) > 0:\n",
    "        text_elements.append(f\"[{tags_str}]\")\n",
    "    # get lyrics\n",
    "    text = re.sub(r\"\\[(.*?)\\:(.*?)\\]\", _clean_inline_tags, text)\n",
    "    if len(text) > 0 and not inference:\n",
    "        # augment tags inside text:\n",
    "        text = re.sub(r\"\\[(.*?)\\:(.*?)\\]\", _augment_inline_tags, text)\n",
    "        if random.random() >= 0.95:\n",
    "            text = text.lower()\n",
    "        if random.random() >= 0.95:\n",
    "            text = re.sub(r\"\\n+\", \" \", text)\n",
    "    if len(text) > 0:\n",
    "        text_elements.append(text.strip())\n",
    "    # get control tags\n",
    "    for n in range(len(text_elements)):\n",
    "        text_elements[n] = text_elements[n].replace(\"{\", \"\").replace(\"}\", \"\")\n",
    "    if (inference or random.random() >= 0.1) and enable_control_tags:\n",
    "        control_tags = _get_control_tags(duration_s, sample_vocal_start_s, do_augment=not inference)\n",
    "        if control_tags is not None:\n",
    "            text_elements = [control_tags] + text_elements\n",
    "\n",
    "    if inference or random.random() >= 0.5:\n",
    "        text = \"\\n\\n\".join(text_elements)\n",
    "    else:\n",
    "        text = \"\"\n",
    "        for t in text_elements:\n",
    "            text += random.choice([\" \", \"\\n\", \"\\n\\n\"]) + t\n",
    "    text = text.strip()\n",
    "    return text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 185,
   "metadata": {},
   "outputs": [],
   "source": [
    "tags = [\"2014\", \"Pop Punk\", \"Emo-Pop\", \"Midwest Emo\", \"Emo\", \"Emo Revival\", \"introspective\", \"melancholic\", \"energetic\"]\n",
    "\n",
    "text = \"\"\"\n",
    "[verse]\n",
    "Embrace the random\n",
    "Avoid the void\n",
    "Stuck in the sky\n",
    "Like a smiley asteroid\n",
    "See you on the other side\n",
    "Clouds knocking on my door\n",
    "Dream to sleep tonight\n",
    "What am I flying for?\n",
    "(four, four, four, four)\n",
    "\n",
    "[chorus]\n",
    "Hey\n",
    "I'm on holiday\n",
    "From the ground\n",
    "Up in the clouds\n",
    "Hey\n",
    "I'm on holiday\n",
    "From the ground\n",
    "Up in the clouds\n",
    "\n",
    "[instrumental break]\n",
    "\n",
    "[bridge]\n",
    "Stars in my wet eyes\n",
    "(Droplets of condensation)\n",
    "Taste the water vapor\n",
    "(Just like a conversation)\n",
    "\n",
    "[chorus]\n",
    "Hey\n",
    "I'm on holiday\n",
    "From the ground\n",
    "Up in the clouds\n",
    "Hey\n",
    "I'm on holiday\n",
    "From the ground\n",
    "Up in the clouds\n",
    "\n",
    "[bridge]\n",
    "Stars in my wet eyes\n",
    "(Droplets of condensation)\n",
    "Taste the water vapor\n",
    "(Just like a conversation)\n",
    "\n",
    "[chorus]\n",
    "Hey\n",
    "I'm on holiday\n",
    "From the ground\n",
    "Up in the clouds\n",
    "Hey\n",
    "I'm on holiday\n",
    "From the ground\n",
    "Up in the clouds\n",
    "\"\"\"\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 196,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{duration:240;min_duration:120;min_duration:240;max_duration:240;max_duration:360;max_duration:480;vocals:normal;vocals:intro}\n",
      "\n",
      "[Emo-Pop, introspective, Pop Punk, Emo Revival, 2014, energetic, Emo, melancholic, Midwest Emo]\n",
      "\n",
      "[verse]\n",
      "Embrace the random\n",
      "Avoid the void\n",
      "Stuck in the sky\n",
      "Like a smiley asteroid\n",
      "See you on the other side\n",
      "Clouds knocking on my door\n",
      "Dream to sleep tonight\n",
      "What am I flying for?\n",
      "(four, four, four, four)\n",
      "\n",
      "[chorus]\n",
      "Hey\n",
      "I'm on holiday\n",
      "From the ground\n",
      "Up in the clouds\n",
      "Hey\n",
      "I'm on holiday\n",
      "From the ground\n",
      "Up in the clouds\n",
      "\n",
      "[instrumental break]\n",
      "\n",
      "[bridge]\n",
      "Stars in my wet eyes\n",
      "(Droplets of condensation)\n",
      "Taste the water vapor\n",
      "(Just like a conversation)\n",
      "\n",
      "[chorus]\n",
      "Hey\n",
      "I'm on holiday\n",
      "From the ground\n",
      "Up in the clouds\n",
      "Hey\n",
      "I'm on holiday\n",
      "From the ground\n",
      "Up in the clouds\n",
      "\n",
      "[bridge]\n",
      "Stars in my wet eyes\n",
      "(Droplets of condensation)\n",
      "Taste the water vapor\n",
      "(Just like a conversation)\n",
      "\n",
      "[chorus]\n",
      "Hey\n",
      "I'm on holiday\n",
      "From the ground\n",
      "Up in the clouds\n",
      "Hey\n",
      "I'm on holiday\n",
      "From the ground\n",
      "Up in the clouds\n"
     ]
    }
   ],
   "source": [
    "text_out = build_text(tags, text, 240, 10, False, False, enable_control_tags=True)\n",
    "print(text_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_env",
   "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}