feat: create a table

dev/inference/wandb-backend.ipynb

@@ -46,7 +46,8 @@
     "batch_size = 8\n",
     "num_images = 128\n",
     "top_k = 8\n",
-    "text_normalizer = TextNormalizer() if normalize_text else None"
+    "text_normalizer = TextNormalizer() if normalize_text else None\n",
+    "padding_item = 'NONE'"
    ]
   },
   {
@@ -95,8 +96,8 @@
     "    samples = []\n",
     "    for row in reader:\n",
     "        samples.append(row)\n",
-    "    # make list multiple of batch_size by adding \"empty\"\n",
-    "    samples_to_add = [{'Caption':'empty', 'Theme':'empty'}] * (-len(samples) % batch_size)\n",
+    "    # make list multiple of batch_size by adding elements\n",
+    "    samples_to_add = [{'Caption':padding_item, 'Theme':padding_item}] * (-len(samples) % batch_size)\n",
     "    samples.extend(samples_to_add)\n",
     "    # reshape\n",
     "    samples = [samples[i:i+batch_size] for i in range(0, len(samples), batch_size)]"
@@ -388,7 +389,6 @@
     "    def p_clip(inputs):\n",
     "        logits = clip(**inputs).logits_per_image\n",
     "        return logits\n",
-    "        scores = jax.nn.softmax(logits, axis=0).squeeze()        \n",
     "    \n",
     "    functions_pmapped = False"
    ]
@@ -649,7 +649,8 @@
    "outputs": [],
    "source": [
     "results = []\n",
-    "columns = ['Caption', 'Theme'] + [f'Image {i+1}' for i in range(top_k)] + [f'Score {i+1}' for i in range(top_k)]"
+    "columns = ['Caption', 'Theme'] + [f'Image {i+1}' for i in range(top_k)] + [f'Score {i+1}' for i in range(top_k)]\n",
+    "logits = jax.device_get(logits)"
    ]
   },
   {
@@ -660,12 +661,23 @@
    "outputs": [],
    "source": [
     "for i, (idx, scores, sample) in enumerate(zip(top_idx, logits, batch)):\n",
+    "    if sample['Caption'] == padding_item: continue\n",
     "    cur_images = [images[x] for x in images_per_prompt_indices + i]\n",
     "    top_images = [wandb.Image(cur_images[x]) for x in idx]\n",
-    "    top_scores = [logits[x] for x in idx]\n",
+    "    top_scores = [scores[x] for x in idx]\n",
     "    results.append([sample['Caption'], sample['Theme']] + top_images + top_scores)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4bf40461-99d3-4d36-b7cc-e0129a3c9053",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "table = wandb.Table(columns=columns, data=results)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,