notebooks/sample.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipywidgets as widgets\n",
    "import torch\n",
    "\n",
    "from discrete_flow_matching_pytorch.model import DiscreteFlowMatchingNet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DiscreteFlowMatchingNet(\n",
       "  (input_projection): Embedding(50259, 768)\n",
       "  (embed_timestep): Sequential(\n",
       "    (0): Embedding(1024, 768)\n",
       "    (1): Unsqueeze()\n",
       "  )\n",
       "  (blocks): ModuleList(\n",
       "    (0-5): 6 x Sequential(\n",
       "      (0): Transpose()\n",
       "      (1): Conv1d(768, 768, kernel_size=(31,), stride=(1,), padding=same)\n",
       "      (2): Transpose()\n",
       "      (3): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (4): GELU(approximate='none')\n",
       "      (5): Linear(in_features=768, out_features=768, bias=True)\n",
       "      (6): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (7): GELU(approximate='none')\n",
       "    )\n",
       "  )\n",
       "  (timestep_embedding_norms): ModuleList(\n",
       "    (0-5): 6 x LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "  )\n",
       "  (output_projection): Linear(in_features=768, out_features=50259, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = DiscreteFlowMatchingNet.load_from_checkpoint(\n",
    "    \"../flow-matching-tiny-stories/1pkkmee9/checkpoints/epoch=0-step=3000.ckpt\"\n",
    ")\n",
    "model.freeze()\n",
    "model.eval()\n",
    "model.to(dtype=torch.bfloat16, device=\"cuda:0\")\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "step_tokens: list[list[str]] = []\n",
    "step_texts: list[str] = []\n",
    "with torch.inference_mode():\n",
    "    for t, samples in model.sample(\n",
    "        num_samples=1,\n",
    "        sequence_length=128,\n",
    "        num_sampling_steps=128,\n",
    "        stochasticity=5.0,\n",
    "        yield_intermediate=True,\n",
    "    ):\n",
    "        step_tokens.append(model.tokenizer.batch_decode(samples[0]))\n",
    "        step_texts.append(model.tokenizer.decode(samples[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "30ad5d91ccda466da693e68c74289aaa",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(IntSlider(value=64, description='step', max=128), Output()), _dom_classes=('widget-inter…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def get_step_text(step):\n",
    "    rows = []\n",
    "    column_size = 16\n",
    "    for i in range(0, len(step_tokens[step]), column_size):\n",
    "        rows.append(step_tokens[step][i : i + column_size])\n",
    "\n",
    "    # Create html table\n",
    "    html_table = '<table style=\"border: 1px solid black\">'\n",
    "    for row in rows:\n",
    "        html_table += \"<tr>\"\n",
    "        for token in row:\n",
    "            style = \"border: 1px solid black; width: 50px; text-align: center;\"\n",
    "            if token == \"[MASK]\":\n",
    "                style += \"background-color: #cccccc;\"\n",
    "            else:\n",
    "                style += \"background-color: #eeeeee;\"\n",
    "            html_table += f'<td style=\"{style}\">{token}</td>'\n",
    "        html_table += \"</tr>\"\n",
    "    html_table += \"</table>\"\n",
    "\n",
    "    return widgets.HTML(html_table)\n",
    "\n",
    "\n",
    "interact_widget = widgets.interact(get_step_text, step=(0, len(step_tokens) - 1, 1))\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}