|
| 1 | +import copy |
| 2 | +import re |
| 3 | +from abc import ABC, abstractmethod |
| 4 | +from dataclasses import dataclass |
| 5 | +from typing import Callable, Iterable, List, Literal, Optional, Union |
| 6 | + |
| 7 | +from graphgen.bases.datatypes import Chunk |
| 8 | +from graphgen.utils import logger |
| 9 | + |
| 10 | + |
| 11 | +@dataclass |
| 12 | +class BaseSplitter(ABC): |
| 13 | + """ |
| 14 | + Abstract base class for splitting text into smaller chunks. |
| 15 | + """ |
| 16 | + |
| 17 | + chunk_size: int = 1024 |
| 18 | + chunk_overlap: int = 100 |
| 19 | + length_function: Callable[[str], int] = len |
| 20 | + keep_separator: bool = False |
| 21 | + add_start_index: bool = False |
| 22 | + strip_whitespace: bool = True |
| 23 | + |
| 24 | + @abstractmethod |
| 25 | + def split_text(self, text: str) -> List[str]: |
| 26 | + """ |
| 27 | + Split the input text into smaller chunks. |
| 28 | +
|
| 29 | + :param text: The input text to be split. |
| 30 | + :return: A list of text chunks. |
| 31 | + """ |
| 32 | + |
| 33 | + def create_chunks( |
| 34 | + self, texts: List[str], metadatas: Optional[List[dict]] = None |
| 35 | + ) -> List[Chunk]: |
| 36 | + """Create chunks from a list of texts.""" |
| 37 | + _metadatas = metadatas or [{}] * len(texts) |
| 38 | + chunks = [] |
| 39 | + for i, text in enumerate(texts): |
| 40 | + index = 0 |
| 41 | + previous_chunk_len = 0 |
| 42 | + for chunk in self.split_text(text): |
| 43 | + metadata = copy.deepcopy(_metadatas[i]) |
| 44 | + if self.add_start_index: |
| 45 | + offset = index + previous_chunk_len - self.chunk_overlap |
| 46 | + index = text.find(chunk, max(0, offset)) |
| 47 | + metadata["start_index"] = index |
| 48 | + previous_chunk_len = len(chunk) |
| 49 | + new_chunk = Chunk(content=chunk, metadata=metadata) |
| 50 | + chunks.append(new_chunk) |
| 51 | + return chunks |
| 52 | + |
| 53 | + def _join_chunks(self, chunks: List[str], separator: str) -> Optional[str]: |
| 54 | + text = separator.join(chunks) |
| 55 | + if self.strip_whitespace: |
| 56 | + text = text.strip() |
| 57 | + if text == "": |
| 58 | + return None |
| 59 | + return text |
| 60 | + |
| 61 | + def _merge_splits(self, splits: Iterable[str], separator: str) -> List[str]: |
| 62 | + # We now want to combine these smaller pieces into medium size chunks to send to the LLM. |
| 63 | + separator_len = self.length_function(separator) |
| 64 | + |
| 65 | + chunks = [] |
| 66 | + current_chunk: List[str] = [] |
| 67 | + total = 0 |
| 68 | + for d in splits: |
| 69 | + _len = self.length_function(d) |
| 70 | + if ( |
| 71 | + total + _len + (separator_len if len(current_chunk) > 0 else 0) |
| 72 | + > self.chunk_size |
| 73 | + ): |
| 74 | + if total > self.chunk_size: |
| 75 | + logger.warning( |
| 76 | + "Created a chunk of size %s, which is longer than the specified %s", |
| 77 | + total, |
| 78 | + self.chunk_size, |
| 79 | + ) |
| 80 | + if len(current_chunk) > 0: |
| 81 | + chunk = self._join_chunks(current_chunk, separator) |
| 82 | + if chunk is not None: |
| 83 | + chunks.append(chunk) |
| 84 | + # Keep on popping if: |
| 85 | + # - we have a larger chunk than in the chunk overlap |
| 86 | + # - or if we still have any chunks and the length is long |
| 87 | + while total > self.chunk_overlap or ( |
| 88 | + total + _len + (separator_len if len(current_chunk) > 0 else 0) |
| 89 | + > self.chunk_size |
| 90 | + and total > 0 |
| 91 | + ): |
| 92 | + total -= self.length_function(current_chunk[0]) + ( |
| 93 | + separator_len if len(current_chunk) > 1 else 0 |
| 94 | + ) |
| 95 | + current_chunk = current_chunk[1:] |
| 96 | + current_chunk.append(d) |
| 97 | + total += _len + (separator_len if len(current_chunk) > 1 else 0) |
| 98 | + chunk = self._join_chunks(current_chunk, separator) |
| 99 | + if chunk is not None: |
| 100 | + chunks.append(chunk) |
| 101 | + return chunks |
| 102 | + |
| 103 | + @staticmethod |
| 104 | + def _split_text_with_regex( |
| 105 | + text: str, separator: str, keep_separator: Union[bool, Literal["start", "end"]] |
| 106 | + ) -> List[str]: |
| 107 | + # Now that we have the separator, split the text |
| 108 | + if separator: |
| 109 | + if keep_separator: |
| 110 | + # The parentheses in the pattern keep the delimiters in the result. |
| 111 | + _splits = re.split(f"({separator})", text) |
| 112 | + splits = ( |
| 113 | + ( |
| 114 | + [ |
| 115 | + _splits[i] + _splits[i + 1] |
| 116 | + for i in range(0, len(_splits) - 1, 2) |
| 117 | + ] |
| 118 | + ) |
| 119 | + if keep_separator == "end" |
| 120 | + else ( |
| 121 | + [_splits[i] + _splits[i + 1] for i in range(1, len(_splits), 2)] |
| 122 | + ) |
| 123 | + ) |
| 124 | + if len(_splits) % 2 == 0: |
| 125 | + splits += _splits[-1:] |
| 126 | + splits = ( |
| 127 | + (splits + [_splits[-1]]) |
| 128 | + if keep_separator == "end" |
| 129 | + else ([_splits[0]] + splits) |
| 130 | + ) |
| 131 | + else: |
| 132 | + splits = re.split(separator, text) |
| 133 | + else: |
| 134 | + splits = list(text) |
| 135 | + return [s for s in splits if s != ""] |
0 commit comments