{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.utils.s3 import read_from_s3\n",
    "from suno_utils.utils.text import read_jsonl, write_jsonl\n",
    "import random\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def summarize_meta(ds_path):\n",
    "    metas = read_from_s3(ds_path, read_f=read_jsonl)\n",
    "    print(\n",
    "        f\"{len(metas):,} tracks with {sum([m['duration_s'] for m in metas])/60/60:,.1f}h total\"\n",
    "    )\n",
    "    print(metas[0].keys())\n",
    "    return metas\n",
    "\n",
    "\n",
    "def sample_meta(metas):\n",
    "    random_index = random.randint(0, len(metas) - 1)\n",
    "    sample = metas[random_index]\n",
    "    for k, v in sample.items():\n",
    "        print(f\"{k}: {v}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"SPLICE\")\n",
    "splice_meta = summarize_meta(\n",
    "    \"s3://suno-data/datasets/bundles/v5/splice/sfx_metas_v0_all.jsonl\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_meta(splice_meta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"POND5_SFX\")\n",
    "pond5_meta = summarize_meta(\n",
    "    \"s3://suno-data/datasets/bundles/v5/pond_sfx/pond5_metas_v0_all.jsonl\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_meta(pond5_meta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"FREESOUND\")\n",
    "freesound_meta = summarize_meta(\n",
    "    \"s3://suno-data/datasets/bundles/v5/freesound/metas_v0_all.jsonl\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_meta(freesound_meta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "final_metas = []\n",
    "\n",
    "for row in freesound_meta:\n",
    "    final_metas.append(\n",
    "        dict(\n",
    "            id=f\"freesound_{row['id']}\",\n",
    "            s3_filepath=row[\"s3_filepath\"],\n",
    "            duration_s=row[\"duration_s\"],\n",
    "            tags=row[\"tags\"],\n",
    "            dataset=\"freesound\",\n",
    "        )\n",
    "    )\n",
    "\n",
    "for row in pond5_meta:\n",
    "    final_metas.append(\n",
    "        dict(\n",
    "            id=f\"pond5_{row['id']}\",\n",
    "            s3_filepath=row[\"s3_filepath\"],\n",
    "            duration_s=row[\"duration_s\"],\n",
    "            tags=row[\"tags\"],\n",
    "            dataset=\"pond5_sfx\",\n",
    "        )\n",
    "    )\n",
    "\n",
    "for row in splice_meta:\n",
    "    final_metas.append(\n",
    "        dict(\n",
    "            id=f\"splice_{row['id']}\",\n",
    "            s3_filepath=row[\"s3_filepath\"],\n",
    "            duration_s=row[\"duration_s\"],\n",
    "            tags=row[\"tags\"],\n",
    "            dataset=\"splice\",\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Combined metas: \")\n",
    "print(\n",
    "    f\"{len(final_metas):,} tracks with {sum([m['duration_s'] for m in final_metas])/60/60:,.1f}h total\"\n",
    ")\n",
    "sample_meta(final_metas)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "random.shuffle(final_metas)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "write_jsonl(final_metas, \"/home/sara/combined_metas.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame.from_dict(final_metas)\n",
    "grouped = df.groupby(\"id\")\n",
    "group_sizes = grouped.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "id_repeats = group_sizes[group_sizes > 1]\n",
    "print(id_repeats)  # this should be empty"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.tasks.dac_vae_fixed_25hz import encode_overlap as encode\n",
    "\n",
    "model_filepath = \"s3://suno-data/minz/models/dac_vae_tuned_25hz.pth\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.tasks.dac_vae_fixed_25hz import preload_models\n",
    "\n",
    "_ = preload_models(\n",
    "    checkpoint_filepath=model_filepath,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "encoded_arrays = []\n",
    "for nn, arr in enumerate(audio_arrays):\n",
    "    encoded_arrays.append(encode([arr], normalize_volume=False)[0])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_train",
   "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": 5
}