Update splitter unit test and fix issue#227

SayaZhang · SayaZhang · commit 50bd4275e9ba · 2024-04-10T12:43:36.000Z
diff --git a/run_tests.sh b/run_tests.sh
@@ -1,2 +1,2 @@
 #!/bin/sh
-python -m unittest discover tests
+python3 -m unittest discover tests
diff --git a/tests/op/extract/split/test_pattern_splitter_op.py b/tests/op/extract/split/test_pattern_splitter_op.py
@@ -6,9 +6,12 @@
 
 class TestPatternSplitter(unittest.TestCase):
     def setUp(self):
-        self.splitter = PatternSplitter("test_splitter")
+        self.splitter = PatternSplitter({}, "test_splitter")
 
     def test_special_function_call(self):
+        """
+        Test special function call.
+        """
         node = Node(name="node1", value_dict={"text": "Hello\n\nWorld"})
 
         output_nodes = self.splitter([node])
@@ -17,7 +20,10 @@ def test_special_function_call(self):
         self.assertEqual(output_nodes[0].value_dict["text"], ["Hello", "World"])
 
     def test_special_function_call_with_custom_splitter(self):
-        splitter = PatternSplitter("test_splitter", splitter=" ")
+        splitter = PatternSplitter(
+            {"separators": " "},
+            "test_splitter",
+        )
         node = Node(name="node1", value_dict={"text": "Hello World"})
 
         output_nodes = splitter([node])
diff --git a/tests/op/extract/split/test_recursive_character_splitter.py b/tests/op/extract/split/test_recursive_character_splitter.py
@@ -8,101 +8,98 @@
 
 class TestRecursiveCharacterSplitter(unittest.TestCase):
     def setUp(self):
-        self.splitter = RecursiveCharacterSplitter("test_splitter", chunk_size=10)
-        self.default_separators = ["\n\n", "\n", " ", ""]
+        self.default_separators = ["\n\n", "\n"]
 
     def test_recursive_splitter(self):
+        splitter = RecursiveCharacterSplitter({"max_chunk_size": 10}, "test_splitter")
         text = "Hello\n\nWorld."
 
-        chunks = self.splitter._recursive_splitter(text, self.default_separators)
+        chunks = splitter._recursive_splitter(text, splitter.default_separators)
 
         self.assertEqual(chunks, ["Hello", "World."])
 
-    def test_recursive_splitter_with_merge_chunk(self):
-        splitter = RecursiveCharacterSplitter("test_splitter", chunk_size=100)
-        text = "Hello\n\nWorld"
+    def test_merge_splits(self):
+        splits = ["Hello", "World"]
+        splitter = RecursiveCharacterSplitter({"max_chunk_size": 20}, "test_splitter")
 
-        chunks = splitter._recursive_splitter(text, self.default_separators)
+        merged = splitter._merge_splits(splits, "\n")
 
-        self.assertEqual(chunks, ["HelloWorld"])
+        self.assertEqual(merged, ["Hello\nWorld"])
 
-    def test_recursive_splitter_with_small_chunk_size(self):
-        splitter = RecursiveCharacterSplitter("test_splitter", chunk_size=1)
-        text = "Hello\n\nWorld"
-        expected_chunks = ["H", "e", "l", "l", "o", "W", "o", "r", "l", "d"]
+    def test_recursive_splitter_with_merge_chunk(self):
+        splitter = RecursiveCharacterSplitter({"max_chunk_size": 20}, "test_splitter")
+        node = Node(name="node1", value_dict={"text": "Hello World"})
 
-        chunks = splitter._recursive_splitter(text, self.default_separators)
+        output_nodes = splitter([node])
 
-        self.assertEqual(chunks, expected_chunks)
+        self.assertEqual(len(output_nodes), 1)
+        self.assertEqual(output_nodes[0].value_dict["text"], ["Hello\\ World"])
 
-    def test_recursive_splitter_with_zero_chunk_size(self):
-        splitter = RecursiveCharacterSplitter("test_splitter", chunk_size=0)
-        text = "Hello\n\nWorld"
+    def test_recursive_splitter_with_small_chunk_size(self):
+        splitter = RecursiveCharacterSplitter(
+            {"max_chunk_size": 1, "chunk_overlap_size": 0}, "test_splitter"
+        )
+        node = Node(name="node1", value_dict={"text": "Hello\n\nWorld"})
         expected_chunks = ["H", "e", "l", "l", "o", "W", "o", "r", "l", "d"]
 
