{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from suno_utils.utils.text import read_jsonl, write_jsonl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\n",
    "    \"/home/christian/code/christian/metadata/v45_splits/ids_keep_sets_v11.json\", \"r\"\n",
    ") as f:\n",
    "    sft_ids = json.load(f)\n",
    "\n",
    "print(len(sft_ids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "discogs_subset 496613\n",
      "genius 419370\n",
      "imslp 76482\n",
      "deezer 135671\n",
      "1055921\n"
     ]
    }
   ],
   "source": [
    "known_sft_ids = set()\n",
    "for dataset_name, ids in sft_ids.items():\n",
    "    print(dataset_name, len(ids))\n",
    "    known_sft_ids.update(ids)\n",
    "print(len(known_sft_ids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "metas_tr = read_jsonl(\"/app2/suno/data/auk_v0/metas_v3_tr.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 56617603/56617603 [00:24<00:00, 2354170.46it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "911801\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "sft_metas_tr = []\n",
    "for meta_tr in tqdm(metas_tr):\n",
    "    if meta_tr[\"id\"] in known_sft_ids:\n",
    "        sft_metas_tr.append(meta_tr)\n",
    "print(len(sft_metas_tr))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'id': 'koHo5-0DVlI',\n",
       " 's3_filepath': 's3://webdataset/bundles/v0/genius/audio/koHo5-0DVlI.opus',\n",
       " 'duration_s': 228.13,\n",
       " 'text': \"[Intro: male vocals; soft; inviting; synthesizer; gentle background sound; inviting; soft touch; Soft synthesizer introduces the song, creating an inviting atmosphere.]\\nYou know you make me wanna...\\nYou know you make me wanna...\\n\\n[Verse 1: male vocals; soothing; emotional; piano; soft drumming; bass; emotional; smooth; Smooth male vocals express yearning, accompanied by gentle piano and bass.]\\nTo start it off, I know you know me\\nTo come to think of it, it was only last week that I\\nHad a dream about us, oh\\nThat's why I'm here, I'm this writing this song\\nTo tell the truth you know I've been hurting all along\\nSomeway let me know, you want me girl\\n\\n[Pre-Chorus: male vocals; expressive; building anticipation; piano; subtle percussion buildup; building anticipation; warm; The vocals build anticipation with expressive tones, layered with light percussion.]\\nEvery time you see me, what do you see?\\nI feel like I'm a poor man and you're the queen\\nOoh baby, you're the only thing that I really need\\nBaby that's why\\n\\n[Chrous]\\nYou make me wanna call you in the middle of the night\\nYou make me wanna hold you till the morning light\\nYou make me wanna love, you make me wanna fall\\nYou make me wanna surrender my soul\\nI know this is a feeling that I just can't fight\\nYou're the first and last thing on my mind\\nYou make me wanna love, you make me wanna fall\\nYou make me wanna surrender my soul\\n\\n[Verse 2: male vocals; tender; relatable; piano; softer percussion; bass underlines emotion; relatable; yearning; Emotional and relatable vocals continue, enhancing the established mood.]\\nWell I know that these feelings won't end, nah nah\\nThey'll get stronger if I see you again\\nBaby I'm tired of being friends, oh\\nI wanna know if you feel the same\\nAnd could you tell me do you feel my pain?\\nDon't leave me in doubt, baby\\n\\n[Pre-Chorus: male vocals; expressive; building anticipation; piano; subtle percussion buildup; warm; building anticipation; Similar build-up occurs, maintaining the emotional connection and engagement.]\\nEverytime you see me, what do you see?\\nI feel like I'm a poor man and you're the queen\\nOoh baby, you're the only thing that I really need\\nBaby that's why\\n\\n[Chorus: male vocals; strong melody; heartfelt delivery; drums; bass; synth layers enhance emotion; passionate; compelling emotion; Passionate delivery by the male vocalist, strengthened by a rich instrumental backdrop, engaging listeners with its heartfelt melody.]\\nYou make me wanna call you in the middle of the night\\nYou make me wanna hold you till the morning light\\nYou make me wanna love, you make me wanna fall\\nYou make me wanna surrender my soul\\nI know this is a feeling that I just can't fight\\nYou're the first and last thing on my mind\\nYou make me wanna love, you make me wanna fall\\nYou make me wanna surrender my soul\\n\\n[Bridge: male vocals; sensual; intimate; piano; soft tambourine; gentle layering of sounds; intimate; sensual vibe; A sensual vibe takes over with gently layered male vocals, piano, and soft tambourine, enhancing intimacy.]\\nI'll take you home real quick and sit you down on the couch\\nPour some Dom Perignon and hit the lights out\\nBaby we can make sweet love\\nThen we'll take it nice and slow\\nI'm gonna touch you like you've never known before\\nWe're gonna make love, ohhh\\n\\n[Chorus: male vocals; strong melody; heartfelt delivery; drums; bass; synth layers enhance emotion; passionate; compelling emotion; A repeat of the strong, passionate delivery that characterizes the song, dominated by drums and bass alongside vocals.]\\nYou make me wanna call you in the middle of the night\\nYou make me wanna hold you till the morning light\\nYou make me wanna love, you make me wanna fall\\nYou make me wanna surrender my soul\\nI know this is a feeling that I just can't fight\\nYou're the first and last thing on my mind\\nYou make me wanna love, you make me wanna fall\\nYou make me wanna surrender my soul\\n\\n[Chorus: male vocals; strong melody; heartfelt delivery; drums; bass; synth layers enhance emotion; passionate; compelling emotion; Strong melody resonates as vocals soar, supported by a vibrant instrumental base.]\\nYou make me wanna call you in the middle of the night\\nYou make me wanna hold you till the morning light\\nYou make me wanna love, you make me wanna fall\\nYou make me wanna surrender my soul\\nI know this is a feeling that I just can't fight\\nYou're the first and last thing on my mind\\nYou make me wanna love, you make me wanna fall\\nYou make me wanna surrender my soul\\n\\n[Outro: male vocals; soft; reflective; synthesizer; fading out with soft piano; reflective; cooling down; The song gently fades out with soft synthesizer and piano, leaving a reflective mood.]\\nYou make me wanna call you in the middle of the night\\nYou make me wanna hold you till the morning light\\nYou make me wanna love, you make me wanna fall\\nYou make me wanna surrender my soul\",\n",
       " 'lang': 'en',\n",
       " 'tags': ['tenor',\n",
       "  'love',\n",
       "  \"A heartfelt pop-R&B track that draws listeners in with its smooth melodies and passionate lyrics. The sonic landscape is enriched by soft piano, warm bass, and gentle percussion, creating an intimate feel throughout. The emotional delivery from the lead vocalist captures the yearning and desire expressed in the lyrics, making it resonate with anyone who's experienced love and longing.\",\n",
       "  'smooth',\n",
       "  'The lyrics delve into themes of love, desire, and emotional vulnerability, offering a glimpse into the complexities of relationships.',\n",
       "  'passionate',\n",
       "  'male vocal',\n",
       "  'drums',\n",
       "  'The male vocals are smooth and expressive, enhancing the romantic themes of the song. They convey deep emotions that switch between yearning and passion, engaging the listener.',\n",
       "  'bass',\n",
       "  'emotional',\n",
       "  'English',\n",
       "  \"The song starts with a gentle introduction before moving into a narrative-driven first verse, building tension through pre-choruses that lead into rich, emotionally-charged choruses. The bridge introduces a sensual interlude, followed by repeated choruses that reinforce the song's central emotional themes, ultimately leading to a soft, reflective ending.\",\n",
       "  'desire',\n",
       "  'synthesizer',\n",
       "  'R&B',\n",
       "  'Pop',\n",
       "  'romantic',\n",
       "  'relationships',\n",
       "  'expressive',\n",
       "  'piano',\n",
       "  'yearning',\n",
       "  '2000s',\n",
       "  'quality: 25',\n",
       "  'spectral_centroid: 1.2',\n",
       "  'bass: 0.5',\n",
       "  'mid: 0.8',\n",
       "  'high: 1.0',\n",
       "  'stereo_width: -0.4',\n",
       "  'spectral_flatness: 0.9',\n",
       "  'crest_factor: -0.8',\n",
       "  'loudness: 0.7'],\n",
       " 'cover_ids': ['Y50cltC4rFo', 'PasfLTmTpXs', 'oHUOAezaGqY'],\n",
       " 'weight': 3.0,\n",
       " 'local_filepath': '/app2/suno/data/raw_audio_opus_v0/koHo5-0DVlI.opus',\n",
       " 'stems': {'Vocals': '/app2/suno/data/sft_stems_12_output_v11/koHo5-0DVlI_Vocals.opus',\n",
       "  'Backing_Vocals': '/app2/suno/data/sft_stems_12_output_v11/koHo5-0DVlI_Backing_Vocals.opus',\n",
       "  'Drums': '/app2/suno/data/sft_stems_12_output_v11/koHo5-0DVlI_Drums.opus',\n",
       "  'Bass': '/app2/suno/data/sft_stems_12_output_v11/koHo5-0DVlI_Bass.opus',\n",
       "  'Guitar': '/app2/suno/data/sft_stems_12_output_v11/koHo5-0DVlI_Guitar.opus',\n",
       "  'Percussion': '/app2/suno/data/sft_stems_12_output_v11/koHo5-0DVlI_Percussion.opus',\n",
       "  'Strings': '/app2/suno/data/sft_stems_12_output_v11/koHo5-0DVlI_Strings.opus',\n",
       "  'Synth': '/app2/suno/data/sft_stems_12_output_v11/koHo5-0DVlI_Synth.opus',\n",
       "  'FX': '/app2/suno/data/sft_stems_12_output_v11/koHo5-0DVlI_FX.opus'}}"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sft_metas_tr[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the sft_metas_tr to a jsonl file\n",
    "write_jsonl(sft_metas_tr, \"/app2/suno/data/dpo/sft/sft_metas_tr_v11.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "911801\n"
     ]
    }
   ],
   "source": [
    "input_jsonl = read_jsonl(\"/app2/suno/data/dpo/sft/sft_metas_tr_v11.jsonl\")\n",
    "\n",
    "print(len(input_jsonl))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "id\n",
      "s3_filepath\n",
      "duration_s\n",
      "text\n",
      "lang\n",
      "tags\n",
      "cover_ids\n",
      "weight\n",
      "local_filepath\n",
      "stems\n"
     ]
    }
   ],
   "source": [
    "for k in input_jsonl[0].keys():\n",
    "    print(k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "import tqdm\n",
    "from collections import defaultdict, Counter\n",
    "from typing import List, Dict\n",
    "from suno_utils.utils.text import write_jsonl, write_json\n",
    "\n",
    "# Constants from the original file\n",
    "SEMANTIC_N_CODEBOOKS = 1\n",
    "COARSE_RATE_HZ = 25\n",
    "SEMANTIC_CODEBOOK_SIZE = 4000\n",
    "SEMANTIC_PAD_TOKEN = SEMANTIC_CODEBOOK_SIZE\n",
    "SEMANTIC_RATE_HZ = 25\n",
    "COARSE_CODEBOOK_SIZE = 2048\n",
    "COARSE_N_CODEBOOKS = 12\n",
    "COARSE_PAD_TOKEN = COARSE_CODEBOOK_SIZE\n",
    "N_TOKENS_AUDIO = 25 * 8 * 60  # max 8 mins of audio\n",
    "\n",
    "\n",
    "def make_sft_dataset(\n",
    "    input_data: List[Dict],\n",
    "    output_data_dir: str,\n",
    "    npz_dir: str = \"/app2/suno/data/dpo/sft/semantic_codes_v11\",\n",
    "    t_data_memmap: int = N_TOKENS_AUDIO,\n",
    "):\n",
    "    \"\"\"\n",
    "    Create SFT dataset from jsonl input data.\n",
    "\n",
    "    Args:\n",
    "        input_data: List of dictionaries from jsonl with keys:\n",
    "            id, s3_filepath, duration_s, text, lang, tags, cover_ids, weight, local_filepath, stems\n",
    "        output_data_dir: Directory to save output files\n",
    "        is_val: Whether this is validation set\n",
    "        npz_dir: Directory containing NPZ files with semantic codes\n",
    "        t_data_memmap: Maximum tokens per audio (default 8 mins)\n",
    "    \"\"\"\n",
    "    print(f\"t_data_memmap is set to: {t_data_memmap}\")\n",
    "    print(f\"Processing {len(input_data)} items\")\n",
    "\n",
    "    # Create output directory if it doesn't exist\n",
    "    os.makedirs(output_data_dir, exist_ok=True)\n",
    "\n",
    "    dset_type = \"sft\"\n",
    "    out_mmap_path = os.path.join(output_data_dir, f\"data_{dset_type}.bin\")\n",
    "    out_metas_path = os.path.join(output_data_dir, f\"meta_{dset_type}.jsonl\")\n",
    "    out_info_filepath = os.path.join(output_data_dir, f\"info_{dset_type}.json\")\n",
    "\n",
    "    def reshift(arr):\n",
    "        \"\"\"Pad and reshape semantic array\"\"\"\n",
    "        sem_start_idx = 0\n",
    "        sem_end_idx = len(arr) - 1\n",
    "        semantic_arr = arr[:, :SEMANTIC_N_CODEBOOKS]\n",
    "\n",
    "        # get array segments\n",
    "        arr_s = semantic_arr[sem_start_idx:sem_end_idx, :].copy()\n",
    "        assert arr_s.max() <= SEMANTIC_PAD_TOKEN\n",
    "        # concat and stack\n",
    "        arr_s = np.pad(\n",
    "            arr_s,\n",
    "            ((0, t_data_memmap - len(arr_s)), (0, 0)),\n",
    "            constant_values=SEMANTIC_PAD_TOKEN,\n",
    "            mode=\"constant\",\n",
    "        )\n",
    "        arr = arr_s\n",
    "        arr = arr.astype(np.uint16)\n",
    "        assert arr.shape == (t_data_memmap, SEMANTIC_N_CODEBOOKS)\n",
    "        return arr\n",
    "\n",
    "    # Initialize mmap\n",
    "    _ = np.memmap(out_mmap_path, dtype=np.uint16, mode=\"w+\", shape=(1,))\n",
    "    n_offs = 0\n",
    "    curr_idx = 0\n",
    "    tot_duration_dict = defaultdict(float)\n",
    "    datasets_info = defaultdict(dict)\n",
    "    total_processed = 0\n",
    "    total_skipped = 0\n",
    "    error_counter = Counter()\n",
    "\n",
    "    # Process each item\n",
    "    for item in tqdm.tqdm(input_data, desc=\"Processing SFT data\"):\n",
    "        try:\n",
    "            # Construct NPZ path\n",
    "            item_id = item[\"id\"]\n",
    "            local_path = os.path.join(npz_dir, f\"{item_id}.npz\")\n",
    "\n",
    "            if not os.path.exists(local_path):\n",
    "                raise ValueError(f\"File does not exist: {local_path}\")\n",
    "\n",
    "            # Load NPZ file\n",
    "            temp_npz = np.load(local_path)\n",
    "\n",
    "            # Try different version keys in order of preference\n",
    "            arr = temp_npz[\"codes\"].reshape(-1, 1)\n",
    "\n",
    "            if arr is None:\n",
    "                raise ValueError(f\"No valid semantic code array found in {local_path}\")\n",
    "\n",
    "            # Validate and crop if needed\n",
    "            if arr.shape[0] > t_data_memmap:\n",
    "                raise ValueError(\n",
    "                    f\"Overflow for {item_id}: {arr.shape[0]} > {t_data_memmap}\"\n",
    "                )\n",
    "\n",
    "            assert arr.shape[0] <= t_data_memmap\n",
    "            assert arr.shape[1] == 1\n",
    "\n",
    "            # Calculate duration\n",
    "            arr_duration = arr.shape[0] / SEMANTIC_RATE_HZ\n",
    "\n",
    "            # Create metadata\n",
    "            add_meta = {\n",
    "                \"dataset\": \"sft\",\n",
    "                \"id\": item_id,\n",
    "                \"s3_filepath\": item.get(\"s3_filepath\", \"\"),\n",
    "                \"duration_s\": item.get(\"duration_s\", arr_duration),\n",
    "                \"text\": item.get(\"text\", \"\"),\n",
    "                \"lang\": item.get(\"lang\", \"\"),\n",
    "                \"tags\": item.get(\"tags\", [])\n",
    "                if isinstance(item.get(\"tags\"), list)\n",
    "                else [item.get(\"tags\", \"\")],\n",
    "                \"cover_ids\": item.get(\"cover_ids\", []),\n",
    "                \"weight\": item.get(\"weight\", 1.0),\n",
    "                \"local_filepath\": item.get(\"local_filepath\", \"\"),\n",
    "                \"stems\": item.get(\"stems\", {}),\n",
    "                \"generated_start_index\": 0,\n",
    "                \"start_s\": 0,\n",
    "                \"vocal_start_s\": None,\n",
    "                \"vocal_end_s\": None,\n",
    "            }\n",
    "\n",
    "            # Process array\n",
    "            arr = reshift(arr)\n",
    "            arr = arr.reshape(\n",
    "                -1,\n",
    "            )\n",
    "\n",
    "            # Write to mmap\n",
    "            out_mm = np.memmap(\n",
    "                out_mmap_path,\n",
    "                dtype=np.uint16,\n",
    "                mode=\"r+\",\n",
    "                shape=(n_offs + arr.size,),\n",
    "            )\n",
    "            out_mm[n_offs : n_offs + arr.size] = arr\n",
    "            out_mm.flush()\n",
    "            del out_mm\n",
    "\n",
    "            # Write metadata\n",
    "            write_jsonl(\n",
    "                [add_meta],\n",
    "                out_metas_path,\n",
    "                do_append=bool(n_offs != 0),\n",
    "            )\n",
    "\n",
    "            # Update info\n",
    "            dataset_name = add_meta[\"dataset\"]\n",
    "            if \"idx_list\" not in datasets_info[dataset_name]:\n",
    "                datasets_info[dataset_name][\"idx_list\"] = [curr_idx]\n",
    "            else:\n",
    "                datasets_info[dataset_name][\"idx_list\"].append(curr_idx)\n",
    "\n",
    "            # Update counters\n",
    "            curr_idx += 1\n",
    "            n_offs += arr.size\n",
    "            tot_duration_dict[\"sft\"] += arr_duration\n",
    "            total_processed += 1\n",
    "\n",
    "        except Exception as e:\n",
    "            error_type = type(e).__name__\n",
    "            error_counter[error_type] += 1\n",
    "            # print(\n",
    "            #     f\"Error processing {item.get('id', 'unknown')}: {error_type}: {str(e)}\"\n",
    "            # )\n",
    "            total_skipped += 1\n",
    "            continue\n",
    "\n",
    "    # Write dataset info\n",
    "    write_json(datasets_info, out_info_filepath)\n",
    "\n",
    "    # Print summary\n",
    "    print(\"\\n\" + \"=\" * 50)\n",
    "    print(f\"Total processed: {total_processed} clips\")\n",
    "    print(f\"Total skipped: {total_skipped} clips\")\n",
    "    print(f\"Success rate: {total_processed/(total_processed + total_skipped)*100:.1f}%\")\n",
    "\n",
    "    for k, v in tot_duration_dict.items():\n",
    "        print(f\"{k}: {v / 60 / 60:.1f} hours\")\n",
    "\n",
    "    if error_counter:\n",
    "        print(\"\\nErrors encountered:\")\n",
    "        for error_type, count in error_counter.most_common():\n",
    "            print(f\"  {error_type}: {count}\")\n",
    "\n",
    "    print(\"Done\")\n",
    "\n",
    "    return {\n",
    "        \"processed\": total_processed,\n",
    "        \"skipped\": total_skipped,\n",
    "        \"total_duration_hours\": sum(tot_duration_dict.values()) / 3600,\n",
    "        \"errors\": dict(error_counter),\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "t_data_memmap is set to: 12000\n",
      "Processing 911801 items\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processing SFT data: 100%|██████████| 911801/911801 [2:26:21<00:00, 103.84it/s]  \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "==================================================\n",
      "Total processed: 897779 clips\n",
      "Total skipped: 14022 clips\n",
      "Success rate: 98.5%\n",
      "sft: 56955.5 hours\n",
      "\n",
      "Errors encountered:\n",
      "  ValueError: 14022\n",
      "Done\n",
      "\n",
      "Final results: {'processed': 897779, 'skipped': 14022, 'total_duration_hours': 56955.474611112586, 'errors': {'ValueError': 14022}}\n"
     ]
    }
   ],
   "source": [
    "# Create the dataset\n",
    "results = make_sft_dataset(\n",
    "    input_data=input_jsonl,\n",
    "    output_data_dir=\"/app2/suno/data/dpo/sft/\",\n",
    "    npz_dir=\"/app2/suno/data/dpo/sft/semantic_codes_v11\",\n",
    "    t_data_memmap=N_TOKENS_AUDIO,\n",
    ")\n",
    "\n",
    "print(f\"\\nFinal results: {results}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'N_TOKENS_AUDIO' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[2], line 10\u001b[0m\n\u001b[1;32m      8\u001b[0m test_metas \u001b[38;5;241m=\u001b[39m read_jsonl(os\u001b[38;5;241m.\u001b[39mpath\u001b[38;5;241m.\u001b[39mjoin(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m/app2/suno/data/dpo/sft/\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmeta_sft.jsonl\u001b[39m\u001b[38;5;124m\"\u001b[39m))\n\u001b[1;32m      9\u001b[0m test_info \u001b[38;5;241m=\u001b[39m read_json(os\u001b[38;5;241m.\u001b[39mpath\u001b[38;5;241m.\u001b[39mjoin(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m/app2/suno/data/dpo/sft/\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124minfo_sft.json\u001b[39m\u001b[38;5;124m\"\u001b[39m))\n\u001b[0;32m---> 10\u001b[0m mm \u001b[38;5;241m=\u001b[39m mm\u001b[38;5;241m.\u001b[39mreshape(\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m, \u001b[43mN_TOKENS_AUDIO\u001b[49m, \u001b[38;5;241m1\u001b[39m)\n\u001b[1;32m     11\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(mm) \u001b[38;5;241m==\u001b[39m \u001b[38;5;28mlen\u001b[39m(test_metas)\n\u001b[1;32m     12\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m mm[:\u001b[38;5;241m100\u001b[39m, :, \u001b[38;5;241m0\u001b[39m]\u001b[38;5;241m.\u001b[39mmin() \u001b[38;5;241m>\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0\u001b[39m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'N_TOKENS_AUDIO' is not defined"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "from suno_utils.utils.text import read_jsonl, read_json\n",
    "\n",
    "mm = np.memmap(\n",
    "    os.path.join(\"/app2/suno/data/dpo/sft/\", f\"data_sft.bin\"), dtype=np.uint16, mode=\"r\"\n",
    ")\n",
    "test_metas = read_jsonl(os.path.join(\"/app2/suno/data/dpo/sft/\", f\"meta_sft.jsonl\"))\n",
    "test_info = read_json(os.path.join(\"/app2/suno/data/dpo/sft/\", f\"info_sft.json\"))\n",
    "mm = mm.reshape(-1, N_TOKENS_AUDIO, 1)\n",
    "assert len(mm) == len(test_metas)\n",
    "assert mm[:100, :, 0].min() >= 0\n",
    "assert mm[:100, :, 0].max() <= 4000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_env_dev",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
