{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "12376f58",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-19T22:09:26.707939Z",
     "start_time": "2023-09-19T22:09:19.829021Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/tony/anaconda3/envs/suno_env/lib/python3.10/site-packages/whisper/timing.py:57: NumbaDeprecationWarning: \u001b[1mThe 'nopython' keyword argument was not supplied to the 'numba.jit' decorator. The implicit default value for this argument is currently False, but it will be changed to True in Numba 0.59.0. See https://numba.readthedocs.io/en/stable/reference/deprecation.html#deprecation-of-object-mode-fall-back-behaviour-when-using-jit for details.\u001b[0m\n",
      "  def backtrace(trace: np.ndarray):\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\"\n",
    "import json\n",
    "\n",
    "from suno_utils.audio import Audio\n",
    "from suno_utils.tasks.mert_v2 import (\n",
    "    encode as mert_encode,\n",
    "    preload_models as preload_mert_models,\n",
    ")\n",
    "from suno_utils.tasks.dac import (\n",
    "    encode as codec_encode,\n",
    "    decode as codec_decode,\n",
    "    preload_models as preload_codec_models,\n",
    ")\n",
    "import numpy as np\n",
    "from suno_utils.tasks.gpt_v2.chirp_v1 import preload_models as preload_chirp_models\n",
    "from suno_utils.tasks.gpt_v2.chirp_v1 import generate_audio\n",
    "from suno_utils.tasks.gpt_v2.generation import generate, load_model\n",
    "from collections import defaultdict\n",
    "\n",
    "from torcheval.metrics import WordErrorRate\n",
    "from torchmetrics.text import CharErrorRate\n",
    "import re\n",
    "import string"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "89676b91",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-19T22:09:26.718805Z",
     "start_time": "2023-09-19T22:09:26.709297Z"
    }
   },
   "outputs": [],
   "source": [
    "output_dir = \"results/test_demuc\"\n",
    "if not os.path.exists(output_dir):\n",
    "    os.makedirs(output_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b6878973",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-19T22:09:26.768498Z",
     "start_time": "2023-09-19T22:09:26.720296Z"
    }
   },
   "outputs": [],
   "source": [
    "# from suno_utils.utils.s3 import _download_s3_file\n",
    "# _download_s3_file(\"s3://suno-data/georg/trained_models/chirp_v1/lid.176.bin\", \"/home/tony/Data/Chirp/lid.176.bin\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "774fc130",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-19T22:09:26.824468Z",
     "start_time": "2023-09-19T22:09:26.769506Z"
    }
   },
   "outputs": [],
   "source": [
    "def clean_text_for_wer(text):\n",
    "    text = text.replace(\"’\", \"'\").replace(\"\\n\", \" \").lower()\n",
    "    text = re.sub(r\"\\[.+?\\]\", \" \", text)\n",
    "    text = re.sub(r\"\\s+\", \" \", text).strip()\n",
    "    text = text.translate(str.maketrans(\"\", \"\", string.punctuation))\n",
    "    text = \"\".join(text)\n",
    "    return text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "51a4ade6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-19T22:09:52.684367Z",
     "start_time": "2023-09-19T22:09:26.825860Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "model loaded: 1720.9M params, 1.411 loss\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at m-a-p/MERT-v1-95M were not used when initializing MERTModel: ['encoder.pos_conv_embed.conv.weight_v', 'encoder.pos_conv_embed.conv.weight_g']\n",
      "- This IS expected if you are initializing MERTModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing MERTModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of MERTModel were not initialized from the model checkpoint at m-a-p/MERT-v1-95M and are newly initialized: ['encoder.pos_conv_embed.conv.parametrizations.weight.original1', 'encoder.pos_conv_embed.conv.parametrizations.weight.original0']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
      "/home/tony/anaconda3/envs/suno_env/lib/python3.10/site-packages/torch/nn/utils/weight_norm.py:30: UserWarning: torch.nn.utils.weight_norm is deprecated in favor of torch.nn.utils.parametrizations.weight_norm.\n",
      "  warnings.warn(\"torch.nn.utils.weight_norm is deprecated in favor of torch.nn.utils.parametrizations.weight_norm.\")\n",
      "Warning : `load_model` does not return WordVectorModel or SupervisedModel any more, but a `FastText` object which is very similar.\n"
     ]
    }
   ],
   "source": [
    "preload_chirp_models(\n",
    "    centroids_filepath=\"/home/mikeys/bundle/2x1k_centroids_mert.npy\",\n",
    "    gpt_ckpt_path=\"/home/tony/Data/Chirp/xl_2.pt\",\n",
    "    codec_ckpt_path=\"/home/mikeys/bundle/d_codec_25x8.pt\",\n",
    "    common_genre_tags_path=\"/home/mikeys/bundle/common_tags.json\",\n",
    "    fasttext_ckpt_path=\"/home/tony/Data/Chirp/lid.176.bin\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "eb6ac3fb",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-19T22:09:52.688411Z",
     "start_time": "2023-09-19T22:09:52.686232Z"
    }
   },
   "outputs": [],
   "source": [
    "exampe_tags = [\"pop\", \"rap\", \"rock\", \"hip-hop\", \"edm\"]\n",
    "languages = [\"english\", \"french\", \"chinese\", \"russian\"] #\"spanish\"\n",
    "themes = [\"life\", \"love\", \"friendship\"]\n",
    "n_var = len(themes)\n",
    "n_last_batch = 4\n",
    "language = \"chinese\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "04fa5e6c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-19T22:31:03.118134Z",
     "start_time": "2023-09-19T22:12:53.992928Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "checking with tag: pop\n",
      "input text [verse]\n",
      "生活啊，多姿多彩如诗画\n",
      "每天都是新的挑战\n",
      "困难和希望交织美好\n",
      "享受每个瞬间，让心充满爱\n",
      "\n",
      "[chorus]\n",
      "我在这个多彩的人生乐园\n",
      "舞动青春的旋律，不畏风云\n",
      "每一天都是独一无二的精彩\n",
      "用心去感受，人生的美好\n",
      "hacking text with intro for non-English\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████| 3185/3185 [01:17<00:00, 41.04it/s]\n",
      "output_file: results/test_demuc/chinese_pop_0_gen0.mp3 already exists and will be overwritten on build\n",
      "output_file: results/test_demuc/chinese_pop_0_gen1.mp3 already exists and will be overwritten on build\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input text [verse]\n",
      "在这个世界上两个人相遇\n",
      "眼中只剩下彼此的倾心\n",
      "一见钟情写下爱的故事\n",
      "心跳加速，情愫沸腾起\n",
      "\n",
      "[chorus]\n",
      "爱情不问东西南北中\n",
      "只要有你，我心满意足\n",
      "手牵手一起走向未来\n",
      "爱永远不变，永不停息\n",
      "\n",
      "(Song ends here)\n",
      "False, 这歌多彩的人声的夜泪无动精臣的旋律不为风雨, 0.434\n",
      "DEMUCED False, 结果多彩的人生了也愿我懂青春的旋律不为风雨, 0.994\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████| 2435/2435 [01:02<00:00, 38.79it/s]\n",
      "output_file: results/test_demuc/chinese_pop_1_gen0.mp3 already exists and will be overwritten on build\n",
      "output_file: results/test_demuc/chinese_pop_1_gen1.mp3 already exists and will be overwritten on build\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input text [verse]\n",
      "朋友是我最珍贵的宝贝\n",
      "相互扶持 青春不会散开\n",
      "一起度过了无数个日夜\n",
      "友情永远 不会改变\n",
      "\n",
      "[chorus]\n",
      "朋友像太阳温暖心房\n",
      "手牵手共同实现梦想\n",
      "不论风雨 我们在一起\n",
      "友谊一生永不散场\n",
      "False, 西南北中只要有你 我心滿意足手牽手一起走向未來, 0.377\n",
      "DEMUCED False, 西南北中只要有你 我心滿意足手牽手一起走向未來, 0.894\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████| 2435/2435 [01:17<00:00, 31.26it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "朋友是我最珍贵的宝贝 相互扶持 青春不会散开 一起度过了无数个日夜 友情永远 不会改变 朋友像太阳温暖心房 手牵手共同实现梦想 不论风雨 我们在一起 友谊一生永不散场\n",
      "又是我最珍贵的宝贝想不服之 青春不会散开一起度过了无数个月夜夜夜友情永永远不会改变朋友像太阳 吻了心放手牵手 共投血 陷弄伤 tensor(1.) tensor(0.5181)\n",
      "朋友是我最珍贵的宝贝相互扶植青春不会散开一起度过了无数个日夜友情永远不会改变朋友像太阳温暖心放手牵手共同实现梦想不论风雨我们在一起 tensor(1.) tensor(0.2410)\n",
      "朋友是我最珍贵的宝贝相误付词情传不会散开脾气度过了无数个月夜有情拥有不会开天朋友像太阳吻了心房手牵手共同视线梦想无论风雨我们在一起 tensor(1.) tensor(0.4337)\n",
      "愛恨又是我最珍貴的寶貝相不符合 情純不會散開一起度過了 無數個日夜有情永遠 不會改變朋友像太陽 溫暖心法手牽手 共同實現夢想不論風與我們在一起有一生 tensor(1.) tensor(0.5904)\n",
      "continuation average wer of 4 is 1.0, cer of 4 is 0.4457831382751465\n",
      "checking with tag: rap\n",
      "input text [verse]\n",
      "人生就像一曲RAP\n",
      "奇幻起落 我在追逐\n",
      "每次挫折都是个机会\n",
      "坚持向前才是根本\n",
      "\n",
      "[chorus]\n",
      "活在当下活出精彩\n",
      "不怕困难 我心自在\n",
      "无论曲折也别低头\n",
      "歌唱人生 梦想无敌\n",
      "hacking text with intro for non-English\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████| 3185/3185 [01:17<00:00, 40.88it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input text [Verse]\n",
      "爱在心中躁动，像冬日的暖阳\n",
      "对你的情感，像海浪般翻涌澎湃\n",
      "每个呼吸都充满你的气息\n",
      "爱情之火在心中燃烧，不止一刻\n",
      "\n",
      "[Chorus]\n",
      "你是我的信仰，让我找到了快乐\n",
      "我会守护你的心，直到永远不变节\n",
      "无论遇到什么挑战和艰难\n",
      "爱的力量会给我力量去战胜\n",
      "False, 自在無論去這也別低頭割傷了, 0.311\n",
      "DEMUCED False, , 2.128\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████| 2435/2435 [01:02<00:00, 38.77it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input text [verse]\n",
      "朋友如兄弟 因缘如镜\n",
      "相互扶持 直到天涯涯\n",
      "时光无情 友情难忘\n",
      "追逐梦想 在成长的道路上\n",
      "\n",
      "[chorus]\n",
      "朋友啊 你是我的明日\n",
      "手牵手 永不分离\n",
      "友情密不可分 在心中永驻\n",
      "这首歌 歌颂我们的友谊\n",
      "False, 能遇到什么挑战很艰难爱的力量会给我力量去战胜, 0.422\n",
      "DEMUCED False, 你都不知道什么条件和缺乏爱的力量会给我力量去暂时, 0.952\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████▌| 2426/2435 [01:16<00:00, 31.71it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "朋友如兄弟 因缘如镜 相互扶持 直到天涯涯 时光无情 友情难忘 追逐梦想 在成长的道路上 朋友啊 你是我的明日 手牵手 永不分离 友情密不可分 在心中永驻 这首歌 歌颂我们的友谊\n",
      "東牛兄弟 陰影人入鏡像 無福之指他天涯涯時光無盡 又盡難忘追逐夢想 在城堂 到道路上朋友啊 你是我的名人抽籤 收用 不分離有情你不可分 在心中用錢這首歌歌 送我們的友誼 tensor(1.) tensor(0.5618)\n",
      "朋友入兄弟 音乐 如今 也像肚子直到天涯涯时光 无情有情 难忘追逐某场在成长的道路上朋友啊 迷失我的名人 手牵手永不分明有情你不可奋在心中 涌向 执手可歌颂我们得要离开 tensor(1.) tensor(0.4719)\n",
      "韓國兄弟因緣如今相互服是遲到天涯時光無情又輕鬧忘追求夢想再成長得到路上朋友啊你是我的名月受牽手永不分離有情你不可分在心中用心這首歌歌送我們都有意 tensor(1.) tensor(0.5618)\n",
      "朋友如兄的音言如计笑不复辞 值得贴牙牙时光无情 又轻难忘追逐梦想在成长的道路上朋友啊你是我的名额手牵手 永不分离有情你不可分在心中 用牵着手歌歌送 我们的友誼 tensor(1.) tensor(0.4270)\n",
      "continuation average wer of 4 is 1.0, cer of 4 is 0.5056179761886597\n",
      "checking with tag: rock\n",
      "input text [Verse]\n",
      "生活不易，如同波涛起伏\n",
      "追寻梦想，考验意志力\n",
      "坚持不懈，努力奋斗着\n",
      "逆风飞翔，展现勇气与力量\n",
      "\n",
      "[Chorus]\n",
      "人生之路，苦尽甘来\n",
      "奋发向前，信念永不崩\n",
      "挑战自我，逆境中成长\n",
      "勇往直前，活出我风采\n",
      "hacking text with intro for non-English\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 90%|████████████████████████████████████████████████████████████████████████████████████████████████▌          | 2873/3185 [01:08<00:07, 41.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input text [verse]\n",
      "我不能停止思念\n",
      "心中只有你的身影\n",
      "爱上你我心悸不已\n",
      "你的存在让我充满勇气\n",
      "\n",
      "[chorus]\n",
      "这份爱如狂风暴雨般热烈\n",
      "我们的心在爱中交织\n",
      "你是我生命中的唯一\n",
      "你的微笑让我感到幸福的喜悦\n",
      "False, 深夜微波 挑戰寂寞一場長長 勇往前前或許我放棄, 0.457\n",
      "DEMUCED False, 深夜吻我叫 time to love誰叫成潮用忘記牽我心有放下, 1.063\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 99%|██████████████████████████████████████████████████████████████████████████████████████████████████████████▎| 2420/2435 [01:02<00:00, 38.74it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input text [verse]\n",
      "朋友，是我一生的财富\n",
      "把你当成我心中最珍贵的宝\n",
      "一起征服人生的高山和大海\n",
      "永远一起，不离不弃，坚不可摧\n",
      "\n",
      "[chorus]\n",
      "友谊如山，永不磨灭\n",
      "心连心，同舟共济\n",
      "无论前路如何坎坷\n",
      "朋友，在一起，走到永远\n",
      "False, 你是我身体中的回忆你都会想让我看到幸福的心愿, 0.377\n",
      "DEMUCED False,  I'm not ready to give you something Don't you agree? I'm not ready to give you something Don't you agree? I'm not ready to give you something, 1.872\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 97%|███████████████████████████████████████████████████████████████████████████████████████████████████████▋   | 2360/2435 [01:13<00:02, 32.01it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "朋友，是我一生的财富 把你当成我心中最珍贵的宝 一起征服人生的高山和大海 永远一起，不离不弃，坚不可摧 友谊如山，永不磨灭 心连心，同舟共济 无论前路如何坎坷 朋友，在一起，走到永远\n",
      "字幕by索兰娅 tensor(1.) tensor(1.)\n",
      "字幕by索兰娅 tensor(1.) tensor(1.)\n",
      "我 你不為什麼 我敢要幸福的喜愛就夠有一種善用不如磨滅 心裂心痛就勾血無論前路如何堪有 tensor(1.) tensor(0.8242)\n",
      "字幕by索兰娅 tensor(1.) tensor(1.)\n",
      "continuation average wer of 4 is 1.0, cer of 4 is 0.9560439586639404\n",
      "checking with tag: hip-hop\n",
      "input text [verse]\n",
      "生活就像一场战斗\n",
      "每天都要付出努力\n",
      "无论风雨还是曙光\n",
      "我坚持走在前方\n",
      "\n",
      "[chorus]\n",
      "生命之舞 不停歇\n",
      "不畏艰辛 直追梦想\n",
      "流淌的節奏 跳动的心\n",
      "这就是我的华语 Hip-Hop 真谛\n",
      "hacking text with intro for non-English\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████| 3185/3185 [01:17<00:00, 40.94it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input text [verse]\n",
      "我在这街头流浪，寻找那失去的爱\n",
      "心痛如刀割，思念无尽的夜晚\n",
      "追逐烟花绽放，找回我内心的光彩\n",
      "唯有你的拥抱，让我感受爱的温暖\n",
      "\n",
      "[chorus]\n",
      "你是我的宿命，爱穿越千山万水\n",
      "心与心永相连，你是我唯一信仰\n",
      "那些甜蜜的回忆，在夜里安慰我的灵魂\n",
      "我们的爱，像烟火一样灿烂绚丽\n",
      "\n",
      "\n",
      "[\n",
      "False, 別為了鐵頭跳動的心這就是我的懷疑一般之真假, 0.403\n",
      "DEMUCED False, 別為了解脫跳動的心就是我的懷疑一般之成全, 1.009\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████| 2435/2435 [01:03<00:00, 38.64it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input text [verse]\n",
      "朋友不分昼夜\n",
      "共度艰难时刻\n",
      "手牵手一起闯荡\n",
      "友情永不会被遗忘\n",
      "\n",
      "[chorus]\n",
      "兄弟姐妹一起走\n",
      "友爱无与伦比\n",
      "我们是挚友一生一世\n",
      "友谊像磁铁永不分离\n",
      "False, 月渐山晚水 心一心永笑脸你是我唯一夕阳那些甜蜜的回忆在夜里安慰我的灵魂我们的, 0.719\n",
      "DEMUCED False, 月間閃完睡 心一心擁笑臉你是我唯一夕陽那些甜蜜的回憶在夜裡安慰我的靈魂我們的, 1.566\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 77%|█████████████████████████████████████████████████████████████████████████████████▉                         | 1865/2435 [00:56<00:17, 33.17it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "朋友不分昼夜 共度艰难时刻 手牵手一起闯荡 友情永不会被遗忘 兄弟姐妹一起走 友爱无与伦比 我们是挚友一生一世 友谊像磁铁永不分离\n",
      "字幕by索兰娅 tensor(1.) tensor(1.)\n",
      "朋友不尊求也共有箭短系可手牽手一起闖蕩眼睛永不回悲遺忘兄弟姐妹一起走有爱会冷漓我们世界有一生一息有影响期间我们不分离 tensor(1.) tensor(0.5846)\n",
      "有不分手也共度皆難是可守堅守一起撞蕩要情永不為陪遇亡終體界面一起走有愛未與人與我們世界有一生一生有一相 今天永不分離 tensor(1.) tensor(0.7077)\n",
      "字幕by索兰娅 tensor(1.) tensor(1.)\n",
      "continuation average wer of 4 is 1.0, cer of 4 is 0.8230769634246826\n",
      "checking with tag: edm\n",
      "input text [verse]\n",
      "人生无常变幻多，睡觉醒来全是梦\n",
      "天地间绽放的光，把我心中激荡\n",
      "\n",
      "[chorus]\n",
      "放声歌唱，寻找自己的轨迹\n",
      "跳动音符，奏响生命的旋律\n",
      "hacking text with intro for non-English\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 97%|███████████████████████████████████████████████████████████████████████████████████████████████████████▎   | 3077/3185 [01:15<00:02, 41.01it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input text [Verse]\n",
      "我迷失在爱的海洋中\n",
      "感觉心跳像是飞翔\n",
      "你的微笑如阳光般灿烂\n",
      "我的世界因你而绚烂\n",
      "\n",
      "[Chorus]\n",
      "爱情的旋律 在心中回响\n",
      "激荡着我灵魂的旋律\n",
      "我们的爱情 像一曲动人的舞曲\n",
      "在这世界上闪耀不灭\n",
      "False, 我心中起扛放生歌唱寻找自己的悔疾跳动音色奏险生命的旋律, 0.587\n",
      "DEMUCED False, 我心中起崗放生歌唱寻找自己的悔悉跳動因此走向生命的旋律, 1.266\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 92%|██████████████████████████████████████████████████████████████████████████████████████████████████▎        | 2237/2435 [00:57<00:05, 39.19it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input text [verse]\n",
      "我们的友情像烈火燃烧\n",
      "一起追逐梦想的光芒\n",
      "无论多远距离阻挡挑战\n",
      "友谊的舞曲继续跳跃\n",
      "\n",
      "[chorus]\n",
      "朋友啊，携手共度时光\n",
      "心与心相连直至永久\n",
      "这感觉让世界充满光荣\n",
      "我们的友谊鼓舞万众\n",
      "False, 即当着破灵魂的旋律我们的爱情下一曲突然流去在这世界上闪耀不灭, 0.562\n",
      "DEMUCED False, 即当折磨灵魂的旋律我们的爱情像一群独人流群在这世界上闪耀的工业, 1.345\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 94%|█████████████████████████████████████████████████████████████████████████████████████████████████████      | 2300/2435 [01:11<00:04, 32.10it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "我们的友情像烈火燃烧 一起追逐梦想的光芒 无论多远距离阻挡挑战 友谊的舞曲继续跳跃 朋友啊，携手共度时光 心与心相连直至永久 这感觉让世界充满光荣 我们的友谊鼓舞万众\n",
      "耶 我們都有情相別忽然少一句 推著夢想逃光芒互動 多遠距離觸發挑戰 有意的過去繼續挑夜 暗秒萬提升空度時光 星與星相列這次永久 照看前要視線充滿光中我們被惹意 呼呼萬中 tensor(1.1250) tensor(0.7831)\n",
      "飛在彩 飛在彩oh yeah 走連起遠朝朝好牛啊 閉手孤獨時光星星 省省時間永久這感覺 到世界窗外光榮我們在 容易孤空中 tensor(1.) tensor(0.8554)\n",
      "字幕by索兰娅 tensor(1.) tensor(1.)\n",
      "唯有無敵的手工都時光千筆線上面只是永久這感覺讓世界充滿光榮我們都有一股萬種 tensor(1.) tensor(0.8675)\n",
      "continuation average wer of 4 is 1.03125, cer of 4 is 0.876505970954895\n"
     ]
    }
   ],
   "source": [
    "tag_wers = defaultdict(list)\n",
    "tag_cers = defaultdict(list)\n",
    "for tag in exampe_tags:\n",
    "    print(f\"checking with tag: {tag}\")\n",
    "    prompts = []\n",
    "    for i in range(n_var):\n",
    "        with open(f\"inputs/{language}_{tag}_{i}.json\", \"r\") as fp:\n",
    "            prompts.append(json.load(fp))\n",
    "\n",
    "    history_cache = defaultdict(dict)\n",
    "    for i, prompt in enumerate(prompts):\n",
    "        # print(i, \"-->\")\n",
    "        print(\"input text\", prompt[\"text\"])\n",
    "        history_arr, raw_arrays, audios, text_tags = generate_audio(\n",
    "            prompt[\"text\"],\n",
    "            text_tags=prompt[\"tag\"] if i == 0 else None,\n",
    "            history_audio=None if i == 0 else history_cache[i - 1][0],\n",
    "            history_text_guess=None if i == 0 else prompts[i - 1][\"text\"],\n",
    "            n_batch=2 if i != len(prompts) - 1 else n_last_batch,\n",
    "            max_gen_duration_s=30 if i != 0 else 40,\n",
    "            return_raw_arrays=True,\n",
    "        )\n",
    "        for j, (raw_array, audio) in enumerate(zip(raw_arrays, audios)):\n",
    "            # print(j)\n",
    "            with open(\n",
    "                os.path.join(\n",
    "                    output_dir, f\"{prompt['lan']}_{prompt['tag']}_{i}_gen{j}.npz\"\n",
    "                ),\n",
    "                \"wb\",\n",
    "            ) as fp:\n",
    "                np.savez(fp, arr=raw_array)\n",
    "            audio.to_mp3(\n",
    "                os.path.join(\n",
    "                    output_dir, f\"{prompt['lan']}_{prompt['tag']}_{i}_gen{j}.mp3\"\n",
    "                )\n",
    "            )\n",
    "            history_cache[i][j] = raw_array\n",
    "\n",
    "    wer = WordErrorRate()\n",
    "    cer = CharErrorRate()\n",
    "    cleaned_prompt = clean_text_for_wer(prompts[i][\"text\"])\n",
    "\n",
    "    from suno_utils.tasks.gpt_v2.chirp_v1 import whisper_en, whisper_multi\n",
    "\n",
    "    cleaned_prompt = clean_text_for_wer(prompts[i][\"text\"])\n",
    "    total_wer = 0\n",
    "    total_cer = 0\n",
    "    print(cleaned_prompt)\n",
    "    for j in range(n_last_batch):\n",
    "        whisper_model = whisper_multi if prompt[\"lan\"] != \"english\" else whisper_en\n",
    "        output_text = whisper_model.transcribe(\n",
    "            os.path.join(output_dir, f\"{prompt['lan']}_{prompt['tag']}_{i}_gen{j}.mp3\"),\n",
    "            condition_on_previous_text=False,\n",
    "        )\n",
    "        cleaned_output = clean_text_for_wer(output_text[\"text\"])\n",
    "        wer.update([cleaned_output], [cleaned_prompt])\n",
    "        wer_value = wer.compute()\n",
    "        wer.reset()\n",
    "        total_wer += wer_value\n",
    "\n",
    "        cer.update([cleaned_output], [cleaned_prompt])\n",
    "        cer_value = cer.compute()\n",
    "        cer.reset()\n",
    "        total_cer += cer_value\n",
    "        tag_wers[tag].append(wer_value)\n",
    "        tag_cers[tag].append(cer_value)\n",
    "        print(cleaned_output, wer_value, cer_value)\n",
    "    print(\n",
    "        f\"continuation average wer of 4 is {total_wer / 4}, cer of 4 is {total_cer / 4}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83db0264",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-19T22:12:39.862430Z",
     "start_time": "2023-09-19T22:12:39.862421Z"
    }
   },
   "outputs": [],
   "source": [
    "test_audio = Audio.from_file(os.path.join(output_dir, \"chinese_rock_2_gen2.mp3\"))\n",
    "test_audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4beae481",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-19T22:12:39.863311Z",
     "start_time": "2023-09-19T22:12:39.863302Z"
    }
   },
   "outputs": [],
   "source": [
    "whisper_multi.transcribe(\n",
    "    os.path.join(output_dir, \"chinese_rock_2_gen2.mp3\"),\n",
    "    condition_on_previous_text=False,\n",
    ")[\"text\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "899816c4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-19T22:12:39.863903Z",
     "start_time": "2023-09-19T22:12:39.863895Z"
    }
   },
   "outputs": [],
   "source": [
    "# import whisper\n",
    "# whisper_test = whisper.load_model(\"medium\")\n",
    "# device = \"cpu\"\n",
    "# _ = whisper_test.to(device)\n",
    "# whisper_test.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12249669",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-19T22:12:39.864723Z",
     "start_time": "2023-09-19T22:12:39.864715Z"
    }
   },
   "outputs": [],
   "source": [
    "whisper_test.transcribe(\n",
    "    os.path.join(output_dir, \"chinese_rock_2_gen2.mp3\"),\n",
    "    condition_on_previous_text=False,\n",
    ")[\"text\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53dfd826",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-19T22:12:39.865429Z",
     "start_time": "2023-09-19T22:12:39.865420Z"
    }
   },
   "outputs": [],
   "source": [
    "# # To load\n",
    "# with open(os.path.join(output_dir, f\"{prompt['lan']}_{prompt['lan']}_{i}_gen{j}.npz\"), \"rb\") as fp:\n",
    "#     check_arr = np.load(fp)[\"arr\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a01c66c0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-19T15:36:16.933210Z",
     "start_time": "2023-09-19T15:36:16.931230Z"
    }
   },
   "source": [
    "# Evaluation on Continuation completeness"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4505cae3",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-19T22:12:39.866109Z",
     "start_time": "2023-09-19T22:12:39.866100Z"
    }
   },
   "outputs": [],
   "source": [
    "tag_wers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fd992a7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-19T22:12:39.866706Z",
     "start_time": "2023-09-19T22:12:39.866698Z"
    }
   },
   "outputs": [],
   "source": [
    "tag_cers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a5b163e",
   "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.12"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
