Merge branch 'main' into nli/performance_improvements

Author: nikos-livathinos

docling_eval/dataset_builders/dataset_builder.py

@@ -4,9 +4,8 @@
 import sys
 from abc import abstractmethod
 from pathlib import Path
-from typing import Iterable, Optional, Union
+from typing import TYPE_CHECKING, Iterable, Optional, Union
 
-import ibm_boto3  # type: ignore
 from docling.utils.utils import chunkify
 from docling_core.types.doc.document import ImageRefMode
 from huggingface_hub import snapshot_download
@@ -29,6 +28,14 @@
 # Get logger
 _log = logging.getLogger(__name__)
 
+if TYPE_CHECKING:
+    try:
+        from ibm_boto3 import client as IbmBoto3Client  # type: ignore
+        from ibm_boto3 import resource as IbmBoto3Resource  # type: ignore
+    except ImportError:
+        IbmBoto3Client = object
+        IbmBoto3Resource = object
+
 
 class HFSource(BaseModel):
     repo_id: str
@@ -45,25 +52,48 @@ class S3Source(BaseModel):
     overwrite_downloads: bool = True
 
     def __init__(self, **data):
+        r""" """
+        # Import guards
+        try:
+            import ibm_boto3
+        except ImportError:
+            raise ImportError(
+                "ibm_boto3 package is missing. Install optional dependencies."
+            )
+
         super().__init__(**data)
         self._cos_resource: ibm_boto3.resource = self.initialize_s3_resource()
         self._cos_client: ibm_boto3.client = self.initialize_s3_client()
 
-    def initialize_s3_client(self) -> ibm_boto3.client:
+    def initialize_s3_client(self) -> "IbmBoto3Client":
         """Initializes boto3 client - s3 instance
         Returns the s3 client
         """
+        # Import guards
+        try:
+            import ibm_boto3
+        except ImportError:
+            raise ImportError(
+                "ibm_boto3 package is missing. Install optional dependencies."
+            )
         return ibm_boto3.client(
             "s3",
             endpoint_url=self.endpoint,
             aws_access_key_id=self.access_key,
             aws_secret_access_key=self.secret_key,
         )
 
-    def initialize_s3_resource(self) -> ibm_boto3.resource:
+    def initialize_s3_resource(self) -> "IbmBoto3Resource":
         """Initializes boto3 resource - s3 instance
         Returns the s3 instance
         """
+        # Import guards
+        try:
+            import ibm_boto3
+        except ImportError:
+            raise ImportError(
+                "ibm_boto3 package is missing. Install optional dependencies."
+            )
 
         return ibm_boto3.resource(
             "s3",

docling_eval/prediction_providers/aws_prediction_provider.py

@@ -4,7 +4,6 @@
 import os
 from typing import Dict, Optional, Set, Tuple
 
-import boto3
 from docling.datamodel.base_models import ConversionStatus
 from docling_core.types.doc.base import BoundingBox, CoordOrigin, Size
 from docling_core.types.doc.document import (
@@ -55,6 +54,13 @@ def __init__(
         true_labels: Optional[Set[DocItemLabel]] = None,
         pred_labels: Optional[Set[DocItemLabel]] = None,
     ):
+        # Import guards
+        try:
+            import boto3
+        except ImportError:
+            raise ImportError(
+                "boto3 package is missing. Install optional packages dependencies."
+            )
         super().__init__(
             do_visualization=do_visualization,
             ignore_missing_predictions=ignore_missing_predictions,

docling_eval/prediction_providers/azure_prediction_provider.py

@@ -5,8 +5,6 @@
 from io import BytesIO
 from typing import Dict, Optional, Set, Tuple
 
-from azure.ai.documentintelligence import DocumentIntelligenceClient
-from azure.ai.documentintelligence.models import AnalyzeOutputOption
 from docling.datamodel.base_models import ConversionStatus
 
 # from docling_core.types import DoclingDocument
@@ -64,6 +62,15 @@ def __init__(
         true_labels: Optional[Set[DocItemLabel]] = None,
         pred_labels: Optional[Set[DocItemLabel]] = None,
     ):  # could be the docling converter options, the remote credentials for MS/Google, etc.
+
+        # Import guards
+        try:
+            from azure.ai.documentintelligence import DocumentIntelligenceClient
+        except ImportError:
+            raise ImportError(
+                "azure.ai.documentintelligence package is missing. Install optional dependencies."
+            )
+
         super().__init__(
             do_visualization=do_visualization,
             ignore_missing_predictions=ignore_missing_predictions,
@@ -404,6 +411,7 @@ def predict(self, record: DatasetRecord) -> DatasetRecordWithPrediction:
         Raises:
             RuntimeError: If ground truth doc is not available or if mime type is unsupported
         """
+        from azure.ai.documentintelligence.models import AnalyzeOutputOption
 
         status = ConversionStatus.SUCCESS
         result_orig = None

docling_eval/prediction_providers/google_prediction_provider.py

@@ -23,8 +23,6 @@
     TextCell,
 )
 from docling_core.types.io import DocumentStream
-from google.cloud import documentai
-from google.oauth2 import service_account
 from google.protobuf.json_format import MessageToDict
 
 from docling_eval.datamodels.dataset_record import (
@@ -458,6 +456,21 @@ def __init__(
         true_labels: Optional[Set[DocItemLabel]] = None,
         pred_labels: Optional[Set[DocItemLabel]] = None,
     ):
+        r""" """
+        # Import guards
+        try:
+            from google.cloud import documentai
+        except ImportError:
+            raise ImportError(
+                "google.cloud package is missing. Install optional dependencies."
+            )
+        try:
+            from google.oauth2 import service_account
+        except ImportError:
+            raise ImportError(
+                "google.oauth2 package is missing. Install optional dependencies."
+            )
+
         super().__init__(
             do_visualization=do_visualization,
             ignore_missing_predictions=ignore_missing_predictions,
@@ -857,6 +870,14 @@ def prediction_format(self) -> PredictionFormats:
     def predict(self, record: DatasetRecord) -> DatasetRecordWithPrediction:
         """For the given document stream (single document), run the API and create the doclingDocument."""
 
+        # Import guards
+        try:
+            from google.cloud import documentai
+        except ImportError:
+            raise ImportError(
+                "google.cloud package is missing. Install optional dependencies."
+            )
+
         status = ConversionStatus.SUCCESS
         assert record.original is not None
 

pyproject.toml

@@ -73,6 +73,7 @@ hyperscalers = [
     'azure-core (>=1.33.0,<2.0.0)',
     'boto3 (>=1.37.8,<2.0.0)',
     'google-cloud-documentai (>=3.2.0,<4.0.0)',
+    'google-oauth (>=1.0.1,<2.0.0)',
     'ibm-cos-sdk (>=2.1.40,<3.0.0)',
 ]
 cvat_tools = [