{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "import yaml\n",
    "from suno_utils.utils.s3 import read_from_s3\n",
    "from suno_utils.utils.text import write_jsonl, read_jsonl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# create a bundle of the slakh dataset\n",
    "# first we need to get all the original files and their metadata\n",
    "# start by getting all the track directories from all the subdirs\n",
    "LOCAL_DATA_DIR = \"/app/suno/data/slakh2100_flac_redux/\"\n",
    "S3_DATA_DIR = \"s3://suno-data/datasets/slakh2100_flac_redux\"\n",
    "\n",
    "track_dirs = glob.glob(os.path.join(LOCAL_DATA_DIR, \"**/*\"))\n",
    "print(f\"Found {len(track_dirs)} track directories\")\n",
    "\n",
    "new_metas = []\n",
    "\n",
    "for track_dir in track_dirs:\n",
    "    if os.path.isdir(track_dir):\n",
    "        # get the metadata\n",
    "        metadata_path = os.path.join(track_dir, \"metadata.yaml\")\n",
    "        with open(metadata_path, \"r\") as f:\n",
    "            metadata = yaml.safe_load(f)\n",
    "        print(metadata)\n",
    "\n",
    "        # remove the local data dir from the track path\n",
    "        track_path = track_dir.replace(LOCAL_DATA_DIR, \"\")\n",
    "        print(track_path)\n",
    "        s3_filepath = f\"{S3_DATA_DIR}/{track_path}/mix.flac\"\n",
    "\n",
    "        # get the mix file            \n",
    "        new_meta = {\n",
    "            \"id\": metadata[\"UUID\"],\n",
    "            \"s3_filepath\": s3_filepath,\n",
    "            \"tags\": [],\n",
    "            \"text\": \"[Instrumental]\",\n",
    "        }\n",
    "        new_metas.append(new_meta)\n",
    "        \n",
    "        for stem_name, stem_metadata in metadata[\"stems\"].items():\n",
    "            stem_filepath = os.path.join(track_dir, f\"{stem_name}.flac\")\n",
    "            s3_filepath = f\"{S3_DATA_DIR}/{track_path}/stems/{stem_name}.flac\"\n",
    "            tags = []\n",
    "            inst_class = stem_metadata.get(\"instrument_class\", \"\")\n",
    "            if inst_class:\n",
    "                tags.append(inst_class)\n",
    "            midi_program_name = stem_metadata.get(\"midi_program_name\", \"\")\n",
    "            if midi_program_name:\n",
    "                tags.append(midi_program_name)\n",
    "            plugin_name = stem_metadata.get(\"plugin_name\", \"\")\n",
    "            if plugin_name:\n",
    "                tags.append(plugin_name.split(\".\")[0].replace(\"_\", \" \"))\n",
    "            new_meta = {\n",
    "                \"id\": f\"{metadata['UUID']}-{stem_name}\",\n",
    "                \"s3_filepath\": s3_filepath,\n",
    "                \"tags\": tags,\n",
    "                \"text\": \"[Instrumental]\",\n",
    "            }\n",
    "            new_metas.append(new_meta)\n",
    "\n",
    "        break\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_metas[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_diff",
   "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.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}