speed up chunking & add separator chunking

Author: rangehow

nano_graphrag/_op.py

@@ -4,6 +4,8 @@
 from typing import Union
 from collections import Counter, defaultdict
 
+import tiktoken
+
 from ._utils import (
     logger,
     clean_str,
@@ -28,26 +30,82 @@
 from .prompt import GRAPH_FIELD_SEP, PROMPTS
 
 
+
+
 def chunking_by_token_size(
-    content: str, overlap_token_size=128, max_token_size=1024, tiktoken_model="gpt-4o"
-):
-    tokens = encode_string_by_tiktoken(content, model_name=tiktoken_model)
-    results = []
-    for index, start in enumerate(
-        range(0, len(tokens), max_token_size - overlap_token_size)
+        tokens_list: list[int], doc_keys,tiktoken_model, overlap_token_size=128, max_token_size=1024,
     ):
-        chunk_content = decode_tokens_by_tiktoken(
-            tokens[start : start + max_token_size], model_name=tiktoken_model
-        )
-        results.append(
-            {
-                "tokens": min(max_token_size, len(tokens) - start),
-                "content": chunk_content.strip(),
-                "chunk_order_index": index,
-            }
-        )
-    return results
+        
+        results=[]
+        for index,tokens in enumerate(tokens_list):
+            chunk_token=[]
+            lengths=[]
+            for start in range(0, len(tokens), max_token_size - overlap_token_size):
+                
+                chunk_token.append(tokens[start : start + max_token_size])
+                lengths.append(min(max_token_size, len(tokens) - start))
+
+            # here somehow tricky, since the whole chunk tokens is list[list[list[int]]] for corpus(doc(chunk)),so it can't be decode entirely
+            chunk_token=tiktoken_model.decode_batch(chunk_token)
+            for i,chunk in enumerate(chunk_token):
+                
+                results.append(
+                    {
+                        "tokens": lengths[i],
+                        "content": chunk.strip(),
+                        "chunk_order_index": i,
+                        "full_doc_id":doc_keys[index],
+                    }
+                )
+
+        return results
+
+def chunking_by_seperators(tokens_list: list[int], doc_keys,tiktoken_model, overlap_token_size=128, max_token_size=1024 ):
+    from nano_graphrag._spliter import SeparatorSplitter
+
+    DEFAULT_SEPERATORS=[
+        # Paragraph separators
+        "\n\n",
+        "\r\n\r\n",
+        # Line breaks
+        "\n",
+        "\r\n",
+        # Sentence ending punctuation
+        "。",  # Chinese period
+        "．",  # Full-width dot
+        ".",  # English period
+        "！",  # Chinese exclamation mark
+        "!",  # English exclamation mark
+        "？",  # Chinese question mark
+        "?",  # English question mark
+        # Whitespace characters
+        " ",  # Space
+        "\t",  # Tab
+        "\u3000",  # Full-width space
+        # Special characters
+        "\u200b",  # Zero-width space (used in some Asian languages)
+    ]
+    
+    splitter=SeparatorSplitter(separators=[tiktoken_model.encode(s) for s in DEFAULT_SEPERATORS],chunk_size=max_token_size,chunk_overlap=overlap_token_size)
+    results=[]
+    for index,tokens in enumerate(tokens_list):
+        chunk_token=splitter.split_tokens(tokens)
+        lengths=[len(c) for c in chunk_token]
+
+        # here somehow tricky, since the whole chunk tokens is list[list[list[int]]] for corpus(doc(chunk)),so it can't be decode entirely
+        chunk_token=tiktoken_model.decode_batch(chunk_token)
+        for i,chunk in enumerate(chunk_token):
+            
+            results.append(
+                {
+                    "tokens": lengths[i],
+                    "content": chunk.strip(),
+                    "chunk_order_index": i,
+                    "full_doc_id":doc_keys[index],
+                }
+            )
 
+    return results
 
 async def _handle_entity_relation_summary(
     entity_or_relation_name: str,

nano_graphrag/_spliter.py

@@ -0,0 +1,121 @@
+from typing import List, Optional, Union, Literal
+
+class SeparatorSplitter:
+    def __init__(
+        self,
+        separators: Optional[List[List[int]]] = None,
+        keep_separator: Union[bool, Literal["start", "end"]] = "end",
+        chunk_size: int = 4000,
+        chunk_overlap: int = 200,
+        length_function: callable = len,
+    ):
+        self._separators = separators or [[10], [13, 10]]  # 默认使用换行符作为分隔符
+        self._keep_separator = keep_separator
+        self._chunk_size = chunk_size
+        self._chunk_overlap = chunk_overlap
+        self._length_function = length_function
+
+    def split_tokens(self, tokens: List[int]) -> List[List[int]]:
+        splits = self._split_tokens_with_separators(tokens)
+        return self._merge_splits(splits)
+
+    def _split_tokens_with_separators(self, tokens: List[int]) -> List[List[int]]:
+        splits = []
+        current_split = []
+        i = 0
+        while i < len(tokens):
+            separator_found = False
+            for separator in self._separators:
+                if tokens[i:i+len(separator)] == separator:
+                    if current_split:
+                        if self._keep_separator == "end":
+                            current_split.extend(separator)
+                            splits.append(current_split)
+                            current_split = []
+                        elif self._keep_separator == "start":
+                            splits.append(current_split)
+                            current_split = separator[:]
+                        else:
+                            splits.append(current_split)
+                            current_split = []
+                    elif self._keep_separator:
+                        current_split.extend(separator)
+                    i += len(separator)
+                    separator_found = True
+                    break
+            if not separator_found:
+                current_split.append(tokens[i])
+                i += 1
+        if current_split:
+            splits.append(current_split)
+        return [s for s in splits if s]
+
+    def _merge_splits(self, splits: List[List[int]]) -> List[List[int]]:
+        merged_splits = []
+        current_split = []
+        current_length = 0
+        separator = [] if self._keep_separator is False else self._separators[-1]
+
+        for split in splits:
+            if self._length_function(current_split) + self._length_function(split) <= self._chunk_size:
+                if current_split and separator:
+                    current_split.extend(separator)
+                current_split.extend(split)
+            else:
+                if current_split:
+                    merged_splits.append(current_split)
+                current_split = split
+            if self._length_function(current_split) >= self._chunk_size:
+                merged_splits.append(current_split)
+                current_split = []
+        if current_split:
+            merged_splits.append(current_split)
+
+        if self._chunk_overlap > 0:
+            return self._enforce_overlap(merged_splits)
+        return merged_splits
+
+    def _enforce_overlap(self, chunks: List[List[int]]) -> List[List[int]]:
+        new_chunks = []
+        for i, chunk in enumerate(chunks):
+            if i == 0:
+                new_chunks.append(chunk)
+            else:
+                overlap_tokens = chunks[i-1][-self._chunk_overlap:]
+                new_chunk = overlap_tokens + chunk
+                if self._length_function(new_chunk) > self._chunk_size:
+                    new_chunk = new_chunk[-self._chunk_size:]
+                new_chunks.append(new_chunk)
+        return new_chunks
+
+# EXAMPLE USAGE
+if __name__ == "__main__":
+    import tiktoken
+    tokenizer = tiktoken.encoding_for_model("gpt-4")
+
+    def tokenize(text: str) -> List[int]:
+        return tokenizer.encode(text)
+
+    def detokenize(tokens: List[int]) -> str:
+        return tokenizer.decode(tokens)
+    
+    # 创建splitter实例
+    splitter = SeparatorSplitter(
+        separators=[tokenize('\n'), tokenize('.')],  # 使用换行符和句号作为分隔符
+        chunk_size=5,
+        chunk_overlap=0,
+        keep_separator="end"
+    )
+
+    # 示例文本
+    text = "This is a sample text. It contains multiple sentences.\nSome sentences are short. Others are longer."
+    tokens = tokenize(text)
+
+    # 分割tokens
+    split_tokens = splitter.split_tokens(tokens)
+
+    print("Split tokens:")
+    for i, token_chunk in enumerate(split_tokens):
+        print(f"Chunk {i + 1}:")
+        print(detokenize(token_chunk))
+        print("---")

nano_graphrag/graphrag.py

@@ -5,6 +5,8 @@
 from functools import partial
 from typing import Callable, Dict, List, Optional, Type, Union, cast
 
+import tiktoken
+
 
 from ._llm import (
     gpt_4o_complete,
@@ -65,7 +67,7 @@ class GraphRAG:
     enable_naive_rag: bool = False
 
     # text chunking
-    chunk_func: Callable[[str, Optional[int], Optional[int], Optional[str]], List[Dict[str, Union[str, int]]]] = chunking_by_token_size
+    chunk_func: Callable[[str,List[str],tiktoken.Encoding, Optional[int], Optional[int], ], List[Dict[str, Union[str, int]]]] = chunking_by_token_size
     chunk_token_size: int = 1200
     chunk_overlap_token_size: int = 100
     tiktoken_model_name: str = "gpt-4o"
@@ -264,20 +266,22 @@ async def ainsert(self, string_or_strings):
 
             # ---------- chunking
             inserting_chunks = {}
-            for doc_key, doc in new_docs.items():
-                chunks = {
-                    compute_mdhash_id(dp["content"], prefix="chunk-"): {
-                        **dp,
-                        "full_doc_id": doc_key,
-                    }
-                    for dp in self.chunk_func(
-                        doc["content"],
-                        overlap_token_size=self.chunk_overlap_token_size,
-                        max_token_size=self.chunk_token_size,
-                        tiktoken_model=self.tiktoken_model_name,
-                    )
-                }
-                inserting_chunks.update(chunks)
+            
+
+            
+            new_docs_list=list(new_docs.items())
+            docs=[new_doc[1]["content"] for new_doc in new_docs_list]
+            doc_keys=[new_doc[0] for new_doc in new_docs_list]
+            
+            
+            ENCODER = tiktoken.encoding_for_model("gpt-4o")
+            tokens=ENCODER.encode_batch(docs,num_threads=16)
+            chunks=self.chunk_func(tokens,overlap_token_size=self.chunk_overlap_token_size,
+                        max_token_size=self.chunk_token_size,doc_keys=doc_keys,tiktoken_model=ENCODER)
+            for chunk in chunks:
+                inserting_chunks.update({compute_mdhash_id(chunk["content"], prefix="chunk-"):chunk})
+            
+
             _add_chunk_keys = await self.text_chunks.filter_keys(
                 list(inserting_chunks.keys())
             )