{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test on dpo data\n",
    "import os\n",
    "import torch\n",
    "import torchaudio\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "def download_audio(s3_filepath: str, example_id: str, tmp_dir: str):\n",
    "    filename = os.path.basename(s3_filepath)\n",
    "    out_filepath = os.path.join(tmp_dir, f\"{example_id}-{filename}\")\n",
    "    # only download the file if its not already downloaded\n",
    "    if not os.path.isfile(out_filepath):\n",
    "        os.system(f\"aws s3 cp {s3_filepath} {out_filepath} > /dev/null 2>&1\")\n",
    "    return out_filepath\n",
    "\n",
    "\n",
    "BASE_S3_DIR = \"s3://suno-data-uploads/studio/uploads\"\n",
    "\n",
    "# load pkl file\n",
    "DATA_PKL_PATH = (\n",
    "    \"/home/tony/Data/Preference/up_v3/interesting_clips_up_u_3_20241216_full.pkl\"\n",
    ")\n",
    "\n",
    "df = pd.read_pickle(DATA_PKL_PATH)\n",
    "print(f\"df shape: {df.shape}\")\n",
    "\n",
    "df_indices = list(range(0, len(df), 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "OUT_DIR = \"/app/suno/christian/data/dpo_diffusion_test_set_1k\"\n",
    "os.makedirs(OUT_DIR, exist_ok=True)\n",
    "\n",
    "# first shuffle the indices\n",
    "np.random.shuffle(df_indices)\n",
    "\n",
    "from concurrent.futures import ThreadPoolExecutor\n",
    "from functools import partial\n",
    "\n",
    "def process_pair(idx, df, OUT_DIR, BASE_S3_DIR):\n",
    "    negative_row = df.iloc[idx]\n",
    "    positive_row = df.iloc[idx+1]\n",
    "    request_id = df.iloc[idx]['request_id']\n",
    "    out_subdir = os.path.join(OUT_DIR, request_id)\n",
    "    os.makedirs(out_subdir, exist_ok=True)\n",
    "\n",
    "    for row in [negative_row, positive_row]:\n",
    "        s3_filepath = f\"{BASE_S3_DIR}/{row['id_x']}.mp3\"\n",
    "        filename = os.path.basename(s3_filepath)\n",
    "        out_filepath = os.path.join(out_subdir, filename)\n",
    "        # only download the file if its not already downloaded\n",
    "        if not os.path.isfile(out_filepath):\n",
    "            os.system(f\"aws s3 cp {s3_filepath} {out_filepath} > /dev/null 2>&1\")\n",
    "\n",
    "# Use ThreadPoolExecutor for parallel downloads\n",
    "with ThreadPoolExecutor(max_workers=8) as executor:\n",
    "    process_fn = partial(process_pair, df=df, OUT_DIR=OUT_DIR, BASE_S3_DIR=BASE_S3_DIR)\n",
    "    list(tqdm(executor.map(process_fn, df_indices[0:1000]), total=1000))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_env",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}