impl text splitter with tiktoken

src/knowledge/core_object/document_object.py

@@ -50,11 +50,7 @@ class DocumentObjectBuilder:
         return self
 
     def build(self) -> DocumentObject:
-        chunk_list = KnowledgeStore().get_chunk_list_writer().create_chunk_list_from_text(
-            self.text,
-            1024 * 4,
-            ".?!\n"
-        )
+        chunk_list = KnowledgeStore().get_chunk_list_writer().create_chunk_list_from_text(self.text)
         doc = DocumentObject(self.meta, self.tags, chunk_list)
         doc_id = doc.calculate_id()
         

src/knowledge/data/writer.py

@@ -1,6 +1,8 @@
 import os
 import hashlib
 import re
+import tiktoken
+import logging
 from typing import Tuple, List
 from .chunk_store import ChunkStore
 from .chunk import ChunkID, PositionFileRange, PositionType
@@ -8,6 +10,131 @@ from ..object import HashValue
 from .tracker import ChunkTracker
 from .chunk_list import ChunkList
 
+def _join_docs(self, docs: List[str], separator: str) -> Optional[str]:
+    text = separator.join(docs)
+    text = text.strip()
+    if text == "":
+        return None
+    else:
+        return text
+
+def _merge_splits(
+        self, 
+        splits: Iterable[str], 
+        separator: str, 
+        chunk_size: int, 
+        chunk_overlap: int,  
+        length_function: Callable[[str], int]
+    ) -> List[str]:
+    # We now want to combine these smaller pieces into medium size
+    # chunks to send to the LLM.
+    separator_len = length_function(separator)
+
+    docs = []
+    current_doc: List[str] = []
+    total = 0
+    for d in splits:
+        _len = length_function(d)
+        if (
+            total + _len + (separator_len if len(current_doc) > 0 else 0)
+            > chunk_size
+        ):
+            if total > chunk_size:
+                logging.warning(
+                    f"Created a chunk of size {total}, "
+                    f"which is longer than the specified {self._chunk_size}"
+                )
+            if len(current_doc) > 0:
+                doc = _join_docs(current_doc, separator)
+                if doc is not None:
+                    docs.append(doc)
+                # Keep on popping if:
+                # - we have a larger chunk than in the chunk overlap
+                # - or if we still have any chunks and the length is long
+                while total > chunk_overlap or (
+                    total + _len + (separator_len if len(current_doc) > 0 else 0)
+                    > chunk_size
+                    and total > 0
+                ):
+                    total -= length_function(current_doc[0]) + (
+                        separator_len if len(current_doc) > 1 else 0
+                    )
+                    current_doc = current_doc[1:]
+        current_doc.append(d)
+        total += _len + (separator_len if len(current_doc) > 1 else 0)
+    doc = _join_docs(current_doc, separator)
+    if doc is not None:
+        docs.append(doc)
+    return docs
+
+
+def _split_text_with_regex(
+    text: str, separator: str, keep_separator: bool
+) -> List[str]:
+    # Now that we have the separator, split the text
+    if separator:
+        if keep_separator:
+            # The parentheses in the pattern keep the delimiters in the result.
+            _splits = re.split(f"({separator})", text)
+            splits = [_splits[i] + _splits[i + 1] for i in range(1, len(_splits), 2)]
+            if len(_splits) % 2 == 0:
+                splits += _splits[-1:]
+            splits = [_splits[0]] + splits
+        else:
+            splits = re.split(separator, text)
+    else:
+        splits = list(text)
+    return [s for s in splits if s != ""]
+
+
+def _split_text(
+        text: str, 
+        separators: List[str], 
+        chunk_size: int, 
+        chunk_overlap: int, 
+        length_function: Callable[[str], int]
+    ) -> List[str]:
+
+    """Split incoming text and return chunks."""
+    final_chunks = []
+    # Get appropriate separator to use
+    separator = separators[-1]
+    new_separators = []
+    for i, _s in enumerate(separators):
+        _separator = re.escape(_s)
+        if _s == "":
+            separator = _s
+            break
+        if re.search(_separator, text):
+            separator = _s
+            new_separators = separators[i + 1 :]
+            break
+
+    keep_separator = True
+    _separator = re.escape(separator)
+    splits = _split_text_with_regex(text, _separator, keep_separator)
+
+    # Now go merging things, recursively splitting longer texts.
+    _good_splits = []
+    _separator = "" if keep_separator else separator
+    for s in splits:
+        if length_function(s) < chunk_size:
+            _good_splits.append(s)
+        else:
+            if _good_splits:
+                merged_text = _merge_splits(_good_splits, _separator, chunk_size, chunk_overlap, length_function)
+                final_chunks.extend(merged_text)
+                _good_splits = []
+            if not new_separators:
+                final_chunks.append(s)
+            else:
+                other_info = _split_text(s, new_separators, chunk_size, chunk_overlap, length_function)
+                final_chunks.extend(other_info)
+    if _good_splits:
+        merged_text = _merge_splits(_good_splits, _separator, chunk_size, chunk_overlap, length_function)
+        final_chunks.extend(merged_text)
+    return final_chunks
+
 class ChunkListWriter:
     def __init__(self, chunk_store: ChunkStore, chunk_tracker: ChunkTracker):
         self.chunk_store = chunk_store
@@ -54,9 +181,24 @@ class ChunkListWriter:
         return ChunkList(chunk_list, file_hash)
 
     def create_chunk_list_from_text(
-        self, text: str, chunk_max_words: int, separator_chars: str = ".,"
+        self, 
+        text: str, 
+        chunk_size: int = 4000, 
+        chunk_overlap: int = 200, 
+        separators: str = ["\n\n", "\n", " ", ""]
     ) -> ChunkList:
-        text_list = self._split_text_list(text, chunk_max_words, separator_chars)
+        enc = tiktoken.encoding_for_model("gpt-3.5-turbo")
+
+        def length_function(text: str) -> int:
+            return len(
+                enc.encode(
+                    text,
+                    allowed_special=set(),
+                    disallowed_special="all",
+                )
+            )
+        
+        text_list = _split_text(text, separators, chunk_size, chunk_overlap, length_function)
         chunk_list = []
         hash_obj = hashlib.sha256()
 
@@ -71,26 +213,3 @@ class ChunkListWriter:
 
         hash = HashValue(hash_obj.digest())
         return ChunkList(chunk_list, hash)
-
-    @staticmethod
-    def _split_text_list(
-        text: str, chunk_max_words: int, separator_chars: str = ".,"
-    ) -> List[str]:
-        sentences = re.split(f"[{separator_chars}]", text)
-        # chunk_list = []
-        # chunk = []
-        # word_count = 0
-        # for sentence in sentences:
-        #     words = sentence.split()
-        #     for word in words:
-        #         if word_count < chunk_max_words:
-        #             chunk.append(word)
-        #             word_count += 1
-        #         else:
-        #             chunk_list.append(" ".join(chunk))
-        #             chunk = [word]
-        #             word_count = 1
-        # if chunk:
-        #     chunk_list.append(" ".join(chunk))
-        # return chunk_list
-        return sentences

test/test_chunk.py

@@ -59,7 +59,7 @@ class TestChunk(unittest.TestCase):
         with open(text_file, "r", encoding="utf-8") as file:
             text = file.read()
 
-        gen.create_chunk_list_from_text(text, 1024)
+        gen.create_chunk_list_from_text(text)
 
 
 if __name__ == "__main__":