{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"2\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.gpt.chirp_v2_5 import (\n",
    "    preload_codec_models,\n",
    "    _get_model_if_needed,\n",
    "    GenerationConfig,\n",
    "    codec_decode_stream_to_full_audio,\n",
    ")\n",
    "from suno_utils.gpt.generation_engine import (\n",
    "    make_prompt,\n",
    "    align_codes,\n",
    "    make_request,\n",
    ")\n",
    "from suno_utils.gpt.engine import Engine, Request, Job\n",
    "from suno_utils.gpt.prompt import Prompt\n",
    "\n",
    "N_BATCH = 2\n",
    "MAX_STREAMS = N_BATCH * 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = preload_codec_models(\"/app/suno/models/chirp_v2/dac_2c_25x12.pt\")\n",
    "gpt_ckpt_path = _get_model_if_needed(\n",
    "    #     \"s3://suno-data/georg/checkpoints/chirp_v2_5/7b_ft.pt\"\n",
    "    #     \"/app/suno/checkpoints/2024-05-09_14-19-24/last_ckpt_infer__fix.pt\"  # 13b\n",
    "    #    \"/app/suno/data/dpo/models/model_30b_150k.pt\"  # 30b\n",
    "    #    \"/app/suno/checkpoints/2024-07-01_14-55-52/last_ckpt_infer.pt\" # 30b 2h 4l\n",
    "    #     \"/app/suno/checkpoints/2024-07-08_15-35-27/last_ckpt_infer.pt\" # 30b ft\n",
    "    #   \"/app/suno/checkpoints/2024-07-11_20-25-42/last_ckpt_infer.pt\" # 30b ft r2\n",
    "    #    \"/app/suno/checkpoints/2024-07-16_03-21-53/10k_ckpt_infer.pt\" # 30b ft r3\n",
    "    #    \"/app/suno/checkpoints/2024-07-20_01-08-38/last_ckpt_infer.pt\"  # 30b ft r4\n",
    "    #    \"/app/suno/data/dpo/models/model_30b_fix_ft2_20k.pt\" # 30b t1\n",
    "    #    \"/app/suno/checkpoints/2024-08-23_13-28-02/last_ckpt_infer.pt\" # 7b delay 1\n",
    "    #    \"/app/suno/checkpoints/2024-08-07_19-38-28/last_ckpt_infer.pt\" # 7b\n",
    "    #    \"/app/suno/checkpoints/2024-07-24_20-16-51/last_ckpt_infer.pt\"  # 30b ft r5 half sem delay\n",
    "    #    \"/app/suno/checkpoints/2024-07-26_15-28-26/last_ckpt_infer.pt\"  # 30b ft r6 32 delay\n",
    "    #   \"/app/suno/checkpoints/2024-07-27_18-21-41/last_ckpt_infer.pt\"  # 30b ft r7 14 delay\n",
    "    #    \"/app/suno/checkpoints/2024-07-31_04-15-07/last_ckpt_infer.pt\"  # 30b ft r8 5 sem delay\n",
    "    #    \"/app/suno/data/dpo/models/model_30b_ft_t1_5k.pt\" # 30b ft t1 2k step\n",
    "    #     \"/app/suno/tmp/2b_plus.pt\"\n",
    "    #     \"/app/suno/tmp/2b_musicfm.pt\"\n",
    "    #     \"/app/suno/tmp/2b.pt\",\n",
    "    #     \"/app/suno/tmp/2b_c25.pt\",\n",
    "    #     \"/app/suno/tmp/2b_s250c1.pt\",\n",
    "    #     \"/app/suno/tmp/2b_s250c5.pt\",\n",
    "    #     \"/app/suno/tmp/2b_s250c25.pt\",\n",
    "    #     \"/app/suno/tmp/2b_noSem_s250c25.pt\",\n",
    "    #    \"/app/suno/data/dpo/models/model_30b_ft_t3.pt\" # 30b t3 v2\n",
    "    \"s3://suno-data/tony/tmp/model_30b_t5_v15_b40.pt\"  # 30b t5\n",
    ")\n",
    "engine = Engine(\n",
    "    gpt_ckpt_path,\n",
    "    \"/app/suno/models/chirp_v2/tokenizer_60k.json\",\n",
    "    max_sequences=MAX_STREAMS,\n",
    "    compile=False,\n",
    ")\n",
    "model = engine.model\n",
    "cfg = model.config\n",
    "tokenizer = engine.tokenizer\n",
    "\n",
    "print(cfg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basic preds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[Verse 1]\n",
    "Almost Heaven, West Virginia\n",
    "Blue Ridge Mountains, Shenandoah River\n",
    "Life is old there, older than the trees\n",
    "Younger than the mountains, growing like a breeze\n",
    "\n",
    "[Chorus]\n",
    "Country roads, take me home\n",
    "To the place I belong\n",
    "West Virginia, mountain mama\n",
    "Take me home, country roads\n",
    "\"\"\"\n",
    "\n",
    "gconf = GenerationConfig(\n",
    "    text=text,\n",
    "    #     text_tags=\"dark epic trailer music, violin, vocalizations, female\",\n",
    "    #     text_tags=\"pop, power ballad, female, violin\",\n",
    "    text_tags=\"bluegrass, female\",\n",
    "    #     text_tags=\"sea shanty, sea shanty, sea shanty\",\n",
    "    cfg_coef=1.3,\n",
    "    cfg_coef_tags=2,\n",
    "    n_repeat_tags=3,\n",
    "    # text_start_control_tags=\"{start;start:0;vocals:intro}\",\n",
    "    # text_end_control_tags=\"{end}\",\n",
    "    cfg_coef_neg_tags=-2,\n",
    "    text_neg_tags=\"repetitive, loop\",\n",
    "    #     temp_semantic=0.85,\n",
    "    #     temp_coarse=0.9,\n",
    "    #     top_k_semantic=None,\n",
    "    #     top_k_coarse=None,\n",
    "    #     top_p_semantic=0.95,\n",
    "    #     top_p_coarse=0.98,\n",
    "    n_batch=1,\n",
    "    min_eos_p=0.1,\n",
    "    #     allow_eos=False,\n",
    "    min_text_offset=0,\n",
    "    eos_pad_duration_s=0,\n",
    "    max_gen_duration_s=60,\n",
    ")\n",
    "\n",
    "request = make_request(\"0\", gconf, cfg, tokenizer)\n",
    "in_arr = request.streams[0].prompt\n",
    "# prompt.visualize(in_arr)\n",
    "# prompt = Prompt(gconf, cfg, tokenizer)\n",
    "# prompt.visualize(in_arr)\n",
    "# for n in range(4):\n",
    "#     print(\"--\", n, \"--\")\n",
    "#     print(tokenizer.decode(request.streams[n].prompt[0][:-1].tolist()))\n",
    "#     print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "requests = [\n",
    "    make_request(f\"{i}\", gconf, engine.model.config, engine.tokenizer)\n",
    "    for i in range(N_BATCH)\n",
    "]\n",
    "# print(requests[0])\n",
    "jobs = engine.run_request(requests, tqdm_enabled=True)\n",
    "\n",
    "for job in jobs:\n",
    "    stream = engine.token_generator(job)\n",
    "    audio = codec_decode_stream_to_full_audio(align_codes(stream, cfg))\n",
    "    audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "requests = [\n",
    "    make_request(f\"{i}\", gconf, engine.model.config, engine.tokenizer)\n",
    "    for i in range(N_BATCH)\n",
    "]\n",
    "# print(requests[0])\n",
    "jobs = engine.run_request(requests, tqdm_enabled=True)\n",
    "\n",
    "for job in jobs:\n",
    "    stream = engine.token_generator(job)\n",
    "    audio = codec_decode_stream_to_full_audio(align_codes(stream, cfg))\n",
    "    audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[Intro]\n",
    "I can feel the heat risin'\n",
    "Everything is on fire\n",
    "Today's a painful re-\n",
    "Minder of why\n",
    "It can only get brighter\n",
    "The further you put it behind ya\n",
    "But right now I'm on the\n",
    "Inside lookin' out, ‘cause—\n",
    "\n",
    "[Chorus]\n",
    "I'm standin' in the flames\n",
    "And it’s a beautiful kind of pain\n",
    "Settin' fire to yesterday\n",
    "To find the light, find the light, find the light\n",
    "Standin' in the flames\n",
    "And it’s a beautiful kind of pain\n",
    "Settin' fire to yesterday\n",
    "Find the light, find the light, find the light\n",
    "\n",
    "[Verse 1]\n",
    "Yesterday was the tornado warning, today's like the morning after\n",
    "Your world is torn in half, you wake in its wake\n",
    "To start the mourning process and rebuilding, you're still a work in progress\n",
    "Today's a whole new chapter, it's like an enormous asthma\n",
    "Thunderstorm has passed ya, you weathered it and poked its\n",
    "Eye out with the thornbush that ya used to smell the roses\n",
    "Stopped to inhale, can't even tell your nose is stuffed\n",
    "So focused on the bright side then you floor the gas pedal\n",
    "And hit the corner faster, more assertive, never looking back\n",
    "May hit the curb, but every day's a new learning curve as ya\n",
    "Steer through life, sometimes you might not wanna swerve but you have to\n",
    "To avert a disaster, lucky no permanent damage\n",
    "‘Cause they hurt you so ba\n",
    "\"\"\"\n",
    "\n",
    "gconf = GenerationConfig(\n",
    "    text=text,\n",
    "    text_tags=\"Rap, In English, USA, Midwest Rap, Detroit Rap, Alternative, Hip-Hop\",\n",
    "    cfg_coef=1.3,\n",
    "    cfg_coef_tags=2,\n",
    "    n_repeat_tags=3,\n",
    "    #     temp_semantic=0.8,\n",
    "    #     temp_coarse=1,\n",
    "    #     top_k_semantic=1000,\n",
    "    #     top_k_coarse=100,\n",
    "    #     top_p_semantic=None,\n",
    "    #     top_p_coarse=None,\n",
    "    #     text_start_control_tags=\"{start;start:0;vocals:intro}\",\n",
    "    #     text_end_control_tags=\"{end}\",\n",
    "    #     cfg_coef_neg_tags=-2,\n",
    "    #     text_neg_tags=\"repetitive, loop\",\n",
    "    n_batch=1,\n",
    "    min_eos_p=0.1,\n",
    "    min_text_offset=0,\n",
    "    eos_pad_duration_s=0,\n",
    "    max_gen_duration_s=60 * 2,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "requests = [\n",
    "    make_request(f\"{i}\", gconf, engine.model.config, engine.tokenizer)\n",
    "    for i in range(N_BATCH)\n",
    "]\n",
    "# print(requests[0])\n",
    "jobs = engine.run_request(requests, tqdm_enabled=True)\n",
    "\n",
    "for job in jobs:\n",
    "    stream = engine.token_generator(job)\n",
    "    audio = codec_decode_stream_to_full_audio(align_codes(stream, cfg))\n",
    "    audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Custom Input Array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from suno_utils.audio import Audio\n",
    "from suno_utils.gpt.chirp_v2_5 import (\n",
    "    codec_encode,\n",
    "    codec_decode,\n",
    "    preload_semantic_models,\n",
    "    semantic_encode,\n",
    ")\n",
    "\n",
    "_ = preload_semantic_models(\n",
    "    checkpoint_filepath=\"/app/suno/models/chirp_v2/mert_25.pt\",\n",
    "    centroids_filepath=\"/app/suno/models/chirp_v2/mert_25_2x4k.npy\",\n",
    "    device=\"cpu\",\n",
    ")\n",
    "\n",
    "# # MusicFM\n",
    "# from suno_utils.tasks.musicfm_v2 import (\n",
    "#     preload_models as preload_semantic_models,\n",
    "#     encode as semantic_encode\n",
    "# )\n",
    "# _ = preload_semantic_models(\n",
    "#     centroids_filepath=\"s3://suno-data/minz/models/musicfm_concat_centroids_2s_8192.npy\",\n",
    "# )\n",
    "\n",
    "\n",
    "def process_audio(audio, cfg):\n",
    "    audio = audio.normalize_volume(-16)\n",
    "    sem_arr = semantic_encode(audio, device=\"cpu\")\n",
    "    coarse_arr = codec_encode(audio)\n",
    "    n_frames = min(sem_arr.shape[0], coarse_arr.shape[0])\n",
    "    sem_arr = sem_arr[:n_frames, : cfg.semantic_n_codebooks]\n",
    "    coarse_arr = coarse_arr[:n_frames, : cfg.coarse_n_codebooks]\n",
    "\n",
    "    a_arr = np.concatenate([sem_arr, coarse_arr], axis=-1)\n",
    "    return a_arr\n",
    "\n",
    "\n",
    "def load_audio(fp):\n",
    "    return Audio.from_file(fp, n_channels=2, sample_rate=48_000, byte_width=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Continue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio = Audio.from_file(\n",
    "    \"audios/50cent.wav\", n_channels=2, sample_rate=48_000, byte_width=2\n",
    ")  # 50 cent\n",
    "audio = audio.get_segment(from_s=24.99, to_s=45)\n",
    "# audio = Audio.from_file(\"tmp/eminem_doctor.wav\", n_channels=2, sample_rate=48_000, byte_width=2) # eminem doctor\n",
    "history_text = \"\"\"\n",
    "Welcome to the candy shop\n",
    "\n",
    "Yeah, uh-huh\n",
    "So seductive\n",
    "\n",
    "I take you to the candy shop\n",
    "I let you lick the lollipop\n",
    "Go 'head, girl, don't you stop\n",
    "\"\"\"\n",
    "text = \"\"\"\n",
    "Almost Heaven, West Virginia\n",
    "Blue Ridge Mountains, Shenandoah River\n",
    "Life is old there, older than the trees\n",
    "Younger than the mountains, growing like a breeze\n",
    "\n",
    "Country roads, take me home\n",
    "To the place I belong\n",
    "West Virginia, mountain mama\n",
    "Take me home, country roads\n",
    "\"\"\"\n",
    "in_history_arr = process_audio(audio, cfg)\n",
    "audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# codec_decode(codec_encode(audio)).play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gconf = GenerationConfig(\n",
    "    text=text,\n",
    "    #     text_tags=\"bluegrass female\",\n",
    "    history_arr=in_history_arr,\n",
    "    history_text=history_text,\n",
    "    n_batch=1,\n",
    "    cfg_coef=1.2,\n",
    "    cfg_coef_tags=0,\n",
    "    max_gen_duration_s=40,\n",
    "    min_text_offset=0,\n",
    "    eos_pad_duration_s=0,\n",
    ")\n",
    "\n",
    "request = make_request(\"0\", gconf, cfg, tokenizer)\n",
    "in_arr = request.streams[0].prompt\n",
    "\n",
    "prompt = Prompt(gconf, cfg, tokenizer)\n",
    "prompt.visualize(in_arr)\n",
    "requests = [\n",
    "    make_request(f\"{i}\", gconf.modify(n_batch=1), engine.model.config, engine.tokenizer)\n",
    "    for i in range(N_BATCH)\n",
    "]\n",
    "for job in engine.run_request(requests, tqdm_enabled=True):\n",
    "    stream = engine.token_generator(job)\n",
    "\n",
    "    audio = codec_decode_stream_to_full_audio(align_codes(stream, cfg))\n",
    "    audio.play()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Oracle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio = load_audio(\"audios/halo.wav\").get_segment(from_s=49.99, to_s=70)\n",
    "text = \"\"\"\n",
    "It's like I've been awakened\n",
    "Every rule, I had you breakin'\n",
    "The risk that I'm takin'\n",
    "I'm never gonna shut you out\n",
    "\n",
    "Everywhere I'm lookin' now\n",
    "I'm surrounded by your embrace\n",
    "Baby, I can see your halo\n",
    "\"\"\"\n",
    "in_oracle_arr = process_audio(audio, cfg)\n",
    "# in_oracle_arr[:,:1] = cfg.semantic_mask_token\n",
    "audio.normalize_volume().play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# codec_decode(codec_encode(audio)[:,:6]).normalize_volume().play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gconf = GenerationConfig(\n",
    "    text=text,\n",
    "    #     text_tags=\"bluegrass female\",\n",
    "    oracle_arr=in_oracle_arr[:, :1],\n",
    "    n_batch=1,\n",
    "    cfg_coef=1,  # .3,\n",
    "    cfg_coef_tags=0,\n",
    "    cfg_coef_neg_tags=0,\n",
    "    max_gen_duration_s=20,\n",
    "    min_text_offset=0,\n",
    "    eos_pad_duration_s=0,\n",
    ")\n",
    "\n",
    "request = make_request(\"0\", gconf, cfg, tokenizer)\n",
    "in_arr = request.streams[0].prompt\n",
    "\n",
    "prompt = Prompt(gconf, cfg, tokenizer)\n",
    "\n",
    "requests = [\n",
    "    make_request(f\"{i}\", gconf.modify(n_batch=1), engine.model.config, engine.tokenizer)\n",
    "    for i in range(N_BATCH)\n",
    "]\n",
    "for job in engine.run_request(requests, tqdm_enabled=True):\n",
    "    stream = engine.token_generator(job)\n",
    "\n",
    "    audio = codec_decode_stream_to_full_audio(align_codes(stream, cfg))\n",
    "    audio.play()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Covers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# audio = Audio.from_file(\"tmp/georg_country.mp3\", n_channels=2).get_segment(to_s=60.01)\n",
    "# audio = Audio.from_file(\"tmp/georg_bday.mp3\", n_channels=2).get_segment(to_s=60.01)\n",
    "# audio = load_audio(\"audios/halo.wav\").get_segment(from_s=49.99, to_s=70)\n",
    "audio = load_audio(\"audios/MurphysLawSchool_short_super.mp3\")\n",
    "audio.play()\n",
    "in_cover_arr = process_audio(audio, cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_npz = np.load(\"audios/4bd01278-b5a7-491d-8fa7-58cbefa09873.npz\")\n",
    "in_cover_arr = test_npz[\"v3.0_raw\"]\n",
    "print(in_cover_arr.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# text = \"\"\"\n",
    "# Happy birthday to you\n",
    "# Happy birthday to you\n",
    "# Happy birthday dear Margu\n",
    "# Happy birthday to you\n",
    "\n",
    "# Happy birthday to you\n",
    "# Happy birthday to you\n",
    "# Happy birthday dear Keenan\n",
    "# Happy birthday to you\n",
    "\n",
    "# Happy birthday to you\n",
    "# Happy birthday to you\n",
    "# Happy birthday dear Mikey\n",
    "# Happy birthday to you\n",
    "\n",
    "# Happy birthday to you\n",
    "# Happy birthday to you\n",
    "# Happy birthday dear Suno\n",
    "# Happy birthday to you\n",
    "\n",
    "# Happy birthday to you\n",
    "# Happy birthday to you\n",
    "# I don't know what I'm gonna do next\n",
    "# Happy birthday to you\n",
    "# \"\"\"\n",
    "\n",
    "# text = \"\"\"\n",
    "# [Verse 1]\n",
    "# Almost Heaven, West Virginia\n",
    "# Blue Ridge Mountains, Shenandoah River\n",
    "# Life is old there, older than the trees\n",
    "# Younger than the mountains, growing like a breeze\n",
    "\n",
    "# [Chorus]\n",
    "# Country roads, take me home\n",
    "# To the place I belong\n",
    "# West Virginia, mountain mama\n",
    "# Take me home, country roads\n",
    "\n",
    "# [Verse 2]\n",
    "# All my memories gather 'round her\n",
    "# Miner's lady, stranger to blue water\n",
    "# Dark and dusty, painted on the sky\n",
    "# Misty taste of moonshine, teardrop in my eye\n",
    "# \"\"\"\n",
    "\n",
    "# text = \"\"\"\n",
    "# been awakened\n",
    "# Every rule, I had you breakin'\n",
    "# The risk that I'm takin'\n",
    "# I'm never gonna shut you out\n",
    "\n",
    "# [Chorus]\n",
    "# Everywhere I'm lookin' now\n",
    "# I'm surrounded by your embrace\n",
    "# Baby, I can see your halo\n",
    "# \"\"\"\n",
    "\n",
    "text = \"\"\"\n",
    "feel the breeze that's calling me\n",
    " golden skies above the sea\n",
    " endless summer chasing dreams \n",
    "yeah you and i we werere wild and free \n",
    "\n",
    "we chillin out just riding the vibe under the stars you by my side it feels so right the night is ours you and me are shooting stars \n",
    "\n",
    "we'll never fade the nights into days a summer will'll always stay \n",
    "\n",
    "on the sand dunes we run we fly \n",
    "burning bright beneath the moonlit sky \n",
    "hold me close don't let this go \n",
    "\n",
    "drifting closer side by side feel us fading with the tide \n",
    "every glance every touch \n",
    "makes me feel like it's enough \n",
    "\n",
    "we were chillining out just riding the vibe under the stars you by my side feels so right the night is ours you and me i' shooting stars \n",
    "\n",
    "we'll never fade the nights into days our summer will always stay \n",
    "\n",
    "on the sand dunes we run we fly \n",
    "building bright beneath the moonlit sky every come don let this go \n",
    "the d on the sand dunes we'll be writing our story \n",
    "\n",
    "every moment is a memory of glory on the sand dun run we fly \n",
    "when the winds blow in we sway we sway to the shifting sands we sway we sway \n",
    "\n",
    "we'll never fade the nights and the days a summer will always stay the moonit sky hearts \n",
    "don't let this go on the sand dun well be writing our story every moment is a memory of glory on the sand dun above we fly on dun\n",
    "\"\"\"\n",
    "\n",
    "gconf = GenerationConfig(\n",
    "    text=text,\n",
    "    text_tags=\"female, k-pop\",\n",
    "    cover_arr=in_cover_arr[:3000, :],\n",
    "    n_batch=1,\n",
    "    cfg_coef=1.0,\n",
    "    cfg_coef_tags=2,\n",
    "    max_gen_duration_s=60,\n",
    "    min_text_offset=0,\n",
    "    eos_pad_duration_s=0,\n",
    "    n_repeat_tags=3,\n",
    ")\n",
    "\n",
    "request = make_request(\"0\", gconf, cfg, tokenizer)\n",
    "in_arr = request.streams[0].prompt\n",
    "\n",
    "prompt = Prompt(gconf, cfg, tokenizer)\n",
    "prompt.visualize(in_arr, compress=True)\n",
    "\n",
    "assert N_BATCH * 4 <= 16\n",
    "requests = [\n",
    "    make_request(f\"{i}\", gconf.modify(n_batch=1), engine.model.config, engine.tokenizer)\n",
    "    for i in range(N_BATCH)\n",
    "]\n",
    "for job in engine.run_request(requests, tqdm_enabled=True):\n",
    "    stream = engine.token_generator(job)\n",
    "\n",
    "    audio = codec_decode_stream_to_full_audio(align_codes(stream, cfg))\n",
    "    audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_npz = np.load(\"/home/tony/dpo_batch_iter_0.npz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for k in test_npz.keys():\n",
    "    print(k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# concat generated to prompt and visualize\n",
    "import numpy as np\n",
    "\n",
    "# in_arr = requests[0].streams[0].prompt\n",
    "# generated = np.stack([t.cpu().numpy() for t in job.generated_tokens]).T\n",
    "# print(generated.shape)\n",
    "# arr = np.concatenate([in_arr, generated], axis=1)\n",
    "\n",
    "arr = test_npz[\"X\"][0, :]\n",
    "print(arr.shape)\n",
    "prompt.visualize(arr)\n",
    "prompt.visualize(arr, compress=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "in_arr = requests[0].streams[0].prompt\n",
    "generated = np.stack([t.cpu().numpy() for t in job.generated_tokens]).T\n",
    "print(generated.shape)\n",
    "arr = np.concatenate([in_arr, generated], axis=1)\n",
    "print(arr.shape)\n",
    "prompt.visualize(arr)\n",
    "prompt.visualize(arr, compress=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Artist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# audio = Audio.from_file(\"audios/georg_country.mp3\", n_channels=2).get_segment(to_s=5.01)\n",
    "audio = Audio.from_file(\n",
    "    \"audios/Erlkoning_tony.mp3\", n_channels=2\n",
    ")  # .get_segment(to_s=5.01)\n",
    "# audio = load_audio(\"tmp/halo.wav\").get_segment(from_s=49.99, to_s=70)\n",
    "audio.play()\n",
    "in_artist_arr = process_audio(audio, cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cover_audio = Audio.from_file(\n",
    "    \"audios/erlkonig_origin.mp3\", n_channels=2\n",
    ")  # .get_segment(to_s=5.01)\n",
    "# # audio = load_audio(\"tmp/halo.wav\").get_segment(from_s=49.99, to_s=70)\n",
    "cover_audio.play()\n",
    "cover_audio_arr = process_audio(cover_audio, cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[Verse]\n",
    "Week grind over, feelin' heavyweight,\n",
    "Pancakes stackin', syrup up on the plate.\n",
    "Phone on mute, stress just evaporate,\n",
    "Netflix loopin', Sunday can't wait.\n",
    "\n",
    "[Verse 2]\n",
    "Sunshine creepin' through a cozy quilt,\n",
    "Slidin' in them slippers, like silk was built.\n",
    "No alarm buzzin', tranquil tilt,\n",
    "Chillin' on the porch, coffee cup spilt.\n",
    "\n",
    "[Chorus]\n",
    "Sunday funday, slow-motion sway,\n",
    "Kickin' back, relaxin' the lazy day.\n",
    "No rush, no fuss, come what may,\n",
    "Vibin' simple, let worries decay.\n",
    "\n",
    "[Verse 3]\n",
    "Neighbors hollerin’, barbecue scent,\n",
    "Feet propped up, bills ain't even rent.\n",
    "Emails ignored, on no response bent,\n",
    "Dreamin' big dreams 'til the day is spent.\n",
    "\n",
    "[Verse 4]\n",
    "Game on the screen, chips in the bowl,\n",
    "Friends gather 'round, hearts feelin’ whole.\n",
    "Laughter echoes, fillin' up the soul,\n",
    "Sunday grooves, under control.\n",
    "\n",
    "[Bridge]\n",
    "City slows down, nightfall arises,\n",
    "Stars in the sky, cosmic surprises.\n",
    "Monday creeps near, but no compromises,\n",
    "Sunday chill, calm compromises.\n",
    "\"\"\"\n",
    "\n",
    "gconf = GenerationConfig(\n",
    "    text=text,\n",
    "    text_tags=\"classical lieder, piano, male vocal\",\n",
    "    artist_arr=in_artist_arr[:3000, :],\n",
    "    n_batch=1,\n",
    "    cfg_coef=1.5,\n",
    "    cfg_coef_tags=4,\n",
    "    max_gen_duration_s=40,\n",
    "    min_text_offset=0,\n",
    "    eos_pad_duration_s=0,\n",
    ")\n",
    "\n",
    "request = make_request(\"0\", gconf, cfg, tokenizer)\n",
    "in_arr = request.streams[0].prompt\n",
    "\n",
    "prompt = Prompt(gconf, cfg, tokenizer)\n",
    "\n",
    "assert N_BATCH * 4 <= 16\n",
    "requests = [\n",
    "    make_request(f\"{i}\", gconf.modify(n_batch=1), engine.model.config, engine.tokenizer)\n",
    "    for i in range(N_BATCH)\n",
    "]\n",
    "for job in engine.run_request(requests, tqdm_enabled=True):\n",
    "    stream = engine.token_generator(job)\n",
    "\n",
    "    audio = codec_decode_stream_to_full_audio(align_codes(stream, cfg))\n",
    "    audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# text = \"\"\"\n",
    "# [Verse 1]\n",
    "# Almost Heaven, West Virginia\n",
    "# Blue Ridge Mountains, Shenandoah River\n",
    "# Life is old there, older than the trees\n",
    "# Younger than the mountains, growing like a breeze\n",
    "\n",
    "# [Chorus]\n",
    "# Country roads, take me home\n",
    "# To the place I belong\n",
    "# West Virginia, mountain mama\n",
    "# Take me home, country roads\n",
    "\n",
    "# [Verse 2]\n",
    "# All my memories gather 'round her\n",
    "# Miner's lady, stranger to blue water\n",
    "# Dark and dusty, painted on the sky\n",
    "# Misty taste of moonshine, teardrop in my eye\n",
    "# \"\"\"\n",
    "\n",
    "# text = \"\"\"\n",
    "# [Verse 1]\n",
    "# Remember those walls I built?\n",
    "# Well, baby, they're tumblin' down\n",
    "# And they didn't even put up a fight\n",
    "# They didn't even make a sound\n",
    "\n",
    "# [Chorus]\n",
    "# Everywhere I'm lookin' now\n",
    "# I'm surrounded by your embrace\n",
    "# Baby, I can see your halo\n",
    "# You know you're my savin' grace\n",
    "\n",
    "# [Verse 2]\n",
    "# Hit me like a ray of sun\n",
    "# Burnin' through my darkest night\n",
    "# You're the only one that I want\n",
    "# Think I'm addicted to your light\n",
    "# \"\"\"\n",
    "\n",
    "\n",
    "text = \"\"\"\n",
    "Wer reitet so spät durch Nacht und Wind?\n",
    "Es ist der Vater mit seinem Kind\n",
    "Er hat den Knaben wohl in dem Arm\n",
    "Er fasst ihn sicher, er hält ihn warm\n",
    "\n",
    "Mein Sohn\n",
    "Was birgst du so bang dein Gesicht?\n",
    "Siehst, Vater, du den Erlkönig nicht?\n",
    "Den Erlenkönig mit Kron' und Schweif?\n",
    "\n",
    "Mein Sohn, es ist ein Nebelstreif\n",
    "Du liebes Kind, komm, geh mit mir\n",
    "Gar schöne Spiele spiel' ich mit dir\n",
    "Manch' bunte Blumen sind an dem Strand\n",
    "Meine Mutter hat manch' gülden Gewand\n",
    "\n",
    "Mein Vater, mein Vater, und hörest du nicht\n",
    "Was Erlenkönig mir leise verspricht?\n",
    "Sei ruhig, bleibe ruhig, mein Kind\n",
    "In dürren Blättern säuselt der Wind\n",
    "\"\"\"\n",
    "\n",
    "gconf = GenerationConfig(\n",
    "    text=text,\n",
    "    text_tags=\"male vocal, piano\",\n",
    "    artist_arr=in_artist_arr[: 25 * 35, :],\n",
    "    cover_arr=cover_audio_arr[25 * 30 : 25 * 120, :],\n",
    "    n_batch=1,\n",
    "    cfg_coef=1.1,\n",
    "    cfg_coef_tags=2.5,\n",
    "    max_gen_duration_s=60,\n",
    "    min_text_offset=0,\n",
    "    eos_pad_duration_s=0,\n",
    ")\n",
    "\n",
    "request = make_request(\"0\", gconf, cfg, tokenizer)\n",
    "in_arr = request.streams[0].prompt\n",
    "\n",
    "prompt = Prompt(gconf, cfg, tokenizer)\n",
    "\n",
    "assert N_BATCH * 4 <= 16\n",
    "requests = [\n",
    "    make_request(f\"{i}\", gconf.modify(n_batch=1), engine.model.config, engine.tokenizer)\n",
    "    for i in range(N_BATCH)\n",
    "]\n",
    "for job in engine.run_request(requests, tqdm_enabled=True):\n",
    "    stream = engine.token_generator(job)\n",
    "\n",
    "    audio = codec_decode_stream_to_full_audio(align_codes(stream, cfg))\n",
    "    audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "request = make_request(\"0\", gconf, cfg, tokenizer)\n",
    "in_arr = request.streams[0].prompt\n",
    "\n",
    "prompt = Prompt(gconf, cfg, tokenizer)\n",
    "\n",
    "assert N_BATCH * 4 <= 16\n",
    "requests = [\n",
    "    make_request(f\"{i}\", gconf.modify(n_batch=1), engine.model.config, engine.tokenizer)\n",
    "    for i in range(N_BATCH)\n",
    "]\n",
    "for job in engine.run_request(requests, tqdm_enabled=True):\n",
    "    stream = engine.token_generator(job)\n",
    "\n",
    "    audio = codec_decode_stream_to_full_audio(align_codes(stream, cfg))\n",
    "    audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "been awakened\n",
    "Every rule, I had you breakin'\n",
    "The risk that I'm takin'\n",
    "I'm never gonna shut you out\n",
    "\n",
    "[Chorus]\n",
    "Everywhere I'm lookin' now\n",
    "I'm surrounded by your embrace\n",
    "Baby, I can see your halo\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gconf = GenerationConfig(\n",
    "    text=text,\n",
    "    text_tags=\"bluegrass, male vocal\",\n",
    "    artist_arr=in_artist_arr,\n",
    "    # cover_arr=in_cover_arr,\n",
    "    n_batch=1,\n",
    "    cfg_coef=1.5,\n",
    "    cfg_coef_tags=2,\n",
    "    max_gen_duration_s=40,\n",
    "    min_text_offset=0,\n",
    "    eos_pad_duration_s=0,\n",
    ")\n",
    "\n",
    "request = make_request(\"0\", gconf, cfg, tokenizer)\n",
    "in_arr = request.streams[0].prompt\n",
    "\n",
    "prompt = Prompt(gconf, cfg, tokenizer)\n",
    "\n",
    "assert N_BATCH * 4 <= 16\n",
    "requests = [\n",
    "    make_request(f\"{i}\", gconf.modify(n_batch=1), engine.model.config, engine.tokenizer)\n",
    "    for i in range(N_BATCH)\n",
    "]\n",
    "for job in engine.run_request(requests, tqdm_enabled=True):\n",
    "    stream = engine.token_generator(job)\n",
    "\n",
    "    audio = codec_decode_stream_to_full_audio(align_codes(stream, cfg))\n",
    "    audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Infill"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# audio = Audio.from_file(\"tmp/georg_country.mp3\", n_channels=2).get_segment(to_s=5.01)\n",
    "audio = load_audio(\"/home/sara/samples/synthwave_house.mp3\").get_segment(\n",
    "    from_s=30, to_s=120\n",
    ")\n",
    "audio.play()\n",
    "in_infill_arr = process_audio(audio, cfg)\n",
    "# in_history_arr = in_infill_arr[:25*10,:].copy()\n",
    "# in_future_arr = in_infill_arr[-25*10:,:].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gconf = GenerationConfig(\n",
    "    text=\"[instrumental]\",\n",
    "    text_tags=\"rock\",\n",
    "    # history_arr=in_history_arr,\n",
    "    # future_arr=in_future_arr,\n",
    "    cover_arr=in_infill_arr,\n",
    "    n_batch=1,\n",
    "    cfg_coef=1.0,  # no text cfg for dpo stream\n",
    "    min_eos_p=0.1,\n",
    "    eos_pad_duration_s=0,\n",
    "    cfg_coef_tags=2,\n",
    "    cfg_coef_neg_tags=-1,\n",
    "    text_neg_tags=\"repetitive, loop\",\n",
    "    n_repeat_tags=1,\n",
    "    use_whisper=False,\n",
    "    text_start_control_tags=\"{start:0} \",\n",
    "    min_text_offset=0,  # this is required for 30b\n",
    "    random_seed=42,\n",
    "    max_gen_duration_s=240,\n",
    ")\n",
    "\n",
    "request = make_request(\"0\", gconf, cfg, tokenizer)\n",
    "in_arr = request.streams[0].prompt\n",
    "\n",
    "prompt = Prompt(gconf, cfg, tokenizer)\n",
    "prompt.visualize(in_arr)\n",
    "assert N_BATCH * 4 <= 16\n",
    "requests = [\n",
    "    make_request(f\"{i}\", gconf.modify(n_batch=1), engine.model.config, engine.tokenizer)\n",
    "    for i in range(N_BATCH)\n",
    "]\n",
    "for job in engine.run_request(requests, tqdm_enabled=True):\n",
    "    stream = engine.token_generator(job)\n",
    "    audio = codec_decode(\n",
    "        np.concatenate(\n",
    "            [\n",
    "                # in_history_arr[:,-12:],\n",
    "                torch.stack(list(align_codes(stream, cfg)))\n",
    "                .detach()\n",
    "                .cpu()\n",
    "                .numpy()[:, -12:],\n",
    "                # in_future_arr[:,-12:],\n",
    "            ],\n",
    "            axis=0,\n",
    "        )\n",
    "    )\n",
    "    audio.play()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Playground"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import numpy as np\n",
    "# in_arr = requests[0].streams[0].prompt\n",
    "# generated = np.stack([t.cpu().numpy() for t in job.generated_tokens]).T\n",
    "# print(generated.shape)\n",
    "# arr = np.concatenate([in_arr, generated], axis=1)\n",
    "# prompt.visualize(arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_jsonl(filepath, parse_idx_set=None):\n",
    "    data = []\n",
    "    with open(filepath) as f:\n",
    "        line_idx = 0\n",
    "        for line in f:\n",
    "            line = line.strip()\n",
    "            if len(line) == 0:\n",
    "                continue\n",
    "            if parse_idx_set is not None and line_idx not in parse_idx_set:\n",
    "                data.append(None)\n",
    "                line_idx += 1\n",
    "                continue\n",
    "            m = json.loads(line)\n",
    "            data.append(m)\n",
    "            line_idx += 1\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_loss = read_jsonl(\n",
    "    \"/app/suno/data/dpo/2b_before_recode_v0/eval_val_loss/loss.jsonl\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(eval_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_loss[1193]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "[x for x in eval_loss if x[\"idx\"] == 311]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "\n",
    "c = Counter()\n",
    "for i in eval_loss:\n",
    "    c[i[\"idx\"]] += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(c))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "gpt_ckpt_path = (\n",
    "    \"/app/suno/checkpoints/2024-07-01_14-55-52/last_ckpt_infer.pt\"  # 30b 4 heads\n",
    ")\n",
    "ckpt = torch.load(gpt_ckpt_path, map_location=\"cpu\", mmap=True)\n",
    "print(ckpt.keys())\n",
    "\n",
    "# ckpt[\"model_args\"][\"global_every_n_layers\"] = 8\n",
    "\n",
    "# torch.save(ckpt, \"/app/suno/victor/30b_g8.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for key in ckpt[\"model\"].keys():\n",
    "#     print(key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "silent_audio = Audio.from_silence(sample_rate=48000, duration_s=1.04, n_channels=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "silent_audio_array = process_audio(silent_audio, cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "silent_audio_array.T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls /app/suno/data/dpo/30b_t1_v7"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"/app/suno/data/dpo/30b_t1_v7/2b_test_bt4_cached_loss.json\", \"r\") as fp:\n",
    "    test_2b_loss_bt4 = json.load(fp)\n",
    "with open(\"/app/suno/data/dpo/30b_t1_v7/2b_test_bt2_cached_loss.json\", \"r\") as fp:\n",
    "    test_2b_loss_bt2 = json.load(fp)\n",
    "with open(\"/app/suno/data/dpo/30b_t1_v7/2b_test_bt8_cached_loss.json\", \"r\") as fp:\n",
    "    test_2b_loss_bt8 = json.load(fp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pos_df_bt8 = pd.DataFrame.from_dict(test_2b_loss_bt8[\"train\"]).transpose()\n",
    "pos_df_bt4 = pd.DataFrame.from_dict(test_2b_loss_bt4[\"train\"]).transpose()\n",
    "pos_df_bt2 = pd.DataFrame.from_dict(test_2b_loss_bt2[\"train\"]).transpose()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pos_df_bt4[\"coarse_0\"].hist(bins=np.linspace(0, 10, 100), label=\"bt4\", alpha=0.5)\n",
    "pos_df_bt2[\"coarse_0\"].hist(bins=np.linspace(0, 10, 100), label=\"bt2\", alpha=0.5)\n",
    "pos_df_bt8[\"coarse_0\"].hist(bins=np.linspace(0, 10, 100), label=\"bt8\", alpha=0.5)\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "a = torch.tensor([3.7168])  # .to(torch.bfloat16)\n",
    "b = torch.tensor([4.2809])  # .to(torch.bfloat16)\n",
    "c = torch.tensor([3.7168])  # .to(torch.bfloat16)\n",
    "d = torch.tensor([4.2809])  # .to(torch.bfloat16)\n",
    "print(a, a.dtype, a.device)\n",
    "print(b, b.dtype, b.device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = a - c\n",
    "y = b - d\n",
    "print(x, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dict = {\"negative_tags\": \"\"}\n",
    "neg_tag = \"fuck\"\n",
    "result = test_dict.get(\"negative_tags\", neg_tag)\n",
    "print(result, len(result))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import torch\n",
    "# sd1 = torch.load(\"/app/suno/checkpoints/2024-09-10_03-55-08/last_ckpt_infer.pt\", map_location=\"cpu\")\n",
    "# sd2 = torch.load(\"/app/suno/checkpoints/2024-09-26_12-26-01/infer_2_5k.pt\", map_location=\"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for k, v in sd1[\"model\"].items():\n",
    "#     if (sd1[\"model\"][k].shape != sd2[\"model\"][k].shape):\n",
    "#         print(k)\n",
    "# for k, v in sd2[\"model\"].items():\n",
    "#     if (sd1[\"model\"][k].shape != sd2[\"model\"][k].shape):\n",
    "#         print(k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for k, v in sd1[\"model_args\"].items():\n",
    "#     # print(k)\n",
    "#     try:\n",
    "#         if (sd1[\"model_args\"][k] != sd2[\"model_args\"][k]):\n",
    "#             print(k, sd1[\"model_args\"][k], sd2[\"model_args\"][k])\n",
    "#     except:\n",
    "#         print(k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for k, v in sd1[\"model\"].items():\n",
    "    print(k)"
   ]
  },
  {
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
