diff --git a/notebooks/stablelm-alpha.ipynb b/notebooks/stablelm-alpha.ipynb
index 2de9e1d..417f067 100644
--- a/notebooks/stablelm-alpha.ipynb
+++ b/notebooks/stablelm-alpha.ipynb
@@ -54,14 +54,20 @@
         "\n",
         "from IPython.display import Markdown, display\n",
         "def hr(): display(Markdown('---'))\n",
-        "def cprint(msg: str, color: str = \"blue\", **kwargs) -> str:\n",
-        "    if color == \"blue\": print(\"\\033[34m\" + msg + \"\\033[0m\", **kwargs)\n",
-        "    elif color == \"red\": print(\"\\033[31m\" + msg + \"\\033[0m\", **kwargs)\n",
-        "    elif color == \"green\": print(\"\\033[32m\" + msg + \"\\033[0m\", **kwargs)\n",
-        "    elif color == \"yellow\": print(\"\\033[33m\" + msg + \"\\033[0m\", **kwargs)\n",
-        "    elif color == \"purple\": print(\"\\033[35m\" + msg + \"\\033[0m\", **kwargs)\n",
-        "    elif color == \"cyan\": print(\"\\033[36m\" + msg + \"\\033[0m\", **kwargs)\n",
-        "    else: raise ValueError(f\"Invalid info color: `{color}`\")"
+        "def cprint(msg: str, color: str = \"blue\", **kwargs) -> None:\n",
+        "    color_codes = {\n",
+        "        \"blue\": \"\\033[34m\",\n",
+        "        \"red\": \"\\033[31m\",\n",
+        "        \"green\": \"\\033[32m\",\n",
+        "        \"yellow\": \"\\033[33m\",\n",
+        "        \"purple\": \"\\033[35m\",\n",
+        "        \"cyan\": \"\\033[36m\",\n",
+        "    }\n",
+        "    \n",
+        "    if color not in color_codes:\n",
+        "        raise ValueError(f\"Invalid info color: `{color}`\")\n",
+        "    \n",
+        "    print(color_codes[color] + msg + \"\\033[0m\", **kwargs)"
       ]
     },
     {