{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load run config\n",
    "import torch\n",
    "import json\n",
    "config_path = \"/home/christian/code/neon/sunoDiff/config/25hz_v45_infill_shared_flow.json\"\n",
    "with open(config_path, \"r\") as f:\n",
    "    run_config = json.load(f)\n",
    "\n",
    "print(run_config)\n",
    "run_config[\"data\"][\"mode\"] = \"pretrain\"\n",
    "run_config[\"data\"][\"patch_size\"] = 1\n",
    "run_config[\"data\"][\"shard_data\"] = False\n",
    "run_config[\"data\"][\"always_pad_semantic\"] = False\n",
    "run_config[\"data\"][\"always_skip_semantic\"] = False\n",
    "run_config[\"data\"][\"respell_augment_prob\"] = 0.0\n",
    "run_config[\"data\"][\"semantic_dropout\"] = 0.1\n",
    "\n",
    "\n",
    "run_config[\"data\"][\"dataset_dir\"] = \"/app/suno/data/diffusion_mix/dac_vae_tuned_25hz\"\n",
    "run_config[\"data\"][\"train_vae_memmap_filename\"] = \"data_vae_tr.bin\"\n",
    "run_config[\"data\"][\"train_semantic_memmap_filename\"] = \"data_semantic_tr.bin\"\n",
    "run_config[\"data\"][\"train_metas_filename\"] = \"metas_tr_extended_apr30.jsonl\"\n",
    "run_config[\"data\"][\"val_vae_memmap_filename\"] = \"data_vae_val.bin\"\n",
    "run_config[\"data\"][\"val_semantic_memmap_filename\"] = \"data_semantic_val.bin\"\n",
    "run_config[\"data\"][\"val_metas_filename\"] = \"metas_val_extended_apr30.jsonl\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.insert(0, \"/home/christian/code/neon/sunoDiff/\")\n",
    "\n",
    "from dataset import GeneralMemmapMapDataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_train = GeneralMemmapMapDataset(\n",
    "    mode=run_config[\"data\"][\"mode\"],\n",
    "    dataset_dir=(\n",
    "        run_config[\"data\"][\"shard_data_dir\"]\n",
    "        if run_config[\"data\"][\"shard_data\"]\n",
    "        else run_config[\"data\"][\"dataset_dir\"]\n",
    "    ),\n",
    "    vae_memmap_filename=run_config[\"data\"][\"train_vae_memmap_filename\"],\n",
    "    semantic_memmap_filename=run_config[\"data\"][\"train_semantic_memmap_filename\"],\n",
    "    metas_filename=run_config[\"data\"][\"train_metas_filename\"],\n",
    "    #info_filename=run_config[\"data\"][\"val_info_filename\"],\n",
    "    vae_dim=run_config[\"data\"][\"vae_dim\"],\n",
    "    vae_n_tokens=run_config[\"model\"][\"block_size\"] * run_config[\"data\"][\"patch_size\"],\n",
    "    semantic_n_tokens=run_config[\"model\"][\"cond_semantic_len\"],\n",
    "    cond_text_len=run_config[\"model\"][\"cond_text_len\"],\n",
    "    vae_scale_factor=run_config[\"data\"][\"vae_scale_factor\"],\n",
    "    semantic_pad_token=run_config[\"model\"][\"cond_semantic_n_vocab\"] - 1,\n",
    "    ctx_len=run_config[\"model\"][\"ctx_len\"],\n",
    "    aligned_text_prob=run_config[\"data\"][\"aligned_text_prob\"],\n",
    "    #patch_size=run_config[\"data\"][\"patch_size\"],\n",
    "    foreign_weight=run_config[\"data\"][\"foreign_weight\"],\n",
    "    is_training=True,\n",
    "    scale_vae_ctx=run_config[\"data\"][\"scale_vae_ctx\"],\n",
    "    always_pad_semantic=run_config[\"data\"][\"always_pad_semantic\"],\n",
    "    semantic_rate_hz=run_config[\"data\"][\"semantic_rate_hz\"],\n",
    "    prev_vae_ctx=run_config[\"data\"][\"prev_vae_ctx\"],\n",
    "    infill_vae_ctx=run_config[\"data\"][\"infill_vae_ctx\"],\n",
    "    noise_ctx=run_config[\"data\"][\"noise_ctx\"],\n",
    "    always_skip_semantic=run_config[\"data\"][\"always_skip_semantic\"],\n",
    "    semantic_skip_factors=run_config[\"data\"][\"semantic_skip_factors\"],\n",
    "    respell_augment_prob=run_config[\"data\"][\"respell_augment_prob\"],\n",
    "    shared_ctx=run_config[\"model\"][\"shared_ctx\"],\n",
    "    semantic_dropout=run_config[\"data\"][\"semantic_dropout\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataloader_train = torch.utils.data.DataLoader(\n",
    "    dataset_train,\n",
    "    batch_size=run_config[\"training\"][\"batch_size\"],\n",
    "    shuffle=True,\n",
    "    #num_workers=run_config[\"data\"][\"num_workers\"],\n",
    "    num_workers=4,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "pbar = tqdm(dataloader_train, total=len(dataloader_train))\n",
    "for bidx, batch in enumerate(pbar):\n",
    "    \n",
    "    if bidx > 100: \n",
    "        break"
   ]
  },
  {
   "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
}