Refactor tutorial to use dataclass for configuration (jax-ml#119)

* Refactor tutorial to use dataclass for configuration

* Imeplemented KV caching (WIP)

* final changes

Author: Aatman09

bonsai/models/unet/tests/UNet_segmentation_example.ipynb

@@ -629,7 +629,7 @@
     "    return state, loss\n",
     "\n",
     "\n",
-    "print(\"🚀 Starting training from checkpoint...\")\n",
+    "print(\"Starting training from checkpoint...\")\n",
     "train_loader, vis_loader = load_dataset()\n",
     "num_epochs = 100\n",
     "state = train_state\n",

bonsai/models/unet/tests/UNet_segmentation_example.md

@@ -256,7 +256,7 @@ def train_step(state: TrainState, other_vars: nnx.State, batch: tuple[jax.Array,
     return state, loss
 
 
-print("🚀 Starting training from checkpoint...")
+print("Starting training from checkpoint...")
 train_loader, vis_loader = load_dataset()
 num_epochs = 100
 state = train_state