refactor: refact split_chunks

Author: ChenZiHong-Gavin

graphgen/graphgen.py

@@ -5,7 +5,6 @@
 from typing import Dict, cast
 
 import gradio as gr
-from tqdm.asyncio import tqdm as tqdm_async
 
 from graphgen.bases.base_storage import StorageNameSpace
 from graphgen.bases.datatypes import Chunk
@@ -18,21 +17,20 @@
     TraverseStrategy,
 )
 from graphgen.operators import (
+    chunk_documents,
     extract_kg,
     generate_cot,
     judge_statement,
     quiz,
     read_files,
     search_all,
-    split_chunks,
     traverse_graph_for_aggregated,
     traverse_graph_for_atomic,
     traverse_graph_for_multi_hop,
 )
 from graphgen.utils import (
     async_to_sync_method,
     compute_content_hash,
-    detect_main_language,
     format_generation_results,
     logger,
 )
@@ -110,7 +108,6 @@ async def insert(self):
         """
         insert chunks into the graph
         """
-
         input_file = self.config["read"]["input_file"]
 
         # Step 1: Read files
@@ -138,33 +135,7 @@ async def insert(self):
             return
         logger.info("[New Docs] inserting %d docs", len(new_docs))
 
-        cur_index = 1
-        doc_number = len(new_docs)
-        async for doc_key, doc in tqdm_async(
-            new_docs.items(), desc="[1/4]Chunking documents", unit="doc"
-        ):
-            doc_language = detect_main_language(doc["content"])
-            text_chunks = split_chunks(
-                doc["content"],
-                language=doc_language,
-                chunk_size=self.config["split"]["chunk_size"],
-                chunk_overlap=self.config["split"]["chunk_overlap"],
-            )
-
-            chunks = {
-                compute_content_hash(txt, prefix="chunk-"): {
-                    "content": txt,
-                    "full_doc_id": doc_key,
-                    "length": len(self.tokenizer_instance.encode_string(txt)),
-                    "language": doc_language,
-                }
-                for txt in text_chunks
-            }
-            inserting_chunks.update(chunks)
-
-            if self.progress_bar is not None:
-                self.progress_bar(cur_index / doc_number, f"Chunking {doc_key}")
-                cur_index += 1
+        inserting_chunks = await chunk_documents(new_docs)
 
         _add_chunk_keys = await self.text_chunks_storage.filter_keys(
             list(inserting_chunks.keys())
@@ -246,7 +217,7 @@ async def search(self):
                         ]
                     )
                 # TODO: fix insert after search
-                await self.async_insert()
+                await self.insert()
 
     @async_to_sync_method
     async def quiz(self):

graphgen/operators/__init__.py

@@ -5,7 +5,7 @@
 from .judge import judge_statement
 from .quiz import quiz
 from .read import read_files
-from .split import split_chunks
+from .split import chunk_documents
 from .traverse_graph import (
     traverse_graph_for_aggregated,
     traverse_graph_for_atomic,

graphgen/operators/split/__init__.py

@@ -1 +1 @@
-from .split_chunks import split_chunks
+from .split_chunks import chunk_documents

graphgen/operators/split/split_chunks.py

@@ -1,7 +1,10 @@
 from functools import lru_cache
 from typing import Union
 
+from tqdm.asyncio import tqdm as tqdm_async
+
 from graphgen.models import ChineseRecursiveTextSplitter, RecursiveCharacterSplitter
+from graphgen.utils import compute_content_hash, detect_main_language
 
 _MAPPING = {
     "en": RecursiveCharacterSplitter,
@@ -26,3 +29,44 @@ def split_chunks(text: str, language: str = "en", **kwargs) -> list:
         )
     splitter = _get_splitter(language, frozenset(kwargs.items()))
     return splitter.split_text(text)
+
+
+async def chunk_documents(
+    new_docs: dict,
+    chunk_size: int = 1024,
+    chunk_overlap: int = 100,
+    tokenizer_instance=None,
+    progress_bar=None,
+) -> dict:
+    inserting_chunks = {}
+    cur_index = 1
+    doc_number = len(new_docs)
+    async for doc_key, doc in tqdm_async(
+        new_docs.items(), desc="[1/4]Chunking documents", unit="doc"
+    ):
+        doc_language = detect_main_language(doc["content"])
+        text_chunks = split_chunks(
+            doc["content"],
+            language=doc_language,
+            chunk_size=chunk_size,
+            chunk_overlap=chunk_overlap,
+        )
+
+        chunks = {
+            compute_content_hash(txt, prefix="chunk-"): {
+                "content": txt,
+                "full_doc_id": doc_key,
+                "length": len(tokenizer_instance.encode_string(txt))
+                if tokenizer_instance
+                else len(txt),
+                "language": doc_language,
+            }
+            for txt in text_chunks
+        }
+        inserting_chunks.update(chunks)
+
+        if progress_bar is not None:
+            progress_bar(cur_index / doc_number, f"Chunking {doc_key}")
+            cur_index += 1
+
+    return inserting_chunks