{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab0f4499",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import glob\n",
    "import json\n",
    "import numpy as np\n",
    "from typing import Iterable, List\n",
    "from suno_utils.utils.text import read_jsonl, write_jsonl\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d75c929",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ----------------------------\n",
    "# Config\n",
    "# ----------------------------\n",
    "# List all source dataset directories here (arbitrary length).\n",
    "BASE_DIRS: List[str] = [\n",
    "    #\"/app/suno/data/dpo/diffv2_v1_t18/\",\n",
    "    \"/app2/suno/data/dpo/diff3_carp_t1_v1/\",\n",
    "    \"/app2/suno/data/christian/outputs/v3-distill-data-ctx-t2/memmaps/t2_labels_0p6\"\n",
    "]\n",
    "\n",
    "OUT_DIR = \"/app2/suno/data/christian/outputs/t2_merge\"\n",
    "SPLITS = [\"tr\", \"val\"]  # do both tr and val\n",
    "\n",
    "# Memmap “kinds” you want to merge: key → dtype\n",
    "# (Extend here if you add more arrays in the future.)\n",
    "MEMMAP_SPECS = {\n",
    "    \"vae\": np.float16,\n",
    "    \"semantic\": np.uint16,\n",
    "}\n",
    "\n",
    "# Filenames:\n",
    "#   metas:     metas_{split}.jsonl\n",
    "#   memmaps:   data_{key}_{split}.bin\n",
    "# ----------------------------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c723e032",
   "metadata": {},
   "outputs": [],
   "source": [
    "merge()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e2aeeef",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def assert_all_exist(dirs: List[str], path_fn) -> bool:\n",
    "    \"\"\"Return True if every dir has path_fn(dir) existing; otherwise warn and return False.\"\"\"\n",
    "    missing = []\n",
    "    for d in dirs:\n",
    "        p = path_fn(d)\n",
    "        if not os.path.exists(p):\n",
    "            missing.append(p)\n",
    "    if missing:\n",
    "        print(f\"[WARN] Missing files:\\n  - \" + \"\\n  - \".join(missing))\n",
    "        return False\n",
    "    return True\n",
    "\n",
    "\n",
    "def merge_jsonl(dirs: List[str], split: str, out_dir: str):\n",
    "    \"\"\"\n",
    "    Merge metas_{split}.jsonl in order of dirs and write to out_dir.\n",
    "    Uses read_jsonl/write_jsonl for simplicity; if your metas are huge,\n",
    "    swap to streaming (shown in comment below).\n",
    "    \"\"\"\n",
    "    out_path = os.path.join(out_dir, f\"metas_{split}.jsonl\")\n",
    "    all_metas = []\n",
    "    for d in dirs:\n",
    "        p = os.path.join(d, f\"metas_{split}.jsonl\")\n",
    "        metas = read_jsonl(p, progress=True)\n",
    "        all_metas.extend(metas)\n",
    "\n",
    "    write_jsonl(all_metas, out_path)\n",
    "    print(f\"[OK] Wrote metas: {out_path}  (items: {len(all_metas)})\")\n",
    "\n",
    "    # --- STREAMING alternative (if metas are too large for RAM) ---\n",
    "    # with open(out_path, \"w\", encoding=\"utf-8\") as fout:\n",
    "    #     for d in dirs:\n",
    "    #         p = os.path.join(d, f\"metas_{split}.jsonl\")\n",
    "    #         with open(p, \"r\", encoding=\"utf-8\") as fin:\n",
    "    #             for line in fin:\n",
    "    #                 fout.write(line)\n",
    "\n",
    "\n",
    "def memmap_len(path: str, dtype: np.dtype) -> int:\n",
    "    \"\"\"Return element count of a 1D memmap file.\"\"\"\n",
    "    mm = np.memmap(path, dtype=dtype, mode=\"r\")\n",
    "    n = mm.shape[0]\n",
    "    del mm\n",
    "    return n\n",
    "\n",
    "\n",
    "def merge_memmaps(dirs: List[str], split: str, out_dir: str, key: str, dtype: np.dtype):\n",
    "    \"\"\"\n",
    "    Merge data_{key}_{split}.bin files across dirs into a single 1D memmap,\n",
    "    writing sequentially to avoid large RAM usage.\n",
    "    \"\"\"\n",
    "    in_paths = [os.path.join(d, f\"data_{key}_{split}.bin\") for d in dirs]\n",
    "    # Verify all exist\n",
    "    for p in in_paths:\n",
    "        if not os.path.exists(p):\n",
    "            raise FileNotFoundError(f\"Missing memmap for '{key}' split '{split}': {p}\")\n",
    "\n",
    "    # Determine total length\n",
    "    lengths = [memmap_len(p, dtype) for p in in_paths]\n",
    "    total_len = sum(lengths)\n",
    "    print(f\"[INFO] Merging '{key}' ({dtype}) split '{split}': total elements = {total_len}\")\n",
    "\n",
    "    out_path = os.path.join(out_dir, f\"data_{key}_{split}.bin\")\n",
    "\n",
    "    # Create destination memmap and fill slices\n",
    "    dest = np.memmap(out_path, dtype=dtype, mode=\"w+\", shape=(total_len,))\n",
    "    offset = 0\n",
    "    for p, n in zip(in_paths, lengths):\n",
    "        src = np.memmap(p, dtype=dtype, mode=\"r\")\n",
    "        dest[offset : offset + n] = src[:]  # vectorized copy\n",
    "        offset += n\n",
    "        del src\n",
    "    dest.flush()\n",
    "    del dest\n",
    "    print(f\"[OK] Wrote memmap: {out_path}  (elements: {total_len})\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29f5586f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def merge():\n",
    "    os.makedirs(OUT_DIR, exist_ok=True)\n",
    "\n",
    "    # Process each split independently; skip if any required input is missing for that split.\n",
    "    for split in SPLITS:\n",
    "        print(f\"\\n=== Split: {split} ===\")\n",
    "\n",
    "        # Check metas\n",
    "        have_metas = assert_all_exist(\n",
    "            BASE_DIRS, lambda d: os.path.join(d, f\"metas_{split}.jsonl\")\n",
    "        )\n",
    "\n",
    "        # Check all memmap kinds\n",
    "        have_memmaps = True\n",
    "        for key, dtype in MEMMAP_SPECS.items():\n",
    "            ok = assert_all_exist(\n",
    "                BASE_DIRS, lambda d, k=key: os.path.join(d, f\"data_{k}_{split}.bin\")\n",
    "            )\n",
    "            have_memmaps = have_memmaps and ok\n",
    "\n",
    "        if not (have_metas and have_memmaps):\n",
    "            print(f\"[WARN] Skipping split '{split}' due to missing inputs.\")\n",
    "            continue\n",
    "\n",
    "        # Merge metas\n",
    "        merge_jsonl(BASE_DIRS, split, OUT_DIR)\n",
    "\n",
    "        # Merge each memmap kind\n",
    "        for key, dtype in MEMMAP_SPECS.items():\n",
    "            merge_memmaps(BASE_DIRS, split, OUT_DIR, key, dtype)\n",
    "\n",
    "    print(\"\\n[DONE]\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6d45d7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tempfile\n",
    "from suno_utils.utils.text import read_jsonl, write_jsonl\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8ea1d58",
   "metadata": {},
   "outputs": [],
   "source": [
    "# read the source jsonl file\n",
    "work_items = read_jsonl(\n",
    "    \"/home/christian/code/christian/metadata/sft/interesting_clips_bluejay_t1_20250811_public_only_lang_balanced_30k.jsonl\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3aeaa0bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "for work_item in work_items:\n",
    "    item_id = work_item[\"id\"]\n",
    "\n",
    "    with tempfile.TemporaryDirectory() as td:\n",
    "        # copy semantic codes to s3 also\n",
    "        semantic_codes_path = os.path.join(td, f\"{item_id}_semantic.npz\")\n",
    "        np.savez(\n",
    "            semantic_codes_path,\n",
    "            semantic_codes=cropped_semantic_codes.cpu().numpy(),\n",
    "        )\n",
    "\n",
    "        s3_filepath = os.path.join(\n",
    "            self.output_path,\n",
    "            f\"{item_id}\",\n",
    "            f\"{item_id}_semantic.npz\",\n",
    "        )\n",
    "        s3_client.upload_file(\n",
    "            semantic_codes_path,\n",
    "            \"suno-data\",\n",
    "            s3_filepath,\n",
    "            ExtraArgs={\n",
    "                \"ContentType\": \"application/octet-stream\",\n",
    "            },\n",
    "        )"
   ]
  }
 ],
 "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": 5
}