{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b15bc7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "893c33d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "extreme_dir = \"/app2/suno/data/extreme_music/\"\n",
    "extreme_meta_file = os.path.join(extreme_dir, \"extreme_music_tracks.jsonl\")\n",
    "assert os.path.exists(extreme_meta_file)\n",
    "\n",
    "extreme_metas = []\n",
    "for line in tqdm(open(extreme_meta_file)):\n",
    "    extreme_metas.append(json.loads(line))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96269871",
   "metadata": {},
   "outputs": [],
   "source": [
    "!head /app2/suno/data/auk_v0/metas_v3_tr.jsonl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09be69df",
   "metadata": {},
   "outputs": [],
   "source": [
    "extreme_metas[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6dd3b3e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def transform_extreme_music_record(record):\n",
    "    \"\"\"Transform extreme music format to SunoGPT metadata format\"\"\"\n",
    "\n",
    "    # Extract basic info\n",
    "    track_id = str(record[\"id\"])  # Convert to string for consistency\n",
    "    title = record.get(\"title\", \"\")\n",
    "    description = record.get(\"description\", \"\")\n",
    "    duration = record[\"sounds\"][0][\"duration\"] if record.get(\"sounds\") else 160  # Default duration\n",
    "\n",
    "    # Build tags from various fields\n",
    "    tags = []\n",
    "\n",
    "    # Add genre/subgenre\n",
    "    if record.get(\"genre\"):\n",
    "        for g in record[\"genre\"]:\n",
    "            tags.append(g[\"label\"])\n",
    "    if record.get(\"subgenre\"):\n",
    "        for sg in record[\"subgenre\"]:\n",
    "            tags.append(sg[\"label\"])\n",
    "\n",
    "    # Add instruments\n",
    "    if record.get(\"instruments\"):\n",
    "        for inst in record[\"instruments\"]:\n",
    "            tags.append(inst[\"label\"].lower())\n",
    "\n",
    "    # Add moods\n",
    "    if record.get(\"moods\"):\n",
    "        for mood in record[\"moods\"]:\n",
    "            tags.append(mood[\"label\"].lower())\n",
    "\n",
    "    # Add keywords\n",
    "    if record.get(\"keywords\"):\n",
    "        for keyword in record[\"keywords\"]:\n",
    "            tags.append(keyword[\"label\"].lower())\n",
    "\n",
    "    # Add tempo info\n",
    "    if record.get(\"tempo\"):\n",
    "        tags.append(f\"tempo: {record['tempo'].lower()}\")\n",
    "    if record.get(\"bpm\"):\n",
    "        tags.append(f\"bpm: {record['bpm']}\")\n",
    "\n",
    "    # Add music key\n",
    "    if record.get(\"music_key\"):\n",
    "        tags.append(f\"key: {record['music_key'].lower()}\")\n",
    "\n",
    "    # Add era if available\n",
    "    if record.get(\"eras\"):\n",
    "        for era in record[\"eras\"]:\n",
    "            tags.append(era[\"label\"])\n",
    "\n",
    "    if description:\n",
    "        tags.append(f\"{description}\")\n",
    "\n",
    "    # Find the full version audio file path\n",
    "    local_filepath = None\n",
    "\n",
    "    # If no full version found, try to construct a default path\n",
    "    if not local_filepath:\n",
    "        local_filepath = f\"/app2/suno/data/extreme_music/audio/{track_id}/Full Version.mp3\"\n",
    "\n",
    "    if not os.path.exists(local_filepath):\n",
    "        return None\n",
    "\n",
    "    # Build the SunoGPT metadata record\n",
    "    suno_record = {\n",
    "        \"id\": f\"extreme_{track_id}\",  # Prefix to avoid ID conflicts\n",
    "        \"local_filepath\": local_filepath,\n",
    "        \"duration_s\": float(duration),\n",
    "        \"weight\": 1.0,  # Default weight\n",
    "        \"tags\": tags[:20],  # Limit to 20 tags to avoid overloading\n",
    "        \"lang\": \"en\",  # Default to English\n",
    "    }\n",
    "\n",
    "    if record.get(\"lyrics\"):\n",
    "        suno_record[\"text\"] = record[\"lyrics\"]\n",
    "\n",
    "    if record.get(\"stems\"):\n",
    "        stem_dict = {}\n",
    "        for stem in record[\"stems\"]:\n",
    "            stem_path = f\"/app2/suno/data/extreme_music/audio/{track_id}/{stem}.mp3\"\n",
    "            stem_dict[stem] = stem_path\n",
    "\n",
    "        if all(os.path.exists(path) for path in stem_dict.values()):\n",
    "            suno_record[\"stems\"] = stem_dict\n",
    "\n",
    "    return suno_record\n",
    "\n",
    "\n",
    "# Test the transformation with a sample record\n",
    "print(\"Testing transformation function...\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eeaf1da5",
   "metadata": {},
   "outputs": [],
   "source": [
    "transform_extreme_music_record(extreme_metas[-10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a2392f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Process all extreme music records\n",
    "print(f\"Processing {len(extreme_metas)} extreme music records...\")\n",
    "\n",
    "from multiprocessing import Pool\n",
    "import multiprocessing as mp\n",
    "\n",
    "\n",
    "def process_record_wrapper(record):\n",
    "    try:\n",
    "        return transform_extreme_music_record(record)\n",
    "    except Exception as e:\n",
    "        print(f\"Error processing record: {e}\")\n",
    "        return None\n",
    "\n",
    "\n",
    "with Pool(processes=mp.cpu_count()) as pool:\n",
    "    results = list(tqdm(pool.imap(process_record_wrapper, extreme_metas), total=len(extreme_metas)))\n",
    "\n",
    "suno_records = [record for record in results if record is not None]\n",
    "print(f\"Successfully processed {len(suno_records)} records\")\n",
    "# Count how many records have stems\n",
    "records_with_stems = sum(1 for record in suno_records if \"stems\" in record)\n",
    "print(f\"Records with stems: {records_with_stems} out of {len(suno_records)}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8498aea6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# split extreme metas into train and val\n",
    "import random\n",
    "\n",
    "# Split suno_records into train and val\n",
    "random.seed(42)\n",
    "random.shuffle(suno_records)\n",
    "\n",
    "split_ratio = 0.95\n",
    "split_index = int(len(suno_records) * split_ratio)\n",
    "\n",
    "suno_records_train = suno_records[:split_index]\n",
    "suno_records_val = suno_records[split_index:]\n",
    "\n",
    "print(\n",
    "    f\"Split {len(suno_records)} records into {len(suno_records_train)} train and {len(suno_records_val)} val\"\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c7093e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "!sudo cp /app2/suno/data/auk_v0/metas_v3_val.jsonl /app2/suno/data/auk_v0/metas_v4_val.jsonl\n",
    "!sudo cp /app2/suno/data/auk_v0/metas_v3_tr.jsonl /app2/suno/data/auk_v0/metas_v4_tr.jsonl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fb1ec83",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Append extreme music records to existing v4 files\n",
    "import json\n",
    "\n",
    "# Append to train file\n",
    "train_file = \"/app2/suno/data/auk_v0/metas_v4_tr.jsonl\"\n",
    "with open(train_file, \"a\") as f:\n",
    "    for record in suno_records_train:\n",
    "        f.write(json.dumps(record) + \"\\n\")\n",
    "\n",
    "# Append to val file\n",
    "val_file = \"/app2/suno/data/auk_v0/metas_v4_val.jsonl\"\n",
    "with open(val_file, \"a\") as f:\n",
    "    for record in suno_records_val:\n",
    "        f.write(json.dumps(record) + \"\\n\")\n",
    "\n",
    "print(f\"Appended {len(suno_records_train)} extreme records to train file\")\n",
    "print(f\"Appended {len(suno_records_val)} extreme records to val file\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eadaa562",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}