From 6555f4e88e3c093dc2edcf067738cdc30b1d9fe4 Mon Sep 17 00:00:00 2001
From: Igor Ilic <igorilic03@gmail.com>
Date: Sun, 27 Oct 2024 22:33:20 +0100
Subject: [PATCH] fix: Resolve chunking issue for notebook

Add cleaning of local data to resolve chunking issue with repeated notebook use

Fix
---
 notebooks/cognee_demo.ipynb | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/notebooks/cognee_demo.ipynb b/notebooks/cognee_demo.ipynb
index ba5a89c86..8b8124490 100644
--- a/notebooks/cognee_demo.ipynb
+++ b/notebooks/cognee_demo.ipynb
@@ -577,6 +577,7 @@
                 "\n",
                 "import cognee\n",
                 "\n",
+                "await cognee.prune.prune_data()\n",
                 "await cognee.prune.prune_system(metadata=True)"
             ]
         },
@@ -639,7 +640,8 @@
                 "    chunks_into_graph, \\\n",
                 "    source_documents_to_chunks, \\\n",
                 "    check_permissions_on_documents, \\\n",
-                "    classify_documents\n",
+                "    classify_documents, \\\n",
+                "    chunk_naive_llm_classifier\n",
                 "from cognee.tasks.summarization import summarize_text\n",
                 "\n",
                 "async def run_cognify_pipeline(dataset: Dataset, user: User = None):\n",
@@ -667,6 +669,10 @@
                 "                    summarization_model = cognee_config.summarization_model,\n",
                 "                    collection_name = \"summaries\",\n",
                 "                ),\n",
+                "             Task(\n",
+                "                    chunk_naive_llm_classifier,\n",
+                "                    classification_model = cognee_config.classification_model,\n",
+                "            ),\n",
                 "            Task(chunk_remove_disconnected), # Remove the obsolete document chunks.\n",
                 "        ]\n",
                 "\n",
@@ -876,7 +882,7 @@
     ],
     "metadata": {
         "kernelspec": {
-            "display_name": "cognee-bGi0WgSG-py3.9",
+            "display_name": ".venv",
             "language": "python",
             "name": "python3"
         },
@@ -890,7 +896,7 @@
             "name": "python",
             "nbconvert_exporter": "python",
             "pygments_lexer": "ipython3",
-            "version": "3.9.5"
+            "version": "3.9.6"
         }
     },
     "nbformat": 4,