add full support for split

Author: sebastianMindee

mindee/v2/__init__.py

@@ -1,15 +1,15 @@
+from mindee.v2.parsing import SplitParameters
 from mindee.v2.parsing.inference.base_inference import BaseInference
 from mindee.v2.parsing.inference.base_inference_response import (
     BaseInferenceResponse,
     TypeInferenceResponse,
 )
-from mindee.v2.parsing.inference.split.split_inference import SplitInference
 from mindee.v2.parsing.inference.split.split_response import SplitResponse
 
 __all__ = [
     "BaseInference",
     "BaseInferenceResponse",
     "TypeInferenceResponse",
-    "SplitInference",
     "SplitResponse",
+    "SplitParameters",
 ]

mindee/v2/parsing/__init__.py

@@ -1,15 +1,15 @@
+from mindee.v2.parsing.inference import SplitParameters
 from mindee.v2.parsing.inference.base_inference import BaseInference
 from mindee.v2.parsing.inference.base_inference_response import (
     BaseInferenceResponse,
     TypeInferenceResponse,
 )
-from mindee.v2.parsing.inference.split.split_inference import SplitInference
 from mindee.v2.parsing.inference.split.split_response import SplitResponse
 
 __all__ = [
     "BaseInference",
     "BaseInferenceResponse",
     "TypeInferenceResponse",
-    "SplitInference",
     "SplitResponse",
+    "SplitParameters",
 ]

mindee/v2/parsing/inference/__init__.py

@@ -3,13 +3,19 @@
     BaseInferenceResponse,
     TypeInferenceResponse,
 )
+from mindee.v2.parsing.inference.split.split import Split
 from mindee.v2.parsing.inference.split.split_inference import SplitInference
+from mindee.v2.parsing.inference.split.split_parameters import SplitParameters
 from mindee.v2.parsing.inference.split.split_response import SplitResponse
+from mindee.v2.parsing.inference.split.split_result import SplitResult
 
 __all__ = [
     "BaseInference",
     "BaseInferenceResponse",
     "TypeInferenceResponse",
+    "Split",
     "SplitInference",
+    "SplitParameters",
     "SplitResponse",
+    "SplitResult",
 ]

mindee/v2/parsing/inference/split/split.py

@@ -0,0 +1,20 @@
+from typing import List
+
+from mindee.parsing.common.string_dict import StringDict
+
+
+class Split:
+    """Split inference result."""
+
+    page_range: List[int]
+    """Page range of the split inference."""
+    document_type: str
+    """Document type of the split inference."""
+
+    def __init__(self, server_response: StringDict):
+        self.page_range = server_response["page_range"]
+        self.document_type = server_response["document_type"]
+
+    def __str__(self) -> str:
+        page_range = ",".join([str(page_index) for page_index in self.page_range])
+        return f":Page Range: {page_range}\n:Document Type: {self.document_type}"

mindee/v2/parsing/inference/split/split_inference.py

@@ -1,20 +1,19 @@
-from typing import Any
-
 from mindee.parsing.common.string_dict import StringDict
 from mindee.v2.parsing.inference.base_inference import BaseInference
+from mindee.v2.parsing.inference.split.split_result import SplitResult
 
 
 class SplitInference(BaseInference):
     """Split inference result."""
 
-    result: Any
+    result: SplitResult
     """Result of a split inference."""
     _slug: str = "split"
     """Slug of the endpoint."""
 
     def __init__(self, raw_response: StringDict) -> None:
         super().__init__(raw_response)
-        self.result = raw_response["result"]
+        self.result = SplitResult(raw_response["result"])
 
     def __str__(self) -> str:
         return f"Inference\n#########\n{self.model}\n{self.file}\n{self.result}\n"

mindee/v2/parsing/inference/split/split_result.py

@@ -0,0 +1,17 @@
+from typing import List
+
+from mindee.parsing.common.string_dict import StringDict
+from mindee.v2.parsing.inference.split.split import Split
+
+
+class SplitResult:
+    """Split result info."""
+
+    split: List[Split]
+
+    def __init__(self, raw_response: StringDict) -> None:
+        self.split = [Split(split) for split in raw_response["split"]]
+
+    def __str__(self) -> str:
+        out_str = f"Splits\n======{self.split}"
+        return out_str

tests/v2/parsing/test_split_integration.py

