diff --git a/src/knowledge/core_object/document_object.py b/src/knowledge/core_object/document_object.py index 672b4a3..b69f59f 100644 --- a/src/knowledge/core_object/document_object.py +++ b/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() diff --git a/src/knowledge/data/writer.py b/src/knowledge/data/writer.py index 93b9ecc..0a87ac1 100644 --- a/src/knowledge/data/writer.py +++ b/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() @@ -70,27 +212,4 @@ class ChunkListWriter: self.chunk_store.put_chunk(chunk_id, chunk_bytes) 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 \ No newline at end of file + return ChunkList(chunk_list, hash) \ No newline at end of file diff --git a/test/test_chunk.py b/test/test_chunk.py index f1f496b..0771f03 100644 --- a/test/test_chunk.py +++ b/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__":