✨ add inference options (#354)

Author: ianardee

mindee/input/inference_parameters.py

@@ -11,7 +11,16 @@ class InferenceParameters:
     model_id: str
     """ID of the model, required."""
     rag: bool = False
-    """If set to `True`, will enable Retrieval-Augmented Generation."""
+    """Use Retrieval-Augmented Generation during inference."""
+    raw_text: bool = False
+    """Extract the entire text from the document as strings, and fill the ``raw_text`` attribute."""
+    polygon: bool = False
+    """Calculate bounding box polygons for values, and fill the ``locations`` attribute of fields"""
+    confidence: bool = False
+    """
+    Calculate confidence scores for values, and fill the ``confidence`` attribute of fields.
+    Useful for automation.
+    """
     alias: Optional[str] = None
     """Use an alias to link the file to your own DB. If empty, no alias will be used."""
     webhook_ids: Optional[List[str]] = None

mindee/mindee_http/mindee_api_v2.py

@@ -84,6 +84,12 @@ def req_post_inference_enqueue(
 
         if params.rag:
             data["rag"] = "true"
+        if params.raw_text:
+            data["raw_text"] = "true"
+        if params.confidence:
+            data["confidence"] = "true"
+        if params.polygon:
+            data["polygon"] = "true"
         if params.webhook_ids and len(params.webhook_ids) > 0:
             data["webhook_ids"] = ",".join(params.webhook_ids)
         if params.alias and len(params.alias):

mindee/parsing/v2/__init__.py

@@ -1,5 +1,6 @@
 from mindee.parsing.v2.error_response import ErrorResponse
 from mindee.parsing.v2.inference import Inference
+from mindee.parsing.v2.inference_active_options import InferenceActiveOptions
 from mindee.parsing.v2.inference_file import InferenceFile
 from mindee.parsing.v2.inference_model import InferenceModel
 from mindee.parsing.v2.inference_response import InferenceResponse
@@ -8,6 +9,7 @@
 
 __all__ = [
     "Inference",
+    "InferenceActiveOptions",
     "InferenceFile",
     "InferenceModel",
     "InferenceResponse",

mindee/parsing/v2/inference.py

@@ -1,6 +1,5 @@
-from typing import Optional
-
 from mindee.parsing.common.string_dict import StringDict
+from mindee.parsing.v2.inference_active_options import InferenceActiveOptions
 from mindee.parsing.v2.inference_file import InferenceFile
 from mindee.parsing.v2.inference_model import InferenceModel
 from mindee.parsing.v2.inference_result import InferenceResult
@@ -9,25 +8,29 @@
 class Inference:
     """Inference object for a V2 API return."""
 
+    id: str
+    """ID of the inference."""
     model: InferenceModel
     """Model info for the inference."""
     file: InferenceFile
     """File info for the inference."""
     result: InferenceResult
     """Result of the inference."""
-    id: Optional[str]
-    """ID of the inference."""
+    active_options: InferenceActiveOptions
+    """Active options for the inference."""
 
     def __init__(self, raw_response: StringDict):
+        self.id = raw_response["id"]
         self.model = InferenceModel(raw_response["model"])
         self.file = InferenceFile(raw_response["file"])
         self.result = InferenceResult(raw_response["result"])
-        self.id = raw_response["id"] if "id" in raw_response else None
+        self.active_options = InferenceActiveOptions(raw_response["active_options"])
 
     def __str__(self) -> str:
         return (
-            f"Inference\n#########\n"
-            f"{self.model}\n\n"
-            f"{self.file}"
-            f"{self.result}\n"
+            f"Inference\n#########"
+            f"\n{self.model}"
+            f"\n\n{self.file}"
+            f"\n\n{self.active_options}"
+            f"\n\n{self.result}\n"
         )

mindee/parsing/v2/inference_active_options.py

@@ -0,0 +1,25 @@
+from mindee.parsing.common.string_dict import StringDict
+
+
+class InferenceActiveOptions:
+    """Active options for the inference."""
+
+    raw_text: bool
+    polygon: bool
+    confidence: bool
+    rag: bool
+
+    def __init__(self, raw_response: StringDict):
+        self.raw_text = raw_response["raw_text"]
+        self.polygon = raw_response["polygon"]
+        self.confidence = raw_response["confidence"]
+        self.rag = raw_response["rag"]
+
+    def __str__(self) -> str:
+        return (
+            f"Active Options\n=============="
+            f"\n:Raw Text: {self.raw_text}"
+            f"\n:Polygon: {self.polygon}"
+            f"\n:Confidence: {self.confidence}"
+            f"\n:RAG: {self.rag}"
+        )

mindee/parsing/v2/inference_result.py

@@ -19,5 +19,5 @@ def __init__(self, raw_response: StringDict) -> None:
             self.raw_text = RawText(raw_response["raw_text"])
 
     def __str__(self) -> str:
-        out_str = f"\n\nFields\n======{self.fields}"
+        out_str = f"Fields\n======{self.fields}"
         return out_str

mindee/parsing/v2/raw_text.py

@@ -12,3 +12,6 @@ class RawText:
 
     def __init__(self, raw_response: StringDict):
         self.pages = [RawTextPage(page) for page in raw_response.get("pages", [])]
+
+    def __str__(self) -> str:
+        return "\n\n".join([page.content for page in self.pages])

tests/data

@@ -37,24 +37,41 @@ def test_parse_file_empty_multiple_pages_must_succeed(
     file & model metadata.
     """
     input_path: Path = FILE_TYPES_DIR / "pdf" / "multipage_cut-2.pdf"
-    assert input_path.exists(), f"sample file missing: {input_path}"
 
-    input_doc = PathInput(input_path)
-    options = InferenceParameters(findoc_model_id)
+    input_source = PathInput(input_path)
+    params = InferenceParameters(
+        model_id=findoc_model_id,
+        rag=False,
+        raw_text=True,
+        polygon=False,
+        confidence=False,
+        webhook_ids=[],
+        alias="py_integration_empty_multiple",
+    )
 
     response: InferenceResponse = v2_client.enqueue_and_get_inference(
-        input_doc, options
+        input_source, params
     )
 
     assert response is not None
     assert response.inference is not None
 
     assert response.inference.file is not None
     assert response.inference.file.name == "multipage_cut-2.pdf"
+    assert response.inference.file.page_count == 2
 
     assert response.inference.model is not None
     assert response.inference.model.id == findoc_model_id
 
+    assert response.inference.active_options is not None
+    assert response.inference.active_options.rag is False
+    assert response.inference.active_options.raw_text is True
+    assert response.inference.active_options.polygon is False
+    assert response.inference.active_options.confidence is False
+
+    assert response.inference.result.raw_text is not None
+    assert len(response.inference.result.raw_text.pages) == 2
+
 
 @pytest.mark.integration
 @pytest.mark.v2
@@ -65,28 +82,46 @@ def test_parse_file_filled_single_page_must_succeed(
     Upload a filled single-page JPEG and verify that common fields are present.
     """
     input_path: Path = PRODUCT_DATA_DIR / "financial_document" / "default_sample.jpg"
-    assert input_path.exists(), f"sample file missing: {input_path}"
 
-    input_doc = PathInput(input_path)
-    options = InferenceParameters(findoc_model_id)
+    input_source = PathInput(input_path)
+    params = InferenceParameters(
+        model_id=findoc_model_id,
+        rag=False,
+        raw_text=False,
+        polygon=False,
+        confidence=False,
+        webhook_ids=[],
+        alias="py_integration_filled_single",
+    )
 
     response: InferenceResponse = v2_client.enqueue_and_get_inference(
-        input_doc, options
+        input_source, params
     )
 
     assert response is not None
     assert response.inference is not None
 
     assert response.inference.file is not None
     assert response.inference.file.name == "default_sample.jpg"
+    assert response.inference.file.page_count == 1
 
     assert response.inference.model is not None
     assert response.inference.model.id == findoc_model_id
 
+    assert response.inference.active_options is not None
+    assert response.inference.active_options.rag is False
+    assert response.inference.active_options.raw_text is False
+    assert response.inference.active_options.polygon is False
+    assert response.inference.active_options.confidence is False
+
+    assert response.inference.result.raw_text is None
+
     assert response.inference.result is not None
     supplier_name = response.inference.result.fields["supplier_name"]
     assert supplier_name is not None
     assert supplier_name.value == "John Smith"
+    assert supplier_name.confidence is None
+    assert len(supplier_name.locations) == 0
 
 
 @pytest.mark.integration
@@ -96,13 +131,12 @@ def test_invalid_uuid_must_throw_error_422(v2_client: ClientV2) -> None:
     Using an invalid model identifier must trigger a 422 HTTP error.
     """
     input_path: Path = FILE_TYPES_DIR / "pdf" / "multipage_cut-2.pdf"
-    assert input_path.exists()
 
-    input_doc = PathInput(input_path)
-    options = InferenceParameters("INVALID MODEL ID")
+    input_source = PathInput(input_path)
+    params = InferenceParameters(model_id="INVALID MODEL ID")
 
     with pytest.raises(MindeeHTTPErrorV2) as exc_info:
-        v2_client.enqueue_inference(input_doc, options)
+        v2_client.enqueue_inference(input_source, params)
 
     exc: MindeeHTTPErrorV2 = exc_info.value
     assert exc.status == 422
@@ -119,10 +153,18 @@ def test_url_input_source_must_not_raise_errors(
     """
     url = os.getenv("MINDEE_V2_SE_TESTS_BLANK_PDF_URL")
 
-    input_doc = UrlInputSource(url)
-    options = InferenceParameters(findoc_model_id)
+    input_source = UrlInputSource(url)
+    params = InferenceParameters(
+        model_id=findoc_model_id,
+        rag=False,
+        raw_text=False,
+        polygon=False,
+        confidence=False,
+        webhook_ids=[],
+        alias="py_integration_url_source",
+    )
     response: InferenceResponse = v2_client.enqueue_and_get_inference(
-        input_doc, options
+        input_source, params
     )
     assert response is not None
     assert response.inference is not None