diff --git a/docs/examples/advanced_chunking_with_merging.ipynb b/docs/examples/advanced_chunking_with_merging.ipynb
index 9bd55c01..5c4d0446 100644
--- a/docs/examples/advanced_chunking_with_merging.ipynb
+++ b/docs/examples/advanced_chunking_with_merging.ipynb
@@ -40,7 +40,7 @@
     "from dataclasses import dataclass\n",
     "from pathlib import Path\n",
     "from tempfile import mkdtemp\n",
-    "from typing import Any, Iterator, Optional\n",
+    "from typing import Any, Iterator, Optional, Union\n",
     "\n",
     "import lancedb\n",
     "import semchunk\n",
@@ -94,7 +94,7 @@
     "    tokenizer: Any\n",
     "    max_tokens: PositiveInt\n",
     "\n",
-    "    def _count_tokens(self, text: Optional[list[str]]):\n",
+    "    def _count_tokens(self, text: Optional[Union[str, list[str]]]):\n",
     "        if text is None:\n",
     "            return 0\n",
     "        elif isinstance(text, list):\n",