{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4924a372-61e0-4806-8820-a716b01787e8",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "VQBASE(\n",
       "  (encoder): Encoder(\n",
       "    (model): Sequential(\n",
       "      (0): Conv2d(3, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (1): ResnetBlock(\n",
       "        (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "        (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (2): ResnetBlock(\n",
       "        (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "        (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (3): Downsample(\n",
       "        (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2))\n",
       "      )\n",
       "      (4): ResnetBlock(\n",
       "        (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "        (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (5): ResnetBlock(\n",
       "        (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "        (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (6): Downsample(\n",
       "        (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2))\n",
       "      )\n",
       "      (7): ResnetBlock(\n",
       "        (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "        (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (nin_shortcut): Conv2d(128, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "      )\n",
       "      (8): ResnetBlock(\n",
       "        (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
       "        (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (9): Downsample(\n",
       "        (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2))\n",
       "      )\n",
       "      (10): ResnetBlock(\n",
       "        (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
       "        (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (nin_shortcut): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "      )\n",
       "      (11): ResnetBlock(\n",
       "        (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (12): Downsample(\n",
       "        (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2))\n",
       "      )\n",
       "      (13): ResnetBlock(\n",
       "        (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (14): AttnBlock(\n",
       "        (norm): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (q): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (k): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (v): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (proj_out): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "      )\n",
       "      (15): ResnetBlock(\n",
       "        (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (16): AttnBlock(\n",
       "        (norm): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (q): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (k): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (v): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (proj_out): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "      )\n",
       "      (17): ResnetBlock(\n",
       "        (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (18): AttnBlock(\n",
       "        (norm): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (q): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (k): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (v): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (proj_out): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "      )\n",
       "      (19): ResnetBlock(\n",
       "        (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (20): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "      (21): Swish()\n",
       "      (22): Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    )\n",
       "  )\n",
       "  (decoder): Decoder(\n",
       "    (model): Sequential(\n",
       "      (0): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (1): ResnetBlock(\n",
       "        (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (2): AttnBlock(\n",
       "        (norm): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (q): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (k): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (v): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (proj_out): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "      )\n",
       "      (3): ResnetBlock(\n",
       "        (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (4): ResnetBlock(\n",
       "        (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (5): AttnBlock(\n",
       "        (norm): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (q): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (k): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (v): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (proj_out): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "      )\n",
       "      (6): ResnetBlock(\n",
       "        (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (7): AttnBlock(\n",
       "        (norm): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (q): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (k): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (v): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (proj_out): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "      )\n",
       "      (8): ResnetBlock(\n",
       "        (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (9): AttnBlock(\n",
       "        (norm): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (q): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (k): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (v): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (proj_out): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "      )\n",
       "      (10): Upsample(\n",
       "        (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (11): ResnetBlock(\n",
       "        (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (12): ResnetBlock(\n",
       "        (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (13): ResnetBlock(\n",
       "        (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (14): Upsample(\n",
       "        (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (15): ResnetBlock(\n",
       "        (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (conv1): Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (nin_shortcut): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "      )\n",
       "      (16): ResnetBlock(\n",
       "        (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
       "        (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (17): ResnetBlock(\n",
       "        (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
       "        (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (18): Upsample(\n",
       "        (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (19): ResnetBlock(\n",
       "        (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
       "        (conv1): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (nin_shortcut): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "      )\n",
       "      (20): ResnetBlock(\n",
       "        (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "        (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (21): ResnetBlock(\n",
       "        (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "        (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (22): Upsample(\n",
       "        (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (23): ResnetBlock(\n",
       "        (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "        (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (24): ResnetBlock(\n",
       "        (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "        (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (25): ResnetBlock(\n",
       "        (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "        (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (26): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "      (27): Swish()\n",
       "      (28): Conv2d(128, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    )\n",
       "  )\n",
       "  (quantize): Codebook(\n",
       "    (embedding): Embedding(8192, 256)\n",
       "  )\n",
       "  (quant_conv): Sequential(\n",
       "    (0): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "    (1): SyncBatchNorm(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       "  (post_quant_conv): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       ")"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "import yaml\n",
    "import time\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "from vqvae import VQBASE\n",
    "import webdataset as wds\n",
    "import albumentations as A\n",
    "from albumentations.pytorch import ToTensorV2\n",
    "from torchvision.utils import make_grid, save_image\n",
    "device = \"cuda\"\n",
    "\n",
    "# Load configuration for the model\n",
    "with open(\"make_a_scene/img_config.yaml\", 'r') as file:\n",
    "    params = yaml.safe_load(file)[\"model\"]\n",
    "    del params[\"_target_\"]\n",
    "\n",
    "\n",
    "# Initialize and load the second model in bfloat16\n",
    "vq_vae = VQBASE(**params).to(device)\n",
    "vq_vae.load_state_dict(torch.load(\"make_a_scene/checkpoint_63.0.pt\", map_location=device)[\"model\"])\n",
    "vq_vae = vq_vae.to(dtype=torch.bfloat16)\n",
    "vq_vae.eval().requires_grad_(False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e14eaa4d-d7c3-4f59-b385-1ab45331c150",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers.tokenization_utils_fast import PreTrainedTokenizerFast"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "bfe41742-65d6-446b-b405-f21261ab3e68",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/alandao/.local/lib/python3.11/site-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
      "  warnings.warn(\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "llama_tokenizer = PreTrainedTokenizerFast.from_pretrained(\"NousResearch/Meta-Llama-3-8B-Instruct\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8f26a6bf-f095-41c1-bc3b-fd2381867708",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Embedding(8192, 256)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Number of to-add vocabs = number of codes in codebook of vq-vae\n",
    "vq_vae.quantize.embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e474eef7-affa-4991-a77a-5ef010147223",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'<|begin_of_text|>': 128000,\n",
       " '<|end_of_text|>': 128001,\n",
       " '<|reserved_special_token_0|>': 128002,\n",
       " '<|reserved_special_token_1|>': 128003,\n",
       " '<|reserved_special_token_2|>': 128004,\n",
       " '<|reserved_special_token_3|>': 128005,\n",
       " '<|start_header_id|>': 128006,\n",
       " '<|end_header_id|>': 128007,\n",
       " '<|reserved_special_token_4|>': 128008,\n",
       " '<|eot_id|>': 128009,\n",
       " '<|reserved_special_token_5|>': 128010,\n",
       " '<|reserved_special_token_6|>': 128011,\n",
       " '<|reserved_special_token_7|>': 128012,\n",
       " '<|reserved_special_token_8|>': 128013,\n",
       " '<|reserved_special_token_9|>': 128014,\n",
       " '<|reserved_special_token_10|>': 128015,\n",
       " '<|reserved_special_token_11|>': 128016,\n",
       " '<|reserved_special_token_12|>': 128017,\n",
       " '<|reserved_special_token_13|>': 128018,\n",
       " '<|reserved_special_token_14|>': 128019,\n",
       " '<|reserved_special_token_15|>': 128020,\n",
       " '<|reserved_special_token_16|>': 128021,\n",
       " '<|reserved_special_token_17|>': 128022,\n",
       " '<|reserved_special_token_18|>': 128023,\n",
       " '<|reserved_special_token_19|>': 128024,\n",
       " '<|reserved_special_token_20|>': 128025,\n",
       " '<|reserved_special_token_21|>': 128026,\n",
       " '<|reserved_special_token_22|>': 128027,\n",
       " '<|reserved_special_token_23|>': 128028,\n",
       " '<|reserved_special_token_24|>': 128029,\n",
       " '<|reserved_special_token_25|>': 128030,\n",
       " '<|reserved_special_token_26|>': 128031,\n",
       " '<|reserved_special_token_27|>': 128032,\n",
       " '<|reserved_special_token_28|>': 128033,\n",
       " '<|reserved_special_token_29|>': 128034,\n",
       " '<|reserved_special_token_30|>': 128035,\n",
       " '<|reserved_special_token_31|>': 128036,\n",
       " '<|reserved_special_token_32|>': 128037,\n",
       " '<|reserved_special_token_33|>': 128038,\n",
       " '<|reserved_special_token_34|>': 128039,\n",
       " '<|reserved_special_token_35|>': 128040,\n",
       " '<|reserved_special_token_36|>': 128041,\n",
       " '<|reserved_special_token_37|>': 128042,\n",
       " '<|reserved_special_token_38|>': 128043,\n",
       " '<|reserved_special_token_39|>': 128044,\n",
       " '<|reserved_special_token_40|>': 128045,\n",
       " '<|reserved_special_token_41|>': 128046,\n",
       " '<|reserved_special_token_42|>': 128047,\n",
       " '<|reserved_special_token_43|>': 128048,\n",
       " '<|reserved_special_token_44|>': 128049,\n",
       " '<|reserved_special_token_45|>': 128050,\n",
       " '<|reserved_special_token_46|>': 128051,\n",
       " '<|reserved_special_token_47|>': 128052,\n",
       " '<|reserved_special_token_48|>': 128053,\n",
       " '<|reserved_special_token_49|>': 128054,\n",
       " '<|reserved_special_token_50|>': 128055,\n",
       " '<|reserved_special_token_51|>': 128056,\n",
       " '<|reserved_special_token_52|>': 128057,\n",
       " '<|reserved_special_token_53|>': 128058,\n",
       " '<|reserved_special_token_54|>': 128059,\n",
       " '<|reserved_special_token_55|>': 128060,\n",
       " '<|reserved_special_token_56|>': 128061,\n",
       " '<|reserved_special_token_57|>': 128062,\n",
       " '<|reserved_special_token_58|>': 128063,\n",
       " '<|reserved_special_token_59|>': 128064,\n",
       " '<|reserved_special_token_60|>': 128065,\n",
       " '<|reserved_special_token_61|>': 128066,\n",
       " '<|reserved_special_token_62|>': 128067,\n",
       " '<|reserved_special_token_63|>': 128068,\n",
       " '<|reserved_special_token_64|>': 128069,\n",
       " '<|reserved_special_token_65|>': 128070,\n",
       " '<|reserved_special_token_66|>': 128071,\n",
       " '<|reserved_special_token_67|>': 128072,\n",
       " '<|reserved_special_token_68|>': 128073,\n",
       " '<|reserved_special_token_69|>': 128074,\n",
       " '<|reserved_special_token_70|>': 128075,\n",
       " '<|reserved_special_token_71|>': 128076,\n",
       " '<|reserved_special_token_72|>': 128077,\n",
       " '<|reserved_special_token_73|>': 128078,\n",
       " '<|reserved_special_token_74|>': 128079,\n",
       " '<|reserved_special_token_75|>': 128080,\n",
       " '<|reserved_special_token_76|>': 128081,\n",
       " '<|reserved_special_token_77|>': 128082,\n",
       " '<|reserved_special_token_78|>': 128083,\n",
       " '<|reserved_special_token_79|>': 128084,\n",
       " '<|reserved_special_token_80|>': 128085,\n",
       " '<|reserved_special_token_81|>': 128086,\n",
       " '<|reserved_special_token_82|>': 128087,\n",
       " '<|reserved_special_token_83|>': 128088,\n",
       " '<|reserved_special_token_84|>': 128089,\n",
       " '<|reserved_special_token_85|>': 128090,\n",
       " '<|reserved_special_token_86|>': 128091,\n",
       " '<|reserved_special_token_87|>': 128092,\n",
       " '<|reserved_special_token_88|>': 128093,\n",
       " '<|reserved_special_token_89|>': 128094,\n",
       " '<|reserved_special_token_90|>': 128095,\n",
       " '<|reserved_special_token_91|>': 128096,\n",
       " '<|reserved_special_token_92|>': 128097,\n",
       " '<|reserved_special_token_93|>': 128098,\n",
       " '<|reserved_special_token_94|>': 128099,\n",
       " '<|reserved_special_token_95|>': 128100,\n",
       " '<|reserved_special_token_96|>': 128101,\n",
       " '<|reserved_special_token_97|>': 128102,\n",
       " '<|reserved_special_token_98|>': 128103,\n",
       " '<|reserved_special_token_99|>': 128104,\n",
       " '<|reserved_special_token_100|>': 128105,\n",
       " '<|reserved_special_token_101|>': 128106,\n",
       " '<|reserved_special_token_102|>': 128107,\n",
       " '<|reserved_special_token_103|>': 128108,\n",
       " '<|reserved_special_token_104|>': 128109,\n",
       " '<|reserved_special_token_105|>': 128110,\n",
       " '<|reserved_special_token_106|>': 128111,\n",
       " '<|reserved_special_token_107|>': 128112,\n",
       " '<|reserved_special_token_108|>': 128113,\n",
       " '<|reserved_special_token_109|>': 128114,\n",
       " '<|reserved_special_token_110|>': 128115,\n",
       " '<|reserved_special_token_111|>': 128116,\n",
       " '<|reserved_special_token_112|>': 128117,\n",
       " '<|reserved_special_token_113|>': 128118,\n",
       " '<|reserved_special_token_114|>': 128119,\n",
       " '<|reserved_special_token_115|>': 128120,\n",
       " '<|reserved_special_token_116|>': 128121,\n",
       " '<|reserved_special_token_117|>': 128122,\n",
       " '<|reserved_special_token_118|>': 128123,\n",
       " '<|reserved_special_token_119|>': 128124,\n",
       " '<|reserved_special_token_120|>': 128125,\n",
       " '<|reserved_special_token_121|>': 128126,\n",
       " '<|reserved_special_token_122|>': 128127,\n",
       " '<|reserved_special_token_123|>': 128128,\n",
       " '<|reserved_special_token_124|>': 128129,\n",
       " '<|reserved_special_token_125|>': 128130,\n",
       " '<|reserved_special_token_126|>': 128131,\n",
       " '<|reserved_special_token_127|>': 128132,\n",
       " '<|reserved_special_token_128|>': 128133,\n",
       " '<|reserved_special_token_129|>': 128134,\n",
       " '<|reserved_special_token_130|>': 128135,\n",
       " '<|reserved_special_token_131|>': 128136,\n",
       " '<|reserved_special_token_132|>': 128137,\n",
       " '<|reserved_special_token_133|>': 128138,\n",
       " '<|reserved_special_token_134|>': 128139,\n",
       " '<|reserved_special_token_135|>': 128140,\n",
       " '<|reserved_special_token_136|>': 128141,\n",
       " '<|reserved_special_token_137|>': 128142,\n",
       " '<|reserved_special_token_138|>': 128143,\n",
       " '<|reserved_special_token_139|>': 128144,\n",
       " '<|reserved_special_token_140|>': 128145,\n",
       " '<|reserved_special_token_141|>': 128146,\n",
       " '<|reserved_special_token_142|>': 128147,\n",
       " '<|reserved_special_token_143|>': 128148,\n",
       " '<|reserved_special_token_144|>': 128149,\n",
       " '<|reserved_special_token_145|>': 128150,\n",
       " '<|reserved_special_token_146|>': 128151,\n",
       " '<|reserved_special_token_147|>': 128152,\n",
       " '<|reserved_special_token_148|>': 128153,\n",
       " '<|reserved_special_token_149|>': 128154,\n",
       " '<|reserved_special_token_150|>': 128155,\n",
       " '<|reserved_special_token_151|>': 128156,\n",
       " '<|reserved_special_token_152|>': 128157,\n",
       " '<|reserved_special_token_153|>': 128158,\n",
       " '<|reserved_special_token_154|>': 128159,\n",
       " '<|reserved_special_token_155|>': 128160,\n",
       " '<|reserved_special_token_156|>': 128161,\n",
       " '<|reserved_special_token_157|>': 128162,\n",
       " '<|reserved_special_token_158|>': 128163,\n",
       " '<|reserved_special_token_159|>': 128164,\n",
       " '<|reserved_special_token_160|>': 128165,\n",
       " '<|reserved_special_token_161|>': 128166,\n",
       " '<|reserved_special_token_162|>': 128167,\n",
       " '<|reserved_special_token_163|>': 128168,\n",
       " '<|reserved_special_token_164|>': 128169,\n",
       " '<|reserved_special_token_165|>': 128170,\n",
       " '<|reserved_special_token_166|>': 128171,\n",
       " '<|reserved_special_token_167|>': 128172,\n",
       " '<|reserved_special_token_168|>': 128173,\n",
       " '<|reserved_special_token_169|>': 128174,\n",
       " '<|reserved_special_token_170|>': 128175,\n",
       " '<|reserved_special_token_171|>': 128176,\n",
       " '<|reserved_special_token_172|>': 128177,\n",
       " '<|reserved_special_token_173|>': 128178,\n",
       " '<|reserved_special_token_174|>': 128179,\n",
       " '<|reserved_special_token_175|>': 128180,\n",
       " '<|reserved_special_token_176|>': 128181,\n",
       " '<|reserved_special_token_177|>': 128182,\n",
       " '<|reserved_special_token_178|>': 128183,\n",
       " '<|reserved_special_token_179|>': 128184,\n",
       " '<|reserved_special_token_180|>': 128185,\n",
       " '<|reserved_special_token_181|>': 128186,\n",
       " '<|reserved_special_token_182|>': 128187,\n",
       " '<|reserved_special_token_183|>': 128188,\n",
       " '<|reserved_special_token_184|>': 128189,\n",
       " '<|reserved_special_token_185|>': 128190,\n",
       " '<|reserved_special_token_186|>': 128191,\n",
       " '<|reserved_special_token_187|>': 128192,\n",
       " '<|reserved_special_token_188|>': 128193,\n",
       " '<|reserved_special_token_189|>': 128194,\n",
       " '<|reserved_special_token_190|>': 128195,\n",
       " '<|reserved_special_token_191|>': 128196,\n",
       " '<|reserved_special_token_192|>': 128197,\n",
       " '<|reserved_special_token_193|>': 128198,\n",
       " '<|reserved_special_token_194|>': 128199,\n",
       " '<|reserved_special_token_195|>': 128200,\n",
       " '<|reserved_special_token_196|>': 128201,\n",
       " '<|reserved_special_token_197|>': 128202,\n",
       " '<|reserved_special_token_198|>': 128203,\n",
       " '<|reserved_special_token_199|>': 128204,\n",
       " '<|reserved_special_token_200|>': 128205,\n",
       " '<|reserved_special_token_201|>': 128206,\n",
       " '<|reserved_special_token_202|>': 128207,\n",
       " '<|reserved_special_token_203|>': 128208,\n",
       " '<|reserved_special_token_204|>': 128209,\n",
       " '<|reserved_special_token_205|>': 128210,\n",
       " '<|reserved_special_token_206|>': 128211,\n",
       " '<|reserved_special_token_207|>': 128212,\n",
       " '<|reserved_special_token_208|>': 128213,\n",
       " '<|reserved_special_token_209|>': 128214,\n",
       " '<|reserved_special_token_210|>': 128215,\n",
       " '<|reserved_special_token_211|>': 128216,\n",
       " '<|reserved_special_token_212|>': 128217,\n",
       " '<|reserved_special_token_213|>': 128218,\n",
       " '<|reserved_special_token_214|>': 128219,\n",
       " '<|reserved_special_token_215|>': 128220,\n",
       " '<|reserved_special_token_216|>': 128221,\n",
       " '<|reserved_special_token_217|>': 128222,\n",
       " '<|reserved_special_token_218|>': 128223,\n",
       " '<|reserved_special_token_219|>': 128224,\n",
       " '<|reserved_special_token_220|>': 128225,\n",
       " '<|reserved_special_token_221|>': 128226,\n",
       " '<|reserved_special_token_222|>': 128227,\n",
       " '<|reserved_special_token_223|>': 128228,\n",
       " '<|reserved_special_token_224|>': 128229,\n",
       " '<|reserved_special_token_225|>': 128230,\n",
       " '<|reserved_special_token_226|>': 128231,\n",
       " '<|reserved_special_token_227|>': 128232,\n",
       " '<|reserved_special_token_228|>': 128233,\n",
       " '<|reserved_special_token_229|>': 128234,\n",
       " '<|reserved_special_token_230|>': 128235,\n",
       " '<|reserved_special_token_231|>': 128236,\n",
       " '<|reserved_special_token_232|>': 128237,\n",
       " '<|reserved_special_token_233|>': 128238,\n",
       " '<|reserved_special_token_234|>': 128239,\n",
       " '<|reserved_special_token_235|>': 128240,\n",
       " '<|reserved_special_token_236|>': 128241,\n",
       " '<|reserved_special_token_237|>': 128242,\n",
       " '<|reserved_special_token_238|>': 128243,\n",
       " '<|reserved_special_token_239|>': 128244,\n",
       " '<|reserved_special_token_240|>': 128245,\n",
       " '<|reserved_special_token_241|>': 128246,\n",
       " '<|reserved_special_token_242|>': 128247,\n",
       " '<|reserved_special_token_243|>': 128248,\n",
       " '<|reserved_special_token_244|>': 128249,\n",
       " '<|reserved_special_token_245|>': 128250,\n",
       " '<|reserved_special_token_246|>': 128251,\n",
       " '<|reserved_special_token_247|>': 128252,\n",
       " '<|reserved_special_token_248|>': 128253,\n",
       " '<|reserved_special_token_249|>': 128254,\n",
       " '<|reserved_special_token_250|>': 128255,\n",
       " '<|img_start|>': 128256,\n",
       " '<|img_end|>': 128257,\n",
       " '<dis_0000>': 128258,\n",
       " '<dis_0001>': 128259,\n",
       " '<dis_0002>': 128260,\n",
       " '<dis_0003>': 128261,\n",
       " '<dis_0004>': 128262,\n",
       " '<dis_0005>': 128263,\n",
       " '<dis_0006>': 128264,\n",
       " '<dis_0007>': 128265,\n",
       " '<dis_0008>': 128266,\n",
       " '<dis_0009>': 128267,\n",
       " '<dis_0010>': 128268,\n",
       " '<dis_0011>': 128269,\n",
       " '<dis_0012>': 128270,\n",
       " '<dis_0013>': 128271,\n",
       " '<dis_0014>': 128272,\n",
       " '<dis_0015>': 128273,\n",
       " '<dis_0016>': 128274,\n",
       " '<dis_0017>': 128275,\n",
       " '<dis_0018>': 128276,\n",
       " '<dis_0019>': 128277,\n",
       " '<dis_0020>': 128278,\n",
       " '<dis_0021>': 128279,\n",
       " '<dis_0022>': 128280,\n",
       " '<dis_0023>': 128281,\n",
       " '<dis_0024>': 128282,\n",
       " '<dis_0025>': 128283,\n",
       " '<dis_0026>': 128284,\n",
       " '<dis_0027>': 128285,\n",
       " '<dis_0028>': 128286,\n",
       " '<dis_0029>': 128287,\n",
       " '<dis_0030>': 128288,\n",
       " '<dis_0031>': 128289,\n",
       " '<dis_0032>': 128290,\n",
       " '<dis_0033>': 128291,\n",
       " '<dis_0034>': 128292,\n",
       " '<dis_0035>': 128293,\n",
       " '<dis_0036>': 128294,\n",
       " '<dis_0037>': 128295,\n",
       " '<dis_0038>': 128296,\n",
       " '<dis_0039>': 128297,\n",
       " '<dis_0040>': 128298,\n",
       " '<dis_0041>': 128299,\n",
       " '<dis_0042>': 128300,\n",
       " '<dis_0043>': 128301,\n",
       " '<dis_0044>': 128302,\n",
       " '<dis_0045>': 128303,\n",
       " '<dis_0046>': 128304,\n",
       " '<dis_0047>': 128305,\n",
       " '<dis_0048>': 128306,\n",
       " '<dis_0049>': 128307,\n",
       " '<dis_0050>': 128308,\n",
       " '<dis_0051>': 128309,\n",
       " '<dis_0052>': 128310,\n",
       " '<dis_0053>': 128311,\n",
       " '<dis_0054>': 128312,\n",
       " '<dis_0055>': 128313,\n",
       " '<dis_0056>': 128314,\n",
       " '<dis_0057>': 128315,\n",
       " '<dis_0058>': 128316,\n",
       " '<dis_0059>': 128317,\n",
       " '<dis_0060>': 128318,\n",
       " '<dis_0061>': 128319,\n",
       " '<dis_0062>': 128320,\n",
       " '<dis_0063>': 128321,\n",
       " '<dis_0064>': 128322,\n",
       " '<dis_0065>': 128323,\n",
       " '<dis_0066>': 128324,\n",
       " '<dis_0067>': 128325,\n",
       " '<dis_0068>': 128326,\n",
       " '<dis_0069>': 128327,\n",
       " '<dis_0070>': 128328,\n",
       " '<dis_0071>': 128329,\n",
       " '<dis_0072>': 128330,\n",
       " '<dis_0073>': 128331,\n",
       " '<dis_0074>': 128332,\n",
       " '<dis_0075>': 128333,\n",
       " '<dis_0076>': 128334,\n",
       " '<dis_0077>': 128335,\n",
       " '<dis_0078>': 128336,\n",
       " '<dis_0079>': 128337,\n",
       " '<dis_0080>': 128338,\n",
       " '<dis_0081>': 128339,\n",
       " '<dis_0082>': 128340,\n",
       " '<dis_0083>': 128341,\n",
       " '<dis_0084>': 128342,\n",
       " '<dis_0085>': 128343,\n",
       " '<dis_0086>': 128344,\n",
       " '<dis_0087>': 128345,\n",
       " '<dis_0088>': 128346,\n",
       " '<dis_0089>': 128347,\n",
       " '<dis_0090>': 128348,\n",
       " '<dis_0091>': 128349,\n",
       " '<dis_0092>': 128350,\n",
       " '<dis_0093>': 128351,\n",
       " '<dis_0094>': 128352,\n",
       " '<dis_0095>': 128353,\n",
       " '<dis_0096>': 128354,\n",
       " '<dis_0097>': 128355,\n",
       " '<dis_0098>': 128356,\n",
       " '<dis_0099>': 128357,\n",
       " '<dis_0100>': 128358,\n",
       " '<dis_0101>': 128359,\n",
       " '<dis_0102>': 128360,\n",
       " '<dis_0103>': 128361,\n",
       " '<dis_0104>': 128362,\n",
       " '<dis_0105>': 128363,\n",
       " '<dis_0106>': 128364,\n",
       " '<dis_0107>': 128365,\n",
       " '<dis_0108>': 128366,\n",
       " '<dis_0109>': 128367,\n",
       " '<dis_0110>': 128368,\n",
       " '<dis_0111>': 128369,\n",
       " '<dis_0112>': 128370,\n",
       " '<dis_0113>': 128371,\n",
       " '<dis_0114>': 128372,\n",
       " '<dis_0115>': 128373,\n",
       " '<dis_0116>': 128374,\n",
       " '<dis_0117>': 128375,\n",
       " '<dis_0118>': 128376,\n",
       " '<dis_0119>': 128377,\n",
       " '<dis_0120>': 128378,\n",
       " '<dis_0121>': 128379,\n",
       " '<dis_0122>': 128380,\n",
       " '<dis_0123>': 128381,\n",
       " '<dis_0124>': 128382,\n",
       " '<dis_0125>': 128383,\n",
       " '<dis_0126>': 128384,\n",
       " '<dis_0127>': 128385,\n",
       " '<dis_0128>': 128386,\n",
       " '<dis_0129>': 128387,\n",
       " '<dis_0130>': 128388,\n",
       " '<dis_0131>': 128389,\n",
       " '<dis_0132>': 128390,\n",
       " '<dis_0133>': 128391,\n",
       " '<dis_0134>': 128392,\n",
       " '<dis_0135>': 128393,\n",
       " '<dis_0136>': 128394,\n",
       " '<dis_0137>': 128395,\n",
       " '<dis_0138>': 128396,\n",
       " '<dis_0139>': 128397,\n",
       " '<dis_0140>': 128398,\n",
       " '<dis_0141>': 128399,\n",
       " '<dis_0142>': 128400,\n",
       " '<dis_0143>': 128401,\n",
       " '<dis_0144>': 128402,\n",
       " '<dis_0145>': 128403,\n",
       " '<dis_0146>': 128404,\n",
       " '<dis_0147>': 128405,\n",
       " '<dis_0148>': 128406,\n",
       " '<dis_0149>': 128407,\n",
       " '<dis_0150>': 128408,\n",
       " '<dis_0151>': 128409,\n",
       " '<dis_0152>': 128410,\n",
       " '<dis_0153>': 128411,\n",
       " '<dis_0154>': 128412,\n",
       " '<dis_0155>': 128413,\n",
       " '<dis_0156>': 128414,\n",
       " '<dis_0157>': 128415,\n",
       " '<dis_0158>': 128416,\n",
       " '<dis_0159>': 128417,\n",
       " '<dis_0160>': 128418,\n",
       " '<dis_0161>': 128419,\n",
       " '<dis_0162>': 128420,\n",
       " '<dis_0163>': 128421,\n",
       " '<dis_0164>': 128422,\n",
       " '<dis_0165>': 128423,\n",
       " '<dis_0166>': 128424,\n",
       " '<dis_0167>': 128425,\n",
       " '<dis_0168>': 128426,\n",
       " '<dis_0169>': 128427,\n",
       " '<dis_0170>': 128428,\n",
       " '<dis_0171>': 128429,\n",
       " '<dis_0172>': 128430,\n",
       " '<dis_0173>': 128431,\n",
       " '<dis_0174>': 128432,\n",
       " '<dis_0175>': 128433,\n",
       " '<dis_0176>': 128434,\n",
       " '<dis_0177>': 128435,\n",
       " '<dis_0178>': 128436,\n",
       " '<dis_0179>': 128437,\n",
       " '<dis_0180>': 128438,\n",
       " '<dis_0181>': 128439,\n",
       " '<dis_0182>': 128440,\n",
       " '<dis_0183>': 128441,\n",
       " '<dis_0184>': 128442,\n",
       " '<dis_0185>': 128443,\n",
       " '<dis_0186>': 128444,\n",
       " '<dis_0187>': 128445,\n",
       " '<dis_0188>': 128446,\n",
       " '<dis_0189>': 128447,\n",
       " '<dis_0190>': 128448,\n",
       " '<dis_0191>': 128449,\n",
       " '<dis_0192>': 128450,\n",
       " '<dis_0193>': 128451,\n",
       " '<dis_0194>': 128452,\n",
       " '<dis_0195>': 128453,\n",
       " '<dis_0196>': 128454,\n",
       " '<dis_0197>': 128455,\n",
       " '<dis_0198>': 128456,\n",
       " '<dis_0199>': 128457,\n",
       " '<dis_0200>': 128458,\n",
       " '<dis_0201>': 128459,\n",
       " '<dis_0202>': 128460,\n",
       " '<dis_0203>': 128461,\n",
       " '<dis_0204>': 128462,\n",
       " '<dis_0205>': 128463,\n",
       " '<dis_0206>': 128464,\n",
       " '<dis_0207>': 128465,\n",
       " '<dis_0208>': 128466,\n",
       " '<dis_0209>': 128467,\n",
       " '<dis_0210>': 128468,\n",
       " '<dis_0211>': 128469,\n",
       " '<dis_0212>': 128470,\n",
       " '<dis_0213>': 128471,\n",
       " '<dis_0214>': 128472,\n",
       " '<dis_0215>': 128473,\n",
       " '<dis_0216>': 128474,\n",
       " '<dis_0217>': 128475,\n",
       " '<dis_0218>': 128476,\n",
       " '<dis_0219>': 128477,\n",
       " '<dis_0220>': 128478,\n",
       " '<dis_0221>': 128479,\n",
       " '<dis_0222>': 128480,\n",
       " '<dis_0223>': 128481,\n",
       " '<dis_0224>': 128482,\n",
       " '<dis_0225>': 128483,\n",
       " '<dis_0226>': 128484,\n",
       " '<dis_0227>': 128485,\n",
       " '<dis_0228>': 128486,\n",
       " '<dis_0229>': 128487,\n",
       " '<dis_0230>': 128488,\n",
       " '<dis_0231>': 128489,\n",
       " '<dis_0232>': 128490,\n",
       " '<dis_0233>': 128491,\n",
       " '<dis_0234>': 128492,\n",
       " '<dis_0235>': 128493,\n",
       " '<dis_0236>': 128494,\n",
       " '<dis_0237>': 128495,\n",
       " '<dis_0238>': 128496,\n",
       " '<dis_0239>': 128497,\n",
       " '<dis_0240>': 128498,\n",
       " '<dis_0241>': 128499,\n",
       " '<dis_0242>': 128500,\n",
       " '<dis_0243>': 128501,\n",
       " '<dis_0244>': 128502,\n",
       " '<dis_0245>': 128503,\n",
       " '<dis_0246>': 128504,\n",
       " '<dis_0247>': 128505,\n",
       " '<dis_0248>': 128506,\n",
       " '<dis_0249>': 128507,\n",
       " '<dis_0250>': 128508,\n",
       " '<dis_0251>': 128509,\n",
       " '<dis_0252>': 128510,\n",
       " '<dis_0253>': 128511,\n",
       " '<dis_0254>': 128512,\n",
       " '<dis_0255>': 128513,\n",
       " '<dis_0256>': 128514,\n",
       " '<dis_0257>': 128515,\n",
       " '<dis_0258>': 128516,\n",
       " '<dis_0259>': 128517,\n",
       " '<dis_0260>': 128518,\n",
       " '<dis_0261>': 128519,\n",
       " '<dis_0262>': 128520,\n",
       " '<dis_0263>': 128521,\n",
       " '<dis_0264>': 128522,\n",
       " '<dis_0265>': 128523,\n",
       " '<dis_0266>': 128524,\n",
       " '<dis_0267>': 128525,\n",
       " '<dis_0268>': 128526,\n",
       " '<dis_0269>': 128527,\n",
       " '<dis_0270>': 128528,\n",
       " '<dis_0271>': 128529,\n",
       " '<dis_0272>': 128530,\n",
       " '<dis_0273>': 128531,\n",
       " '<dis_0274>': 128532,\n",
       " '<dis_0275>': 128533,\n",
       " '<dis_0276>': 128534,\n",
       " '<dis_0277>': 128535,\n",
       " '<dis_0278>': 128536,\n",
       " '<dis_0279>': 128537,\n",
       " '<dis_0280>': 128538,\n",
       " '<dis_0281>': 128539,\n",
       " '<dis_0282>': 128540,\n",
       " '<dis_0283>': 128541,\n",
       " '<dis_0284>': 128542,\n",
       " '<dis_0285>': 128543,\n",
       " '<dis_0286>': 128544,\n",
       " '<dis_0287>': 128545,\n",
       " '<dis_0288>': 128546,\n",
       " '<dis_0289>': 128547,\n",
       " '<dis_0290>': 128548,\n",
       " '<dis_0291>': 128549,\n",
       " '<dis_0292>': 128550,\n",
       " '<dis_0293>': 128551,\n",
       " '<dis_0294>': 128552,\n",
       " '<dis_0295>': 128553,\n",
       " '<dis_0296>': 128554,\n",
       " '<dis_0297>': 128555,\n",
       " '<dis_0298>': 128556,\n",
       " '<dis_0299>': 128557,\n",
       " '<dis_0300>': 128558,\n",
       " '<dis_0301>': 128559,\n",
       " '<dis_0302>': 128560,\n",
       " '<dis_0303>': 128561,\n",
       " '<dis_0304>': 128562,\n",
       " '<dis_0305>': 128563,\n",
       " '<dis_0306>': 128564,\n",
       " '<dis_0307>': 128565,\n",
       " '<dis_0308>': 128566,\n",
       " '<dis_0309>': 128567,\n",
       " '<dis_0310>': 128568,\n",
       " '<dis_0311>': 128569,\n",
       " '<dis_0312>': 128570,\n",
       " '<dis_0313>': 128571,\n",
       " '<dis_0314>': 128572,\n",
       " '<dis_0315>': 128573,\n",
       " '<dis_0316>': 128574,\n",
       " '<dis_0317>': 128575,\n",
       " '<dis_0318>': 128576,\n",
       " '<dis_0319>': 128577,\n",
       " '<dis_0320>': 128578,\n",
       " '<dis_0321>': 128579,\n",
       " '<dis_0322>': 128580,\n",
       " '<dis_0323>': 128581,\n",
       " '<dis_0324>': 128582,\n",
       " '<dis_0325>': 128583,\n",
       " '<dis_0326>': 128584,\n",
       " '<dis_0327>': 128585,\n",
       " '<dis_0328>': 128586,\n",
       " '<dis_0329>': 128587,\n",
       " '<dis_0330>': 128588,\n",
       " '<dis_0331>': 128589,\n",
       " '<dis_0332>': 128590,\n",
       " '<dis_0333>': 128591,\n",
       " '<dis_0334>': 128592,\n",
       " '<dis_0335>': 128593,\n",
       " '<dis_0336>': 128594,\n",
       " '<dis_0337>': 128595,\n",
       " '<dis_0338>': 128596,\n",
       " '<dis_0339>': 128597,\n",
       " '<dis_0340>': 128598,\n",
       " '<dis_0341>': 128599,\n",
       " '<dis_0342>': 128600,\n",
       " '<dis_0343>': 128601,\n",
       " '<dis_0344>': 128602,\n",
       " '<dis_0345>': 128603,\n",
       " '<dis_0346>': 128604,\n",
       " '<dis_0347>': 128605,\n",
       " '<dis_0348>': 128606,\n",
       " '<dis_0349>': 128607,\n",
       " '<dis_0350>': 128608,\n",
       " '<dis_0351>': 128609,\n",
       " '<dis_0352>': 128610,\n",
       " '<dis_0353>': 128611,\n",
       " '<dis_0354>': 128612,\n",
       " '<dis_0355>': 128613,\n",
       " '<dis_0356>': 128614,\n",
       " '<dis_0357>': 128615,\n",
       " '<dis_0358>': 128616,\n",
       " '<dis_0359>': 128617,\n",
       " '<dis_0360>': 128618,\n",
       " '<dis_0361>': 128619,\n",
       " '<dis_0362>': 128620,\n",
       " '<dis_0363>': 128621,\n",
       " '<dis_0364>': 128622,\n",
       " '<dis_0365>': 128623,\n",
       " '<dis_0366>': 128624,\n",
       " '<dis_0367>': 128625,\n",
       " '<dis_0368>': 128626,\n",
       " '<dis_0369>': 128627,\n",
       " '<dis_0370>': 128628,\n",
       " '<dis_0371>': 128629,\n",
       " '<dis_0372>': 128630,\n",
       " '<dis_0373>': 128631,\n",
       " '<dis_0374>': 128632,\n",
       " '<dis_0375>': 128633,\n",
       " '<dis_0376>': 128634,\n",
       " '<dis_0377>': 128635,\n",
       " '<dis_0378>': 128636,\n",
       " '<dis_0379>': 128637,\n",
       " '<dis_0380>': 128638,\n",
       " '<dis_0381>': 128639,\n",
       " '<dis_0382>': 128640,\n",
       " '<dis_0383>': 128641,\n",
       " '<dis_0384>': 128642,\n",
       " '<dis_0385>': 128643,\n",
       " '<dis_0386>': 128644,\n",
       " '<dis_0387>': 128645,\n",
       " '<dis_0388>': 128646,\n",
       " '<dis_0389>': 128647,\n",
       " '<dis_0390>': 128648,\n",
       " '<dis_0391>': 128649,\n",
       " '<dis_0392>': 128650,\n",
       " '<dis_0393>': 128651,\n",
       " '<dis_0394>': 128652,\n",
       " '<dis_0395>': 128653,\n",
       " '<dis_0396>': 128654,\n",
       " '<dis_0397>': 128655,\n",
       " '<dis_0398>': 128656,\n",
       " '<dis_0399>': 128657,\n",
       " '<dis_0400>': 128658,\n",
       " '<dis_0401>': 128659,\n",
       " '<dis_0402>': 128660,\n",
       " '<dis_0403>': 128661,\n",
       " '<dis_0404>': 128662,\n",
       " '<dis_0405>': 128663,\n",
       " '<dis_0406>': 128664,\n",
       " '<dis_0407>': 128665,\n",
       " '<dis_0408>': 128666,\n",
       " '<dis_0409>': 128667,\n",
       " '<dis_0410>': 128668,\n",
       " '<dis_0411>': 128669,\n",
       " '<dis_0412>': 128670,\n",
       " '<dis_0413>': 128671,\n",
       " '<dis_0414>': 128672,\n",
       " '<dis_0415>': 128673,\n",
       " '<dis_0416>': 128674,\n",
       " '<dis_0417>': 128675,\n",
       " '<dis_0418>': 128676,\n",
       " '<dis_0419>': 128677,\n",
       " '<dis_0420>': 128678,\n",
       " '<dis_0421>': 128679,\n",
       " '<dis_0422>': 128680,\n",
       " '<dis_0423>': 128681,\n",
       " '<dis_0424>': 128682,\n",
       " '<dis_0425>': 128683,\n",
       " '<dis_0426>': 128684,\n",
       " '<dis_0427>': 128685,\n",
       " '<dis_0428>': 128686,\n",
       " '<dis_0429>': 128687,\n",
       " '<dis_0430>': 128688,\n",
       " '<dis_0431>': 128689,\n",
       " '<dis_0432>': 128690,\n",
       " '<dis_0433>': 128691,\n",
       " '<dis_0434>': 128692,\n",
       " '<dis_0435>': 128693,\n",
       " '<dis_0436>': 128694,\n",
       " '<dis_0437>': 128695,\n",
       " '<dis_0438>': 128696,\n",
       " '<dis_0439>': 128697,\n",
       " '<dis_0440>': 128698,\n",
       " '<dis_0441>': 128699,\n",
       " '<dis_0442>': 128700,\n",
       " '<dis_0443>': 128701,\n",
       " '<dis_0444>': 128702,\n",
       " '<dis_0445>': 128703,\n",
       " '<dis_0446>': 128704,\n",
       " '<dis_0447>': 128705,\n",
       " '<dis_0448>': 128706,\n",
       " '<dis_0449>': 128707,\n",
       " '<dis_0450>': 128708,\n",
       " '<dis_0451>': 128709,\n",
       " '<dis_0452>': 128710,\n",
       " '<dis_0453>': 128711,\n",
       " '<dis_0454>': 128712,\n",
       " '<dis_0455>': 128713,\n",
       " '<dis_0456>': 128714,\n",
       " '<dis_0457>': 128715,\n",
       " '<dis_0458>': 128716,\n",
       " '<dis_0459>': 128717,\n",
       " '<dis_0460>': 128718,\n",
       " '<dis_0461>': 128719,\n",
       " '<dis_0462>': 128720,\n",
       " '<dis_0463>': 128721,\n",
       " '<dis_0464>': 128722,\n",
       " '<dis_0465>': 128723,\n",
       " '<dis_0466>': 128724,\n",
       " '<dis_0467>': 128725,\n",
       " '<dis_0468>': 128726,\n",
       " '<dis_0469>': 128727,\n",
       " '<dis_0470>': 128728,\n",
       " '<dis_0471>': 128729,\n",
       " '<dis_0472>': 128730,\n",
       " '<dis_0473>': 128731,\n",
       " '<dis_0474>': 128732,\n",
       " '<dis_0475>': 128733,\n",
       " '<dis_0476>': 128734,\n",
       " '<dis_0477>': 128735,\n",
       " '<dis_0478>': 128736,\n",
       " '<dis_0479>': 128737,\n",
       " '<dis_0480>': 128738,\n",
       " '<dis_0481>': 128739,\n",
       " '<dis_0482>': 128740,\n",
       " '<dis_0483>': 128741,\n",
       " '<dis_0484>': 128742,\n",
       " '<dis_0485>': 128743,\n",
       " '<dis_0486>': 128744,\n",
       " '<dis_0487>': 128745,\n",
       " '<dis_0488>': 128746,\n",
       " '<dis_0489>': 128747,\n",
       " '<dis_0490>': 128748,\n",
       " '<dis_0491>': 128749,\n",
       " '<dis_0492>': 128750,\n",
       " '<dis_0493>': 128751,\n",
       " '<dis_0494>': 128752,\n",
       " '<dis_0495>': 128753,\n",
       " '<dis_0496>': 128754,\n",
       " '<dis_0497>': 128755,\n",
       " '<dis_0498>': 128756,\n",
       " '<dis_0499>': 128757,\n",
       " '<dis_0500>': 128758,\n",
       " '<dis_0501>': 128759,\n",
       " '<dis_0502>': 128760,\n",
       " '<dis_0503>': 128761,\n",
       " '<dis_0504>': 128762,\n",
       " '<dis_0505>': 128763,\n",
       " '<dis_0506>': 128764,\n",
       " '<dis_0507>': 128765,\n",
       " '<dis_0508>': 128766,\n",
       " '<dis_0509>': 128767,\n",
       " '<dis_0510>': 128768,\n",
       " '<dis_0511>': 128769,\n",
       " '<dis_0512>': 128770,\n",
       " '<dis_0513>': 128771,\n",
       " '<dis_0514>': 128772,\n",
       " '<dis_0515>': 128773,\n",
       " '<dis_0516>': 128774,\n",
       " '<dis_0517>': 128775,\n",
       " '<dis_0518>': 128776,\n",
       " '<dis_0519>': 128777,\n",
       " '<dis_0520>': 128778,\n",
       " '<dis_0521>': 128779,\n",
       " '<dis_0522>': 128780,\n",
       " '<dis_0523>': 128781,\n",
       " '<dis_0524>': 128782,\n",
       " '<dis_0525>': 128783,\n",
       " '<dis_0526>': 128784,\n",
       " '<dis_0527>': 128785,\n",
       " '<dis_0528>': 128786,\n",
       " '<dis_0529>': 128787,\n",
       " '<dis_0530>': 128788,\n",
       " '<dis_0531>': 128789,\n",
       " '<dis_0532>': 128790,\n",
       " '<dis_0533>': 128791,\n",
       " '<dis_0534>': 128792,\n",
       " '<dis_0535>': 128793,\n",
       " '<dis_0536>': 128794,\n",
       " '<dis_0537>': 128795,\n",
       " '<dis_0538>': 128796,\n",
       " '<dis_0539>': 128797,\n",
       " '<dis_0540>': 128798,\n",
       " '<dis_0541>': 128799,\n",
       " '<dis_0542>': 128800,\n",
       " '<dis_0543>': 128801,\n",
       " '<dis_0544>': 128802,\n",
       " '<dis_0545>': 128803,\n",
       " '<dis_0546>': 128804,\n",
       " '<dis_0547>': 128805,\n",
       " '<dis_0548>': 128806,\n",
       " '<dis_0549>': 128807,\n",
       " '<dis_0550>': 128808,\n",
       " '<dis_0551>': 128809,\n",
       " '<dis_0552>': 128810,\n",
       " '<dis_0553>': 128811,\n",
       " '<dis_0554>': 128812,\n",
       " '<dis_0555>': 128813,\n",
       " '<dis_0556>': 128814,\n",
       " '<dis_0557>': 128815,\n",
       " '<dis_0558>': 128816,\n",
       " '<dis_0559>': 128817,\n",
       " '<dis_0560>': 128818,\n",
       " '<dis_0561>': 128819,\n",
       " '<dis_0562>': 128820,\n",
       " '<dis_0563>': 128821,\n",
       " '<dis_0564>': 128822,\n",
       " '<dis_0565>': 128823,\n",
       " '<dis_0566>': 128824,\n",
       " '<dis_0567>': 128825,\n",
       " '<dis_0568>': 128826,\n",
       " '<dis_0569>': 128827,\n",
       " '<dis_0570>': 128828,\n",
       " '<dis_0571>': 128829,\n",
       " '<dis_0572>': 128830,\n",
       " '<dis_0573>': 128831,\n",
       " '<dis_0574>': 128832,\n",
       " '<dis_0575>': 128833,\n",
       " '<dis_0576>': 128834,\n",
       " '<dis_0577>': 128835,\n",
       " '<dis_0578>': 128836,\n",
       " '<dis_0579>': 128837,\n",
       " '<dis_0580>': 128838,\n",
       " '<dis_0581>': 128839,\n",
       " '<dis_0582>': 128840,\n",
       " '<dis_0583>': 128841,\n",
       " '<dis_0584>': 128842,\n",
       " '<dis_0585>': 128843,\n",
       " '<dis_0586>': 128844,\n",
       " '<dis_0587>': 128845,\n",
       " '<dis_0588>': 128846,\n",
       " '<dis_0589>': 128847,\n",
       " '<dis_0590>': 128848,\n",
       " '<dis_0591>': 128849,\n",
       " '<dis_0592>': 128850,\n",
       " '<dis_0593>': 128851,\n",
       " '<dis_0594>': 128852,\n",
       " '<dis_0595>': 128853,\n",
       " '<dis_0596>': 128854,\n",
       " '<dis_0597>': 128855,\n",
       " '<dis_0598>': 128856,\n",
       " '<dis_0599>': 128857,\n",
       " '<dis_0600>': 128858,\n",
       " '<dis_0601>': 128859,\n",
       " '<dis_0602>': 128860,\n",
       " '<dis_0603>': 128861,\n",
       " '<dis_0604>': 128862,\n",
       " '<dis_0605>': 128863,\n",
       " '<dis_0606>': 128864,\n",
       " '<dis_0607>': 128865,\n",
       " '<dis_0608>': 128866,\n",
       " '<dis_0609>': 128867,\n",
       " '<dis_0610>': 128868,\n",
       " '<dis_0611>': 128869,\n",
       " '<dis_0612>': 128870,\n",
       " '<dis_0613>': 128871,\n",
       " '<dis_0614>': 128872,\n",
       " '<dis_0615>': 128873,\n",
       " '<dis_0616>': 128874,\n",
       " '<dis_0617>': 128875,\n",
       " '<dis_0618>': 128876,\n",
       " '<dis_0619>': 128877,\n",
       " '<dis_0620>': 128878,\n",
       " '<dis_0621>': 128879,\n",
       " '<dis_0622>': 128880,\n",
       " '<dis_0623>': 128881,\n",
       " '<dis_0624>': 128882,\n",
       " '<dis_0625>': 128883,\n",
       " '<dis_0626>': 128884,\n",
       " '<dis_0627>': 128885,\n",
       " '<dis_0628>': 128886,\n",
       " '<dis_0629>': 128887,\n",
       " '<dis_0630>': 128888,\n",
       " '<dis_0631>': 128889,\n",
       " '<dis_0632>': 128890,\n",
       " '<dis_0633>': 128891,\n",
       " '<dis_0634>': 128892,\n",
       " '<dis_0635>': 128893,\n",
       " '<dis_0636>': 128894,\n",
       " '<dis_0637>': 128895,\n",
       " '<dis_0638>': 128896,\n",
       " '<dis_0639>': 128897,\n",
       " '<dis_0640>': 128898,\n",
       " '<dis_0641>': 128899,\n",
       " '<dis_0642>': 128900,\n",
       " '<dis_0643>': 128901,\n",
       " '<dis_0644>': 128902,\n",
       " '<dis_0645>': 128903,\n",
       " '<dis_0646>': 128904,\n",
       " '<dis_0647>': 128905,\n",
       " '<dis_0648>': 128906,\n",
       " '<dis_0649>': 128907,\n",
       " '<dis_0650>': 128908,\n",
       " '<dis_0651>': 128909,\n",
       " '<dis_0652>': 128910,\n",
       " '<dis_0653>': 128911,\n",
       " '<dis_0654>': 128912,\n",
       " '<dis_0655>': 128913,\n",
       " '<dis_0656>': 128914,\n",
       " '<dis_0657>': 128915,\n",
       " '<dis_0658>': 128916,\n",
       " '<dis_0659>': 128917,\n",
       " '<dis_0660>': 128918,\n",
       " '<dis_0661>': 128919,\n",
       " '<dis_0662>': 128920,\n",
       " '<dis_0663>': 128921,\n",
       " '<dis_0664>': 128922,\n",
       " '<dis_0665>': 128923,\n",
       " '<dis_0666>': 128924,\n",
       " '<dis_0667>': 128925,\n",
       " '<dis_0668>': 128926,\n",
       " '<dis_0669>': 128927,\n",
       " '<dis_0670>': 128928,\n",
       " '<dis_0671>': 128929,\n",
       " '<dis_0672>': 128930,\n",
       " '<dis_0673>': 128931,\n",
       " '<dis_0674>': 128932,\n",
       " '<dis_0675>': 128933,\n",
       " '<dis_0676>': 128934,\n",
       " '<dis_0677>': 128935,\n",
       " '<dis_0678>': 128936,\n",
       " '<dis_0679>': 128937,\n",
       " '<dis_0680>': 128938,\n",
       " '<dis_0681>': 128939,\n",
       " '<dis_0682>': 128940,\n",
       " '<dis_0683>': 128941,\n",
       " '<dis_0684>': 128942,\n",
       " '<dis_0685>': 128943,\n",
       " '<dis_0686>': 128944,\n",
       " '<dis_0687>': 128945,\n",
       " '<dis_0688>': 128946,\n",
       " '<dis_0689>': 128947,\n",
       " '<dis_0690>': 128948,\n",
       " '<dis_0691>': 128949,\n",
       " '<dis_0692>': 128950,\n",
       " '<dis_0693>': 128951,\n",
       " '<dis_0694>': 128952,\n",
       " '<dis_0695>': 128953,\n",
       " '<dis_0696>': 128954,\n",
       " '<dis_0697>': 128955,\n",
       " '<dis_0698>': 128956,\n",
       " '<dis_0699>': 128957,\n",
       " '<dis_0700>': 128958,\n",
       " '<dis_0701>': 128959,\n",
       " '<dis_0702>': 128960,\n",
       " '<dis_0703>': 128961,\n",
       " '<dis_0704>': 128962,\n",
       " '<dis_0705>': 128963,\n",
       " '<dis_0706>': 128964,\n",
       " '<dis_0707>': 128965,\n",
       " '<dis_0708>': 128966,\n",
       " '<dis_0709>': 128967,\n",
       " '<dis_0710>': 128968,\n",
       " '<dis_0711>': 128969,\n",
       " '<dis_0712>': 128970,\n",
       " '<dis_0713>': 128971,\n",
       " '<dis_0714>': 128972,\n",
       " '<dis_0715>': 128973,\n",
       " '<dis_0716>': 128974,\n",
       " '<dis_0717>': 128975,\n",
       " '<dis_0718>': 128976,\n",
       " '<dis_0719>': 128977,\n",
       " '<dis_0720>': 128978,\n",
       " '<dis_0721>': 128979,\n",
       " '<dis_0722>': 128980,\n",
       " '<dis_0723>': 128981,\n",
       " '<dis_0724>': 128982,\n",
       " '<dis_0725>': 128983,\n",
       " '<dis_0726>': 128984,\n",
       " '<dis_0727>': 128985,\n",
       " '<dis_0728>': 128986,\n",
       " '<dis_0729>': 128987,\n",
       " '<dis_0730>': 128988,\n",
       " '<dis_0731>': 128989,\n",
       " '<dis_0732>': 128990,\n",
       " '<dis_0733>': 128991,\n",
       " '<dis_0734>': 128992,\n",
       " '<dis_0735>': 128993,\n",
       " '<dis_0736>': 128994,\n",
       " '<dis_0737>': 128995,\n",
       " '<dis_0738>': 128996,\n",
       " '<dis_0739>': 128997,\n",
       " '<dis_0740>': 128998,\n",
       " '<dis_0741>': 128999,\n",
       " ...}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "llama_tokenizer.add_tokens(\"<|img_start|>\",special_tokens=True)\n",
    "llama_tokenizer.add_tokens(\"<|img_end|>\",special_tokens=True)\n",
    "for img_token in range(0, 8192):\n",
    "    padded_token = f\"<dis_{img_token:04}>\"  # This pads the img_token with zeros to ensure it is 4 digits long.\n",
    "    llama_tokenizer.add_tokens(padded_token)\n",
    "llama_tokenizer.get_added_vocab()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3ca6861c-4182-4d23-984f-92e0b5ba22f2",
   "metadata": {},
   "outputs": [
    {
     "ename": "AssertionError",
     "evalue": "Key <|img_start|> is not a special token",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAssertionError\u001b[0m                            Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[8], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mllama_tokenizer\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43madd_special_tokens\u001b[49m\u001b[43m(\u001b[49m\u001b[43m{\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43m<|img_start|>\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43masdf\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m}\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.local/lib/python3.11/site-packages/transformers/tokenization_utils_base.py:947\u001b[0m, in \u001b[0;36mSpecialTokensMixin.add_special_tokens\u001b[0;34m(self, special_tokens_dict, replace_additional_special_tokens)\u001b[0m\n\u001b[1;32m    945\u001b[0m added_tokens \u001b[38;5;241m=\u001b[39m []\n\u001b[1;32m    946\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m key, value \u001b[38;5;129;01min\u001b[39;00m special_tokens_dict\u001b[38;5;241m.\u001b[39mitems():\n\u001b[0;32m--> 947\u001b[0m     \u001b[38;5;28;01massert\u001b[39;00m key \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mSPECIAL_TOKENS_ATTRIBUTES, \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mKey \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mkey\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m is not a special token\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    949\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mverbose:\n\u001b[1;32m    950\u001b[0m         logger\u001b[38;5;241m.\u001b[39minfo(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mAssigning \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mvalue\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m to the \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mkey\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m key of the tokenizer\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "\u001b[0;31mAssertionError\u001b[0m: Key <|img_start|> is not a special token"
     ]
    }
   ],
   "source": [
    "llama_tokenizer.regist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aaf0a3d1-dd5d-402b-9802-b692921c9079",
   "metadata": {},
   "outputs": [],
   "source": [
    "llama_tokenizer.add_special_tokens({\"image start\":\"<|img_start|>\"})\n",
    "llama_tokenizer.add_special_tokens(\"<|img_end|>\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "a3e9b926-2cd4-4870-b219-7f92cd9295c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We need to pad the beginning position of vqvae since discrete token is range(0,8192)\n",
    "pad_idx_vqvae = llama_tokenizer.vocab['<dis_0000>']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17f19396-b4f2-4764-a0ed-fa0bcb98ca43",
   "metadata": {},
   "source": [
    "## Process some sample dataset from LAION\n",
    "\n",
    "Our target here is to generate a simple text only dataset that is suitable to be used by traditional tokenizers, but also suitable for representing the images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "7998066b-6fe7-4ca8-8ede-3870f37f3725",
   "metadata": {},
   "outputs": [],
   "source": [
    "import albumentations as A\n",
    "from albumentations.pytorch import ToTensorV2\n",
    "import torch\n",
    "import webdataset as wds\n",
    "\n",
    "def process_data(data):\n",
    "    pretransforms = A.Compose([\n",
    "        A.SmallestMaxSize(512),\n",
    "        A.CenterCrop(512, 512, always_apply=True),\n",
    "        ToTensorV2()\n",
    "    ])\n",
    "    data[\"jpg\"] = pretransforms(image=data[\"jpg\"])[\"image\"]\n",
    "    # Convert image to bfloat16\n",
    "    data[\"jpg\"] = data[\"jpg\"].to(torch.bfloat16)\n",
    "    return data\n",
    "\n",
    "url = \"file:make_a_scene/00000.tar\"\n",
    "dataset = wds.WebDataset(url).decode(\"rgb\").map(process_data).to_tuple(\"jpg\", \"txt\")\n",
    "\n",
    "def collate(batch):\n",
    "    images = torch.stack([i[0] for i in batch], dim=0)\n",
    "    captions = [i[1] for i in batch]\n",
    "    return [images, captions]\n",
    "\n",
    "loader = torch.utils.data.DataLoader(dataset, batch_size=1, collate_fn=collate)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "8919fa2e-4e07-48d9-876f-c033a0fd1ab8",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_text_dataset = [] "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "c97bf893-aeda-43e4-b703-bdb425dcda3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "counter = 0\n",
    "for image,description in dataset:\n",
    "    counter += 1\n",
    "    if counter > 10:\n",
    "        break\n",
    "    tokenized_image_text = []\n",
    "\n",
    "    # Get the tokens for the image part, do not forget to pad the position of codebook\n",
    "    discrete_tokens_padded = list(vq_vae.encode(image.to(device).unsqueeze(0))[2] + pad_idx_vqvae)\n",
    "\n",
    "    # Get the tokens of the image description\n",
    "    describe_text = f\"\"\"DESCRIPTION:\n",
    "    {description}\n",
    "    IMAGE:\n",
    "    \"\"\"\n",
    "\n",
    "    describe_text_tokens = llama_tokenizer.encode(describe_text)\n",
    "\n",
    "    pos_img_start = llama_tokenizer.vocab['<|img_start|>']\n",
    "    pos_img_end = llama_tokenizer.vocab['<|img_end|>']\n",
    "    # Combine the tokens of image and text\n",
    "    tokenized_image_text = describe_text_tokens + [pos_img_start] + discrete_tokens_padded + [pos_img_end]\n",
    "\n",
    "    # Reconstruct the text\n",
    "    recontructed_text = llama_tokenizer.decode(tokenized_image_text) \n",
    "    image_text_dataset.append(recontructed_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "320c91a9-e0c0-4786-8eaa-424fa5e8e41e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<|begin_of_text|>DESCRIPTION:\n",
      "    No Choc Easter Gifts for Babies First Easter Shoes\n",
      "    IMAGE:\n",
      "    <|img_start|><dis_1540><dis_7914><dis_5946><dis_8063><dis_4710><dis_2294><dis_7914><dis_5679><dis_7136><dis_7254><dis_6156><dis_5146><dis_5449><dis_4710><dis_7874><dis_5146><dis_5146><dis_4710><dis_1595><dis_1022><dis_2921><dis_0711><dis_4710><dis_0726><dis_4145><dis_0513><dis_6015><dis_1313><dis_6156><dis_7914><dis_0658><dis_1292><dis_1022><dis_7874><dis_7150><dis_7150><dis_5146><dis_4008><dis_4008><dis_7874><dis_5864><dis_1635><dis_4998><dis_4178><dis_7466><dis_5232><dis_0915><dis_4879><dis_2460><dis_7626><dis_5320><dis_4998><dis_0280><dis_0915><dis_4920><dis_4831><dis_7914><dis_7601><dis_4920><dis_0252><dis_5146><dis_5146><dis_4632><dis_2294><dis_1833><dis_4710><dis_5146><dis_5146><dis_5146><dis_2119><dis_7150><dis_4710><dis_5146><dis_4145><dis_1191><dis_2119><dis_5481><dis_5246><dis_4368><dis_7427><dis_7855><dis_0122><dis_3258><dis_4300><dis_4024><dis_2997><dis_5300><dis_4019><dis_7544><dis_5414><dis_6021><dis_2495><dis_5798><dis_5146><dis_5146><dis_7365><dis_0971><dis_6516><dis_7728><dis_2119><dis_0711><dis_0711><dis_4512><dis_5146><dis_4998><dis_6002><dis_6544><dis_4920><dis_5507><dis_5801><dis_1352><dis_1555><dis_3703><dis_2799><dis_1328><dis_3258><dis_2170><dis_1603><dis_3524><dis_5684><dis_5891><dis_6768><dis_7246><dis_3831><dis_2997><dis_4539><dis_4512><dis_1168><dis_0726><dis_2119><dis_2119><dis_2119><dis_6998><dis_2366><dis_4178><dis_0398><dis_6002><dis_1680><dis_4920><dis_5507><dis_6708><dis_1348><dis_6228><dis_7065><dis_6358><dis_4558><dis_0124><dis_1518><dis_2077><dis_3240><dis_7336><dis_3825><dis_0780><dis_0429><dis_4028><dis_5414><dis_4618><dis_5798><dis_4710><dis_6115><dis_0356><dis_4710><dis_5146><dis_5384><dis_5607><dis_1080><dis_3668><dis_5966><dis_3396><dis_5062><dis_5067><dis_1661><dis_1555><dis_4552><dis_0573><dis_6980><dis_6704><dis_2066><dis_5229><dis_3792><dis_1358><dis_4060><dis_5728><dis_1556><dis_7141><dis_2429><dis_0698><dis_0339><dis_6364><dis_5146><dis_1052><dis_5449><dis_0320><dis_0256><dis_7728><dis_7728><dis_7728><dis_4258><dis_5744><dis_6950><dis_7675><dis_4751><dis_6512><dis_1555><dis_1555><dis_7087><dis_1326><dis_5318><dis_8167><dis_0445><dis_3831><dis_2936><dis_7929><dis_7720><dis_3104><dis_4887><dis_1874><dis_1874><dis_7741><dis_5519><dis_8084><dis_0690><dis_5146><dis_2294><dis_0658><dis_0711><dis_4035><dis_2967><dis_4035><dis_6653><dis_7828><dis_6557><dis_4019><dis_2702><dis_2803><dis_5832><dis_4552><dis_5113><dis_7563><dis_6780><dis_0531><dis_7648><dis_4409><dis_5152><dis_3039><dis_0203><dis_5841><dis_4637><dis_0382><dis_0938><dis_7945><dis_0784><dis_1602><dis_6501><dis_0931><dis_0356><dis_0510><dis_4035><dis_1249><dis_0972><dis_5888><dis_1525><dis_2986><dis_3643><dis_3643><dis_0348><dis_2072><dis_3652><dis_7366><dis_1401><dis_1862><dis_3247><dis_3875><dis_0348><dis_0348><dis_1660><dis_0491><dis_4235><dis_2368><dis_3683><dis_3296><dis_7020><dis_0847><dis_1253><dis_4925><dis_8063><dis_4450><dis_6881><dis_1313><dis_6334><dis_7684><dis_8014><dis_0333><dis_1842><dis_2803><dis_3643><dis_2375><dis_2415><dis_1225><dis_5073><dis_7932><dis_3063><dis_3354><dis_6317><dis_4285><dis_6886><dis_2720><dis_4269><dis_8045><dis_1236><dis_1790><dis_7524><dis_5500><dis_0998><dis_7197><dis_0718><dis_5832><dis_5165><dis_5744><dis_6933><dis_7143><dis_3879><dis_3765><dis_2082><dis_1488><dis_8104><dis_7025><dis_7446><dis_6287><dis_2190><dis_7098><dis_6914><dis_4101><dis_4210><dis_4343><dis_2886><dis_2415><dis_2409><dis_5305><dis_3794><dis_7703><dis_4051><dis_6108><dis_3357><dis_7912><dis_4481><dis_2961><dis_0247><dis_6829><dis_5779><dis_3869><dis_5537><dis_4590><dis_0372><dis_4906><dis_1488><dis_2864><dis_4154><dis_1425><dis_3574><dis_1060><dis_5108><dis_0780><dis_3701><dis_7704><dis_4135><dis_5867><dis_6039><dis_7286><dis_4382><dis_1471><dis_4596><dis_4546><dis_4303><dis_7867><dis_4271><dis_1887><dis_5555><dis_8053><dis_2748><dis_1014><dis_6212><dis_4560><dis_1233><dis_6793><dis_4893><dis_2409><dis_7950><dis_2184><dis_4348><dis_5397><dis_1800><dis_1660><dis_1823><dis_3165><dis_2749><dis_7507><dis_0376><dis_3948><dis_6228><dis_1741><dis_2217><dis_2993><dis_5297><dis_7914><dis_0670><dis_5892><dis_7556><dis_4441><dis_6364><dis_4560><dis_0187><dis_2563><dis_2617><dis_5790><dis_0254><dis_2184><dis_4028><dis_0333><dis_5450><dis_2172><dis_6228><dis_7902><dis_1069><dis_1526><dis_0423><dis_7691><dis_0417><dis_1003><dis_4348><dis_5675><dis_2234><dis_4857><dis_6840><dis_3450><dis_5388><dis_4022><dis_6933><dis_7753><dis_3323><dis_6523><dis_1599><dis_5039><dis_6119><dis_1175><dis_0733><dis_5143><dis_1233><dis_1988><dis_3707><dis_0262><dis_4098><dis_3254><dis_2859><dis_1947><dis_2674><dis_0989><dis_4287><dis_3500><dis_1555><dis_6281><dis_5925><dis_6768><dis_1643><dis_8094><dis_2735><dis_2839><dis_1546><dis_3831><dis_6933><dis_1589><dis_5777><dis_6557><dis_0093><dis_3500><dis_5963><dis_6638><dis_5581><dis_4125><dis_1255><dis_4925><dis_1799><dis_4925><dis_4976><dis_3786><dis_1048><dis_6716><dis_7325><dis_6789><dis_2930><dis_7527><dis_5462><dis_0059><dis_0532><dis_2281><dis_0334><dis_7606><dis_2894><dis_2731><dis_3916><dis_0947><dis_0784><dis_1121><dis_1249><dis_6533><dis_3006><dis_7815><dis_3010><dis_6540><dis_6277><dis_7573><dis_5679><dis_4609><dis_4287><dis_7527><dis_3346><dis_0274><dis_0482><dis_2941><dis_7498><dis_5832><dis_7264><dis_1569><dis_7902><dis_0399><dis_4277><dis_1621><dis_3846><dis_0076><dis_4018><dis_0984><dis_4163><dis_3126><dis_2041><dis_3032><dis_4088><dis_4088><dis_1175><dis_6634><dis_6593><dis_5784><dis_1045><dis_5777><dis_5594><dis_0376><dis_0644><dis_0670><dis_2144><dis_8017><dis_0538><dis_3875><dis_3810><dis_2705><dis_3643><dis_4893><dis_6352><dis_4580><dis_1100><dis_4564><dis_4906><dis_0427><dis_1271><dis_1789><dis_2984><dis_5800><dis_1530><dis_6946><dis_1121><dis_0615><dis_1045><dis_6987><dis_1045><dis_8055><dis_1220><dis_3336><dis_2584><dis_7556><dis_8082><dis_7534><dis_0349><dis_6488><dis_2066><dis_0301><dis_4287><dis_7286><dis_3745><dis_5745><dis_3810><dis_6581><dis_0503><dis_6552><dis_8163><dis_3875><dis_1741><dis_3188><dis_2184><dis_6638><dis_5454><dis_7669><dis_1400><dis_1814><dis_5477><dis_2965><dis_5047><dis_4645><dis_0968><dis_6948><dis_5225><dis_2046><dis_5580><dis_0561><dis_6191><dis_3875><dis_0417><dis_2223><dis_1425><dis_0347><dis_1555><dis_3584><dis_1556><dis_3164><dis_7417><dis_3148><dis_0123><dis_0123><dis_1741><dis_2012><dis_3786><dis_7148><dis_5777><dis_7399><dis_3425><dis_1887><dis_3680><dis_0029><dis_0570><dis_4219><dis_1832><dis_6659><dis_2080><dis_7866><dis_0080><dis_4433><dis_0576><dis_0160><dis_3825><dis_0059><dis_4019><dis_7189><dis_4667><dis_3089><dis_4623><dis_2145><dis_0820><dis_7211><dis_6108><dis_0711><dis_0711><dis_3259><dis_4724><dis_3487><dis_4088><dis_0356><dis_3912><dis_7819><dis_7408><dis_3680><dis_1664><dis_0151><dis_1453><dis_5381><dis_4409><dis_4101><dis_1818><dis_0348><dis_0733><dis_7527><dis_4109><dis_4019><dis_1660><dis_2715><dis_3652><dis_2942><dis_1181><dis_5518><dis_6422><dis_0711><dis_1559><dis_4946><dis_0202><dis_3786><dis_2899><dis_6715><dis_4831><dis_2981><dis_4197><dis_5481><dis_5094><dis_5751><dis_0254><dis_0318><dis_7839><dis_2092><dis_8107><dis_1069><dis_1415><dis_5772><dis_0301><dis_2930><dis_3766><dis_1493><dis_1837><dis_2803><dis_0803><dis_6564><dis_6269><dis_6948><dis_6948><dis_6015><dis_1463><dis_6015><dis_0873><dis_2230><dis_3749><dis_7383><dis_2350><dis_7225><dis_3850><dis_5381><dis_6529><dis_1183><dis_7413><dis_7211><dis_6747><dis_7928><dis_0483><dis_6950><dis_6693><dis_0967><dis_0613><dis_3643><dis_2581><dis_6950><dis_0930><dis_1135><dis_8151><dis_3992><dis_7169><dis_0510><dis_1177><dis_3764><dis_0252><dis_0599><dis_5547><dis_0905><dis_0051><dis_7228><dis_7367><dis_1179><dis_0683><dis_3645><dis_4197><dis_5956><dis_6529><dis_2834><dis_0733><dis_0320><dis_6063><dis_2045><dis_4906><dis_4462><dis_2982><dis_2026><dis_5250><dis_3805><dis_5998><dis_0820><dis_0107><dis_5062><dis_6334><dis_4161><dis_4632><dis_7989><dis_6015><dis_0318><dis_5384><dis_4691><dis_7624><dis_5027><dis_5381><dis_6197><dis_7080><dis_4603><dis_1356><dis_1472><dis_7148><dis_5287><dis_4603><dis_2158><dis_5381><dis_8017><dis_3145><dis_7417><dis_5791><dis_2763><dis_4451><dis_6119><dis_6334><dis_7989><dis_1670><dis_4738><dis_4363><dis_4321><dis_4321><dis_2682><dis_4805><dis_1670><dis_1080><dis_7174><dis_6969><dis_8183><dis_0683><dis_7899><dis_4867><dis_6593><dis_7815><dis_6496><dis_1175><dis_7080><dis_1957><dis_4867><dis_5650><dis_1421><dis_5018><dis_8149><dis_4007><dis_1336><dis_0318><dis_4363><dis_7969><dis_7969><dis_7188><dis_7227><dis_0174><dis_2056><dis_7066><dis_5056><dis_0103><dis_2918><dis_0711><dis_0174><dis_7664><dis_7283><dis_5393><dis_0280><dis_7080><dis_4867><dis_1339><dis_1313><dis_1052><dis_4603><dis_7613><dis_8035><dis_7763><dis_7276><dis_3539><dis_3298><dis_3664><dis_0599><dis_7078><dis_3135><dis_5056><dis_4321><dis_4800><dis_1318><dis_7188><dis_3457><dis_3307><dis_6496><dis_6998><dis_6998><dis_0318><dis_5016><dis_1950><dis_4044><dis_4040><dis_2825><dis_4867><dis_6197><dis_1175><dis_6277><dis_3307><dis_5777><dis_6158><dis_2158><dis_0561><dis_0546><dis_1069><dis_5217><dis_5259><dis_4738><dis_1559><dis_7890><dis_6796><dis_2237><dis_4363><dis_3711><dis_1341><dis_7344><dis_3395><dis_4465><dis_3764><dis_7515><dis_0854><dis_4920><dis_7619><dis_1271><dis_7932><dis_1630><dis_2825><dis_6540><dis_1175><dis_5454><dis_5481><dis_1472><dis_6651><dis_4970><dis_3488><dis_0680><dis_0853><dis_3627><dis_7080><dis_2394><dis_6286><dis_1991><dis_0854><dis_3847><dis_7613><dis_1039><dis_7989><dis_4738><dis_6194><dis_2727><dis_1846><dis_7890><dis_7844><dis_1241><dis_6081><dis_2276><dis_4060><dis_0808><dis_5228><dis_7652><dis_1643><dis_6121><dis_1694><dis_5311><dis_8188><dis_0511><dis_0340><dis_6621><dis_5553><dis_0432><dis_4450><dis_4161><dis_5047><dis_3645><dis_3148><dis_6814><dis_2089><dis_4161><dis_2682><dis_3217><dis_5056><dis_2009><dis_4465><dis_5016><dis_2004><dis_3395><dis_4161><dis_2276><dis_0658><dis_7772><dis_0628><dis_2479><dis_7075><dis_2235><dis_7554><dis_6251><dis_3272><dis_1469><dis_0820><dis_5194><dis_3378><dis_6156><dis_6036><dis_7821><dis_6119><dis_2230><dis_7788><dis_4619><dis_3893><dis_1016><dis_1554><dis_5449><dis_0892><dis_0356><dis_1168><dis_0658><dis_2294><dis_7258><dis_0356><dis_3879><dis_5062><dis_5543><dis_3766><dis_3015><dis_6787><dis_4323><dis_4931><dis_2599><dis_1121><dis_7601><dis_1914><|img_end|>\n"
     ]
    }
   ],
   "source": [
    "print(image_text_dataset[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "85037f9a-91aa-4991-ae41-251e926343b9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['<|begin_of_text|>',\n",
       " 'DESCRIPTION',\n",
       " ':Ċ',\n",
       " 'ĠĠĠ',\n",
       " 'ĠNo',\n",
       " 'ĠCh',\n",
       " 'oc',\n",
       " 'ĠEaster',\n",
       " 'ĠGifts',\n",
       " 'Ġfor',\n",
       " 'ĠBabies',\n",
       " 'ĠFirst',\n",
       " 'ĠEaster',\n",
       " 'ĠShoes',\n",
       " 'Ċ',\n",
       " 'ĠĠĠ',\n",
       " 'ĠIMAGE',\n",
       " ':Ċ',\n",
       " 'ĠĠĠĠ',\n",
       " '<|img_start|>',\n",
       " '<dis_1540>',\n",
       " '<dis_7914>',\n",
       " '<dis_5946>',\n",
       " '<dis_8063>',\n",
       " '<dis_4710>',\n",
       " '<dis_2294>',\n",
       " '<dis_7914>',\n",
       " '<dis_5679>',\n",
       " '<dis_7136>',\n",
       " '<dis_7254>',\n",
       " '<dis_6156>',\n",
       " '<dis_5146>',\n",
       " '<dis_5449>',\n",
       " '<dis_4710>',\n",
       " '<dis_7874>',\n",
       " '<dis_5146>',\n",
       " '<dis_5146>',\n",
       " '<dis_4710>',\n",
       " '<dis_1595>',\n",
       " '<dis_1022>',\n",
       " '<dis_2921>',\n",
       " '<dis_0711>',\n",
       " '<dis_4710>',\n",
       " '<dis_0726>',\n",
       " '<dis_4145>',\n",
       " '<dis_0513>',\n",
       " '<dis_6015>',\n",
       " '<dis_1313>',\n",
       " '<dis_6156>',\n",
       " '<dis_7914>',\n",
       " '<dis_0658>',\n",
       " '<dis_1292>',\n",
       " '<dis_1022>',\n",
       " '<dis_7874>',\n",
       " '<dis_7150>',\n",
       " '<dis_7150>',\n",
       " '<dis_5146>',\n",
       " '<dis_4008>',\n",
       " '<dis_4008>',\n",
       " '<dis_7874>',\n",
       " '<dis_5864>',\n",
       " '<dis_1635>',\n",
       " '<dis_4998>',\n",
       " '<dis_4178>',\n",
       " '<dis_7466>',\n",
       " '<dis_5232>',\n",
       " '<dis_0915>',\n",
       " '<dis_4879>',\n",
       " '<dis_2460>',\n",
       " '<dis_7626>',\n",
       " '<dis_5320>',\n",
       " '<dis_4998>',\n",
       " '<dis_0280>',\n",
       " '<dis_0915>',\n",
       " '<dis_4920>',\n",
       " '<dis_4831>',\n",
       " '<dis_7914>',\n",
       " '<dis_7601>',\n",
       " '<dis_4920>',\n",
       " '<dis_0252>',\n",
       " '<dis_5146>',\n",
       " '<dis_5146>',\n",
       " '<dis_4632>',\n",
       " '<dis_2294>',\n",
       " '<dis_1833>',\n",
       " '<dis_4710>',\n",
       " '<dis_5146>',\n",
       " '<dis_5146>',\n",
       " '<dis_5146>',\n",
       " '<dis_2119>',\n",
       " '<dis_7150>',\n",
       " '<dis_4710>',\n",
       " '<dis_5146>',\n",
       " '<dis_4145>',\n",
       " '<dis_1191>',\n",
       " '<dis_2119>',\n",
       " '<dis_5481>',\n",
       " '<dis_5246>',\n",
       " '<dis_4368>',\n",
       " '<dis_7427>',\n",
       " '<dis_7855>',\n",
       " '<dis_0122>',\n",
       " '<dis_3258>',\n",
       " '<dis_4300>',\n",
       " '<dis_4024>',\n",
       " '<dis_2997>',\n",
       " '<dis_5300>',\n",
       " '<dis_4019>',\n",
       " '<dis_7544>',\n",
       " '<dis_5414>',\n",
       " '<dis_6021>',\n",
       " '<dis_2495>',\n",
       " '<dis_5798>',\n",
       " '<dis_5146>',\n",
       " '<dis_5146>',\n",
       " '<dis_7365>',\n",
       " '<dis_0971>',\n",
       " '<dis_6516>',\n",
       " '<dis_7728>',\n",
       " '<dis_2119>',\n",
       " '<dis_0711>',\n",
       " '<dis_0711>',\n",
       " '<dis_4512>',\n",
       " '<dis_5146>',\n",
       " '<dis_4998>',\n",
       " '<dis_6002>',\n",
       " '<dis_6544>',\n",
       " '<dis_4920>',\n",
       " '<dis_5507>',\n",
       " '<dis_5801>',\n",
       " '<dis_1352>',\n",
       " '<dis_1555>',\n",
       " '<dis_3703>',\n",
       " '<dis_2799>',\n",
       " '<dis_1328>',\n",
       " '<dis_3258>',\n",
       " '<dis_2170>',\n",
       " '<dis_1603>',\n",
       " '<dis_3524>',\n",
       " '<dis_5684>',\n",
       " '<dis_5891>',\n",
       " '<dis_6768>',\n",
       " '<dis_7246>',\n",
       " '<dis_3831>',\n",
       " '<dis_2997>',\n",
       " '<dis_4539>',\n",
       " '<dis_4512>',\n",
       " '<dis_1168>',\n",
       " '<dis_0726>',\n",
       " '<dis_2119>',\n",
       " '<dis_2119>',\n",
       " '<dis_2119>',\n",
       " '<dis_6998>',\n",
       " '<dis_2366>',\n",
       " '<dis_4178>',\n",
       " '<dis_0398>',\n",
       " '<dis_6002>',\n",
       " '<dis_1680>',\n",
       " '<dis_4920>',\n",
       " '<dis_5507>',\n",
       " '<dis_6708>',\n",
       " '<dis_1348>',\n",
       " '<dis_6228>',\n",
       " '<dis_7065>',\n",
       " '<dis_6358>',\n",
       " '<dis_4558>',\n",
       " '<dis_0124>',\n",
       " '<dis_1518>',\n",
       " '<dis_2077>',\n",
       " '<dis_3240>',\n",
       " '<dis_7336>',\n",
       " '<dis_3825>',\n",
       " '<dis_0780>',\n",
       " '<dis_0429>',\n",
       " '<dis_4028>',\n",
       " '<dis_5414>',\n",
       " '<dis_4618>',\n",
       " '<dis_5798>',\n",
       " '<dis_4710>',\n",
       " '<dis_6115>',\n",
       " '<dis_0356>',\n",
       " '<dis_4710>',\n",
       " '<dis_5146>',\n",
       " '<dis_5384>',\n",
       " '<dis_5607>',\n",
       " '<dis_1080>',\n",
       " '<dis_3668>',\n",
       " '<dis_5966>',\n",
       " '<dis_3396>',\n",
       " '<dis_5062>',\n",
       " '<dis_5067>',\n",
       " '<dis_1661>',\n",
       " '<dis_1555>',\n",
       " '<dis_4552>',\n",
       " '<dis_0573>',\n",
       " '<dis_6980>',\n",
       " '<dis_6704>',\n",
       " '<dis_2066>',\n",
       " '<dis_5229>',\n",
       " '<dis_3792>',\n",
       " '<dis_1358>',\n",
       " '<dis_4060>',\n",
       " '<dis_5728>',\n",
       " '<dis_1556>',\n",
       " '<dis_7141>',\n",
       " '<dis_2429>',\n",
       " '<dis_0698>',\n",
       " '<dis_0339>',\n",
       " '<dis_6364>',\n",
       " '<dis_5146>',\n",
       " '<dis_1052>',\n",
       " '<dis_5449>',\n",
       " '<dis_0320>',\n",
       " '<dis_0256>',\n",
       " '<dis_7728>',\n",
       " '<dis_7728>',\n",
       " '<dis_7728>',\n",
       " '<dis_4258>',\n",
       " '<dis_5744>',\n",
       " '<dis_6950>',\n",
       " '<dis_7675>',\n",
       " '<dis_4751>',\n",
       " '<dis_6512>',\n",
       " '<dis_1555>',\n",
       " '<dis_1555>',\n",
       " '<dis_7087>',\n",
       " '<dis_1326>',\n",
       " '<dis_5318>',\n",
       " '<dis_8167>',\n",
       " '<dis_0445>',\n",
       " '<dis_3831>',\n",
       " '<dis_2936>',\n",
       " '<dis_7929>',\n",
       " '<dis_7720>',\n",
       " '<dis_3104>',\n",
       " '<dis_4887>',\n",
       " '<dis_1874>',\n",
       " '<dis_1874>',\n",
       " '<dis_7741>',\n",
       " '<dis_5519>',\n",
       " '<dis_8084>',\n",
       " '<dis_0690>',\n",
       " '<dis_5146>',\n",
       " '<dis_2294>',\n",
       " '<dis_0658>',\n",
       " '<dis_0711>',\n",
       " '<dis_4035>',\n",
       " '<dis_2967>',\n",
       " '<dis_4035>',\n",
       " '<dis_6653>',\n",
       " '<dis_7828>',\n",
       " '<dis_6557>',\n",
       " '<dis_4019>',\n",
       " '<dis_2702>',\n",
       " '<dis_2803>',\n",
       " '<dis_5832>',\n",
       " '<dis_4552>',\n",
       " '<dis_5113>',\n",
       " '<dis_7563>',\n",
       " '<dis_6780>',\n",
       " '<dis_0531>',\n",
       " '<dis_7648>',\n",
       " '<dis_4409>',\n",
       " '<dis_5152>',\n",
       " '<dis_3039>',\n",
       " '<dis_0203>',\n",
       " '<dis_5841>',\n",
       " '<dis_4637>',\n",
       " '<dis_0382>',\n",
       " '<dis_0938>',\n",
       " '<dis_7945>',\n",
       " '<dis_0784>',\n",
       " '<dis_1602>',\n",
       " '<dis_6501>',\n",
       " '<dis_0931>',\n",
       " '<dis_0356>',\n",
       " '<dis_0510>',\n",
       " '<dis_4035>',\n",
       " '<dis_1249>',\n",
       " '<dis_0972>',\n",
       " '<dis_5888>',\n",
       " '<dis_1525>',\n",
       " '<dis_2986>',\n",
       " '<dis_3643>',\n",
       " '<dis_3643>',\n",
       " '<dis_0348>',\n",
       " '<dis_2072>',\n",
       " '<dis_3652>',\n",
       " '<dis_7366>',\n",
       " '<dis_1401>',\n",
       " '<dis_1862>',\n",
       " '<dis_3247>',\n",
       " '<dis_3875>',\n",
       " '<dis_0348>',\n",
       " '<dis_0348>',\n",
       " '<dis_1660>',\n",
       " '<dis_0491>',\n",
       " '<dis_4235>',\n",
       " '<dis_2368>',\n",
       " '<dis_3683>',\n",
       " '<dis_3296>',\n",
       " '<dis_7020>',\n",
       " '<dis_0847>',\n",
       " '<dis_1253>',\n",
       " '<dis_4925>',\n",
       " '<dis_8063>',\n",
       " '<dis_4450>',\n",
       " '<dis_6881>',\n",
       " '<dis_1313>',\n",
       " '<dis_6334>',\n",
       " '<dis_7684>',\n",
       " '<dis_8014>',\n",
       " '<dis_0333>',\n",
       " '<dis_1842>',\n",
       " '<dis_2803>',\n",
       " '<dis_3643>',\n",
       " '<dis_2375>',\n",
       " '<dis_2415>',\n",
       " '<dis_1225>',\n",
       " '<dis_5073>',\n",
       " '<dis_7932>',\n",
       " '<dis_3063>',\n",
       " '<dis_3354>',\n",
       " '<dis_6317>',\n",
       " '<dis_4285>',\n",
       " '<dis_6886>',\n",
       " '<dis_2720>',\n",
       " '<dis_4269>',\n",
       " '<dis_8045>',\n",
       " '<dis_1236>',\n",
       " '<dis_1790>',\n",
       " '<dis_7524>',\n",
       " '<dis_5500>',\n",
       " '<dis_0998>',\n",
       " '<dis_7197>',\n",
       " '<dis_0718>',\n",
       " '<dis_5832>',\n",
       " '<dis_5165>',\n",
       " '<dis_5744>',\n",
       " '<dis_6933>',\n",
       " '<dis_7143>',\n",
       " '<dis_3879>',\n",
       " '<dis_3765>',\n",
       " '<dis_2082>',\n",
       " '<dis_1488>',\n",
       " '<dis_8104>',\n",
       " '<dis_7025>',\n",
       " '<dis_7446>',\n",
       " '<dis_6287>',\n",
       " '<dis_2190>',\n",
       " '<dis_7098>',\n",
       " '<dis_6914>',\n",
       " '<dis_4101>',\n",
       " '<dis_4210>',\n",
       " '<dis_4343>',\n",
       " '<dis_2886>',\n",
       " '<dis_2415>',\n",
       " '<dis_2409>',\n",
       " '<dis_5305>',\n",
       " '<dis_3794>',\n",
       " '<dis_7703>',\n",
       " '<dis_4051>',\n",
       " '<dis_6108>',\n",
       " '<dis_3357>',\n",
       " '<dis_7912>',\n",
       " '<dis_4481>',\n",
       " '<dis_2961>',\n",
       " '<dis_0247>',\n",
       " '<dis_6829>',\n",
       " '<dis_5779>',\n",
       " '<dis_3869>',\n",
       " '<dis_5537>',\n",
       " '<dis_4590>',\n",
       " '<dis_0372>',\n",
       " '<dis_4906>',\n",
       " '<dis_1488>',\n",
       " '<dis_2864>',\n",
       " '<dis_4154>',\n",
       " '<dis_1425>',\n",
       " '<dis_3574>',\n",
       " '<dis_1060>',\n",
       " '<dis_5108>',\n",
       " '<dis_0780>',\n",
       " '<dis_3701>',\n",
       " '<dis_7704>',\n",
       " '<dis_4135>',\n",
       " '<dis_5867>',\n",
       " '<dis_6039>',\n",
       " '<dis_7286>',\n",
       " '<dis_4382>',\n",
       " '<dis_1471>',\n",
       " '<dis_4596>',\n",
       " '<dis_4546>',\n",
       " '<dis_4303>',\n",
       " '<dis_7867>',\n",
       " '<dis_4271>',\n",
       " '<dis_1887>',\n",
       " '<dis_5555>',\n",
       " '<dis_8053>',\n",
       " '<dis_2748>',\n",
       " '<dis_1014>',\n",
       " '<dis_6212>',\n",
       " '<dis_4560>',\n",
       " '<dis_1233>',\n",
       " '<dis_6793>',\n",
       " '<dis_4893>',\n",
       " '<dis_2409>',\n",
       " '<dis_7950>',\n",
       " '<dis_2184>',\n",
       " '<dis_4348>',\n",
       " '<dis_5397>',\n",
       " '<dis_1800>',\n",
       " '<dis_1660>',\n",
       " '<dis_1823>',\n",
       " '<dis_3165>',\n",
       " '<dis_2749>',\n",
       " '<dis_7507>',\n",
       " '<dis_0376>',\n",
       " '<dis_3948>',\n",
       " '<dis_6228>',\n",
       " '<dis_1741>',\n",
       " '<dis_2217>',\n",
       " '<dis_2993>',\n",
       " '<dis_5297>',\n",
       " '<dis_7914>',\n",
       " '<dis_0670>',\n",
       " '<dis_5892>',\n",
       " '<dis_7556>',\n",
       " '<dis_4441>',\n",
       " '<dis_6364>',\n",
       " '<dis_4560>',\n",
       " '<dis_0187>',\n",
       " '<dis_2563>',\n",
       " '<dis_2617>',\n",
       " '<dis_5790>',\n",
       " '<dis_0254>',\n",
       " '<dis_2184>',\n",
       " '<dis_4028>',\n",
       " '<dis_0333>',\n",
       " '<dis_5450>',\n",
       " '<dis_2172>',\n",
       " '<dis_6228>',\n",
       " '<dis_7902>',\n",
       " '<dis_1069>',\n",
       " '<dis_1526>',\n",
       " '<dis_0423>',\n",
       " '<dis_7691>',\n",
       " '<dis_0417>',\n",
       " '<dis_1003>',\n",
       " '<dis_4348>',\n",
       " '<dis_5675>',\n",
       " '<dis_2234>',\n",
       " '<dis_4857>',\n",
       " '<dis_6840>',\n",
       " '<dis_3450>',\n",
       " '<dis_5388>',\n",
       " '<dis_4022>',\n",
       " '<dis_6933>',\n",
       " '<dis_7753>',\n",
       " '<dis_3323>',\n",
       " '<dis_6523>',\n",
       " '<dis_1599>',\n",
       " '<dis_5039>',\n",
       " '<dis_6119>',\n",
       " '<dis_1175>',\n",
       " '<dis_0733>',\n",
       " '<dis_5143>',\n",
       " '<dis_1233>',\n",
       " '<dis_1988>',\n",
       " '<dis_3707>',\n",
       " '<dis_0262>',\n",
       " '<dis_4098>',\n",
       " '<dis_3254>',\n",
       " '<dis_2859>',\n",
       " '<dis_1947>',\n",
       " '<dis_2674>',\n",
       " '<dis_0989>',\n",
       " '<dis_4287>',\n",
       " '<dis_3500>',\n",
       " '<dis_1555>',\n",
       " '<dis_6281>',\n",
       " '<dis_5925>',\n",
       " '<dis_6768>',\n",
       " '<dis_1643>',\n",
       " '<dis_8094>',\n",
       " '<dis_2735>',\n",
       " '<dis_2839>',\n",
       " '<dis_1546>',\n",
       " '<dis_3831>',\n",
       " '<dis_6933>',\n",
       " '<dis_1589>',\n",
       " '<dis_5777>',\n",
       " '<dis_6557>',\n",
       " '<dis_0093>',\n",
       " '<dis_3500>',\n",
       " '<dis_5963>',\n",
       " '<dis_6638>',\n",
       " '<dis_5581>',\n",
       " '<dis_4125>',\n",
       " '<dis_1255>',\n",
       " '<dis_4925>',\n",
       " '<dis_1799>',\n",
       " '<dis_4925>',\n",
       " '<dis_4976>',\n",
       " '<dis_3786>',\n",
       " '<dis_1048>',\n",
       " '<dis_6716>',\n",
       " '<dis_7325>',\n",
       " '<dis_6789>',\n",
       " '<dis_2930>',\n",
       " '<dis_7527>',\n",
       " '<dis_5462>',\n",
       " '<dis_0059>',\n",
       " '<dis_0532>',\n",
       " '<dis_2281>',\n",
       " '<dis_0334>',\n",
       " '<dis_7606>',\n",
       " '<dis_2894>',\n",
       " '<dis_2731>',\n",
       " '<dis_3916>',\n",
       " '<dis_0947>',\n",
       " '<dis_0784>',\n",
       " '<dis_1121>',\n",
       " '<dis_1249>',\n",
       " '<dis_6533>',\n",
       " '<dis_3006>',\n",
       " '<dis_7815>',\n",
       " '<dis_3010>',\n",
       " '<dis_6540>',\n",
       " '<dis_6277>',\n",
       " '<dis_7573>',\n",
       " '<dis_5679>',\n",
       " '<dis_4609>',\n",
       " '<dis_4287>',\n",
       " '<dis_7527>',\n",
       " '<dis_3346>',\n",
       " '<dis_0274>',\n",
       " '<dis_0482>',\n",
       " '<dis_2941>',\n",
       " '<dis_7498>',\n",
       " '<dis_5832>',\n",
       " '<dis_7264>',\n",
       " '<dis_1569>',\n",
       " '<dis_7902>',\n",
       " '<dis_0399>',\n",
       " '<dis_4277>',\n",
       " '<dis_1621>',\n",
       " '<dis_3846>',\n",
       " '<dis_0076>',\n",
       " '<dis_4018>',\n",
       " '<dis_0984>',\n",
       " '<dis_4163>',\n",
       " '<dis_3126>',\n",
       " '<dis_2041>',\n",
       " '<dis_3032>',\n",
       " '<dis_4088>',\n",
       " '<dis_4088>',\n",
       " '<dis_1175>',\n",
       " '<dis_6634>',\n",
       " '<dis_6593>',\n",
       " '<dis_5784>',\n",
       " '<dis_1045>',\n",
       " '<dis_5777>',\n",
       " '<dis_5594>',\n",
       " '<dis_0376>',\n",
       " '<dis_0644>',\n",
       " '<dis_0670>',\n",
       " '<dis_2144>',\n",
       " '<dis_8017>',\n",
       " '<dis_0538>',\n",
       " '<dis_3875>',\n",
       " '<dis_3810>',\n",
       " '<dis_2705>',\n",
       " '<dis_3643>',\n",
       " '<dis_4893>',\n",
       " '<dis_6352>',\n",
       " '<dis_4580>',\n",
       " '<dis_1100>',\n",
       " '<dis_4564>',\n",
       " '<dis_4906>',\n",
       " '<dis_0427>',\n",
       " '<dis_1271>',\n",
       " '<dis_1789>',\n",
       " '<dis_2984>',\n",
       " '<dis_5800>',\n",
       " '<dis_1530>',\n",
       " '<dis_6946>',\n",
       " '<dis_1121>',\n",
       " '<dis_0615>',\n",
       " '<dis_1045>',\n",
       " '<dis_6987>',\n",
       " '<dis_1045>',\n",
       " '<dis_8055>',\n",
       " '<dis_1220>',\n",
       " '<dis_3336>',\n",
       " '<dis_2584>',\n",
       " '<dis_7556>',\n",
       " '<dis_8082>',\n",
       " '<dis_7534>',\n",
       " '<dis_0349>',\n",
       " '<dis_6488>',\n",
       " '<dis_2066>',\n",
       " '<dis_0301>',\n",
       " '<dis_4287>',\n",
       " '<dis_7286>',\n",
       " '<dis_3745>',\n",
       " '<dis_5745>',\n",
       " '<dis_3810>',\n",
       " '<dis_6581>',\n",
       " '<dis_0503>',\n",
       " '<dis_6552>',\n",
       " '<dis_8163>',\n",
       " '<dis_3875>',\n",
       " '<dis_1741>',\n",
       " '<dis_3188>',\n",
       " '<dis_2184>',\n",
       " '<dis_6638>',\n",
       " '<dis_5454>',\n",
       " '<dis_7669>',\n",
       " '<dis_1400>',\n",
       " '<dis_1814>',\n",
       " '<dis_5477>',\n",
       " '<dis_2965>',\n",
       " '<dis_5047>',\n",
       " '<dis_4645>',\n",
       " '<dis_0968>',\n",
       " '<dis_6948>',\n",
       " '<dis_5225>',\n",
       " '<dis_2046>',\n",
       " '<dis_5580>',\n",
       " '<dis_0561>',\n",
       " '<dis_6191>',\n",
       " '<dis_3875>',\n",
       " '<dis_0417>',\n",
       " '<dis_2223>',\n",
       " '<dis_1425>',\n",
       " '<dis_0347>',\n",
       " '<dis_1555>',\n",
       " '<dis_3584>',\n",
       " '<dis_1556>',\n",
       " '<dis_3164>',\n",
       " '<dis_7417>',\n",
       " '<dis_3148>',\n",
       " '<dis_0123>',\n",
       " '<dis_0123>',\n",
       " '<dis_1741>',\n",
       " '<dis_2012>',\n",
       " '<dis_3786>',\n",
       " '<dis_7148>',\n",
       " '<dis_5777>',\n",
       " '<dis_7399>',\n",
       " '<dis_3425>',\n",
       " '<dis_1887>',\n",
       " '<dis_3680>',\n",
       " '<dis_0029>',\n",
       " '<dis_0570>',\n",
       " '<dis_4219>',\n",
       " '<dis_1832>',\n",
       " '<dis_6659>',\n",
       " '<dis_2080>',\n",
       " '<dis_7866>',\n",
       " '<dis_0080>',\n",
       " '<dis_4433>',\n",
       " '<dis_0576>',\n",
       " '<dis_0160>',\n",
       " '<dis_3825>',\n",
       " '<dis_0059>',\n",
       " '<dis_4019>',\n",
       " '<dis_7189>',\n",
       " '<dis_4667>',\n",
       " '<dis_3089>',\n",
       " '<dis_4623>',\n",
       " '<dis_2145>',\n",
       " '<dis_0820>',\n",
       " '<dis_7211>',\n",
       " '<dis_6108>',\n",
       " '<dis_0711>',\n",
       " '<dis_0711>',\n",
       " '<dis_3259>',\n",
       " '<dis_4724>',\n",
       " '<dis_3487>',\n",
       " '<dis_4088>',\n",
       " '<dis_0356>',\n",
       " '<dis_3912>',\n",
       " '<dis_7819>',\n",
       " '<dis_7408>',\n",
       " '<dis_3680>',\n",
       " '<dis_1664>',\n",
       " '<dis_0151>',\n",
       " '<dis_1453>',\n",
       " '<dis_5381>',\n",
       " '<dis_4409>',\n",
       " '<dis_4101>',\n",
       " '<dis_1818>',\n",
       " '<dis_0348>',\n",
       " '<dis_0733>',\n",
       " '<dis_7527>',\n",
       " '<dis_4109>',\n",
       " '<dis_4019>',\n",
       " '<dis_1660>',\n",
       " '<dis_2715>',\n",
       " '<dis_3652>',\n",
       " '<dis_2942>',\n",
       " '<dis_1181>',\n",
       " '<dis_5518>',\n",
       " '<dis_6422>',\n",
       " '<dis_0711>',\n",
       " '<dis_1559>',\n",
       " '<dis_4946>',\n",
       " '<dis_0202>',\n",
       " '<dis_3786>',\n",
       " '<dis_2899>',\n",
       " '<dis_6715>',\n",
       " '<dis_4831>',\n",
       " '<dis_2981>',\n",
       " '<dis_4197>',\n",
       " '<dis_5481>',\n",
       " '<dis_5094>',\n",
       " '<dis_5751>',\n",
       " '<dis_0254>',\n",
       " '<dis_0318>',\n",
       " '<dis_7839>',\n",
       " '<dis_2092>',\n",
       " '<dis_8107>',\n",
       " '<dis_1069>',\n",
       " '<dis_1415>',\n",
       " '<dis_5772>',\n",
       " '<dis_0301>',\n",
       " '<dis_2930>',\n",
       " '<dis_3766>',\n",
       " '<dis_1493>',\n",
       " '<dis_1837>',\n",
       " '<dis_2803>',\n",
       " '<dis_0803>',\n",
       " '<dis_6564>',\n",
       " '<dis_6269>',\n",
       " '<dis_6948>',\n",
       " '<dis_6948>',\n",
       " '<dis_6015>',\n",
       " '<dis_1463>',\n",
       " '<dis_6015>',\n",
       " '<dis_0873>',\n",
       " '<dis_2230>',\n",
       " '<dis_3749>',\n",
       " '<dis_7383>',\n",
       " '<dis_2350>',\n",
       " '<dis_7225>',\n",
       " '<dis_3850>',\n",
       " '<dis_5381>',\n",
       " '<dis_6529>',\n",
       " '<dis_1183>',\n",
       " '<dis_7413>',\n",
       " '<dis_7211>',\n",
       " '<dis_6747>',\n",
       " '<dis_7928>',\n",
       " '<dis_0483>',\n",
       " '<dis_6950>',\n",
       " '<dis_6693>',\n",
       " '<dis_0967>',\n",
       " '<dis_0613>',\n",
       " '<dis_3643>',\n",
       " '<dis_2581>',\n",
       " '<dis_6950>',\n",
       " '<dis_0930>',\n",
       " '<dis_1135>',\n",
       " '<dis_8151>',\n",
       " '<dis_3992>',\n",
       " '<dis_7169>',\n",
       " '<dis_0510>',\n",
       " '<dis_1177>',\n",
       " '<dis_3764>',\n",
       " '<dis_0252>',\n",
       " '<dis_0599>',\n",
       " '<dis_5547>',\n",
       " '<dis_0905>',\n",
       " '<dis_0051>',\n",
       " '<dis_7228>',\n",
       " '<dis_7367>',\n",
       " '<dis_1179>',\n",
       " '<dis_0683>',\n",
       " '<dis_3645>',\n",
       " '<dis_4197>',\n",
       " '<dis_5956>',\n",
       " '<dis_6529>',\n",
       " '<dis_2834>',\n",
       " '<dis_0733>',\n",
       " '<dis_0320>',\n",
       " '<dis_6063>',\n",
       " '<dis_2045>',\n",
       " '<dis_4906>',\n",
       " '<dis_4462>',\n",
       " '<dis_2982>',\n",
       " '<dis_2026>',\n",
       " '<dis_5250>',\n",
       " '<dis_3805>',\n",
       " '<dis_5998>',\n",
       " '<dis_0820>',\n",
       " '<dis_0107>',\n",
       " '<dis_5062>',\n",
       " '<dis_6334>',\n",
       " '<dis_4161>',\n",
       " '<dis_4632>',\n",
       " '<dis_7989>',\n",
       " '<dis_6015>',\n",
       " '<dis_0318>',\n",
       " '<dis_5384>',\n",
       " '<dis_4691>',\n",
       " '<dis_7624>',\n",
       " '<dis_5027>',\n",
       " '<dis_5381>',\n",
       " '<dis_6197>',\n",
       " '<dis_7080>',\n",
       " '<dis_4603>',\n",
       " '<dis_1356>',\n",
       " '<dis_1472>',\n",
       " '<dis_7148>',\n",
       " '<dis_5287>',\n",
       " '<dis_4603>',\n",
       " '<dis_2158>',\n",
       " '<dis_5381>',\n",
       " '<dis_8017>',\n",
       " '<dis_3145>',\n",
       " '<dis_7417>',\n",
       " '<dis_5791>',\n",
       " '<dis_2763>',\n",
       " '<dis_4451>',\n",
       " '<dis_6119>',\n",
       " '<dis_6334>',\n",
       " '<dis_7989>',\n",
       " '<dis_1670>',\n",
       " '<dis_4738>',\n",
       " '<dis_4363>',\n",
       " '<dis_4321>',\n",
       " '<dis_4321>',\n",
       " '<dis_2682>',\n",
       " '<dis_4805>',\n",
       " '<dis_1670>',\n",
       " '<dis_1080>',\n",
       " '<dis_7174>',\n",
       " '<dis_6969>',\n",
       " '<dis_8183>',\n",
       " '<dis_0683>',\n",
       " '<dis_7899>',\n",
       " '<dis_4867>',\n",
       " '<dis_6593>',\n",
       " '<dis_7815>',\n",
       " '<dis_6496>',\n",
       " '<dis_1175>',\n",
       " '<dis_7080>',\n",
       " '<dis_1957>',\n",
       " '<dis_4867>',\n",
       " '<dis_5650>',\n",
       " '<dis_1421>',\n",
       " '<dis_5018>',\n",
       " '<dis_8149>',\n",
       " '<dis_4007>',\n",
       " '<dis_1336>',\n",
       " '<dis_0318>',\n",
       " '<dis_4363>',\n",
       " '<dis_7969>',\n",
       " '<dis_7969>',\n",
       " '<dis_7188>',\n",
       " '<dis_7227>',\n",
       " '<dis_0174>',\n",
       " '<dis_2056>',\n",
       " '<dis_7066>',\n",
       " '<dis_5056>',\n",
       " '<dis_0103>',\n",
       " '<dis_2918>',\n",
       " '<dis_0711>',\n",
       " '<dis_0174>',\n",
       " '<dis_7664>',\n",
       " '<dis_7283>',\n",
       " '<dis_5393>',\n",
       " '<dis_0280>',\n",
       " '<dis_7080>',\n",
       " '<dis_4867>',\n",
       " '<dis_1339>',\n",
       " '<dis_1313>',\n",
       " '<dis_1052>',\n",
       " '<dis_4603>',\n",
       " '<dis_7613>',\n",
       " '<dis_8035>',\n",
       " '<dis_7763>',\n",
       " '<dis_7276>',\n",
       " '<dis_3539>',\n",
       " '<dis_3298>',\n",
       " '<dis_3664>',\n",
       " '<dis_0599>',\n",
       " '<dis_7078>',\n",
       " '<dis_3135>',\n",
       " '<dis_5056>',\n",
       " '<dis_4321>',\n",
       " '<dis_4800>',\n",
       " '<dis_1318>',\n",
       " '<dis_7188>',\n",
       " '<dis_3457>',\n",
       " '<dis_3307>',\n",
       " '<dis_6496>',\n",
       " '<dis_6998>',\n",
       " '<dis_6998>',\n",
       " '<dis_0318>',\n",
       " '<dis_5016>',\n",
       " '<dis_1950>',\n",
       " '<dis_4044>',\n",
       " '<dis_4040>',\n",
       " '<dis_2825>',\n",
       " '<dis_4867>',\n",
       " '<dis_6197>',\n",
       " '<dis_1175>',\n",
       " '<dis_6277>',\n",
       " '<dis_3307>',\n",
       " '<dis_5777>',\n",
       " '<dis_6158>',\n",
       " '<dis_2158>',\n",
       " '<dis_0561>',\n",
       " '<dis_0546>',\n",
       " '<dis_1069>',\n",
       " '<dis_5217>',\n",
       " '<dis_5259>',\n",
       " '<dis_4738>',\n",
       " '<dis_1559>',\n",
       " '<dis_7890>',\n",
       " '<dis_6796>',\n",
       " '<dis_2237>',\n",
       " '<dis_4363>',\n",
       " '<dis_3711>',\n",
       " '<dis_1341>',\n",
       " '<dis_7344>',\n",
       " '<dis_3395>',\n",
       " '<dis_4465>',\n",
       " '<dis_3764>',\n",
       " '<dis_7515>',\n",
       " '<dis_0854>',\n",
       " '<dis_4920>',\n",
       " '<dis_7619>',\n",
       " '<dis_1271>',\n",
       " '<dis_7932>',\n",
       " '<dis_1630>',\n",
       " '<dis_2825>',\n",
       " '<dis_6540>',\n",
       " '<dis_1175>',\n",
       " '<dis_5454>',\n",
       " '<dis_5481>',\n",
       " '<dis_1472>',\n",
       " '<dis_6651>',\n",
       " '<dis_4970>',\n",
       " '<dis_3488>',\n",
       " '<dis_0680>',\n",
       " '<dis_0853>',\n",
       " '<dis_3627>',\n",
       " '<dis_7080>',\n",
       " '<dis_2394>',\n",
       " '<dis_6286>',\n",
       " '<dis_1991>',\n",
       " '<dis_0854>',\n",
       " '<dis_3847>',\n",
       " '<dis_7613>',\n",
       " '<dis_1039>',\n",
       " '<dis_7989>',\n",
       " '<dis_4738>',\n",
       " '<dis_6194>',\n",
       " '<dis_2727>',\n",
       " '<dis_1846>',\n",
       " '<dis_7890>',\n",
       " '<dis_7844>',\n",
       " '<dis_1241>',\n",
       " '<dis_6081>',\n",
       " '<dis_2276>',\n",
       " '<dis_4060>',\n",
       " '<dis_0808>',\n",
       " '<dis_5228>',\n",
       " '<dis_7652>',\n",
       " '<dis_1643>',\n",
       " '<dis_6121>',\n",
       " '<dis_1694>',\n",
       " '<dis_5311>',\n",
       " '<dis_8188>',\n",
       " '<dis_0511>',\n",
       " '<dis_0340>',\n",
       " '<dis_6621>',\n",
       " '<dis_5553>',\n",
       " '<dis_0432>',\n",
       " '<dis_4450>',\n",
       " '<dis_4161>',\n",
       " '<dis_5047>',\n",
       " '<dis_3645>',\n",
       " '<dis_3148>',\n",
       " '<dis_6814>',\n",
       " '<dis_2089>',\n",
       " '<dis_4161>',\n",
       " '<dis_2682>',\n",
       " '<dis_3217>',\n",
       " '<dis_5056>',\n",
       " '<dis_2009>',\n",
       " '<dis_4465>',\n",
       " '<dis_5016>',\n",
       " '<dis_2004>',\n",
       " '<dis_3395>',\n",
       " '<dis_4161>',\n",
       " ...]"
      ]
     },
     "execution_count": 75,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "llama_tokenizer.tokenize(image_text_dataset[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4495b402-95be-48c2-8829-f240ac883969",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}