* Prepend [bos] to image encodings, rename to "labels".

model/data-pipeline.ipynb

@@ -161,7 +161,8 @@
    "source": [
     "# Setting padding=\"max_length\" as we need fixed length inputs for jitted functions\n",
     "max_length = 256   # Read from data_args.max_source_length\n",
-    "tokenizer = BartTokenizer.from_pretrained('facebook/bart-large-cnn')"
+    "tokenizer = BartTokenizer.from_pretrained('facebook/bart-large-cnn')\n",
+    "image_bos = 16384   # Max token is 16383 in our VQGAN configuration"
    ]
   },
   {
@@ -178,7 +179,7 @@
     "        inputs, max_length=max_length, padding=\"max_length\", truncation=True, return_tensors=\"np\"\n",
     "    )\n",
     "\n",
-    "    model_inputs[\"eval_encoding\"] = [eval(indices) for indices in examples['encoding']]\n",
+    "    model_inputs[\"labels\"] = [[image_bos] + eval(indices) for indices in examples['encoding']]\n",
     "\n",
     "    return model_inputs"
    ]
@@ -192,10 +193,10 @@
    "source": [
     "num_workers = 48     # We have 96 processors in the TPU\n",
     "column_names = dataset.column_names\n",
-    "dataset = dataset.map(preprocess_function,\n",
-    "                      remove_columns=column_names,\n",
-    "                      batched=True,\n",
-    "                      num_proc=48\n",
+    "input_dataset = dataset.map(preprocess_function,\n",
+    "                            remove_columns=column_names,\n",
+    "                            batched=True,\n",
+    "                            num_proc=48\n",
     ")"
    ]
   },
@@ -240,7 +241,7 @@
      "text": [
       "INFO:absl:Starting the local TPU driver.\n",
       "INFO:absl:Unable to initialize backend 'tpu_driver': Not found: Unable to find driver in registry given worker: local://\n",
-      "INFO:absl:Unable to initialize backend 'gpu': Not found: Could not find registered platform with name: \"cuda\". Available platform names are: Interpreter TPU Host\n"
+      "INFO:absl:Unable to initialize backend 'gpu': Not found: Could not find registered platform with name: \"cuda\". Available platform names are: Host TPU Interpreter\n"
      ]
     }
    ],
@@ -257,7 +258,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "loader = data_loader(rng, dataset, batch_size=super_batch_size)"
+    "loader = data_loader(rng, input_dataset, batch_size=super_batch_size)"
    ]
   },
   {
@@ -279,7 +280,7 @@
     {
      "data": {
       "text/plain": [
-       "dict_keys(['attention_mask', 'eval_encoding', 'input_ids'])"
+       "dict_keys(['attention_mask', 'input_ids', 'labels'])"
       ]
      },
      "execution_count": 13,
@@ -309,7 +310,7 @@
     }
    ],
    "source": [
-    "len(superbatch[\"eval_encoding\"])"
+    "len(superbatch[\"labels\"])"
    ]
   },
   {
@@ -321,7 +322,7 @@
     {
      "data": {
       "text/plain": [
-       "(8, 64, 256)"
+       "(8, 64, 257)"
       ]
      },
      "execution_count": 15,
@@ -330,15 +331,33 @@
     }
    ],
    "source": [
-    "superbatch[\"eval_encoding\"].shape"
+    "superbatch[\"labels\"].shape"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6800153b",
+   "metadata": {},
+   "source": [
+    "Any image sequence should begin with `image_bos`:"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 16,
    "id": "cfe23a71",
    "metadata": {},
    "outputs": [],
+   "source": [
+    "assert superbatch[\"labels\"][1][5][0].item() == image_bos"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0fb899b4",
+   "metadata": {},
+   "outputs": [],
    "source": []
   }
  ],