Modify _get_sub_docs to use Custom Separator (#254)

Move _get_sub_docs to private function

llama_parse/base.py

@@ -25,20 +25,7 @@
 # if passing as bytes or a buffer, must provide the file_name in extra_info
 FileInput = Union[str, bytes, BufferedIOBase]
 
-
-def _get_sub_docs(docs: List[Document]) -> List[Document]:
-    """Split docs into pages, by separator."""
-    sub_docs = []
-    for doc in docs:
-        doc_chunks = doc.text.split("\n---\n")
-        for doc_chunk in doc_chunks:
-            sub_doc = Document(
-                text=doc_chunk,
-                metadata=deepcopy(doc.metadata),
-            )
-            sub_docs.append(sub_doc)
-
-    return sub_docs
+_DEFAULT_SEPARATOR = "\n---\n"
 
 
 class LlamaParse(BasePydanticReader):
@@ -132,7 +119,7 @@ class LlamaParse(BasePydanticReader):
     )
     split_by_page: bool = Field(
         default=True,
-        description="Whether to split by page (NOTE: using a predefined separator `\n---\n`)",
+        description="Whether to split by page using the page separator",
     )
     vendor_multimodal_api_key: Optional[str] = Field(
         default=None,
@@ -318,7 +305,7 @@ async def _aload_data(
                 )
             ]
             if self.split_by_page:
-                return _get_sub_docs(docs)
+                return self._get_sub_docs(docs)
             else:
                 return docs
 
@@ -492,3 +479,18 @@ def get_images(self, json_result: List[dict], download_path: str) -> List[dict]:
                 return []
             else:
                 raise e
+
+    def _get_sub_docs(self, docs: List[Document]) -> List[Document]:
+        """Split docs into pages, by separator."""
+        sub_docs = []
+        separator = self.page_separator or _DEFAULT_SEPARATOR
+        for doc in docs:
+            doc_chunks = doc.text.split(separator)
+            for doc_chunk in doc_chunks:
+                sub_doc = Document(
+                    text=doc_chunk,
+                    metadata=deepcopy(doc.metadata),
+                )
+                sub_docs.append(sub_doc)
+
+        return sub_docs