tutorials/multimodal/Multimodal Data Preparation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Multimodal Dataset Preparation\n",
    "\n",
    "The first step of pre-training any deep learning model is data preparation. This notebook will walk you through the 5 stages of data preparation for training a multimodal model:\n",
    "1. Download your Data\n",
    "2. Extract Images and Text\n",
    "3. Re-organize to ensure uniform text-image pairs\n",
    "4. Precache Encodings\n",
    "5. Generate Metadata required for training\n"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "88adf24c9f52084f"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "693d0fcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\n",
    "\n",
    "Instructions for setting up Colab are as follows:\n",
    "1. Open a new Python 3 notebook.\n",
    "2. Import this notebook from GitHub (File -> Upload Notebook -> \"GITHUB\" tab -> copy/paste GitHub URL)\n",
    "3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select \"GPU\" for hardware accelerator)\n",
    "4. Run this cell to set up dependencies.\n",
    "\"\"\"\n",
    "# If you're using Google Colab and not running locally, run this cell.\n",
    "\n",
    "## Install dependencies\n",
    "! pip install img2dataset\n",
    "\n",
    "### Install NeMo\n",
    "BRANCH = 'main'\n",
    "!python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[all]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "For both running this notebook locally and in a nemo container:\n",
    "We need to downgrade opencv version to resolve this issue: https://github.com/opencv/opencv-python/issues/884\n",
    "\"\"\"\n",
    "! pip uninstall -y opencv-python-headless\n",
    "! pip install opencv-python==4.8.0.74\n",
    "\"\"\"\n",
    "We need to downgrade dask version to prevent a repartition error.\n",
    "\"\"\"\n",
    "! pip install dask==2024.1.1"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "bb0c8d61cdb92704"
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c06f3527",
   "metadata": {},
   "source": [
    "\n",
    "This notebook will show you how to prepare an image-text dataset into the [WebDataset](https://github.com/webdataset/webdataset) format. The Webdataset format is required to train all multimodal models in NeMo, such as Stable Diffusion and Imagen. \n",
    "\n",
    "This notebook is designed to demonstrate the different stages of multimodal dataset preparation. It is not meant to be used to process large-scale datasets since many stages are too time-consuming to run without parallelism. For large workloads, we recommend running the multimodal dataset preparation pipeline with the NeMo-Framework-Launcher on multiple processors/GPUs. NeMo-Framework-Launcher packs the same 5 scripts in this notebook into one runnable command and one config file to enable a smooth and a streamlined workflow.\n",
    "\n",
    "Depending on your use case, not all 5 stages need to be run. Please go to [NeMo Multimodal Documentation](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/multimodal/text2img/datasets.html) for an overview of the 5 stages.\n",
    "    \n",
    "We will use a [dummy dataset](https://huggingface.co/datasets/cuichenx/dummy-image-text-dataset) as the dataset example throughout this notebook. This dataset is formatted as a table with one column storing the text captions, and one column storing the URL link to download the corresponding image. This is the same format as most common text-image datasets. The use of this dummy dataset is for demonstration purposes only. **Each user is responsible for checking the content of the dataset and the applicable licenses to determine if it is suitable for the intended use.**\n",
    "\n",
    "Let's first set up some paths."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bef3833e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "LAUNCHER_DIR = \"/opt/NeMo-Framework-Launcher\"  # formerly NeMo-Megatron-Launcher\n",
    "SCRIPT_DIR = os.path.join(LAUNCHER_DIR, \"launcher_scripts/nemo_launcher/collections/dataprep_scripts/multimodal_dataprep\")\n",
    "CONF_DIR = \"conf\"\n",
    "DATA_DIR = \"dummy_data\"\n",
    "os.makedirs(CONF_DIR, exist_ok=True)\n",
    "os.makedirs(DATA_DIR, exist_ok=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "25d76419",
   "metadata": {},
   "source": [
    "\n",
    "## Stage 1: Download Parquet Files from HuggingFace\n",
    ">**Alternative workflows:**\n",
    ">- **If your dataset is not hosted on HuggingFace or your dataset does not contain .parquet files, please move on to Stage 2**\n",
    ">- **If you want to experiment with local image and text files, please see Appendix 1 for a tutorial to create a WebDataset from local images, then move on to Stage 3**\n",
    ">- **If you have a dataset in the WebDataset format already and only want to precache the embeddings, please move on to Stage 4**\n",
    "\n",
    "In this stage, we download the raw data files (.parquet format) from HuggingFace. The parquet files should contain the text captions and the urls to download each image. We then optionally subpartition the parquet file so that the next stage can be parallelized more efficiently.\n",
    "\n",
    "Script: download_parquet.py\n",
    "\n",
    "Arguments:\n",
    "- `dataset_repo_id`: huggingface dataset repo id, in the format of {user_or_company}/{dataset_name}. See [here](https://huggingface.co/datasets?task_categories=task_categories:text-to-image&sort=downloads) for a list of datasets on HuggingFace.\n",
    "- `output_dir`: output of this stage\n",
    "- `parquet_subpartitions`: increase the number of partitions to reduce the image downloading time (next stage) of each task. Useful if the next stage is parallelized over multiple tasks. We will use 3 for this example to keep the run time of subsequent stages short.\n",
    "- `parquet_pattern`: `glob` pattern to use to find the parquet files. Defaults to `*.parquet` (all files that end with the extension)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a06a904b",
   "metadata": {},
   "outputs": [],
   "source": [
    "! python $SCRIPT_DIR/download_parquet.py \\\n",
    "    dataset_repo_id='cuichenx/dummy-image-text-dataset' \\\n",
    "    output_dir=$DATA_DIR/parquet \\\n",
    "    parquet_subpartitions=3 \\\n",
    "    parquet_pattern='*.parquet'"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ebfd2565",
   "metadata": {},
   "source": [
    "**Milestone**: You should now see 3 files in the output directory `$DATA_DIR/parquet/dummy_dataset50000.parquet_parts`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ccc6d7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "! ls $DATA_DIR/parquet/dummy_dataset50000.parquet_parts | wc -l\n",
    "# should output 3"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "56d5dbb9",
   "metadata": {},
   "source": [
    "\n",
    "## Stage 2: Download Images Files\n",
    ">**Alternative workflows:**\n",
    ">- **If your dataset is not hosted on HuggingFace or your dataset does not contain .parquet files, please consult the README page of [img2dataset](https://github.com/rom1504/img2dataset) and modify the command to call img2dataset.**\n",
    "\n",
    "In this stage, we extract the images and texts from the parquet files into the WebDataset format using an open-source tool, [img2dataset](https://github.com/rom1504/img2dataset). \n",
    "This stage will typically benefit from a large degree of parallelism (e.g. thousands of tasks). We will pretend there are 3 tasks running (3 was set in the previous stage), and only work on the first of the 3 parquet subpartitions (i.e. shards) in this notebook.\n",
    "\n",
    "Script: download_images.py\n",
    "\n",
    "Environment variables (automatically set by SLURM if running with NeMo-Framework-Launcher):\n",
    "- `SLURM_ARRAY_TASK_COUNT`: total number of tasks, should be set to the number of parquet files in `$DATA_DIR/parquet/dummy_dataset50000.parquet_parts`. (i.e. `parquet_subpartitions` x `num_parquets_downloaded`)\n",
    "- `SLURM_ARRAY_TASK_ID`: id of the current task (0 <= SLURM_ARRAY_TASK_ID < SLURM_ARRAY_TASK_COUNT)\n",
    "\n",
    "Arguments:\n",
    "- `input_dir`: parquet download dir from the previous stage.\n",
    "- `output_dir`: output of this stage\n",
    "- `parquet_pattern`: see stage 1\n",
    "- `download_num_processes`: number of processes to use. This should be set to number of CPUs in the machine\n",
    "- `download_num_threads`: number of threads to use. This should be tuned to balance cpu usage, internet bandwidth and disk bandwidth. \n",
    "- `img2dataset_additional_arguments`: see [img2dataset](https://github.com/rom1504/img2dataset) for complete list of parameters and [here](https://github.com/rom1504/img2dataset/tree/main/dataset_examples) for some examples. In this example, we use encode_quality=95 for jpeg compression quality, and resize_mode=no to keep the original images on disk. You can also override these arguments to suit your own needs: input_format (default is parquet), caption_col (default is TEXT), url_col (default is URL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72d7e0f8",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# pretend that we're the first task out of 3 tasks\n",
    "! SLURM_ARRAY_TASK_ID=0 SLURM_ARRAY_TASK_COUNT=3 python $SCRIPT_DIR/download_images.py \\\n",
    "    input_dir=$DATA_DIR/parquet \\\n",
    "    output_dir=$DATA_DIR/tarfiles_raw \\\n",
    "    parquet_pattern='*.parquet' \\\n",
    "    download_num_processes=2 \\\n",
    "    download_num_threads=16 \\\n",
    "    \"img2dataset_additional_arguments={{encode_quality:95,resize_mode:10}}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Note: In this dummy dataset, you will likely see a success rate of 1.000 (no failures). However, for read datasets, the success rate will always be much less than 1.000"
   ],
   "id": "eaffa123548d6a5e"
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "27cf740d",
   "metadata": {},
   "source": [
    "**Milestone**: You should now see tar files (along with other files) in the output directory `$DATA_DIR/tarfiles_raw/part.0.parquet`. Inside each tar file, you should see the .jpg image, the corresponding .txt caption, and .json metadata."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1b6477e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "! ls $DATA_DIR/tarfiles_raw/part.0.parquet | head -n 6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f8572f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "! tar -tf $DATA_DIR/tarfiles_raw/part.0.parquet/00000.tar | tail -n 6"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "09fea68d",
   "metadata": {},
   "source": [
    "## Stage 3: Reorganize Tarfiles to Same Number of Image-Text Pairs"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2d0cbfc0",
   "metadata": {},
   "source": [
    "*Note: This stage is required to train multimodal models in NeMo.*\n",
    "\n",
    "In this stage, we reorganize the contents of tar files from the download_images step, so that the tar files are uniform\n",
    "(i.e. each containing an equal number (usually 1000) of training examples (image-text pairs)).\n",
    "The tar files created from the download_images step are not uniform, because there is always a portion of images\n",
    "that fail to download or are no long available.\n",
    "Uniform tar files are important if a sequential sampler is used during training (i.e. not infinite sampler).\n",
    "Uniform tar files are also important for precaching because a sequential sampler is used there.\n",
    "\n",
    "Script: reorganize_tar.py\n",
    "\n",
    "Environment variables (automatically set by SLURM if running with NeMo-Framework-Launcher):\n",
    "- `SLURM_ARRAY_TASK_COUNT`: total number of tasks, should be set to parquet_subpartitions x num_parquets_downloaded\n",
    "- `SLURM_ARRAY_TASK_ID`: id of the current task (0 <= `SLURM_ARRAY_TASK_ID` < `SLURM_ARRAY_TASK_COUNT`)\n",
    "\n",
    "Arguments:\n",
    "- `input_dir`: image download dir from the previous stage.\n",
    "- `output_dir`: output of this stage\n",
    "- `file_ext_in_tar`: target file extensions in each tar file to transfer to the reorganized tar files. In this example, we have .jpg, .txt, and .json in the downloaded tar files, but we will only keep the image and text and discard the .json metadata.\n",
    "- `tar_chunk_size`: number of training examples in each output tar file\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc3d1e60",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "! SLURM_ARRAY_TASK_ID=0 SLURM_ARRAY_TASK_COUNT=1 python $SCRIPT_DIR/reorganize_tar.py \\\n",
    "    input_dir=$DATA_DIR/tarfiles_raw \\\n",
    "    output_dir=$DATA_DIR/tarfiles_reorganized \\\n",
    "    tar_chunk_size=1000 \\\n",
    "    file_ext_in_tar=[.jpg,.txt]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "95c77097",
   "metadata": {},
   "source": [
    "**Milestone**: You should now see tar files (along with other files) in the output directory `$DATA_DIR/tarfiles_reorganized/`. Inside each tar file, you should see exactly 1000 pairs of .jpg image and .txt caption, for a total of 2000 files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4495428f",
   "metadata": {},
   "outputs": [],
   "source": [
    "! tar -tf $DATA_DIR/tarfiles_reorganized/task0000/00001.tar | wc -l\n",
    "# should output 2000"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "09b83dd7",
   "metadata": {},
   "source": [
    "## Stage 4: Precache Encodings\n",
    "\n",
    ">**Alternative workflows:**\n",
    ">- **If you're only testing out the NeMo text2image models and do not care about good training performance, you can skip this step and move on to Stage 5.**\n",
    "\n",
    "### General Format\n",
    "\n",
    "Precaching refers to the offline computation of image and text encodings prior to training a model. This technique\n",
    "is suitable for any model that uses pretrained, frozen encoders during training.\n",
    "By using precached encodings, embeddings for image and text do not need to be recomputed in each epoch,\n",
    "thereby significantly improving training throughput (up to 60% higher).\n",
    "Precached encodings are saved in the format of WebDataset.\n",
    "Each tar file contains one pickle file to store all the modality embeddings for each training example. Optionally,\n",
    "the tar file may also include the original image or text files\n",
    "\n",
    "```\n",
    "t0_r0_0.tar\n",
    "|---- 00000.pickle\n",
    "|---- 00000.jpg (optional)\n",
    "|---- 00000.txt (optional)\n",
    "|---- 00001.pickle\n",
    "|---- 00001.jpg (optional)\n",
    "|---- 00001.txt (optional)\n",
    "...\n",
    "```\n",
    "Each pickle file stores one python dictionary, with key value pairs storing the embedding name and the embedding as a\n",
    "numpy array.\n",
    "\n",
    "### Precaching Config\n",
    "Configuration for precaching can be extensive and intricate for some models. To maintain clarity and ensure an\n",
    "organized workflow, we utilize a separate YAML file for these configurations. The YAML file looks like this:\n",
    "\n",
    "```\n",
    "encodings:\n",
    "  - modality: image\n",
    "    extension: jpg\n",
    "    key: autoencoderkl_image\n",
    "    precision: 16\n",
    "    encoder_config:\n",
    "      cls: nemo.collections.multimodal.models.text_to_image.stable_diffusion.ldm.autoencoder.AutoencoderKL\n",
    "      ... (kwargs to initialize the encoder)\n",
    "  - modality: text\n",
    "    extension: txt\n",
    "    key: clip-vit-large-patch14_text\n",
    "    precision: 32\n",
    "    store_pad_tokens: True\n",
    "    encoder_config:\n",
    "      cls: nemo.collections.multimodal.modules.stable_diffusion.encoders.modules.FrozenCLIPEmbedder\n",
    "      ... (kwargs to initialize the encoder)\n",
    "```\n",
    "\n",
    "In this YAML file, the encodings field specifies a list of embeddings to be saved in the pickle file.\n",
    "Each entry can have the following attributes:\n",
    "\n",
    "\n",
    "- `modality`: either image or text\n",
    "- `extension`: file extension for this modality in the tar file (e.g. 'jpg', 'txt')\n",
    "- `key`: dictionary key for the encoding. It is recommended to follow the format `{model_name}-{model_variant}_{modality}`, if applicable. e.g. `clip-vit-large-patch14_text`\n",
    "- `precision`: precision of the stored tensors (32 or 16)\n",
    "- `store_pad_tokens`: Whether to store the PAD tokens. Not storing PAD tokens can significantly reduce disk usage, but the training script must account for this. Ignored for image modality.\n",
    "- `encoder_config`: This dictionary must contain `cls` which points to the location of the encoder class. The rest of the parameters are treated as kwargs to initiate the encoder class.\n",
    "  - Note: the encoder class must implement an `encode` or `__call__` function. If `store_pad_tokens`, this function must return the encoded tensor. Otherwise, this function must return a tuple of (encoded_tensor, mask). The mask is needed so the script knows which tokens are pad tokens and should be ignored for caching. A mask value of 1 denotes regular tokens, and 0 denotes pad tokens.\n",
    "\n",
    "\n",
    "Note that it is not required to have only one encoding per modality, in the case of multiple encoders.\n",
    "The `encodings` field is designed as a list to account for this. For example, it's possible to have one image embedding\n",
    "and two text embeddings (e.g. one from CLIP and one from T5) and both are used during training. An example is shown below.\n",
    "\n",
    "```\n",
    "encodings:\n",
    "  - modality: image\n",
    "    extension: jpg\n",
    "    key: image_emb\n",
    "    encoder_config:\n",
    "      cls: path.to.ImageEncoder\n",
    "      ...\n",
    "  - modality: text\n",
    "    extension: txt\n",
    "    key: text_emb_1\n",
    "    encoder_config:\n",
    "      cls: path.to.TextEncoder1\n",
    "      ...\n",
    "  - modality: text\n",
    "    extension: txt\n",
    "    key: text_emb_2\n",
    "    encoder_config:\n",
    "      cls: path.to.TextEncoder2\n",
    "      ...\n",
    "```\n",
    "\n",
    "In this tutorial, we will show an example of precaching workflow for Stable Diffusion. "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "27b26036",
   "metadata": {},
   "source": [
    "Let's download an example precaching config file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2bd39f5",
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "! wget https://raw.githubusercontent.com/NVIDIA/NeMo-Framework-Launcher/master/launcher_scripts/conf/data_preparation/multimodal/precache_sd.yaml -P $CONF_DIR/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "986045fe",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from omegaconf import OmegaConf\n",
    "precache_cfg = OmegaConf.load(os.path.join(CONF_DIR, \"precache_sd.yaml\"))\n",
    "del precache_cfg.encodings[1].encoder_config['use_fp16']\n",
    "# visualize the config\n",
    "print(OmegaConf.to_yaml(precache_cfg))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f5bee705",
   "metadata": {},
   "source": [
    "There are a few things to note about this config file:\n",
    "- `batch_size_per_GPU`: this should be set to as much as your GPU memory can fit\n",
    "- `save_original_in_tar`: for SD, original images or text are not used during training (if using precached encodings), so we can leave this empty. If you want the original image and text copied into the tar file, you can set this to [image, text]. \n",
    "\n",
    "In this example, we need to download the weights of the image autoencoder from the HuggingFace [Stable Diffusion v1.5 repo](https://huggingface.co/runwayml/stable-diffusion-v1-5/blob/main/vae/diffusion_pytorch_model.bin). Text encoder weights will be downloaded automatically when the model `FrozenCLIPEmbedder` is initialized."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d6804d4",
   "metadata": {},
   "outputs": [],
   "source": "! wget https://huggingface.co/CompVis/stable-diffusion-v1-4/resolve/main/vae/diffusion_pytorch_model.bin"
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "d2e5d6aa",
   "metadata": {},
   "source": [
    "Then, we modify the `from_pretrained` field with the weights file, and save this config as a yaml file to disk.\n",
    "We also adjust the config to use 1 GPU for this tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d509be7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "precache_cfg.encodings[0].encoder_config.from_pretrained = 'diffusion_pytorch_model.bin'\n",
    "precache_cfg.lightning.devices=1\n",
    "# precache_cfg.batch_size_per_GPU=8  # adjust if needed\n",
    "\n",
    "OmegaConf.save(precache_cfg, os.path.join(CONF_DIR, \"precache_sd_example.yaml\"))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "765d281e",
   "metadata": {},
   "source": [
    "Now we can run the precaching script. \n",
    "\n",
    "Script: precache_encodings.py\n",
    "\n",
    "Environment variables (automatically set by SLURM if running with NeMo-Framework-Launcher):\n",
    "- `SLURM_ARRAY_TASK_COUNT`: total number of tasks, should be set to parquet_subpartitions x num_parquets_downloaded\n",
    "- `SLURM_ARRAY_TASK_ID`: id of the current task (0 <= `SLURM_ARRAY_TASK_ID` < `SLURM_ARRAY_TASK_COUNT`)\n",
    "\n",
    "Arguments:\n",
    "- `input_dir`: reorganized tar dir from the previous stage.\n",
    "- `output_dir`: output of this stage\n",
    "- `tar_chunk_size`: number of training examples in each output tar file\n",
    "- `precache_cfg`: precaching config file as describe above\n",
    "\n",
    "This stage will typically benefit from a large degree of parallelism (e.g. thousands of tasks). We will pretend that we are the first task out of 2 tasks. Since the input directory is already 1/3 of the full dataset, the result of this stage in the tutorial will be 1/6 of the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63c9b44c",
   "metadata": {},
   "outputs": [],
   "source": [
    "! SLURM_ARRAY_TASK_ID=0 SLURM_ARRAY_TASK_COUNT=2 python $SCRIPT_DIR/precache_encodings.py \\\n",
    "    input_dir=$DATA_DIR/tarfiles_reorganized \\\n",
    "    output_dir=$DATA_DIR/tarfiles_precached \\\n",
    "    tar_chunk_size=1000 \\\n",
    "    precache_config_path=$CONF_DIR/precache_sd_example.yaml"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "76082b01",
   "metadata": {},
   "source": [
    "**Milestone**: You should now see tar files (along with other files) in the output directory `$DATA_DIR/tarfiles_precached/`. Inside each tar file, you should see exactly 1000 .pickle files storing the image and text embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e93883ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "! tar -tf `ls -d $DATA_DIR/tarfiles_precached/* | head -n 1` | wc -l\n",
    "# should output 1000"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "6394e07f",
   "metadata": {},
   "source": [
    "## Stage 5: Generate `wdinfo` Metadata\n",
    "\n",
    "This stage generates the metadata required by the NeMo multimodal training data pipeline. The metadata contains the chunk size, the total size of dataset, and most importantly, the list of tarfiles.\n",
    "\n",
    "The metadata will only include tarfiles with exactly `tar_chunk_size` (1000 in this tutorial) examples, and ignore/discard incomplete tar files. Incomplete tar files are the last tar file generated by each process which contain the leftover training examples, most likely less than `tar_chunk_size`. \n",
    "\n",
    "In the case of a high degree of parallelism, there can be a significant number of incomplete tarfiles leading to a waste of discarded training examples. Therefore, before creating the metadata file, this stage will also find all the incomplete tar files generated in the previous stage, and combine them in a single-process script so that there is at most one incomplete tarfile throughout the entire dataset.\n",
    "\n",
    "Script: generate_wdinfo.py\n",
    "\n",
    "Arguments:\n",
    "- `input_dir`: output tar dir from stage 3 or stage 4.\n",
    "- `output_wdinfo_path`: output of this stage\n",
    "- `tar_chunk_size`: number of training examples in each output tar file\n",
    "- `file_ext_in_tar`: see explanation in Stage 3. If you performed precaching without copying any files from the source tar (i.e. save_original_in_tar: null in precache_sd_example.yaml) then this should be [.pickle]\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b72a03a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "! python $SCRIPT_DIR/generate_wdinfo.py \\\n",
    "    input_dir=$DATA_DIR/tarfiles_precached \\\n",
    "    output_wdinfo_path=$DATA_DIR/wdinfo.pkl \\\n",
    "    tar_chunk_size=1000 \\\n",
    "    file_ext_in_tar=[.pickle]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a9c8e2ca",
   "metadata": {},
   "source": [
    "**Milestone**: You should now see the wdinfo.pkl file generated. The content of the file is printed above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cc2e2a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "! test -f $DATA_DIR/wdinfo.pkl && echo \"File exists\" || echo \"File does not exist\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Appendix 1: Create a WebDataset from Local Image Text Files\n",
    "\n",
    "\n",
    "If you have image and text files already downloaded, you can quickly convert your dataset to the WebDataset format and proceed with Stages 3-5 of this tutorial without wasting any time on download. Simply follow the steps below.\n",
    "\n",
    "1. Manipulate your dataset so that each text caption is stored in a single file, and shares the same file path as the corresponding image, except the extension. The file path can contain subfolders. An example is shown below\n",
    "\n",
    "    ```bash\n",
    "    > cd dataset\n",
    "    > find . -type f\n",
    "      ./train/n00001234/00010000.jpg    ./train/n00001234/00010000.txt\n",
    "      ./train/n00001234/00010001.jpg    ./train/n00001234/00010001.txt\n",
    "      ./train/n00001234/00010002.jpg    ./train/n00001234/00010002.txt\n",
    "      ./train/n00001234/00010003.jpg    ./train/n00001234/00010003.txt\n",
    "      ./train/n00001234/00010004.jpg    ./train/n00001234/00010004.txt\n",
    "      ./train/n00001235/00010000.jpg    ./train/n00001235/00010000.txt\n",
    "      ./train/n00001235/00010001.jpg    ./train/n00001235/00010001.txt\n",
    "      ...\n",
    "    > cd ..\n",
    "    ```\n",
    "\n",
    "2. Run this command to create of tarball of the folder with sorted file names. It is important for WebDataset to have the image and text files in consecutive blocks on disk, hence the sorting is necessary.\n",
    "\n",
    "```bash\n",
    "> tar --sort=name -cf dataset.tar dataset/\n",
    "```\n",
    "\n",
    "For more information, please visit [Creating a WebDataset](https://github.com/webdataset/webdataset#creating-a-webdataset)\n",
    "\n",
    "After this, you can proceed with Stage 3 of the tutorial.\n",
    "Note: if you can use a script to create folders with exactly `tar_chunk_size` (1000 in the tutorial) image-text pairs, and create  multiple tarfiles each with `tar_chunk_size` pairs of data, then you can skip Stage 3 and proceed with Stage 4 of the tutorial."
   ],
   "id": "217dacb92b870798"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}