@@ -35,5 +35,5 @@ def test_split_blank(v2_client: ClientV2, split_model_id: str):
     )  # Note: do not use blank_1.pdf for this.
     assert response.inference is not None
     assert response.inference.file.name == "default_sample.pdf"
-    assert response.inference.result.get("split")
-    assert len(response.inference.result.get("split")) == 2
+    assert response.inference.result.split
+    assert len(response.inference.result.split) == 2

tests/v2/parsing/test_split_response.py

@@ -1,8 +1,10 @@
 import pytest
 
 from mindee import LocalResponse
+from mindee.v2.parsing.inference.split.split import Split
 from mindee.v2.parsing.inference.split.split_inference import SplitInference
 from mindee.v2.parsing.inference.split.split_response import SplitResponse
+from mindee.v2.parsing.inference.split.split_result import SplitResult
 from tests.utils import V2_UTILITIES_DATA_DIR
 
 
@@ -11,44 +13,33 @@ def test_split_single():
     input_inference = LocalResponse(V2_UTILITIES_DATA_DIR / "split_single.json")
     split_response = input_inference.deserialize_response(SplitResponse)
     assert isinstance(split_response.inference, SplitInference)
-    assert split_response.inference.result.get("split")
-    assert len(split_response.inference.result.get("split")[0].get("page_range")) == 2
-    assert split_response.inference.result.get("split")[0].get("page_range")[0] == 0
-    assert split_response.inference.result.get("split")[0].get("page_range")[1] == 0
-    assert (
-        split_response.inference.result.get("split")[0].get("document_type")
-        == "receipt"
-    )
+    assert split_response.inference.result.split
+    assert len(split_response.inference.result.split[0].page_range) == 2
+    assert split_response.inference.result.split[0].page_range[0] == 0
+    assert split_response.inference.result.split[0].page_range[1] == 0
+    assert split_response.inference.result.split[0].document_type == "receipt"
 
 
 @pytest.mark.v2
 def test_split_multiple():
     input_inference = LocalResponse(V2_UTILITIES_DATA_DIR / "split_multiple.json")
     split_response = input_inference.deserialize_response(SplitResponse)
     assert isinstance(split_response.inference, SplitInference)
-    assert split_response.inference.result.get("split")
-    assert len(split_response.inference.result.get("split")) == 3
-
-    assert len(split_response.inference.result.get("split")[0].get("page_range")) == 2
-    assert split_response.inference.result.get("split")[0].get("page_range")[0] == 0
-    assert split_response.inference.result.get("split")[0].get("page_range")[1] == 0
-    assert (
-        split_response.inference.result.get("split")[0].get("document_type")
-        == "invoice"
-    )
-
-    assert len(split_response.inference.result.get("split")[1].get("page_range")) == 2
-    assert split_response.inference.result.get("split")[1].get("page_range")[0] == 1
-    assert split_response.inference.result.get("split")[1].get("page_range")[1] == 3
-    assert (
-        split_response.inference.result.get("split")[1].get("document_type")
-        == "invoice"
-    )
-
-    assert len(split_response.inference.result.get("split")[2].get("page_range")) == 2
-    assert split_response.inference.result.get("split")[2].get("page_range")[0] == 4
-    assert split_response.inference.result.get("split")[2].get("page_range")[1] == 4
-    assert (
-        split_response.inference.result.get("split")[2].get("document_type")
-        == "invoice"
-    )
+    assert isinstance(split_response.inference.result, SplitResult)
+    assert isinstance(split_response.inference.result.split[0], Split)
+    assert len(split_response.inference.result.split) == 3
+
+    assert len(split_response.inference.result.split[0].page_range) == 2
+    assert split_response.inference.result.split[0].page_range[0] == 0
+    assert split_response.inference.result.split[0].page_range[1] == 0
+    assert split_response.inference.result.split[0].document_type == "invoice"
+
+    assert len(split_response.inference.result.split[1].page_range) == 2
+    assert split_response.inference.result.split[1].page_range[0] == 1
+    assert split_response.inference.result.split[1].page_range[1] == 3
+    assert split_response.inference.result.split[1].document_type == "invoice"
+
+    assert len(split_response.inference.result.split[2].page_range) == 2
+    assert split_response.inference.result.split[2].page_range[0] == 4
+    assert split_response.inference.result.split[2].page_range[1] == 4
+    assert split_response.inference.result.split[2].document_type == "invoice"