{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d0b72877",
   "metadata": {},
   "source": [
    "# VQGAN JAX Encoding for 🤗 Datasets in streaming mode"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba7b31e6",
   "metadata": {},
   "source": [
    "This notebook shows how to pre-encode images to token sequences using JAX, VQGAN and 🤗 Datasets in streaming mode.\n",
    "\n",
    "This example uses our YFCC100M dataset, but it should be easy to adapt to any other image/caption dataset in the huggingface hub."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b59489e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import io\n",
    "\n",
    "import requests\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "\n",
    "import torch\n",
    "import torchvision.transforms as T\n",
    "import torchvision.transforms.functional as TF\n",
    "from torchvision.transforms import InterpolationMode\n",
    "import os\n",
    "\n",
    "import jax\n",
    "from jax import pmap"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7c4c1e6",
   "metadata": {},
   "source": [
    "## Dataset and Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d45a289e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import datasets\n",
    "from datasets import Dataset, load_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f26e4f18",
   "metadata": {},
   "source": [
    "We'll use the `validation` set for testing. Adjust accordingly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28893c3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = load_dataset('dalle-mini/YFCC100M_OpenAI_subset', use_auth_token=True, streaming=True, split='validation')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33861477",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "yfcc100m = Path.home()/'data'/'YFCC100M_OpenAI_subset'\n",
    "yfcc100m_output = yfcc100m/'encoded'      # Output directory for encoded files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e7b71c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 128     # Per device\n",
    "num_workers = 16     # Unused in streaming mode"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0793c26a",
   "metadata": {},
   "source": [
    "### Data preparation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86415769",
   "metadata": {},
   "source": [
    "* Images: we transform them so they are center-cropped and square, all of the same size so we can build batches for TPU/GPU processing.\n",
    "* Captions: we extract a single `caption` column from the source data, by concatenating the cleaned title and description.\n",
    "\n",
    "These transformations are done using the Datasets `map` function. In the case of streaming datasets, transformations will run as needed instead of pre-processing the dataset at once."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fdf1851",
   "metadata": {},
   "source": [
    "This helper function is used to decode images from the bytes retrieved in `streaming` mode."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5bbca804",
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "import io\n",
    "\n",
    "def get_image(byte_stream):\n",
    "    image = Image.open(io.BytesIO(byte_stream))\n",
    "    return image.convert('RGB')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b435290b",
   "metadata": {},
   "source": [
    "Image processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e73dfa3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def center_crop(image, max_size=256):\n",
    "    # Note: we allow upscaling too. We should exclude small images.    \n",
    "    image = TF.resize(image, max_size, interpolation=InterpolationMode.LANCZOS)\n",
    "    image = TF.center_crop(image, output_size=2 * [max_size])\n",
    "    return image\n",
    "\n",
    "preprocess_image = T.Compose([\n",
    "    get_image,\n",
    "    center_crop,\n",
    "    T.ToTensor(),\n",
    "    lambda t: t.permute(1, 2, 0)   # Reorder, we need dimensions last\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e3ac8de",
   "metadata": {},
   "source": [
    "Caption preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aadb4d23",
   "metadata": {},
   "outputs": [],
   "source": [
    "import string\n",
    "\n",
    "def create_caption(title, description):\n",
    "    title = title.strip()\n",
    "    description = description.strip()\n",
    "    if len(title) > 0 and title[-1] not in '.!?': title += '.'\n",
    "    return f'{title} {description}'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c4522b9",
   "metadata": {},
   "source": [
    "And this is the basic transformation function to use in `map`. We don't really need the `key`, but we'll keep it for reference. Since we are returning a new dictionary (as opposed to adding entries to the input), this also removes any metadata columns we don't need."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2566ff68",
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_item(item):\n",
    "    return {\n",
    "        'key': item['key'],\n",
    "        'caption': create_caption(item['title_clean'], item['description_clean']),\n",
    "        'image': preprocess_image(item['img'])\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e519e475",
   "metadata": {},
   "source": [
    "Unlike when using non-streaming datasets, the following operation completes immediately in streaming mode. In streaming mode, `num_proc` is not supported."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10d7750e",
   "metadata": {},
   "outputs": [],
   "source": [
    "prepared_dataset = dataset.map(prepare_item, batched=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8595539",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "item = next(iter(prepared_dataset))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04a6eeb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert(list(item.keys()) == ['key', 'caption', 'image'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40d3115f",
   "metadata": {},
   "outputs": [],
   "source": [
    "item['image'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd844e1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "T.ToPILImage()(item['image'].permute(2, 0, 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44d50a51",
   "metadata": {},
   "source": [
    "### Torch DataLoader"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17a4bbc6",
   "metadata": {},
   "source": [
    "We'll create a PyTorch DataLoader for convenience. This allows us to easily take batches of our desired size.\n",
    "\n",
    "We won't be using parallel processing of the DataLoader for now, as the items will be retrieved on the fly. We could attempt to do it using these recommendations: https://pytorch.org/docs/stable/data.html#multi-process-data-loading. For performance considerations, please refer to this thread: https://discuss.huggingface.co/t/allow-streaming-of-large-datasets-with-image-audio/8062/13"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1c08b7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a296677",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch_dataset = prepared_dataset.with_format(\"torch\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29ab13bc",
   "metadata": {},
   "source": [
    "**Note**: according to my tests, `num_workers` is not compatible with Datasets in streaming mode. Processes deadlock and there's no progress."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2df5e13",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataloader = DataLoader(torch_dataset, batch_size=batch_size * jax.device_count())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c15e3783",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch = next(iter(dataloader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71d027fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch['image'].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a354472b",
   "metadata": {},
   "source": [
    "## VQGAN-JAX model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fcf01d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from vqgan_jax.modeling_flax_vqgan import VQModel"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9daa636d",
   "metadata": {},
   "source": [
    "We'll use a VQGAN trained with Taming Transformers and converted to a JAX model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47a8b818",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model = VQModel.from_pretrained(\"flax-community/vqgan_f16_16384\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62ad01c3",
   "metadata": {},
   "source": [
    "## Encoding"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20357f74",
   "metadata": {},
   "source": [
    "Encoding is really simple using `shard` to automatically distribute \"superbatches\" across devices, and `pmap`. This is all it takes to create our encoding function, that will be jitted on first use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6686b004",
   "metadata": {},
   "outputs": [],
   "source": [
    "from flax.training.common_utils import shard\n",
    "from functools import partial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "322a4619",
   "metadata": {},
   "outputs": [],
   "source": [
    "@partial(jax.pmap, axis_name=\"batch\")\n",
    "def encode(batch):\n",
    "    # Not sure if we should `replicate` params, does not seem to have any effect\n",
    "    _, indices = model.encode(batch)\n",
    "    return indices"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14375a41",
   "metadata": {},
   "source": [
    "### Encoding loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff6c10d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "\n",
    "def encode_captioned_dataset(dataloader, output_dir, save_every=14):\n",
    "    output_dir.mkdir(parents=True, exist_ok=True)\n",
    "        \n",
    "    # Saving strategy:\n",
    "    # - Create a new file every so often to prevent excessive file seeking.\n",
    "    # - Save each batch after processing.\n",
    "    # - Keep the file open until we are done with it.\n",
    "    file = None        \n",
    "    for n, batch in enumerate(tqdm(iter(dataloader))):\n",
    "        if (n % save_every == 0):\n",
    "            if file is not None:\n",
    "                file.close()\n",
    "            split_num = n // save_every\n",
    "            file = open(output_dir/f'split_{split_num:05x}.jsonl', 'w')\n",
    "\n",
    "        images = batch[\"image\"].numpy()\n",
    "        images = shard(images.squeeze())\n",
    "        encoded = encode(images)\n",
    "        encoded = encoded.reshape(-1, encoded.shape[-1])\n",
    "\n",
    "        keys = batch[\"key\"]\n",
    "        captions = batch[\"caption\"]\n",
    "\n",
    "        encoded_as_string = list(map(lambda item: np.array2string(item, separator=',', max_line_width=50000, formatter={'int':lambda x: str(x)}), encoded))\n",
    "        batch_df = pd.DataFrame.from_dict({\"key\": keys, \"caption\": captions, \"encoding\": encoded_as_string})\n",
    "        batch_df.to_json(file, orient='records', lines=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09ff75a3",
   "metadata": {},
   "source": [
    "Create a new file every 318 iterations. This should produce splits of ~500 MB each, when using a total batch size of 1024."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96222bb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "save_every = 318"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7704863d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "28it [01:17,  1.60s/it]"
     ]
    }
   ],
   "source": [
    "encode_captioned_dataset(dataloader, yfcc100m_output, save_every=save_every)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e266a70a",
   "metadata": {},
   "source": [
    "This is ~10-15 slower than local encoding from an SSD. For performance considerations, see the discussion at https://discuss.huggingface.co/t/allow-streaming-of-large-datasets-with-image-audio/8062/13."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8953dd84",
   "metadata": {},
   "source": [
    "----"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "db471c52d602b4f5f40ecaf278e88ccfef85c29d0a1a07185b0d51fc7acf4e26"
  },
  "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}