✨ add inference options

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/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 ActiveOptions
 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: ActiveOptions
+    """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 = ActiveOptions(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 ActiveOptions:
+    """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])