{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.utils.text import read_jsonl, write_jsonl, read_json\n",
    "import random\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def summarize_meta(ds_path):\n",
    "    metas = read_jsonl(ds_path, progress=True)\n",
    "    print(f\"{len(metas):,} tracks with {sum([m['duration_s'] for m in metas])/60/60:,.1f}h total\")\n",
    "    return metas\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": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "## Format SFX for GPT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "metas_sfx = summarize_meta(\"/app2/suno/data/diffusion/sfx/v6/combined_v3_w_extreme_metas_opus_name_norm.jsonl\")\n",
    "id_to_reliable = {m['id']: False for m in metas_sfx}\n",
    "sample_meta(metas_sfx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "reliable = read_json(\"/app2/suno/data/diffusion/sfx/v6/info_all_v0.json\")\n",
    "\n",
    "seen = 0\n",
    "for ds, ids in reliable.items():\n",
    "    for id in ids:\n",
    "        id_to_reliable[id] = True\n",
    "        seen += 1\n",
    "print(seen)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def consolidate_tags_for_row(row):\n",
    "    dur = row[\"duration_s\"]\n",
    "    key = row.get(\"key\", None)\n",
    "    bpm = row.get(\"bpm\", None)\n",
    "    tags = row.get(\"tags\", []).copy()\n",
    "    if key is not None:\n",
    "        key_center = key.split(\" \")[0]\n",
    "        tonality = key.split(\" \")[-1] if len(key.split(\" \")) == 2 else None\n",
    "        tags.append(f\"key: {key_center}\")\n",
    "        if tonality is not None:\n",
    "            tags.append(f\"key: {key_center} {tonality}\")\n",
    "    if bpm is not None:\n",
    "        tags.append(f\"tempo: {bpm}\")\n",
    "    tags.append(f\"duration_s: {dur}\")\n",
    "\n",
    "    return tags\n",
    "\n",
    "def create_gpt_row(meta):\n",
    "    id = meta['id']\n",
    "    suffix = \".opus\"\n",
    "    return dict(\n",
    "        id=id,\n",
    "        s3_filepath=meta['s3_filepath'],\n",
    "        duration_s=meta['duration_s'],\n",
    "        tags = consolidate_tags_for_row(meta),\n",
    "        weight=1.0,\n",
    "        audio_type=\"sfx\",\n",
    "        local_filepath=f\"/app2/suno/data/raw_audio_opus_v0/{id}{suffix}\",\n",
    "        is_reliable=meta[\"is_reliable\"]\n",
    "    )\n",
    "\n",
    "def create_gpt_meta(metadata, reliable_map=None):\n",
    "    if reliable_map is None:\n",
    "        reliable_map = {m['id']: True for m in metadata}\n",
    "\n",
    "    gpt_metas = []\n",
    "    for row in metadata:\n",
    "        gpt_row = create_gpt_row(row)\n",
    "        gpt_metas.append(gpt_row)\n",
    "        #if reliable_map[gpt_row['id']] and os.path.exists(gpt_row['local_filepath']):\n",
    "        #    gpt_metas.append(gpt_row)\n",
    "\n",
    "    return gpt_metas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "sfx_gpt_format = create_gpt_meta(metas_sfx, reliable_map=id_to_reliable)\n",
    "#retained_percent = 100. * len(sfx_gpt_format) / len(metas_sfx)\n",
    "#print(f\"Retained {retained_percent}% of rows.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_meta(sfx_gpt_format)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "write_jsonl(sfx_gpt_format, \"gpt_format_sfx.jsonl\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_clean",
   "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
}