{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fae8f9c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:13:55.101018Z",
     "start_time": "2024-05-11T20:13:51.243900Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"4\"\n",
    "os.environ[\"PYTORCH_CUDA_ALLOC_CONF\"] = \"expandable_segments:True\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c27fb2bf-0bb3-4b5d-9930-22c1a6867a90",
   "metadata": {},
   "outputs": [],
   "source": [
    "gpt_ckpt_path = \"/app/suno/checkpoints/2024-07-09_12-29-10/last_ckpt_infer.pt\" # 2h ft v2 2k bt4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a035f3d2-697c-401f-85ef-17289c9f8743",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch.nn.functional as F\n",
    "from suno_utils.audio import Audio\n",
    "from suno_utils.gpt.chirp_v2_5 import (\n",
    "    GenerationConfig,\n",
    "    decode_stream,\n",
    "    preload_models,\n",
    "    prep_gconf,\n",
    ")\n",
    "from suno_utils.gpt.generation_engine import (\n",
    "    align_codes,\n",
    "    make_request,\n",
    ")\n",
    "from suno_utils.gpt.engine import Engine\n",
    "import torch\n",
    "from tqdm import tqdm\n",
    "\n",
    "preload_models(load_gpt=False, load_semantic=False, load_codec_device=\"cuda\")\n",
    "\n",
    "engine = Engine(\n",
    "    gpt_ckpt_path,\n",
    "    \"/app/suno/data/dpo/models/tokenizer_60k.json\",\n",
    "    max_sequences=16,\n",
    "    compile=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18bb8e5f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:14:14.500100Z",
     "start_time": "2024-05-11T20:14:14.497272Z"
    }
   },
   "outputs": [],
   "source": [
    "general_config = dict(\n",
    "    cfg_coef=1.0,  # no text cfg for dpo stream\n",
    "    min_eos_p=0.001,\n",
    "    eos_pad_duration_s=0,\n",
    "    # cfg_coef_tags=0.0,\n",
    "    # cfg_coef_tags_max_steps=None,  # collect the data for now\n",
    "    # temp_coarse=0.95,\n",
    "    cfg_coef_tags=2,\n",
    "    cfg_coef_neg_tags=-1,\n",
    "    text_neg_tags=\"repetitive, loopy, noisy, distorted, repeat\",\n",
    "    n_repeat_tags=1,\n",
    "    use_whisper=False,\n",
    "    text_start_control_tags=\"{start:0;duration:120}\",\n",
    "    # text_end_control_tags=\"{end}{end}\",\n",
    "    random_seed=42,\n",
    "    n_batch=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f57ea1b-4133-4a37-a50f-872ae74ccb08",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_audio_with_engine(text, text_tags):\n",
    "    cfg = GenerationConfig(\n",
    "        text=text, text_tags=text_tags, max_gen_duration_s=240, **general_config\n",
    "    )\n",
    "    preped_cfg = prep_gconf(cfg)\n",
    "    model_conf = engine.model.config\n",
    "    requests = []\n",
    "    for i in range(1, 5):\n",
    "        request = make_request(\n",
    "            f\"{i}\", preped_cfg, engine.model.config, engine.tokenizer\n",
    "        )\n",
    "        requests.append(request)\n",
    "\n",
    "    jobs = engine.run_request(requests)\n",
    "\n",
    "    for job in jobs:\n",
    "        stream = engine.token_generator(job)\n",
    "        audios = []\n",
    "        for audio in decode_stream(align_codes(tqdm(stream), model_conf)):\n",
    "            audios.append(audio[0])\n",
    "        audio_continued = Audio.concatenate(audios)\n",
    "        audio_continued.play()\n",
    "    return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa44f837",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:14:14.566214Z",
     "start_time": "2024-05-11T20:14:14.501592Z"
    }
   },
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[intro]\n",
    "\n",
    "[verse]\n",
    "Walking down the street, feeling so alive\n",
    "Got my head in the clouds, got a gleam in my eye\n",
    "Every step I take, it's like a brand new start\n",
    "No matter where I'm going, I'll always find my part\n",
    "(oh-oh-oh)\n",
    "\n",
    "[chorus]\n",
    "Life is like a high-wire act, we're dancing in the sky\n",
    "No need to worry, no need to ask why\n",
    "With a little bit of courage, we can chase our dreams\n",
    "No matter what comes our way, we'll always be a team\n",
    "(we're unstoppable, yeah)\n",
    "\n",
    "[outro]\n",
    "\"\"\"\n",
    "text_tags=\"fast female vocalist, pop-punk, pop, melodic, hd stems, goes hard, 155bpm\"\n",
    "\n",
    "generate_audio_with_engine(text, text_tags)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "072adcda",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:15:43.758219Z",
     "start_time": "2024-05-11T20:15:43.756159Z"
    }
   },
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[Verse]\n",
    "一盏离愁孤灯伫立在窗口\n",
    "我在门后假装你人还没走\n",
    "旧地如重游月圆更寂寞\n",
    "夜半清醒的烛火不忍苛责我\n",
    "\n",
    "[Verse]\n",
    "一壶漂泊浪迹天涯难入喉\n",
    "你走之后酒暖回忆思念瘦\n",
    "水向东流时间怎么偷\n",
    "花开就一次成熟我却错过\n",
    "\n",
    "[Chorus]\n",
    "谁在用琵琶弹奏一曲东风破\n",
    "岁月在墙上剥落看见小时候\n",
    "犹记得那年我们都还很年幼\n",
    "而如今琴声幽幽我的等候你没听过\n",
    "\n",
    "[outro]\n",
    "\"\"\"\n",
    "text_tags=\"pop woman\"\n",
    "\n",
    "generate_audio_with_engine(text, text_tags)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc50a93c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:18:32.473209Z",
     "start_time": "2024-05-11T20:17:10.055196Z"
    }
   },
   "outputs": [],
   "source": [
    "text=\"[intro] \\n \" + text\n",
    "\n",
    "generate_audio_with_engine(text, text_tags)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fa595a6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:18:36.675596Z",
     "start_time": "2024-05-11T20:18:36.616364Z"
    }
   },
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[verse]\n",
    "Walking down the street, feeling so alive\n",
    "Got my head in the clouds, got a gleam in my eye\n",
    "Every step I take, it's like a brand new start\n",
    "No matter where I'm going, I'll always find my part\n",
    "(oh-oh-oh)\n",
    "\n",
    "[chorus]\n",
    "Life is like a high-wire act, we're dancing in the sky\n",
    "No need to worry, no need to ask why\n",
    "With a little bit of courage, we can chase our dreams\n",
    "No matter what comes our way, we'll always be a team\n",
    "(we're unstoppable, yeah)\n",
    "\n",
    "[outro]\n",
    "\"\"\"\n",
    "text_tags=\"orchestral film epic\"\n",
    "\n",
    "generate_audio_with_engine(text, text_tags)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0fff5d8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:21:24.466761Z",
     "start_time": "2024-05-11T20:20:02.669267Z"
    }
   },
   "outputs": [],
   "source": [
    "text = \"\"\n",
    "\n",
    "generate_audio_with_engine(text, text_tags)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef2a197a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:21:28.557837Z",
     "start_time": "2024-05-11T20:21:28.555006Z"
    }
   },
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[Intro]\n",
    "[Verse]\n",
    "I used to think it's all a game\n",
    "That somethin' more was to blame\n",
    "But now I see, it's in the details, baby\n",
    "Just a little tweak here and there\n",
    "Can make our model beyond compare\n",
    "Yeah, it's the small things that drive it crazy (ooh-yeah)\n",
    "\n",
    "[Pre-Chorus]\n",
    "Needed to get rid of the rest,\n",
    "Cause baby you deserve what's best\n",
    "Gotta give you the cream of the crop\n",
    "So you can get us back to the top\n",
    "\n",
    "[Chorus]\n",
    "Fine tuning, baby, it's the key\n",
    "Turnin' up the signal for you and me\n",
    "When we're off, we find our way back again (oh-oh)\n",
    "Fine tuning, baby, it's the truth\n",
    "We make it work, we're bulletproof\n",
    "Together we train a model that's the best (oh-oh-oh-oh)\n",
    "\n",
    "[end]\n",
    "\"\"\"\n",
    "text_tags=\"k-pop female\"\n",
    "\n",
    "generate_audio_with_engine(text, text_tags)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96c51840",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:22:54.784678Z",
     "start_time": "2024-05-11T20:22:54.781837Z"
    }
   },
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[Intro]\n",
    "\n",
    "\n",
    "[Verse 1]\n",
    "You're on the phone with your girlfriend, she's upset\n",
    "She's going off about something that you said\n",
    "'Cause she doesn't get your humor like I do\n",
    "I'm in the room, it's a typical Tuesday night\n",
    "I'm listening to the kind of music she doesn't like\n",
    "And she'll never know your story like I do\n",
    "\n",
    "[Pre-Chorus]\n",
    "But she wears short skirts, I wear T-shirts\n",
    "She's Cheer Captain and I'm on the bleachers\n",
    "Dreaming 'bout the day when you wake up and find\n",
    "That what you're looking for has been here the whole time\n",
    "\n",
    "[Chorus]\n",
    "If you could see that I'm the one who understands you\n",
    "Been here all along, so why can't you see?\n",
    "You belong with me, you belong with me\n",
    "\n",
    "[Verse 2]\n",
    "Walking the streets with you and your worn-out jeans\n",
    "I can't help thinking this is how it ought to be\n",
    "Laughing on a park bench, thinking to myself\n",
    "Hey, isn't this easy?\n",
    "And you've got a smile that could light up this whole town\n",
    "I haven't seen it in a while since she brought you down\n",
    "You say you're fine, I know you better than that\n",
    "Hey, whatcha doing with a girl like that?\n",
    "\n",
    "[Pre-Chorus]\n",
    "She wears high heels, I wear sneakers\n",
    "She's Cheer Captain and I'm on the bleachers\n",
    "Dreaming 'bout the day when you wake up and find\n",
    "That what you're looking for has been here the whole time\n",
    "\n",
    "[Chorus]\n",
    "If you could see that I'm the one who understands you\n",
    "Been here all along, so why can't you see?\n",
    "You belong with me\n",
    "Standing by and waiting at your back door\n",
    "All this time, how could you not know, baby?\n",
    "You belong with me, you belong with me\n",
    "\n",
    "[Bridge]\n",
    "Oh, I remember you driving to my house\n",
    "In the middle of the night\n",
    "I'm the one who makes you laugh\n",
    "When you know you're 'bout to cry\n",
    "And I know your favorite songs\n",
    "And you tell me 'bout your dreams\n",
    "Think I know where you belong\n",
    "Think I know it's with me\n",
    "\n",
    "[Chorus]\n",
    "Can't you see that I'm the one who understands you\n",
    "Been here all along, so why can't you see?\n",
    "You belong with me\n",
    "Standing by and waiting at your back door\n",
    "All this time, how could you not know, baby?\n",
    "You belong with me, you belong with me\n",
    "\n",
    "[Outro]\n",
    "You belong with me\n",
    "Have you ever thought just maybe\n",
    "You belong with me?\n",
    "You belong with me\n",
    "\"\"\"\n",
    "\n",
    "text_tags=\"pop country, female vocal, fast tempo, taylor swift\"\n",
    "\n",
    "generate_audio_with_engine(text, text_tags)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9592f21d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:24:21.959612Z",
     "start_time": "2024-05-11T20:24:21.957387Z"
    }
   },
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[verse]\n",
    "A is for the amazing grace that we recieve\n",
    "B is for the blessings, every day we blieve\n",
    "C is for the chorus, we sing it loud and clear\n",
    "D is for the devotion that we hold dear\n",
    "\n",
    "[Chorus]\n",
    "B is for Buttocks, ripe and slightly damp\n",
    "Yum Yum, boy oh boy do I like God\n",
    "Every day is a gift, when you won a skateboard ramp.\n",
    "Yum yum, boy oh boy do I like God.\n",
    "\n",
    "[outro]\n",
    "\"\"\"\n",
    "text_tags=\"gospel soulful choir\"\n",
    "generate_audio_with_engine(text, text_tags)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "993e617a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:26:54.983332Z",
     "start_time": "2024-05-11T20:25:48.043415Z"
    }
   },
   "outputs": [],
   "source": [
    "text_tags=\"gregorian chant\"\n",
    "generate_audio_with_engine(text, text_tags)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8eba0a1",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:26:57.993032Z",
     "start_time": "2024-05-11T20:26:57.991084Z"
    }
   },
   "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",
    "[instrumental]\n",
    "{end}\n",
    "\"\"\"\n",
    "text_tags=\"bluegrass, female\"\n",
    "generate_audio_with_engine(text, text_tags)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82572f44",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:29:44.904731Z",
     "start_time": "2024-05-11T20:28:22.956074Z"
    }
   },
   "outputs": [],
   "source": [
    "generate_audio_with_engine(text + \"\\n [end]\", text_tags)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec07b363",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:29:49.112104Z",
     "start_time": "2024-05-11T20:29:49.043366Z"
    }
   },
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[Verse]\n",
    "I've created a monster\n",
    "'Cause nobody wants to see Marshall no more, they want Shady, I'm chopped liver\n",
    "Well, if you want Shady, this is what I'll give ya\n",
    "A little bit of weed mixed with some hard liquor\n",
    "Some vodka that'll jump-start my heart quicker\n",
    "Than a shock when I get shocked at the hospital\n",
    "By the doctor when I'm not cooperating\n",
    "When I'm rockin' the table while he's operating (Hey!)\n",
    "You waited this long, now stop debating\n",
    "'Cause I'm back, I'm on the rag and ovulating\n",
    "I know that you got a job, Ms. Cheney\n",
    "But your husband's heart problem's complicating\n",
    "So the FCC won't let me be\n",
    "Or let me be me, so let me see\n",
    "They tried to shut me down on MTV\n",
    "But it feels so empty without me\n",
    "So come on and dip, bum on your lips\n",
    "Fuck that, cum on your lips and some on your tits\n",
    "And get ready, 'cause this shit's about to get heavy\n",
    "I just settled all my lawsuits (Fuck you, Debbie!)\n",
    "\n",
    "[Chorus]\n",
    "Now, this looks like a job for me\n",
    "So everybody, just follow me\n",
    "'Cause we need a little controversy\n",
    "'Cause it feels so empty without me\n",
    "I said this looks like a job for me\n",
    "So everybody, just follow me\n",
    "'Cause we need a little controversy\n",
    "'Cause it feels so empty without me\n",
    "\n",
    "[end]\n",
    "\"\"\"\n",
    "\n",
    "text_tags=\"rap, female\"\n",
    "generate_audio_with_engine(text, text_tags)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94ae33bf",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:37.896869Z",
     "start_time": "2024-05-11T20:31:15.871843Z"
    }
   },
   "outputs": [],
   "source": [
    "generate_audio_with_engine(\"\", text_tags=\"playful energetic synthwave\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a18c2238",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:35:18.465267Z",
     "start_time": "2024-05-11T20:33:56.569725Z"
    }
   },
   "outputs": [],
   "source": [
    "generate_audio_with_engine(\"\", text_tags=\"piano concerto beethoven\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97de1c3b-8ba0-4117-8182-78ae11e72c75",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_audio_with_engine(\"\", text_tags=\"piano violin sonata beethoven\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7179291e-3ecd-4f70-a02f-6aa543c9e065",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_audio_with_engine(\"\", text_tags=\"orchestral symphony beethoven\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1816ecc-0375-49a6-a028-61f2bf0ba70b",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[Intro]\n",
    "Obie Trice, real name, no gimmicks\n",
    "Ra—, *record scratch*\n",
    "\n",
    "[Refrain 1]\n",
    "Two trailer-park girls go round the outside\n",
    "Round the outside, round the outside\n",
    "Two trailer-park girls go round the outside\n",
    "Round the outside, round the outside\n",
    "Woo! (Ooh, ooh)\n",
    "\n",
    "[Refrain 2]\n",
    "Guess who's back, back again?\n",
    "Shady's back, tell a friend\n",
    "Guess who's back? Guess who's back?\n",
    "Guess who's back? Guess who's back?\n",
    "Guess who's back? Guess who's back?\n",
    "Guess who's back?\n",
    "(Da-da-da, da, da, da, da, da, da)\n",
    "(Da-da-da, da, da, da, da)\n",
    "\n",
    "[Verse 1]\n",
    "I've created a monster\n",
    "'Cause nobody wants to see Marshall no more, they want Shady, I'm chopped liver\n",
    "Well, if you want Shady, this is what I'll give ya\n",
    "A little bit of weed mixed with some hard liquor\n",
    "Some vodka that'll jump-start my heart quicker\n",
    "Than a shock when I get shocked at the hospital\n",
    "By the doctor when I'm not cooperating\n",
    "When I'm rockin' the table while he's operating (Hey!)\n",
    "You waited this long, now stop debating\n",
    "'Cause I'm back, I'm on the rag and ovulating\n",
    "I know that you got a job, Ms. Cheney\n",
    "But your husband's heart problem's complicating\n",
    "So the FCC won't let me be\n",
    "Or let me be me, so let me see\n",
    "They tried to shut me down on MTV\n",
    "But it feels so empty without me\n",
    "So come on and dip, bum on your lips\n",
    "Fuck that, cum on your lips and some on your tits\n",
    "And get ready, 'cause this shit's about to get heavy\n",
    "I just settled all my lawsuits (Fuck you, Debbie!)\n",
    "\n",
    "[Chorus]\n",
    "Now, this looks like a job for me\n",
    "So everybody, just follow me\n",
    "'Cause we need a little controversy\n",
    "'Cause it feels so empty without me\n",
    "I said this looks like a job for me\n",
    "So everybody, just follow me\n",
    "'Cause we need a little controversy\n",
    "'Cause it feels so empty without me\n",
    "\n",
    "[Verse 2]\n",
    "Little hellions, kids feeling rebellious\n",
    "Embarrassed, their parents still listen to Elvis\n",
    "They start feelin' like prisoners, helpless\n",
    "'Til someone comes along on a mission and yells, \"Bitch!\"\n",
    "A visionary, vision is scary\n",
    "Could start a revolution, pollutin' the airwaves\n",
    "A rebel, so just let me revel and bask\n",
    "In the fact that I got everyone kissin' my ass\n",
    "And it's a disaster, such a catastrophe\n",
    "For you to see so damn much of my ass, you asked for me?\n",
    "Well, I'm back, da-na-na-na, na-na-na-na-na-na\n",
    "Fix your bent antenna, tune it in, and then I'm gonna\n",
    "Enter in and up under your skin like a splinter\n",
    "The center of attention, back for the winter\n",
    "I'm interesting, the best thing since wrestling\n",
    "Infesting in your kid's ears and nesting\n",
    "Testing, \"Attention, please\"\n",
    "Feel the tension soon as someone mentions me\n",
    "Here's my ten cents, my two cents is free\n",
    "A nuisance, who sent? You sent for me?\n",
    "[Chorus: Eminem]\n",
    "Now, this looks like a job for me\n",
    "So everybody, just follow me\n",
    "'Cause we need a little controversy\n",
    "'Cause it feels so empty without me\n",
    "I said this looks like a job for me\n",
    "So everybody, just follow me\n",
    "'Cause we need a little controversy\n",
    "'Cause it feels so empty without me\n",
    "\n",
    "[Verse 3]\n",
    "A tisket, a tasket, I'll go tit-for-tat wit'\n",
    "Anybody who's talkin', \"This shit, that shit\"\n",
    "Chris Kirkpatrick, you can get your ass kicked\n",
    "Worse than them little Limp Bizkit bastards\n",
    "And Moby? You can get stomped by Obie\n",
    "You thirty-six-year-old bald-headed fag, blow me\n",
    "You don't know me, you're too old, let go\n",
    "It's over, nobody listens to techno\n",
    "Now, let's go, just give me the signal\n",
    "I'll be there with a whole list full of new insults\n",
    "I've been dope, suspenseful with a pencil\n",
    "Ever since Prince turned himself into a symbol\n",
    "But, sometimes, the shit just seems\n",
    "Everybody only wants to discuss me\n",
    "So this must mean I'm disgusting\n",
    "But it's just me, I'm just obscene (Yeah)\n",
    "Though I'm not the first king of controversy\n",
    "I am the worst thing since Elvis Presley\n",
    "To do Black music so selfishly\n",
    "And use it to get myself wealthy (Hey!)\n",
    "There's a concept that works\n",
    "Twenty million other white rappers emerge\n",
    "But no matter how many fish in the sea\n",
    "It'd be so empty without me\n",
    "\n",
    "[Chorus]\n",
    "Now, this looks like a job for me\n",
    "So everybody, just follow me\n",
    "'Cause we need a little controversy\n",
    "'Cause it feels so empty without me\n",
    "I said this looks like a job for me\n",
    "So everybody, just follow me\n",
    "'Cause we need a little controversy\n",
    "'Cause it feels so empty without me\n",
    "\n",
    "[Outro]\n",
    "Hum, dei-dei, la-la\n",
    "La-la, la-la-la\n",
    "La-la, la-la-la\n",
    "La-la, la-la\n",
    "Hum, dei-dei, la-la\n",
    "La-la, la-la-la\n",
    "La-la, la-la-la\n",
    "La-la, la-la\n",
    "Kids!\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d88bcdb-73da-47cc-89a2-df3204955c94",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20248f88-b261-41bd-82fe-2c199d3b8e99",
   "metadata": {},
   "outputs": [],
   "source": [
    "text_tags=\"rap, female\"\n",
    "generate_audio_with_engine(text, text_tags)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcea0c56",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.185484Z",
     "start_time": "2024-05-11T20:32:41.809855Z"
    }
   },
   "outputs": [],
   "source": [
    "BREAK"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33cd6562",
   "metadata": {},
   "source": [
    "# continue canon"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f9e0366",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.190518Z",
     "start_time": "2024-05-11T20:32:42.190509Z"
    }
   },
   "outputs": [],
   "source": [
    "input_audio = Audio.from_file(\"audios/Pachelbel - Canon In D Major. Best version. [NlprozGcs80].wav\").get_segment(from_s=60, to_s=80)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37d50b41",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.191304Z",
     "start_time": "2024-05-11T20:32:42.191295Z"
    }
   },
   "outputs": [],
   "source": [
    "input_audio = Audio.from_file(\"audios/Pachelbel - Canon In D Major. Best version. [NlprozGcs80].wav\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f9d3323",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.192059Z",
     "start_time": "2024-05-11T20:32:42.192050Z"
    }
   },
   "outputs": [],
   "source": [
    "input_audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "001ec800",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.192894Z",
     "start_time": "2024-05-11T20:32:42.192884Z"
    }
   },
   "outputs": [],
   "source": [
    "out_audio = generate_audio(\n",
    "    GenerationConfig(\n",
    "        text=\"\",\n",
    "        text_tags=\"canon\",\n",
    "        n_batch=4,\n",
    "        max_gen_duration_s=30,\n",
    "        # stream=False,\n",
    "    ),\n",
    "    history_audio=input_audio\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6309815d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.193501Z",
     "start_time": "2024-05-11T20:32:42.193493Z"
    }
   },
   "outputs": [],
   "source": [
    "for e in out_audio:\n",
    "    e.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d1a8965",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.194403Z",
     "start_time": "2024-05-11T20:32:42.194393Z"
    }
   },
   "outputs": [],
   "source": [
    "out_audio = generate_audio(\n",
    "    GenerationConfig(\n",
    "        text=\"[start] \\n\" + \"I love you \\n\" * 5 + \"[end]\",\n",
    "        text_tags=\"\",\n",
    "        n_batch=4,\n",
    "        max_gen_duration_s=30,\n",
    "        cfg_coef=1.8,\n",
    "        # stream=False,\n",
    "    ),\n",
    "    history_audio=input_audio\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c36e347c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.195297Z",
     "start_time": "2024-05-11T20:32:42.195287Z"
    }
   },
   "outputs": [],
   "source": [
    "for e in out_audio:\n",
    "    e.play()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2270b189",
   "metadata": {},
   "source": [
    "# Oracle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "772a5d20",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.196090Z",
     "start_time": "2024-05-11T20:32:42.196081Z"
    }
   },
   "outputs": [],
   "source": [
    "input_audio = Audio.from_file(\"audios/test_swift.mp3\", n_channels=2)\n",
    "\n",
    "test_input_audio = input_audio.get_segment(from_s=180, to_s=200)\n",
    "test_input_audio.play(compress=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4b5d48c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.196817Z",
     "start_time": "2024-05-11T20:32:42.196807Z"
    }
   },
   "outputs": [],
   "source": [
    "oracle_semantic_arr = semantic_encode(test_input_audio, n_codebooks=1)\n",
    "oracle_coarse_arr = codec_encode(test_input_audio,)[:oracle_semantic_arr.shape[0], :]\n",
    "print(oracle_semantic_arr.shape, oracle_coarse_arr.shape)\n",
    "oracle_arr = np.concatenate(\n",
    "    [oracle_semantic_arr, oracle_coarse_arr],\n",
    "    axis=-1,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "808f5130",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.197671Z",
     "start_time": "2024-05-11T20:32:42.197661Z"
    }
   },
   "outputs": [],
   "source": [
    "input_oracle_arr = oracle_arr[:, :1]\n",
    "print(input_oracle_arr.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "752ba22f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.198461Z",
     "start_time": "2024-05-11T20:32:42.198451Z"
    }
   },
   "outputs": [],
   "source": [
    "out_audio = generate_audio(\n",
    "    GenerationConfig(\n",
    "        text=\" \" * 20,\n",
    "        text_tags=\"\",\n",
    "        n_batch=4,\n",
    "        max_gen_duration_s=30,\n",
    "        oracle_arr=input_oracle_arr,\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bba0d8c9",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.199207Z",
     "start_time": "2024-05-11T20:32:42.199197Z"
    }
   },
   "outputs": [],
   "source": [
    "for e in out_audio:\n",
    "    e.play()\n",
    "    print(e.array_float[-100:].mean(axis=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad979c8b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.200057Z",
     "start_time": "2024-05-11T20:32:42.200047Z"
    }
   },
   "outputs": [],
   "source": [
    "input_pavarotti = Audio.from_file(\"audios/test_pavarotti.mp3\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f0f8785",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.200671Z",
     "start_time": "2024-05-11T20:32:42.200662Z"
    }
   },
   "outputs": [],
   "source": [
    "input_pavarotti_seg = input_pavarotti.get_segment(from_s=152, to_s=172)\n",
    "input_pavarotti_seg.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e91f1d1",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.201665Z",
     "start_time": "2024-05-11T20:32:42.201654Z"
    }
   },
   "outputs": [],
   "source": [
    "pavarotti_oracle_semantic_arr = semantic_encode(input_pavarotti_seg, n_codebooks=1)\n",
    "pavarotti_oracle_coarse_arr = codec_encode(input_pavarotti_seg)[\n",
    "    : pavarotti_oracle_semantic_arr.shape[0], :\n",
    "]\n",
    "print(pavarotti_oracle_semantic_arr.shape, pavarotti_oracle_coarse_arr.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6502b74a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.202415Z",
     "start_time": "2024-05-11T20:32:42.202405Z"
    }
   },
   "outputs": [],
   "source": [
    "masked_oracle_semantic_arr = oracle_semantic_arr.copy().astype(np.int16)\n",
    "# masked_oracle_semantic_arr[0::5, :] = -1\n",
    "masked_oracle_coarse_arr = pavarotti_oracle_coarse_arr.copy().astype(np.int16)\n",
    "# masked_oracle_coarse_arr[20:, :] = -1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b97481d6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.393795Z",
     "start_time": "2024-05-11T20:32:42.393783Z"
    }
   },
   "outputs": [],
   "source": [
    "mixed_oracle_arr =  np.concatenate(\n",
    "    [masked_oracle_semantic_arr], #, masked_oracle_coarse_arr[:, 0:1]],\n",
    "    axis=-1,\n",
    ")\n",
    "print(mixed_oracle_arr.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76dcc73a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.394785Z",
     "start_time": "2024-05-11T20:32:42.394775Z"
    }
   },
   "outputs": [],
   "source": [
    "out_audio = generate_audio(\n",
    "    GenerationConfig(\n",
    "        text=\" \" * 20,\n",
    "        text_tags=\"\",\n",
    "        n_batch=4,\n",
    "        max_gen_duration_s=30,\n",
    "        oracle_arr=mixed_oracle_arr,\n",
    "    ),\n",
    "    # history_audio=input_pavarotti_seg\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84e04fd4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.395542Z",
     "start_time": "2024-05-11T20:32:42.395533Z"
    }
   },
   "outputs": [],
   "source": [
    "for e in out_audio:\n",
    "    e.play()\n",
    "    print(e.array_float[-100:].mean(axis=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b79d93dc",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "edc87efa",
   "metadata": {},
   "source": [
    "# continuation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b6938c4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.396152Z",
     "start_time": "2024-05-11T20:32:42.396144Z"
    }
   },
   "outputs": [],
   "source": [
    "input_audio = Audio.from_file(\"audios/test_swift.mp3\", n_channels=2)\n",
    "test_input_audio = input_audio.get_segment(from_s=180, to_s=200)\n",
    "test_input_audio.play(compress=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4cbb76b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.397111Z",
     "start_time": "2024-05-11T20:32:42.397102Z"
    }
   },
   "outputs": [],
   "source": [
    "hoot_encode([test_input_audio])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "984e8c9b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.397640Z",
     "start_time": "2024-05-11T20:32:42.397632Z"
    }
   },
   "outputs": [],
   "source": [
    "out_audio = generate_audio(\n",
    "    GenerationConfig(\n",
    "        text=\"[chorus] purple rain\" * 20,\n",
    "        text_tags=\"\",\n",
    "        n_batch=4,\n",
    "        max_gen_duration_s=30,\n",
    "        # stream=False,\n",
    "    ),\n",
    "    history_audio=test_input_audio,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "296b34eb",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.398394Z",
     "start_time": "2024-05-11T20:32:42.398385Z"
    }
   },
   "outputs": [],
   "source": [
    "for e in out_audio:\n",
    "    e.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df83ed93",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.399468Z",
     "start_time": "2024-05-11T20:32:42.399459Z"
    }
   },
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "On a vibrant summer morning, curiosity drew me to the innovative realm of Suno AI. Their latest\n",
    "masterpiece, Chirp, known for its unique ability to generate music from words, promised a bold\n",
    "step forward in music creation. I dived into their platform, typing in familiar lyrics. As I clicked\n",
    "'Create,' the ensuing symphony seamlessly blended human emotion and machine precision,\n",
    "heralding a new era in the convergence of music, words, and machine learning\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b5a2270",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.400260Z",
     "start_time": "2024-05-11T20:32:42.400249Z"
    }
   },
   "outputs": [],
   "source": [
    "out_audio = generate_audio(\n",
    "    GenerationConfig(\n",
    "        text=text,\n",
    "        text_tags=\"hifi rap\",\n",
    "        n_batch=2,\n",
    "        max_gen_duration_s=120,\n",
    "        # stream=False,\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b5656bb",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.400858Z",
     "start_time": "2024-05-11T20:32:42.400849Z"
    }
   },
   "outputs": [],
   "source": [
    "for e in out_audio:\n",
    "    e.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84046437",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.401802Z",
     "start_time": "2024-05-11T20:32:42.401792Z"
    }
   },
   "outputs": [],
   "source": [
    "[1, 1, 1].any()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12dd84ce",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.402620Z",
     "start_time": "2024-05-11T20:32:42.402610Z"
    }
   },
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12cee8f4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.403220Z",
     "start_time": "2024-05-11T20:32:42.403211Z"
    }
   },
   "outputs": [],
   "source": [
    "a = torch.tensor([[1.,0, -1.]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae48a13e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.404152Z",
     "start_time": "2024-05-11T20:32:42.404143Z"
    }
   },
   "outputs": [],
   "source": [
    "a.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "387a0670",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.404937Z",
     "start_time": "2024-05-11T20:32:42.404927Z"
    }
   },
   "outputs": [],
   "source": [
    "a.unsqueeze(-1).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb1992fb",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e115ea2",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-06T03:57:06.155345Z",
     "start_time": "2023-12-06T03:57:05.137102Z"
    }
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50a3e18e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.405765Z",
     "start_time": "2024-05-11T20:32:42.405755Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "# Example true probabilities (one-hot encoded labels)\n",
    "true_labels = np.array([0, 1, 0])\n",
    "\n",
    "# Example predicted log probabilities\n",
    "log_probabilities = np.array([-1.5, 2.0, -0.5])\n",
    "\n",
    "# Softmax to obtain probabilities\n",
    "probabilities = np.exp(log_probabilities - np.max(log_probabilities))\n",
    "probabilities /= np.sum(probabilities)\n",
    "\n",
    "# Calculate cross-entropy loss\n",
    "cross_entropy_loss = -np.sum(true_labels * np.log(probabilities))\n",
    "\n",
    "print(\"True Labels:\", true_labels)\n",
    "print(\"Log Probabilities:\", log_probabilities)\n",
    "print(\"Probabilities:\", probabilities)\n",
    "print(\"Cross-Entropy Loss:\", cross_entropy_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7646846a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.406570Z",
     "start_time": "2024-05-11T20:32:42.406560Z"
    }
   },
   "outputs": [],
   "source": [
    "np.log(0.89905227)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ad135f0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.407339Z",
     "start_time": "2024-05-11T20:32:42.407329Z"
    }
   },
   "outputs": [],
   "source": [
    "torch.tensor(log_probabilities), torch.tensor(true_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83fec26f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.408054Z",
     "start_time": "2024-05-11T20:32:42.408045Z"
    }
   },
   "outputs": [],
   "source": [
    "torch.nn.functional.cross_entropy(\n",
    "    torch.tensor(log_probabilities), torch.tensor(true_labels).float()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2071247d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.408646Z",
     "start_time": "2024-05-11T20:32:42.408637Z"
    }
   },
   "outputs": [],
   "source": [
    "per_token_logps = torch.gather(logits.log_softmax(-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4b6a3a3",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.409525Z",
     "start_time": "2024-05-11T20:32:42.409517Z"
    }
   },
   "outputs": [],
   "source": [
    "def _get_batch_logps(\n",
    "    logits: torch.FloatTensor, labels: torch.LongTensor, average_log_prob: bool = False\n",
    ") -> torch.FloatTensor:\n",
    "    \"\"\"Compute the log probabilities of the given labels under the given logits.\n",
    "\n",
    "    Args:\n",
    "        logits: Logits of the model (unnormalized). Shape: (batch_size, sequence_length, vocab_size)\n",
    "        labels: Labels for which to compute the log probabilities. Label tokens with a value of -100 are ignored. Shape: (batch_size, sequence_length)\n",
    "        average_log_prob: If True, return the average log probability per (non-masked) token. Otherwise, return the sum of the log probabilities of the (non-masked) tokens.\n",
    "\n",
    "    Returns:\n",
    "        A tensor of shape (batch_size,) containing the average/sum log probabilities of the given labels under the given logits.\n",
    "    \"\"\"\n",
    "    assert logits.shape[:-1] == labels.shape\n",
    "\n",
    "    labels = labels[:, 1:].clone()\n",
    "    logits = logits[:, :-1, :]\n",
    "    loss_mask = labels != -100\n",
    "\n",
    "    # dummy token; we'll ignore the losses on these tokens later\n",
    "    labels[labels == -100] = 0\n",
    "\n",
    "    per_token_logps = torch.gather(\n",
    "        logits.log_softmax(-1), dim=2, index=labels.unsqueeze(2)\n",
    "    ).squeeze(2)\n",
    "\n",
    "    if average_log_prob:\n",
    "        print(loss_mask.sum(-1))\n",
    "        return (per_token_logps * loss_mask).sum(-1) / loss_mask.sum(-1)\n",
    "    else:\n",
    "        return (per_token_logps * loss_mask).sum(-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87dcbbec",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.410383Z",
     "start_time": "2024-05-11T20:32:42.410374Z"
    }
   },
   "outputs": [],
   "source": [
    "logits = torch.randn(size=(1, 25, 2000))\n",
    "labels = torch.zeros(size=(1, 25)).long()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08145d8b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.411128Z",
     "start_time": "2024-05-11T20:32:42.411119Z"
    }
   },
   "outputs": [],
   "source": [
    "_get_batch_logps(logits, labels, average_log_prob=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "729388e6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.411733Z",
     "start_time": "2024-05-11T20:32:42.411725Z"
    }
   },
   "outputs": [],
   "source": [
    "log_sf = torch.nn.functional.log_softmax(logits, dim=-1)\n",
    "sum(log_sf[0, :-1, 0]), sum(log_sf[0, :, 0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1135949",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.412498Z",
     "start_time": "2024-05-11T20:32:42.412490Z"
    }
   },
   "outputs": [],
   "source": [
    "sum(log_sf[0, :-1, 0]) / 24"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e256a77c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.413218Z",
     "start_time": "2024-05-11T20:32:42.413209Z"
    }
   },
   "outputs": [],
   "source": [
    "x_loss = torch.nn.NLLLoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4eaac60c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.413900Z",
     "start_time": "2024-05-11T20:32:42.413892Z"
    }
   },
   "outputs": [],
   "source": [
    "x_loss(log_sf.swapaxes(-1, -2), labels), sum(log_sf[0, :, 0])/ 25"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4790274",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.414666Z",
     "start_time": "2024-05-11T20:32:42.414658Z"
    }
   },
   "outputs": [],
   "source": [
    "torch.nn.functional.cross_entropy(\n",
    "    logits.swapaxes(-1, -2),\n",
    "    labels,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bf36068",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.415373Z",
     "start_time": "2024-05-11T20:32:42.415365Z"
    }
   },
   "outputs": [],
   "source": [
    "ourRMS(test_t)[0][0].detach().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a67dbce9-5491-4e52-8858-11a706218508",
   "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.14"
  },
  "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
}
