gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor: Add summarization task to notebook

Browse source 
Added summarization task to jupyter notebook

Refactor #COG-387
This commit is contained in:
Igor Ilic 2024-10-09 17:12:32 +02:00
parent 3322891132
commit eab6bda2d3
1 changed files with 9 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
notebooks/cognee_demo.ipynb
View file

@ -563,6 +563,7 @@
"from cognee.shared.data_models import KnowledgeGraph\n",
"from cognee.modules.data.models import Dataset, Data\n",
"from cognee.modules.data.methods.get_dataset_data import get_dataset_data\n",
"from cognee.modules.cognify.config import get_cognify_config\n",
"from cognee.modules.pipelines.tasks.Task import Task\n",
"from cognee.modules.pipelines import run_tasks\n",
"from cognee.modules.users.models import User\n",
@ -574,6 +575,7 @@
" source_documents_to_chunks, \\\n",
" check_permissions_on_documents, \\\n",
" classify_documents\n",
"from cognee.tasks.summarization import summarize_text\n",
"\n",
"async def run_cognify_pipeline(dataset: Dataset, user: User = None):\n",
" data_documents: list[Data] = await get_dataset_data(dataset_id = dataset.id)\n",
@ -582,6 +584,8 @@
"\n",
" root_node_id = None\n",
"\n",
" cognee_config = get_cognify_config()\n",
"\n",
" tasks = [\n",
" Task(classify_documents),\n",
" Task(check_permissions_on_documents, user = user, permissions = [\"write\"]),\n",
@ -593,6 +597,11 @@
" save_chunks_to_store,\n",
" collection_name = \"chunks\",\n",
" ), \n",
" Task(\n",
" summarize_text,\n",
" summarization_model = cognee_config.summarization_model,\n",
" collection_name = \"summaries\",\n",
" ),\n",
" Task(chunk_remove_disconnected), # Remove the obsolete document chunks.\n",
" ]\n",
"\n",