Add splitter config for recursive splitter

Author: SayaZhang

example/extract/extract_pdf_with_recursive_splitter.ipynb

@@ -70,7 +70,7 @@
     "import pandas as pd\n",
     "import pprint\n",
     "from uniflow.flow.client import ExtractClient, TransformClient\n",
-    "from uniflow.flow.config import TransformOpenAIConfig, ExtractPDFConfig\n",
+    "from uniflow.flow.config import TransformOpenAIConfig, ExtractPDFConfig, SplitterConfig\n",
     "from uniflow.op.model.model_config import OpenAIModelConfig, NougatModelConfig\n",
     "from uniflow.op.prompt import PromptTemplate, Context\n",
     "from uniflow.op.extract.split.splitter_factory import SplitterOpsFactory\n",
@@ -136,26 +136,23 @@
    "cell_type": "code",
    "execution_count": 5,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/home/ubuntu/anaconda3/envs/uniflow/lib/python3.10/site-packages/torch/functional.py:504: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at ../aten/src/ATen/native/TensorShape.cpp:3526.)\n",
-      "  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "data = [\n",
     "    {\"filename\": input_file},\n",
     "]\n",
     "\n",
+    "splitter_config = SplitterConfig(\n",
+    "            max_chunk_size = 1024,\n",
+    "            splitter_func = RECURSIVE_CHARACTER_SPLITTER\n",
+    "        )\n",
+    "splitter_config.chunk_overlap_size = 5\n",
+    "\n",
     "config = ExtractPDFConfig(\n",
     "    model_config=NougatModelConfig(\n",
     "        batch_size = 1 # When batch_size>1, nougat will run on CUDA, otherwise it will run on CPU\n",
     "    ),\n",
-    "    splitter=RECURSIVE_CHARACTER_SPLITTER,\n",
+    "    splitter_config=splitter_config,\n",
     ")\n",
     "nougat_client = ExtractClient(config)"
    ]
@@ -176,7 +173,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "100%|██████████| 1/1 [00:05<00:00,  5.07s/it]\n"
+      "100%|██████████| 1/1 [00:03<00:00,  3.23s/it]\n"
      ]
     }
    ],

uniflow/flow/config.py

@@ -72,7 +72,7 @@ class ExtractPDFConfig(ExtractConfig):
 
     flow_name: str = "ExtractPDFFlow"
     model_config: ModelConfig = field(default_factory=NougatModelConfig)
-    splitter: str = PARAGRAPH_SPLITTER
+    splitter_config: SplitterConfig = field(default_factory=SplitterConfig)
 
 
 @dataclass

uniflow/flow/extract/extract_pdf_flow.py

@@ -6,7 +6,6 @@
 from uniflow.flow.flow import Flow
 from uniflow.node import Node
 from uniflow.op.extract.load.pdf_op import ExtractPDFOp, ProcessPDFOp
-from uniflow.op.extract.split.constants import PARAGRAPH_SPLITTER
 from uniflow.op.extract.split.splitter_factory import SplitterOpsFactory
 from uniflow.op.model.cv.model import CvModel
 
@@ -17,9 +16,7 @@ class ExtractPDFFlow(Flow):
     TAG = EXTRACT
 
     def __init__(
-        self,
-        model_config: Dict[str, Any],
-        splitter: str = PARAGRAPH_SPLITTER,
+        self, model_config: Dict[str, Any], splitter_config: Dict[str, Any]
     ) -> None:
         """Extract PDF Flow Constructor.
 
@@ -35,7 +32,7 @@ def __init__(
             ),
         )
         self._process_pdf_op = ProcessPDFOp(name="process_pdf_op")
-        self._split_op = SplitterOpsFactory.get(splitter)
+        self._split_op = SplitterOpsFactory.get(splitter_config)
 
     def run(self, nodes: Sequence[Node]) -> Sequence[Node]:
         """Run Model Flow.

uniflow/op/extract/split/recursive_character_splitter.py

@@ -2,7 +2,7 @@
 
 import copy
 import re
-from typing import Iterable, List, Optional, Sequence
+from typing import Iterable, List, Sequence
 
 import tiktoken  # Import necessary for token-based splitting
 
@@ -13,15 +13,15 @@
 class RecursiveCharacterSplitter(Op):
     """Recursive character splitter class."""
 
-    default_separators = ["\n\n", "\n", ". ", " ", ""]
+    default_chunk_size = 1024
+    default_chunk_overlap_size = 32
+    default_separators = "\n\n|\n|. |.|, | "
+    default_splitting_mode = "char"
 
     def __init__(
         self,
-        name: str,
-        chunk_size: int = 1024,
-        chunk_overlap_size: int = 0,
-        separators: Optional[List[str]] = None,
-        splitting_mode: str = "char",  # Added parameter for splitting mode
+        splitterConfig: dict[str, any],
+        name: str = "recursive_character_splitter_op",
     ) -> None:
         """Recursive Splitter Op Constructor
 
@@ -36,10 +36,22 @@ def __init__(
             splitting_mode (str): "char" for character count, "token" for token count. Defaults to "char".
         """
         super().__init__(name)
-        self._chunk_size = chunk_size
-        self._chunk_overlap_size = chunk_overlap_size
-        self._separators = separators or self.default_separators
-        self._splitting_mode = splitting_mode  # Track splitting mode
+
+        # 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
+        ).split("|")
+
+        # Set up the splitter configuration for recursive splitting
+        self._chunk_overlap_size = (
+            "chunk_overlap_size" in splitterConfig
+            and splitterConfig["chunk_overlap_size"]
+        ) or self.default_chunk_overlap_size
+        self._splitting_mode = (
+            "splitting_mode" in splitterConfig and splitterConfig["splitting_mode"]
+        ) or self.default_splitting_mode
+
         self._encoder = tiktoken.encoding_for_model(
             "gpt-3.5"
         )  # Setup encoder for token-based splitting