-        chunks = splitter._recursive_splitter(text, self.default_separators)
-
-        self.assertEqual(chunks, expected_chunks)
-
-    def test_recursive_splitter_with_no_separators(self):
-        text = "Hello\n\nWorld"
-        separators = []
+        output_nodes = splitter([node])
 
-        chunks = self.splitter._recursive_splitter(text, separators)
-
-        self.assertEqual(chunks, [])
+        self.assertEqual(len(output_nodes), 1)
+        self.assertEqual(output_nodes[0].value_dict["text"], expected_chunks)
 
     def test_recursive_splitter_with_no_split(self):
-        text = "HelloWorld"
+        splitter = RecursiveCharacterSplitter({"max_chunk_size": 10}, "test_splitter")
+        node = Node(name="node1", value_dict={"text": "HelloWorld"})
 
-        chunks = self.splitter._recursive_splitter(text, self.default_separators)
+        output_nodes = splitter([node])
 
-        self.assertEqual(chunks, ["HelloWorld"])
+        self.assertEqual(len(output_nodes), 1)
+        self.assertEqual(output_nodes[0].value_dict["text"], ["HelloWorld"])
 
     def test_recursive_splitter_with_custom_separators(self):
-        text = "Hello--World."
-        separators = ["-", " "]
+        splitter = RecursiveCharacterSplitter(
+            {"max_chunk_size": 10, "separators": "--"}, "test_splitter"
+        )
+        node = Node(name="node1", value_dict={"text": "Hello--World"})
 
-        chunks = self.splitter._recursive_splitter(text, separators)
+        output_nodes = splitter([node])
 
-        self.assertEqual(chunks, ["Hello", "World."])
+        self.assertEqual(len(output_nodes), 1)
+        self.assertEqual(output_nodes[0].value_dict["text"], ["Hello", "World"])
 
     def test_recursive_splitter_with_large_text_default_chunk(self):
-        text = "Hello\n\nWorld\n\n" * 100
+        splitter = RecursiveCharacterSplitter({"max_chunk_size": 20}, "test_splitter")
+        node = Node(name="node1", value_dict={"text": "Hello\n\nWorld\n\n" * 100})
 
-        chunks = self.splitter._recursive_splitter(text, self.default_separators)
+        output_nodes = splitter([node])
 
-        self.assertEqual(len(chunks), 100)
+        self.assertEqual(len(output_nodes), 1)
+        self.assertEqual(len(output_nodes[0].value_dict["text"]), 100)
 
     def test_recursive_splitter_with_large_text_large_chunk(self):
-        splitter = RecursiveCharacterSplitter("test_splitter", chunk_size=9999)
-        text = "Hello\n\nWorld\n\n" * 100
+        splitter = RecursiveCharacterSplitter({"max_chunk_size": 9999}, "test_splitter")
+        node = Node(name="node1", value_dict={"text": "Hello\n\nWorld\n\n" * 100})
 
-        chunks = splitter._recursive_splitter(text, self.default_separators)
-
-        self.assertEqual(len(chunks), 1)
-        self.assertEqual(chunks, ["HelloWorld" * 100])
-
-    def test_special_function_call(self):
-        node = Node(name="node1", value_dict={"text": "Hello\n\nWorld"})
-        output_nodes = self.splitter([node])
+        output_nodes = splitter([node])
 
         self.assertEqual(len(output_nodes), 1)
-        self.assertEqual(output_nodes[0].value_dict["text"], ["HelloWorld"])
+        self.assertEqual(len(output_nodes[0].value_dict["text"]), 1)
 
     def test_special_function_call_with_multiple_nodes(self):
+        splitter = RecursiveCharacterSplitter({"max_chunk_size": 10}, "test_splitter")
+
         node0 = Node(name="node1", value_dict={"text": "Hello\n\nWorld"})
         node1 = Node(name="node1", value_dict={"text": "Hello\n\nWorld."})
         node2 = Node(name="node1", value_dict={"text": "Hello\n\nWorld\n\n" * 10})
         node3 = Node(name="node1", value_dict={"text": "Hello\n\nWorld.\n\n" * 2})
         expected_texts = [
-            ["HelloWorld"],
+            ["Hello", "World"],
             ["Hello", "World."],
-            ["HelloWorld"] * 10,
+            ["Hello", "World"] * 10,
             ["Hello", "World.", "Hello", "World."],
         ]
 
