feat: cleanup encode_dataset

tools/dataset/{make_dataset.ipynb → encode_dataset.ipynb}

@@ -5,7 +5,7 @@
    "id": "d0b72877",
    "metadata": {},
    "source": [
-    "# VQGAN JAX Encoding for `webdataset`"
+    "# Pre-encoding a dataset for DALLE·mini"
    ]
   },
   {
@@ -15,7 +15,11 @@
    "source": [
     "This notebook shows how to pre-encode images to token sequences using JAX, VQGAN and a dataset in the [`webdataset` format](https://webdataset.github.io/webdataset/).\n",
     "\n",
-    "This example uses a small subset of YFCC100M we created for testing, but it should be easy to adapt to any other image/caption dataset in the `webdataset` format."
+    "Adapt it to your own dataset and image encoder.\n",
+    "\n",
+    "At the end you should have a dataset of pairs:\n",
+    "* a caption defined as a string\n",
+    "* an encoded image defined as a list of int."
    ]
   },
   {
@@ -25,19 +29,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import numpy as np\n",
-    "from tqdm import tqdm\n",
+    "from tqdm.notebook 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 math\n",
     "\n",
     "import webdataset as wds\n",
     "\n",
     "import jax\n",
-    "from jax import pmap"
+    "import braceexpand\n",
+    "from pathlib import Path"
    ]
   },
   {
@@ -45,184 +45,110 @@
    "id": "c7c4c1e6",
    "metadata": {},
    "source": [
-    "## Dataset and Parameters"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "9822850f",
-   "metadata": {},
-   "source": [
-    "The following is the list of shards we'll process. We hardcode the length of data so that we can see nice progress bars using `tqdm`."
+    "## Configuration Parameters"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 3,
    "id": "1265dbfe",
    "metadata": {},
    "outputs": [],
    "source": [
-    "shards = 'https://huggingface.co/datasets/dalle-mini/YFCC100M_OpenAI_subset/resolve/main/data/shard-{0000..0008}.tar'\n",
-    "length = 8320"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "7e38fa14",
-   "metadata": {},
-   "source": [
-    "If we are extra cautious or our server is unreliable, we can enable retries by providing a custom `curl` retrieval command:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "4c8c5960",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Enable curl retries to try to work around temporary network / server errors.\n",
-    "# This shouldn't be necessary when using reliable servers.\n",
-    "# shards = f'pipe:curl -s --retry 5 --retry-delay 5 -L {shards} || true'"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "13c6631b",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from pathlib import Path\n",
+    "shards = \"my_images/shard-{0000..0008}.tar\"  # defined using braceexpand format as used by webdataset\n",
+    "encoded_output = Path(\"encoded_data\")  # where we will save our encoded data\n",
     "\n",
-    "# Output directory for encoded files\n",
-    "encoded_output = Path.home()/'data'/'wds'/'encoded'\n",
+    "VQGAN_REPO, VQGAN_COMMIT_ID = (\n",
+    "    \"dalle-mini/vqgan_imagenet_f16_16384\",\n",
+    "    \"85eb5d3b51a1c62a0cc8f4ccdee9882c0d0bd384\",\n",
+    ")\n",
     "\n",
-    "batch_size = 128           # Per device\n",
-    "num_workers = 8            # For parallel processing"
+    "# good defaults for a TPU v3-8\n",
+    "batch_size = 128  # Per device\n",
+    "num_workers = 8  # For parallel processing\n",
+    "total_bs = batch_size * jax.device_count()  # You can use a smaller size while testing\n",
+    "save_frequency = 128  # Number of batches to create a new file (180MB for f16 and 720MB for f8 per file)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
-   "id": "3435fb85",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "bs = batch_size * jax.device_count()    # You can use a smaller size while testing\n",
-    "batches = math.ceil(length / bs)"
+   "execution_count": 5,
+   "id": "cd956ec6-7d98-4d4d-a454-f80fe857eadd",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['XXX/shard-0000.tar',\n",
+       " 'XXX/shard-0001.tar',\n",
+       " 'XXX/shard-0002.tar',\n",
+       " 'XXX/shard-0003.tar',\n",
+       " 'XXX/shard-0004.tar',\n",
+       " 'XXX/shard-0005.tar',\n",
+       " 'XXX/shard-0006.tar',\n",
+       " 'XXX/shard-0007.tar',\n",
+       " 'XXX/shard-0008.tar']"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "shards = list(\n",
+    "    braceexpand.braceexpand(shards)\n",
+    ")  # better display for tqdm with known length"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "88598e4b",
+   "id": "75dba8e2",
    "metadata": {},
    "source": [
-    "Image processing"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "669b35df",
-   "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",
-    "    center_crop,\n",
-    "    T.ToTensor(),\n",
-    "    lambda t: t.permute(1, 2, 0)   # Reorder, we need dimensions last\n",
-    "])"
+    "## Load data"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "a185e90c",
+   "id": "a1e8fb95",
    "metadata": {},
    "source": [
-    "Caption preparation.\n",
-    "\n",
-    "Note that we receive the contents of the `json` structure, which will be replaced by the string we return.\n",
-    "If we want to keep other fields inside `json`, we can add `caption` as a new field."
+    "We load data using `webdataset`."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "423ee10e",
+   "id": "9ef5de9e",
    "metadata": {},
    "outputs": [],
    "source": [
-    "def create_caption(item):\n",
-    "    title = item['title_clean'].strip()\n",
-    "    description = item['description_clean'].strip()\n",
-    "    if len(title) > 0 and title[-1] not in '.!?': title += '.'\n",
-    "    return f'{title} {description}'"
+    "ds = (\n",
+    "    wds.WebDataset(shards, handler=wds.warn_and_continue)\n",
+    "    .decode(\"rgb\", handler=wds.warn_and_continue)\n",
+    "    .to_tuple(\"jpg\", \"txt\")  # assumes image is in `jpg` and caption in `txt`\n",
+    "    .batched(total_bs)  # load in batch per worker (faster)\n",
+    ")"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "8d3a95db",
+   "id": "90981824",
    "metadata": {},
    "source": [
-    "When an error occurs (a download is disconnected, an image cannot be decoded, etc) the process stops with an exception. We can use one of the exception handlers provided by the `webdataset` library, such as `wds.warn_and_continue` or `wds.ignore_and_continue` to ignore the offending entry and keep iterating.\n",
-    "\n",
-    "**IMPORTANT WARNING:** Do not use error handlers to ignore exceptions until you have tested that your processing pipeline works fine. Otherwise, the process will continue trying to find a valid entry, and it will consume your whole dataset without doing any work.\n",
-    "\n",
-    "We can also create our custom exception handler as demonstrated here:"
+    "Note:\n",
+    "* you can also shuffle shards and items using `shardshuffle` and `shuffle` if necessary.\n",
+    "* you may need to resize images in your pipeline (with `map_dict` for example), we assume they are already set to 256x256.\n",
+    "* you can also filter out some items using `select`."
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "369d9719",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# UNUSED - Log exceptions to a file\n",
-    "def ignore_and_log(exn):\n",
-    "    with open('errors.txt', 'a') as f:\n",
-    "        f.write(f'{repr(exn)}\\n')\n",
-    "    return True"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "27de1414",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Or simply use `wds.ignore_and_continue`\n",
-    "exception_handler = wds.warn_and_continue"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "5149b6d5",
+   "cell_type": "markdown",
+   "id": "129c377d",
    "metadata": {},
-   "outputs": [],
    "source": [
-    "dataset = wds.WebDataset(shards,\n",
-    "                         length=batches,              # Hint so `len` is implemented\n",
-    "                         shardshuffle=False,          # Keep same order for encoded files for easier bookkeeping. Set to `True` for training.\n",
-    "                         handler=exception_handler,   # Ignore read errors instead of failing.\n",
-    ")\n",
-    "\n",
-    "dataset = (dataset           \n",
-    "      .decode('pil')                     # decode image with PIL\n",
-    "#       .map_dict(jpg=preprocess_image, json=create_caption, handler=exception_handler)    # Process fields with functions defined above\n",
-    "      .map_dict(jpg=preprocess_image, json=create_caption)    # Process fields with functions defined above\n",
-    "      .to_tuple('__key__', 'jpg', 'json') # filter to keep only key (for reference), image, caption.\n",
-    "      .batched(bs))                      # better to batch in the dataset (but we could also do it in the dataloader) - this arg does not affect speed and we could remove it"
+    "We can now inspect our data."
    ]
   },
   {
@@ -235,7 +161,7 @@
    "outputs": [],
    "source": [
     "%%time\n",
-    "keys, images, captions = next(iter(dataset))"
+    "images, captions = next(iter(ds))"
    ]
   },
   {
@@ -251,54 +177,50 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "c24693c0",
+   "id": "5acfc4d8",
    "metadata": {},
    "outputs": [],
    "source": [
-    "T.ToPILImage()(images[0].permute(2, 0, 1))"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "44d50a51",
-   "metadata": {},
-   "source": [
-    "### Torch DataLoader"
+    "captions[:10]"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "e2df5e13",
+   "id": "c24693c0",
    "metadata": {},
    "outputs": [],
    "source": [
-    "dl = torch.utils.data.DataLoader(dataset, batch_size=None, num_workers=num_workers)"
+    "T.ToPILImage()(images[0].permute(2, 0, 1))"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "a354472b",
+   "id": "3059ffb1",
    "metadata": {},
    "source": [
-    "## VQGAN-JAX model"
+    "Finally we create our dataloader."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "2fcf01d7",
+   "id": "c227c551",
    "metadata": {},
    "outputs": [],
    "source": [
-    "from vqgan_jax.modeling_flax_vqgan import VQModel"
+    "dl = (\n",
+    "    wds.WebLoader(ds, batch_size=None, num_workers=8).unbatched().batched(total_bs)\n",
+    ")  # avoid partial batch at the end of each worker"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "9daa636d",
+   "id": "a354472b",
    "metadata": {},
    "source": [
+    "## Image encoder\n",
+    "\n",
     "We'll use a VQGAN trained with Taming Transformers and converted to a JAX model."
    ]
   },
@@ -311,7 +233,11 @@
    },
    "outputs": [],
    "source": [
-    "model = VQModel.from_pretrained(\"flax-community/vqgan_f16_16384\")"
+    "from vqgan_jax.modeling_flax_vqgan import VQModel\n",
+    "from flax.jax_utils import replicate\n",
+    "\n",
+    "vqgan = VQModel.from_pretrained(\"flax-community/vqgan_f16_16384\")\n",
+    "vqgan_params = replicate(vqgan.params)"
    ]
   },
   {
@@ -327,18 +253,7 @@
    "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"
+    "Encoding is really simple using `shard` to automatically distribute batches across devices and `pmap`."
    ]
   },
   {
@@ -348,21 +263,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from flax.training.common_utils import shard\n",
+    "from functools import partial\n",
+    "\n",
+    "\n",
     "@partial(jax.pmap, axis_name=\"batch\")\n",
-    "def encode(batch):\n",
+    "def p_encode(batch, params):\n",
     "    # Not sure if we should `replicate` params, does not seem to have any effect\n",
-    "    _, indices = model.encode(batch)\n",
+    "    _, indices = vqgan.encode(batch, params=params)\n",
     "    return indices"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "id": "14375a41",
-   "metadata": {},
-   "source": [
-    "### Encoding loop"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -370,49 +281,48 @@
    "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, (keys, images, captions) in enumerate(tqdm(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 = shard(images.numpy().squeeze())\n",
-    "        encoded = encode(images)\n",
+    "def encode_dataset(dataloader, output_dir, save_frequency):\n",
+    "    output_dir.mkdir(parents=True, exist_ok=True)\n",
+    "    all_captions = []\n",
+    "    all_encoding = []\n",
+    "    n_file = 1\n",
+    "    for idx, (images, captions) in enumerate(tqdm(dataloader)):\n",
+    "        images = images.numpy()\n",
+    "        n = len(images) // 8 * 8\n",
+    "        if n != len(images):\n",
+    "            # get the max number of images we can (multiple of 8)\n",
+    "            print(f\"Different sizes {n} vs {len(images)}\")\n",
+    "            images = images[:n]\n",
+    "            captions = captions[:n]\n",
+    "        if not len(captions):\n",
+    "            print(f\"No images/captions in batch...\")\n",
+    "            continue\n",
+    "        images = shard(images)\n",
+    "        encoded = p_encode(images, vqgan_params)\n",
     "        encoded = encoded.reshape(-1, encoded.shape[-1])\n",
+    "        all_captions.extend(captions)\n",
+    "        all_encoding.extend(encoded.tolist())\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"
+    "        # save files\n",
+    "        if (idx + 1) % save_frequency == 0:\n",
+    "            print(f\"Saving file {n_file}\")\n",
+    "            batch_df = pd.DataFrame.from_dict(\n",
+    "                {\"caption\": all_captions, \"encoding\": all_encoding}\n",
+    "            )\n",
+    "            batch_df.to_parquet(f\"{output_dir}/{n_file:03d}.parquet\")\n",
+    "            all_captions = []\n",
+    "            all_encoding = []\n",
+    "            n_file += 1\n",
+    "\n",
+    "    if len(all_captions):\n",
+    "        print(f\"Saving final file {n_file}\")\n",
+    "        batch_df = pd.DataFrame.from_dict(\n",
+    "            {\"caption\": all_captions, \"encoding\": all_encoding}\n",
+    "        )\n",
+    "        batch_df.to_parquet(f\"{output_dir}/{n_file:03d}.parquet\")"
    ]
   },
   {
@@ -422,7 +332,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "encode_captioned_dataset(dl, encoded_output, save_every=save_every)"
+    "encode_dataset(dl, output_dir=encoded_output, save_frequency=save_frequency)"
    ]
   },
   {
@@ -453,7 +363,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.10"
+   "version": "3.9.7"
   }
  },
  "nbformat": 4,