{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-07T20:16:00.889634Z",
     "iopub.status.busy": "2025-03-07T20:16:00.889315Z",
     "iopub.status.idle": "2025-03-07T20:16:08.112803Z",
     "shell.execute_reply": "2025-03-07T20:16:08.112142Z",
     "shell.execute_reply.started": "2025-03-07T20:16:00.889614Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/tony/anaconda3/envs/suno_env_dev/lib/python3.10/site-packages/transformers/tokenization_utils_base.py:1601: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be depracted in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import gc\n",
    "import re\n",
    "import tqdm\n",
    "import math\n",
    "import copy\n",
    "import json\n",
    "import funcy\n",
    "import random\n",
    "import torch\n",
    "import tempfile\n",
    "import fasttext\n",
    "import numpy as np\n",
    "import collections\n",
    "import pandas as pd\n",
    "from joblib import Parallel, delayed\n",
    "\n",
    "from collections import Counter\n",
    "from transformers import PreTrainedTokenizerFast\n",
    "from transformers import BertTokenizerFast\n",
    "from tokenizers import (\n",
    "    decoders,\n",
    "    models,\n",
    "    normalizers,\n",
    "    pre_tokenizers,\n",
    "    processors,\n",
    "    trainers,\n",
    "    Tokenizer,\n",
    ")\n",
    "\n",
    "\n",
    "from suno_utils.audio import Audio\n",
    "from suno_utils.utils.tokenizers import tokenize\n",
    "from suno_utils.utils.lyrics import remove_speakers\n",
    "from suno_utils.utils.display import capture_output\n",
    "from suno_utils.utils.s3 import read_from_s3, check_s3_file_exists, open_from_s3\n",
    "from suno_utils.harvest.youtube.constants.text_lang import BASE_TO_FASTTEXT_REMAP\n",
    "from suno_utils.utils.text import (\n",
    "    write_jsonl,\n",
    "    read_jsonl,\n",
    "    write_json,\n",
    "    read_json,\n",
    "    normalize_whitespace,\n",
    ")\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"true\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-07T20:16:08.115355Z",
     "iopub.status.busy": "2025-03-07T20:16:08.115210Z",
     "iopub.status.idle": "2025-03-07T20:16:08.118277Z",
     "shell.execute_reply": "2025-03-07T20:16:08.117779Z",
     "shell.execute_reply.started": "2025-03-07T20:16:08.115341Z"
    }
   },
   "outputs": [],
   "source": [
    "def _augment_tag(s):\n",
    "    if random.random() >= 0.95:\n",
    "        s = s.upper()\n",
    "    elif random.random() >= 0.95:\n",
    "        s = s.capitalize()\n",
    "    elif random.random() >= 0.9:\n",
    "        s = s.title()\n",
    "    elif random.random() >= 0.9:\n",
    "        s = s.lower()\n",
    "    if random.random() >= 0.5:\n",
    "        s = s.replace(\"-\", \" \").strip()\n",
    "    return s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-07T20:16:08.120153Z",
     "iopub.status.busy": "2025-03-07T20:16:08.120033Z",
     "iopub.status.idle": "2025-03-07T20:16:08.122695Z",
     "shell.execute_reply": "2025-03-07T20:16:08.122207Z",
     "shell.execute_reply.started": "2025-03-07T20:16:08.120140Z"
    }
   },
   "outputs": [],
   "source": [
    "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 \" \""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-07T20:16:08.124597Z",
     "iopub.status.busy": "2025-03-07T20:16:08.124308Z",
     "iopub.status.idle": "2025-03-07T20:16:08.127309Z",
     "shell.execute_reply": "2025-03-07T20:16:08.126822Z",
     "shell.execute_reply.started": "2025-03-07T20:16:08.124583Z"
    }
   },
   "outputs": [],
   "source": [
    "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-07T20:16:08.129194Z",
     "iopub.status.busy": "2025-03-07T20:16:08.129068Z",
     "iopub.status.idle": "2025-03-07T20:16:08.131640Z",
     "shell.execute_reply": "2025-03-07T20:16:08.131164Z",
     "shell.execute_reply.started": "2025-03-07T20:16:08.129182Z"
    }
   },
   "outputs": [],
   "source": [
    "def _augment_case(s):\n",
    "    if random.random() >= 0.95:\n",
    "        s = s.upper()\n",
    "    elif random.random() >= 0.95:\n",
    "        s = s.lower()\n",
    "    return s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-07T20:16:08.133639Z",
     "iopub.status.busy": "2025-03-07T20:16:08.133514Z",
     "iopub.status.idle": "2025-03-07T20:16:08.138677Z",
     "shell.execute_reply": "2025-03-07T20:16:08.138165Z",
     "shell.execute_reply.started": "2025-03-07T20:16:08.133626Z"
    }
   },
   "outputs": [],
   "source": [
    "def make_tokenizer(base_dir: str, metas_dir: str, tmp_dir: str):\n",
    "    tokenizer_train_filepath = os.path.join(tmp_dir, \"tokenizer_train.txt\")\n",
    "    # create text file with all text\n",
    "    # with open(tokenizer_train_filepath, \"w\") as f:\n",
    "    #     for fn in tqdm.tqdm(\n",
    "    #         [\n",
    "    #             \"13b_s8_v29_meta_tr\",\n",
    "    #             \"13b_s32_v29_meta_tr\",\n",
    "    #             \"30b_t6_v33_meta_tr\",\n",
    "    #             # \"clean_covers_v0_metas\",\n",
    "    #             \"clean_genius_v0_metas\",\n",
    "    #             \"clean_karaoke_v0_metas\",\n",
    "    #             \"clean_pond5_v0_metas\",\n",
    "    #             \"clean_deezer_v0_metas\",\n",
    "    #             \"clean_discogs_v0_metas\",\n",
    "    #             \"clean_imslp_v0_metas\",\n",
    "    #             \"clean_musdb_v0_metas\",\n",
    "    #             \"clean_youtube_music_v0_metas\",\n",
    "    #         ]\n",
    "    #     ):\n",
    "    #         metas_filepath = os.path.join(metas_dir, f\"{fn}.jsonl\")\n",
    "\n",
    "    #         if not os.path.isfile(metas_filepath):\n",
    "    #             print(f\"{metas_filepath} not found. Skipping...\")\n",
    "    #             continue\n",
    "\n",
    "    #         for index, m in enumerate(read_jsonl(metas_filepath)):\n",
    "    #             # for preference data, we only keep one of the two text fields\n",
    "    #             if \"_tr\" in fn and index % 2 == 0:\n",
    "    #                 continue\n",
    "    #             text_key = (\n",
    "    #                 \"text\"\n",
    "    #                 if \"private_text\" not in m or random.random() >= 0.2\n",
    "    #                 else \"private_text\"\n",
    "    #             )\n",
    "    #             if text_key in m:\n",
    "    #                 f.write(_simplify_whitespace(_augment_case(m[text_key])) + \"\\n\")\n",
    "    #             text_key = (\n",
    "    #                 \"text_segments\"\n",
    "    #                 if \"private_text_segments\" not in m or random.random() >= 0.2\n",
    "    #                 else \"private_text_segments\"\n",
    "    #             )\n",
    "    #             if text_key in m:\n",
    "    #                 for mm in m[text_key]:\n",
    "    #                     f.write(_simplify_whitespace(_augment_case(mm[\"text\"])) + \"\\n\")\n",
    "    #             text_key = (\n",
    "    #                 \"tags\"\n",
    "    #                 if \"private_tags\" not in m or random.random() >= 0.2\n",
    "    #                 else \"private_tags\"\n",
    "    #             )\n",
    "    #             if text_key in m:\n",
    "    #                 join_char = random.choice([\", \", \"; \", \" \"])\n",
    "    #                 f.write(\n",
    "    #                     _simplify_whitespace(\n",
    "    #                         join_char.join([_augment_tag(t) for t in m[text_key]])\n",
    "    #                     )\n",
    "    #                     + \"\\n\"\n",
    "    #                 )\n",
    "\n",
    "    # min_char_n = 3\n",
    "    # with open(tokenizer_train_filepath) as f:\n",
    "    #     corpus = f.read()\n",
    "    # char_counts = Counter(corpus)\n",
    "    # print(len(char_counts), \"unique chars\")\n",
    "    # trans_table = str.maketrans(\n",
    "    #     {k: \" \" for k, v in char_counts.items() if v < min_char_n}\n",
    "    # )\n",
    "    # corpus = corpus.translate(trans_table)\n",
    "    # char_counts_2 = Counter(corpus)\n",
    "    # print(len(char_counts_2), \"unique chars after filtering\")\n",
    "    # corpus_lines = corpus.split(\"\\n\")\n",
    "    # corpus_lines = [\n",
    "    #     line_clean for line in corpus_lines if len(line_clean := line.strip()) > 0\n",
    "    # ]\n",
    "    random.seed(6006)\n",
    "\n",
    "    tokenizer_train_clean_filepath = os.path.join(tmp_dir, \"tokenizer_train_clean.txt\")\n",
    "\n",
    "    # with open(tokenizer_train_clean_filepath, \"w\") as f:\n",
    "    #     for line in corpus_lines:\n",
    "    #         f.write(line + \"\\n\")\n",
    "    print(\"loading from cached training file\")\n",
    "    with open(tokenizer_train_clean_filepath, \"r\") as f:\n",
    "        corpus_lines = f.readlines()\n",
    "\n",
    "    print(\"Initializing tokenizer\")\n",
    "    # train and save tokenizer\n",
    "    tokenizer = Tokenizer(models.WordPiece(unk_token=\"[UNK]\"))\n",
    "    tokenizer.normalizer = normalizers.BertNormalizer(lowercase=False)\n",
    "    tokenizer.pre_tokenizer = pre_tokenizers.BertPreTokenizer()\n",
    "    special_tokens = [\"[UNK]\", \"[PAD]\"]\n",
    "    trainer = trainers.WordPieceTrainer(\n",
    "        vocab_size=60_000, special_tokens=special_tokens\n",
    "    )\n",
    "    print(f\"Training tokenizer with {len(corpus_lines)} lines\")\n",
    "    tokenizer.train([tokenizer_train_clean_filepath], trainer=trainer)\n",
    "    tokenizer_path = os.path.join(base_dir, \"tokenizer_60k.json\")\n",
    "    tokenizer.save(tokenizer_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-07T20:16:08.140489Z",
     "iopub.status.busy": "2025-03-07T20:16:08.140359Z",
     "iopub.status.idle": "2025-03-07T20:16:08.143511Z",
     "shell.execute_reply": "2025-03-07T20:16:08.143025Z",
     "shell.execute_reply.started": "2025-03-07T20:16:08.140476Z"
    }
   },
   "outputs": [],
   "source": [
    "metas_dir = \"/app/suno/tmp\"\n",
    "tmp_dir = os.path.join(os.getcwd(), \"tmp\")\n",
    "os.makedirs(tmp_dir, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-07T20:16:08.145271Z",
     "iopub.status.busy": "2025-03-07T20:16:08.145154Z",
     "iopub.status.idle": "2025-03-07T20:31:35.962615Z",
     "shell.execute_reply": "2025-03-07T20:31:35.962001Z",
     "shell.execute_reply.started": "2025-03-07T20:16:08.145259Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading from cached training file\n",
      "Initializing tokenizer\n",
      "Training tokenizer with 432822527 lines\n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "make_tokenizer(tmp_dir, metas_dir, tmp_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-07T20:31:35.965594Z",
     "iopub.status.busy": "2025-03-07T20:31:35.965465Z",
     "iopub.status.idle": "2025-03-07T20:31:35.968236Z",
     "shell.execute_reply": "2025-03-07T20:31:35.967799Z",
     "shell.execute_reply.started": "2025-03-07T20:31:35.965581Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "done!\n"
     ]
    }
   ],
   "source": [
    "print(\"done!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}