diff --git a/run.ipynb b/run.ipynb
index 26b4b77..95c5b0f 100644
--- a/run.ipynb
+++ b/run.ipynb
@@ -13,15 +13,6 @@
         "More details can be found in the [README.md](README.md) ."
       ]
     },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "nX78K1Fi-38d"
-      },
-      "source": [
-        "## Clone Project"
-      ]
-    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -30,78 +21,73 @@
       },
       "outputs": [],
       "source": [
+        "#@title Clone Project\n",
+        "\n",
         "!git clone https://github.com/karpathy/llama2.c.git\n",
         "%cd llama2.c"
       ]
     },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "1MB5LGla-8Ln"
-      },
-      "source": [
-        "## Build"
-      ]
-    },
     {
       "cell_type": "code",
       "execution_count": null,
-      "metadata": {
-        "id": "V1EGferJv_7o"
-      },
+      "metadata": {},
       "outputs": [],
       "source": [
+        "#@title Build\n",
+        "\n",
         "!make run"
       ]
     },
     {
-      "cell_type": "markdown",
+      "cell_type": "code",
+      "execution_count": null,
       "metadata": {
-        "id": "MuMpJio8_AKi"
+        "id": "thm0ZBrtSgoC"
       },
+      "outputs": [],
       "source": [
-        "## Run"
+        "#@title Pick Your Model\n",
+        "\n",
+        "#@markdown Choose model\n",
+        "model = \"stories15M\" #@param [\"stories15M\", \"stories42M\", \"stories110M\"]\n",
+        "\n",
+        "download_url = \"\"\n",
+        "\n",
+        "if(model == \"stories15M\"):\n",
+        "  download_url = \"https://huggingface.co/karpathy/tinyllamas/resolve/main/stories15M.bin\"\n",
+        "if(model == \"stories42M\"):\n",
+        "  download_url = \"https://huggingface.co/karpathy/tinyllamas/resolve/main/stories42M.bin\"\n",
+        "if(model == \"stories110M\"):\n",
+        "  download_url = \"https://huggingface.co/karpathy/tinyllamas/resolve/main/stories110M.bin\"\n",
+        "\n",
+        "print(f\"download_url: {download_url}\")\n",
+        "\n",
+        "!wget $download_url\n",
+        "\n",
+        "model_file = model + \".bin\""
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "MRy23xxavJNO"
+        "id": "OgAc3KjuT-NM"
       },
       "outputs": [],
       "source": [
-        "# run stories15M\n",
-        "# !wget https://huggingface.co/karpathy/tinyllamas/resolve/main/stories15M.bin\n",
-        "# !./run stories110M.bin\n",
+        "#@title Generate Stories\n",
         "\n",
-        "# run stories42M\n",
-        "# !wget https://huggingface.co/karpathy/tinyllamas/resolve/main/stories42M.bin\n",
-        "# !./run stories42M.bin\n",
+        "# Generate args\n",
+        "max_token = 256 #@param {type:\"slider\", min:32, max:1024, step:32}\n",
+        "temperature = 0.8 #@param {type:\"slider\", min:0.0, max:1, step:0.05}\n",
+        "top_p = 0.9 #@param {type:\"slider\", min:0.0, max:1.0, step:0.05}\n",
+        "prompt = \"One day, Lily met a Shoggoth\" #@param {type:\"string\"}\n",
         "\n",
-        "# run stories110M\n",
-        "!wget https://huggingface.co/karpathy/tinyllamas/resolve/main/stories110M.bin\n",
-        "!./run stories110M.bin"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "Wi48eJKI_FKO"
-      },
-      "source": [
-        "## Run with args"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "YFOv6U74vSSZ"
-      },
-      "outputs": [],
-      "source": [
-        "!./run stories110M.bin -t 0.8 -n 256 -i \"One day, Lily met a Shoggoth\""
+        "print(f\"model: {model_file}, max_token: {max_token}, temperature: {temperature}, top_p: {top_p}, prompt: {prompt}\")\n",
+        "print(f\"----------------------------\\n\")\n",
+        "\n",
+        "cmd = f'./run {model_file} -t {temperature} -p {top_p} -n {max_token} -i \"{prompt}\"'\n",
+        "!{cmd}"
       ]
     }
   ],