feat: add functions

dev/inference/wandb-backend.ipynb

@@ -2,13 +2,15 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 197,
    "id": "4ff2a984-b8b2-4a69-89cf-0d16da2393c8",
    "metadata": {},
    "outputs": [],
    "source": [
     "import csv\n",
     "import tempfile\n",
+    "from functools import partial\n",
+    "import jax\n",
     "import wandb\n",
     "from dalle_mini.model import CustomFlaxBartForConditionalGeneration\n",
     "from vqgan_jax.modeling_flax_vqgan import VQModel\n",
@@ -42,26 +44,82 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 245,
    "id": "e57797ab-0b3a-4490-be58-03d8d1c23fe9",
    "metadata": {},
    "outputs": [],
    "source": [
     "with open('samples.csv', newline='', encoding='utf8') as f:\n",
-    "    reader = csv.reader(f)\n",
+    "    reader = csv.DictReader(f)\n",
+    "    samples = []\n",
     "    for row in reader:\n",
-    "        breakpoint()"
+    "        samples.append(row)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 246,
+   "id": "f75b2869-fc25-4f56-b937-e97bbb712ede",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "101"
+      ]
+     },
+     "execution_count": 246,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "len(samples)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 248,
+   "id": "2ea0b166-a20c-4d78-bffb-b792ca512d17",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "104"
+      ]
+     },
+     "execution_count": 248,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "samples_to_add = ['empty'] * (-len(samples) % 8)\n",
+    "samples.extend(samples_to_add)\n",
+    "len(samples)"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "a2c629e9-1a82-40c6-a260-ca1780c19a2e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "api = wandb.Api()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 204,
    "id": "3ffb1d09-bd1c-4f57-9ae5-3eda6f7d3a08",
    "metadata": {},
    "outputs": [],
    "source": [
+    "# TODO: iterate on runs\n",
     "wandb_run = wandb_runs[0]\n",
-    "api = wandb.Api()"
+    "functions_pmapped = False"
    ]
   },
   {
@@ -280,27 +338,30 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 207,
    "id": "d1cc9993-1bfc-4ec6-a004-c056189c42ac",
    "metadata": {},
    "outputs": [],
-   "source": []
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "43d2a99b-3501-4b30-b041-0fdeead12380",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "06472541-75f1-44e5-841f-a4a26a0493e3",
-   "metadata": {},
-   "outputs": [],
-   "source": []
+   "source": [
+    "# function to generate encoded images\n",
+    "# we should generate this function only once per run\n",
+    "if not functions_pmapped:\n",
+    "    @partial(jax.pmap, axis_name=\"batch\")\n",
+    "    def p_generate(tokenized_prompt, key, params):\n",
+    "        return model.generate(\n",
+    "            **tokenized_prompt,\n",
+    "            do_sample=True,\n",
+    "            num_beams=1,\n",
+    "            prng_key=key,\n",
+    "            params=params\n",
+    "        )\n",
+    "    \n",
+    "    @partial(jax.pmap, axis_name=\"batch\")\n",
+    "    def p_decode(indices, params):\n",
+    "        return vqgan.decode_code(indices, params=params)\n",
+    "    \n",
+    "    functions_pmapped = False"
+   ]
   },
   {
    "cell_type": "code",
@@ -323,7 +384,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b37c1714-d54b-479e-a9e8-740affc0de2c",
+   "id": "e79ac8f2-adc2-4a16-970c-dadcceadd566",
    "metadata": {},
    "outputs": [],
    "source": []