-        output_nodes = self.splitter([node0, node1, node2, node3])
+        output_nodes = splitter([node0, node1, node2, node3])
         output_texts = [node.value_dict["text"] for node in output_nodes]
 
         self.assertEqual(output_texts, expected_texts)
diff --git a/tests/op/extract/split/test_splitter_factory.py b/tests/op/extract/split/test_splitter_factory.py
@@ -1,44 +1,22 @@
 import unittest
 
-from uniflow.op.extract.split.constants import (
-    MARKDOWN_HEADER_SPLITTER,
-    PARAGRAPH_SPLITTER,
-    RECURSIVE_CHARACTER_SPLITTER,
-)
-from uniflow.op.extract.split.markdown_header_splitter import MarkdownHeaderSplitter
+from uniflow.op.extract.split.constants import PARAGRAPH_SPLITTER
 from uniflow.op.extract.split.pattern_splitter_op import PatternSplitter
-from uniflow.op.extract.split.recursive_character_splitter import (
-    RecursiveCharacterSplitter,
-)
 from uniflow.op.extract.split.splitter_factory import SplitterOpsFactory
 
 
 class TestSplitterOpsFactory(unittest.TestCase):
-    def setUp(self):
-        self.paragraph_splitter = SplitterOpsFactory.get(PARAGRAPH_SPLITTER)
-        self.markdown_header_splitter = SplitterOpsFactory.get(MARKDOWN_HEADER_SPLITTER)
-        self.recursive_character_splitter = SplitterOpsFactory.get(
-            RECURSIVE_CHARACTER_SPLITTER
-        )
+    def test_get_with_valid_config(self):
+        config = {"splitter_func": PARAGRAPH_SPLITTER}
+        splitter = SplitterOpsFactory.get(config)
+        self.assertIsInstance(splitter, PatternSplitter)
 
-    def test_get(self):
-        self.assertTrue(isinstance(self.paragraph_splitter, PatternSplitter))
-        self.assertTrue(
-            isinstance(self.markdown_header_splitter, MarkdownHeaderSplitter)
-        )
-        self.assertTrue(
-            isinstance(self.recursive_character_splitter, RecursiveCharacterSplitter)
-        )
-
-    def test_get_with_invalid_name(self):
+    def test_get_with_invalid_config(self):
+        config = {"splitter_func": "invalid"}
         with self.assertRaises(ValueError):
-            SplitterOpsFactory.get("")
-
-    def test_list(self):
-        excepted_splitters = [
-            PARAGRAPH_SPLITTER,
-            MARKDOWN_HEADER_SPLITTER,
-            RECURSIVE_CHARACTER_SPLITTER,
-        ]
+            SplitterOpsFactory.get(config)
 
-        self.assertEqual(SplitterOpsFactory.list(), excepted_splitters)
+    def test_get_with_empty_config(self):
+        config = {}
+        with self.assertRaises(KeyError):
+            SplitterOpsFactory.get(config)
diff --git a/uniflow/op/extract/split/pattern_splitter_op.py b/uniflow/op/extract/split/pattern_splitter_op.py
@@ -11,6 +11,9 @@
 class PatternSplitter(Op):
     """Pattern Splitter Op Class"""
 
+    default_separators = "\n\n|\n"
+    default_min_chunk_size = 1
+
     def __init__(
         self, splitterConfig: dict[str, any], name: str = "paragraph_split_op"
     ) -> None:
