pre-commit checks passed

Author: Sid Mohan

README.md

@@ -27,7 +27,6 @@ DataFog is an open-source DevSecOps platform that lets you scan and redact Perso
 
 ![image](https://github.com/DataFog/datafog-python/assets/61345237/57fba4e5-21cc-458f-ac6a-6fbbb70a8de1)
 
-
 How do you keep:
 
 - Customer PII
@@ -45,7 +44,6 @@ from entering a Generative AI environment in the first place? What you need is a
 
 ![image](https://github.com/DataFog/datafog-python/assets/61345237/91f4634a-8a9f-4621-81bc-09930feda78a)
 
-
 ### There's lots of PII tools out there; why DataFog?
 
 If you look at the landscape of PII detection tools, their very existence was in many cases driven by regulatory requirements (i.e. 'comply with CCPA/GDPR/HIPAA').
@@ -55,11 +53,10 @@ are purpose-built for the problem that they are solving.
 However, Generative AI changes how we think about privacy. There's now a changing set of privacy requirements (new M&A deals, internal discussions means new terms to scan/redact) as well as different and varying document sources to contend with. PII detection is no longer just about compliance, it's an active - and for some, new - internal security threat for CISOs and Eng Leaders to contend with. We want DataFog to be built and driven to meet the needs of the open-source community as they tackle this challenge.
 
 ### Roadmap
-DataFog is an active project with regular weekly releases to production (typically on/around Monday evenings US PT).  Here's a snapshot of our coming roadmap; if you have questions or would like to weigh in, join our discord and let us know what we can do to make the product better!
-
-![image](https://github.com/DataFog/datafog-python/assets/61345237/62964d22-a221-4f1d-a0e6-0cc99de2ba92)
 
+DataFog is an active project with regular weekly releases to production (typically on/around Monday evenings US PT). Here's a snapshot of our coming roadmap; if you have questions or would like to weigh in, join our discord and let us know what we can do to make the product better!
 
+![image](https://github.com/DataFog/datafog-python/assets/61345237/62964d22-a221-4f1d-a0e6-0cc99de2ba92)
 
 ## Installation
 

src/datafog/__about__.py

@@ -1,2 +1 @@
 __version__ = "3.0.0-beta.6"
-

src/datafog/__init__.py

@@ -1,15 +1,18 @@
-from .main import DataFog
-from .config import PipelineOperationType, OperationConfig
+from .config import OperationConfig, PipelineOperationType
 from .donuttransformer import DonutImageProcessor
-from .pii_annotation import PIIAnnotationModel, PIIAnnotationRequest, PIIAnnotationPipeline
+from .main import DataFog
+from .pii_annotation import (
+    PIIAnnotationModel,
+    PIIAnnotationPipeline,
+    PIIAnnotationRequest,
+)
 
 __all__ = [
-    'DataFog',
-    'PipelineOperationType',
-    'OperationConfig',
-    'DonutImageProcessor',
-    'PIIAnnotationModel',
-    'PIIAnnotationRequest',
-    'PIIAnnotationPipeline'
+    "DataFog",
+    "PipelineOperationType",
+    "OperationConfig",
+    "DonutImageProcessor",
+    "PIIAnnotationModel",
+    "PIIAnnotationRequest",
+    "PIIAnnotationPipeline",
 ]
-

src/datafog/config.py

@@ -1,17 +1,21 @@
 from enum import Enum
+from typing import Optional
+
 from pydantic import BaseModel
-from typing import Optional, Dict, List
+
 
 class PipelineOperationType(Enum):
-    READ_IMAGE = 'read_image'
-    PARSE_IMAGE = 'parse_image'
-    TEXT_PII_ANNOTATION = 'text_pii_annotation'
-    TEXT_PII_ANNOTATION_WITH_IMAGE = 'text_pii_annotation_with_image'
+    READ_IMAGE = "read_image"
+    PARSE_IMAGE = "parse_image"
+    TEXT_PII_ANNOTATION = "text_pii_annotation"
+    TEXT_PII_ANNOTATION_WITH_IMAGE = "text_pii_annotation_with_image"
+
 
 class ModelConfig(BaseModel):
     model: str
     processor: str
 
+
 class OperationConfig(BaseModel):
     operation_type: PipelineOperationType
     config: Optional[ModelConfig] = None
@@ -21,16 +25,18 @@ def model_validator(cls, v, values):
         configs = {
             "read_image": ModelConfig(
                 model="naver-clova-ix/donut-base-finetuned-rvlcdip",
-                processor="naver-clova-ix/donut-base-finetuned-rvlcdip"),
+                processor="naver-clova-ix/donut-base-finetuned-rvlcdip",
+            ),
             "parse_image": ModelConfig(
                 model="naver-clova-ix/donut-base-finetuned-cord-v2",
-                processor="naver-clova-ix/donut-base-finetuned-cord-v2")
+                processor="naver-clova-ix/donut-base-finetuned-cord-v2",
+            ),
         }
-        operation_type = values.get('operation_type')
+        operation_type = values.get("operation_type")
         if operation_type and operation_type.value in configs:
             return configs[operation_type.value]
         return v
-    
+
     class Config:
         use_enum_values = True
         validate_assignment = True

src/datafog/donuttransformer.py

@@ -1,11 +1,10 @@
-from pydantic import HttpUrl
-from transformers import DonutProcessor, VisionEncoderDecoderModel
-from PIL import Image
+import re
 import warnings
 from io import BytesIO
-import re
-import torch
-from enum import Enum
+
+from PIL import Image
+from transformers import DonutProcessor, VisionEncoderDecoderModel
+
 from .config import OperationConfig, PipelineOperationType
 
 
@@ -35,9 +34,12 @@ class DonutImageProcessor:
         _process_image(image: Image.Image, operation_type_prompt: str) -> dict:
             A helper method to process the image with the model using a specific operation type prompt.
     """
+
     def __init__(self, operation_type: PipelineOperationType):
         self.operation_type = operation_type
-        model_config = OperationConfig.model_validator(None, {'operation_type': operation_type})
+        model_config = OperationConfig.model_validator(
+            None, {"operation_type": operation_type}
+        )
         self.processor = DonutProcessor.from_pretrained(model_config.processor)
         self.model = VisionEncoderDecoderModel.from_pretrained(model_config.model)
         self.device = "cpu"
@@ -52,20 +54,28 @@ def read_image(file: bytes) -> Image.Image:
                 image = image.convert("RGB")
             return image
         except IOError as e:
-            raise ValueError(f"Unable to read the image file: {e}. Ensure it is a valid image.")
+            raise ValueError(
+                f"Unable to read the image file: {e}. Ensure it is a valid image."
+            )
 
     def classify_image(self, image: Image.Image) -> dict:
         return self._process_image(image, operation_type_prompt="<s_rvlcdip>")
 
     def parse_image(self, image: Image.Image) -> dict:
         return self._process_image(image, operation_type_prompt="<s_cord-v2>")
 
-    def question_image(self, image: Image.Image, question: str = "what is shown in this image?") -> dict:
-        operation_type_prompt = f"<s_docvqa><s_question>{question}</s_question><s_answer>"
+    def question_image(
+        self, image: Image.Image, question: str = "what is shown in this image?"
+    ) -> dict:
+        operation_type_prompt = (
+            f"<s_docvqa><s_question>{question}</s_question><s_answer>"
+        )
         return self._process_image(image, operation_type_prompt=operation_type_prompt)
 
     def _process_image(self, image: Image.Image, operation_type_prompt: str) -> dict:
-        decoder_input_ids = self.processor.tokenizer(operation_type_prompt, add_special_tokens=False, return_tensors="pt").input_ids
+        decoder_input_ids = self.processor.tokenizer(
+            operation_type_prompt, add_special_tokens=False, return_tensors="pt"
+        ).input_ids
         pixel_values = self.processor(image, return_tensors="pt").pixel_values
 
         outputs = self.model.generate(
@@ -82,7 +92,11 @@ def _process_image(self, image: Image.Image, operation_type_prompt: str) -> dict
         )
 
         sequence = self.processor.batch_decode(outputs.sequences)[0]
-        sequence = sequence.replace(self.processor.tokenizer.eos_token, "").replace(self.processor.tokenizer.pad_token, "")
-        sequence = re.sub(r"<.*?>", "", sequence, count=1).strip()  # remove first operation_type start token
+        sequence = sequence.replace(self.processor.tokenizer.eos_token, "").replace(
+            self.processor.tokenizer.pad_token, ""
+        )
+        sequence = re.sub(
+            r"<.*?>", "", sequence, count=1
+        ).strip()  # remove first operation_type start token
 
         return self.processor.token2json(sequence)

src/datafog/main.py

@@ -1,10 +1,9 @@
-from .donuttransformer import DonutImageProcessor,  PipelineOperationType
-from PIL import Image
-from .pii_annotation import PIIAnnotationModel, PIIAnnotationRequest, PIIAnnotationPipeline
-from typing import Any, List, Optional, Tuple
-from enum import Enum
-
-
+from .donuttransformer import DonutImageProcessor, PipelineOperationType
+from .pii_annotation import (
+    PIIAnnotationModel,
+    PIIAnnotationPipeline,
+    PIIAnnotationRequest,
+)
 
 
 class DataFog:
@@ -26,21 +25,25 @@ class DataFog:
         process_image_with_text(file: bytes, text: str) -> dict:
             Processes the image for different operations using the Donut model and text for PII entities using the Spacy model.
     """
-    def __init__(self, operation_type: PipelineOperationType, text_pii_annotation: bool = True, image_processor: bool = False):
+
+    def __init__(
+        self,
+        operation_type: PipelineOperationType,
+        text_pii_annotation: bool = True,
+        image_processor: bool = False,
+    ):
         self.text_pii_annotation = text_pii_annotation
         if text_pii_annotation:
             self.text_annotator = PIIAnnotationModel()
         self.image_processor = None
         if image_processor:
             self.image_processor = DonutImageProcessor(operation_type=operation_type)
 
-
     def process_text(self, text: str) -> list:
         request = PIIAnnotationRequest(text=text)
         workflow = PIIAnnotationPipeline(request=request, model=self.text_annotator)
         entities = workflow.process_request()
         return entities
-    
 
     def process_image(self, file: bytes) -> dict:
         image = DonutImageProcessor.read_image(file)
@@ -49,7 +52,9 @@ def process_image(self, file: bytes) -> dict:
         elif self.image_processor.operation_type == PipelineOperationType.PARSE_IMAGE:
             result = self.image_processor.parse_image(image)
         else:
-            raise ValueError(f"Unsupported operation type: {self.image_processor.operation_type}")
+            raise ValueError(
+                f"Unsupported operation type: {self.image_processor.operation_type}"
+            )
         return result
 
     def annotate_pii_in_images(self, file: bytes) -> dict:
@@ -59,9 +64,8 @@ def annotate_pii_in_images(self, file: bytes) -> dict:
             for key, value in result.items():
                 if isinstance(value, list):
                     for item in value:
-                        if 'nm' in item:
-                            item['entities'] = self.process_text(item['nm'])  # Store entities directly within the item
+                        if "nm" in item:
+                            item["entities"] = self.process_text(
+                                item["nm"]
+                            )  # Store entities directly within the item
         return result  # Return the modified result dictionary
-
-
-

src/datafog/pii_annotation.py

@@ -1,9 +1,9 @@
 # Description: Define the data models for the PII Detection Workflow
 
-from typing import Any, List, Optional, Tuple
+from typing import Any, List, Tuple
+
 import en_spacy_pii_fast
-import requests
-from pydantic import BaseModel, HttpUrl, FilePath, DirectoryPath
+from pydantic import BaseModel, DirectoryPath, FilePath, HttpUrl
 
 
 class PIIAnnotationModel(BaseModel):
@@ -32,8 +32,7 @@ def validate_fields(self):
         if not any([self.text, self.file_path, self.url, self.directory_path]):
             raise ValueError("At least one of the fields must be filled out")
         return True
-    
-    
+
 
 class PIIAnnotationResponse(BaseModel):
     text: str
@@ -57,4 +56,3 @@ def run(self):
 
     class Config:
         arbitrary_types_allowed = True
-

tests/test_donuttransformer.py

@@ -1,41 +1,42 @@
-import unittest
-from PIL import Image
 import pytest
-from PIL import Image
-from datafog import PipelineOperationType, DonutImageProcessor
-
-# class TestDonutImageProcessor(unittest.TestCase):
-#     def setUp(self):
-#         self.processor = DonutImageProcessor(operation_type=PipelineOperationType.PARSE_IMAGE)
-#         with open("/Users/sidmohan/Desktop/v3.0.0/datafog-python/src/datafog/test-invoice.png", "rb") as image_file:
-#             self.image_data = image_file.read()
-
-#     def test_parse_image(self):
-#         image = DonutImageProcessor.read_image(self.image_data)
-#         result = self.processor.parse_image(image)
-#         self.assertIsInstance(result, dict)
-#         self.assertIn('MEDICAL BILLING INVOICE', [item['nm'] for sublist in result.values() if isinstance(sublist, list) for item in sublist])
-#         self.assertIn('12245', [item['price']['unitprice'] for sublist in result.values() if isinstance(sublist, list) for item in sublist if 'unitprice' in item['price']])
-#         self.assertIn('Full Check Up', [item['nm'] for sublist in result.values() if isinstance(sublist, list) for item in sublist])
-#         self.assertIn('Ear & Throat Examination', [item['nm'] for sublist in result.values() if isinstance(sublist, list) for item in sublist])
-
-# if __name__ == '__main__':
-#     unittest.main()
 
+from datafog import DonutImageProcessor, PipelineOperationType
 
 
 @pytest.fixture
 def processor():
     return DonutImageProcessor(operation_type=PipelineOperationType.PARSE_IMAGE)
 
+
 def test_parse_image(processor):
     sample_image_path = "tests/test-invoice.png"
     with open(sample_image_path, "rb") as image_file:
         image_data = image_file.read()
     image = DonutImageProcessor.read_image(image_data)
     result = processor.parse_image(image)
     assert isinstance(result, dict)
-    assert 'MEDICAL BILLING INVOICE' in [item['nm'] for sublist in result.values() if isinstance(sublist, list) for item in sublist]
-    assert '12245' in [item['price']['unitprice'] for sublist in result.values() if isinstance(sublist, list) for item in sublist if 'unitprice' in item['price']]
-    assert 'Full Check Up' in [item['nm'] for sublist in result.values() if isinstance(sublist, list) for item in sublist]
-    assert 'Ear & Throat Examination' in [item['nm'] for sublist in result.values() if isinstance(sublist, list) for item in sublist]
+    assert "MEDICAL BILLING INVOICE" in [
+        item["nm"]
+        for sublist in result.values()
+        if isinstance(sublist, list)
+        for item in sublist
+    ]
+    assert "12245" in [
+        item["price"]["unitprice"]
+        for sublist in result.values()
+        if isinstance(sublist, list)
+        for item in sublist
+        if "unitprice" in item["price"]
+    ]
+    assert "Full Check Up" in [
+        item["nm"]
+        for sublist in result.values()
+        if isinstance(sublist, list)
+        for item in sublist
+    ]
+    assert "Ear & Throat Examination" in [
+        item["nm"]
+        for sublist in result.values()
+        if isinstance(sublist, list)
+        for item in sublist
+    ]