Replace notebooks with the correct versions.

dev/encoding/vqgan-jax-encoding-webdataset.ipynb

@@ -5,7 +5,7 @@
    "id": "d0b72877",
    "metadata": {},
    "source": [
-    "# vqgan-jax-encoding-alamy"
+    "# VQGAN JAX Encoding for `webdataset`"
    ]
   },
   {
@@ -13,12 +13,14 @@
    "id": "ba7b31e6",
    "metadata": {},
    "source": [
-    "Encoding notebook for Alamy dataset."
+    "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."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "id": "3b59489e",
    "metadata": {},
    "outputs": [],
@@ -46,44 +48,83 @@
     "## 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`."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "13c6631b",
+   "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": [
-    "shards = 'https://s3.us-west-1.wasabisys.com/doodlebot-wasabi/datasets/alamy/webdataset/alamy-{000..895}.tar'\n",
-    "\n",
     "# 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'\n",
-    "\n",
-    "length = 44710810    # estimate\n",
-    "\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",
     "\n",
     "# Output directory for encoded files\n",
-    "encoded_output = Path.home()/'data'/'alamy'/'encoded'\n",
+    "encoded_output = Path.home()/'data'/'wds'/'encoded'\n",
     "\n",
     "batch_size = 128           # Per device\n",
-    "num_workers = 8            # Using larger numbers seemed to be less reliable in this case."
+    "num_workers = 8            # For parallel processing"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "id": "3435fb85",
    "metadata": {},
    "outputs": [],
    "source": [
-    "bs = batch_size * jax.device_count()    # Use a smaller size for testing\n",
+    "bs = batch_size * jax.device_count()    # You can use a smaller size while testing\n",
     "batches = math.ceil(length / bs)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "88598e4b",
+   "metadata": {},
+   "source": [
+    "Image processing"
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "id": "669b35df",
    "metadata": {},
    "outputs": [],
@@ -98,97 +139,125 @@
     "    center_crop,\n",
     "    T.ToTensor(),\n",
     "    lambda t: t.permute(1, 2, 0)   # Reorder, we need dimensions last\n",
-    "])\n",
+    "])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a185e90c",
+   "metadata": {},
+   "source": [
+    "Caption preparation.\n",
     "\n",
-    "# Is there a shortcut for this?\n",
-    "def extract_from_json(item):\n",
-    "    item['caption'] = item['json']['caption']\n",
-    "    item['url'] = item['json']['url']\n",
-    "    return item"
+    "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."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
+   "id": "423ee10e",
+   "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}'"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8d3a95db",
+   "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:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
    "id": "369d9719",
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Log exceptions to a hardcoded file\n",
+    "# UNUSED - Log exceptions to a file\n",
     "def ignore_and_log(exn):\n",
     "    with open('errors.txt', 'a') as f:\n",
-    "        f.write(f'{exn}\\n')\n",
-    "    return True\n",
-    "\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 = ignore_and_log\n",
     "exception_handler = wds.warn_and_continue"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "id": "5149b6d5",
    "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\n",
-    "                         handler=exception_handler,   # Ignore read errors instead of failing. See also: `warn_and_continue`\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(extract_from_json)\n",
-    "      .map_dict(jpg=preprocess_image, handler=exception_handler)\n",
-    "      .to_tuple('url', 'jpg', 'caption') # filter to keep only url (for reference), image, caption.\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"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": null,
    "id": "8cac98cb",
    "metadata": {
     "scrolled": true
    },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "CPU times: user 8min 26s, sys: 12.5 s, total: 8min 38s\n",
-      "Wall time: 14.4 s\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "%%time\n",
-    "urls, images, captions = next(iter(dataset))"
+    "keys, images, captions = next(iter(dataset))"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "id": "cd268fbf",
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "torch.Size([1024, 256, 256, 3])"
-      ]
-     },
-     "execution_count": 7,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
     "images.shape"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c24693c0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "T.ToPILImage()(images[0].permute(2, 0, 1))"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "44d50a51",
@@ -199,7 +268,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "id": "e2df5e13",
    "metadata": {},
    "outputs": [],
@@ -217,7 +286,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "id": "2fcf01d7",
    "metadata": {},
    "outputs": [],
@@ -235,20 +304,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": null,
    "id": "47a8b818",
    "metadata": {
     "scrolled": true
    },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Working with z of shape (1, 256, 16, 16) = 65536 dimensions.\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "model = VQModel.from_pretrained(\"flax-community/vqgan_f16_16384\")"
    ]
@@ -271,7 +332,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": null,
    "id": "6686b004",
    "metadata": {},
    "outputs": [],
@@ -282,7 +343,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": null,
    "id": "322a4619",
    "metadata": {},
    "outputs": [],
@@ -304,7 +365,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": null,
    "id": "ff6c10d4",
    "metadata": {},
    "outputs": [],
@@ -320,7 +381,7 @@
     "    # - Save each batch after processing.\n",
     "    # - Keep the file open until we are done with it.\n",
     "    file = None        \n",
-    "    for n, (urls, images, captions) in enumerate(tqdm(dataloader)):\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",
@@ -332,7 +393,7 @@
     "        encoded = encoded.reshape(-1, encoded.shape[-1])\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({\"url\": urls, \"caption\": captions, \"encoding\": encoded_as_string})\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)"
    ]
   },
@@ -346,7 +407,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": null,
    "id": "96222bb4",
    "metadata": {},
    "outputs": [],
@@ -359,15 +420,7 @@
    "execution_count": null,
    "id": "7704863d",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "  2%|█▌                                                           | 1085/43663 [31:58<20:43:42,  1.75s/it]"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "encode_captioned_dataset(dl, encoded_output, save_every=save_every)"
    ]