@@ -22,6 +25,12 @@ def __init__(
         """
         super().__init__(name)
         self._splitter_config = splitterConfig
+        self._separators = (
+            "separators" in splitterConfig and splitterConfig["separators"]
+        ) or self.default_separators
+        self._min_chunk_size = (
+            "min_chunk_size" in splitterConfig and splitterConfig["min_chunk_size"]
+        ) or self.default_min_chunk_size
 
     def __call__(
         self,
@@ -39,8 +48,8 @@ def __call__(
         for node in nodes:
             value_dict = copy.deepcopy(node.value_dict)
             text = value_dict["text"]
-            text = re.split(self._splitter_config["separators"], text)
-            text = [p for p in text if len(p) > self._splitter_config["min_chunk_size"]]
+            text = re.split(self._separators, text)
+            text = [p for p in text if len(p) > self._min_chunk_size]
             output_nodes.append(
                 Node(
                     name=self.unique_name(),
diff --git a/uniflow/op/extract/split/recursive_character_splitter.py b/uniflow/op/extract/split/recursive_character_splitter.py
@@ -4,7 +4,7 @@
 import re
 from typing import Iterable, List, Sequence
 
-import tiktoken  # Import necessary for token-based splitting
+import tiktoken
 
 from uniflow.node import Node
 from uniflow.op.op import Op
@@ -15,7 +15,7 @@ class RecursiveCharacterSplitter(Op):
 
     default_chunk_size = 1024
     default_chunk_overlap_size = 32
-    default_separators = "\n\n|\n|. |.|, | "
+    default_separators = "\n\n|\n|. |.|, | |"
     default_splitting_mode = "char"
 
     def __init__(
@@ -34,14 +34,18 @@ def __init__(
             chunk_overlap_size (int): Overlap in characters between chunks.
             separators (List[str]): Separators to use.
             splitting_mode (str): "char" for character count, "token" for token count. Defaults to "char".
+            keep_separator (bool): Whether to keep the separator. Defaults to True.
+            is_separator_regex (bool): Whether the separator is a regex. Defaults to False.
         """
         super().__init__(name)
 
         # Set up the splitter configuration
         self._chunk_size = splitterConfig["max_chunk_size"] or self.default_chunk_size
         self._separators = (
-            splitterConfig["separators"] or self.default_separators
+            ("separators" in splitterConfig and splitterConfig["separators"])
+            or self.default_separators
         ).split("|")
+        print(f"Separators: {self._separators}")
 
         # Set up the splitter configuration for recursive splitting
         self._chunk_overlap_size = (
@@ -51,6 +55,16 @@ def __init__(
         self._splitting_mode = (
             "splitting_mode" in splitterConfig and splitterConfig["splitting_mode"]
         ) or self.default_splitting_mode
+        self._keep_separator = (
+            True
+            and ("keep_separator" in splitterConfig)
+            and splitterConfig["keep_separator"]
+        )
+        self._is_separator_regex = (
+            ("is_separator_regex" in splitterConfig)
+            and splitterConfig["is_separator_regex"]
+            or False
+        )
 
         self._encoder = tiktoken.encoding_for_model(
             "gpt-3.5"
@@ -118,18 +132,21 @@ def _recursive_splitter(self, text: str, separators: List[str]) -> List[str]:
                 break
 
         # Splited by current separator firstly
-        cur_separator = re.escape(cur_separator)
+        cur_separator = (
+            cur_separator if self._is_separator_regex else re.escape(cur_separator)
+        )
         splits = [s for s in re.split(cur_separator, text) if s != ""]
 
         # Then go merging things, recursively splitting longer texts.
-        _tmp_splits, _separator = [], ""
+        _tmp_splits = []
+        merge_separator = "" if self._keep_separator else _separator
         for s in splits:
             if self._get_length(s) <= self._chunk_size:
                 _tmp_splits.append(s)
             else:
                 # merge splitted texts into a chunk
                 if _tmp_splits:
-                    merged_text = self._merge_splits(_tmp_splits, _separator)
+                    merged_text = self._merge_splits(_tmp_splits, merge_separator)
                     final_chunks.extend(merged_text)
                     # reset tmp_splits
                     _tmp_splits = []
@@ -142,7 +159,7 @@ def _recursive_splitter(self, text: str, separators: List[str]) -> List[str]:
                     final_chunks.extend(other_info)
 
         if _tmp_splits:
-            merged_text = self._merge_splits(_tmp_splits, _separator)
+            merged_text = self._merge_splits(_tmp_splits, merge_separator)
             final_chunks.extend(merged_text)
 
         return final_chunks
@@ -177,6 +194,7 @@ def _merge_splits(self, splits: Iterable[str], separator: str) -> List[str]:
                     doc = separator.join(current_doc).strip()
                     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

Original file line number	Diff line number	Diff line change
`@@ -1,2 +1,2 @@`
`1`	`1`	`#!/bin/sh`
`2`		`-python -m unittest discover tests`
	`2`	`+python3 -m unittest discover tests`