Merge pull request #35 from Clarifai/DEVX-452-Add-Schema-for-Data-Ingestion-Pipeline

[DEVX-452]: Added Schema for Data Ingestion Pipeline Transformations

Author: srikanthbachala20

clarifai_datautils/multimodal/__init__.py

@@ -1,7 +1,7 @@
 from clarifai_datautils.multimodal.pipeline.base import Pipeline
 from clarifai_datautils.multimodal.pipeline.Docx import DocxPartition
 from clarifai_datautils.multimodal.pipeline.Markdown import MarkdownPartition
-from clarifai_datautils.multimodal.pipeline.PDF import PDFPartition
+from clarifai_datautils.multimodal.pipeline.PDF import PDFPartition, PDFPartitionMultimodal
 from clarifai_datautils.multimodal.pipeline.Text import TextPartition
 
 __all__ = [

clarifai_datautils/multimodal/pipeline/base.py

@@ -1,13 +1,29 @@
 import os
 from concurrent.futures import ThreadPoolExecutor
 from typing import List, Type
+from schema import And, Schema
 
 from tqdm import tqdm
 
 from .basetransform import BaseTransform
 from .loaders import MultiModalLoader, TextDataLoader
 
 
+def get_schema() -> Schema:
+  """Initialize the schema for Data Ingestion Pipeline transformations.
+
+        This schema validates:
+
+        - transformations must be a list
+        - First item in the list must be one of the following: PDFPartition, TextPartition, PDFPartitionMultimodal, DocxPartition, MarkdownPartition
+        - Each item in the list must be of BaseTransform instance
+
+        Returns:
+            Schema: The schema for transformations.
+        """
+  return Schema(And(list, lambda x: x[0].__class__.__name__ in ['PDFPartition', 'TextPartition', 'PDFPartitionMultimodal', 'DocxPartition', 'MarkdownPartition'], lambda x: all(isinstance(item, BaseTransform) for item in x)), error="Invalid transformations data.")
+
+
 class Pipeline:
   """Text processing pipeline object from files"""
 
@@ -25,11 +41,8 @@ def __init__(
     """
     self.name = name
     self.transformations = transformations
-    for transform in self.transformations:
-      if not isinstance(transform, BaseTransform):
-        raise ValueError('All transformations should be of type BaseTransform.')
-
-    #TODO: Schema for transformations
+    self.transformation_schema = get_schema()
+    self.transformation_schema.validate(self.transformations)
 
   def run(self,
           files: str = None,

requirements.txt

@@ -3,3 +3,4 @@ llama-index-llms-clarifai==0.1.2
 pi_heif==0.18.0
 markdown==3.7
 python-docx==1.1.2
+schema==0.7.5

tests/pipelines/test_ready_to_use_pipelines.py

@@ -1,12 +1,30 @@
 import os.path as osp
+from typing import List
 
-from clarifai_datautils.multimodal import Pipeline
+import pytest
+from schema import SchemaError
+
+from clarifai_datautils.multimodal import PDFPartition, Pipeline
+from clarifai_datautils.multimodal.pipeline.cleaners import Clean_extra_whitespace
+from clarifai_datautils.multimodal.pipeline.extractors import (ExtractDateTimeTz,
+                                                               ExtractEmailAddress)
 
 PDF_FILE_PATH = osp.abspath(osp.join(osp.dirname(__file__), "assets", "DA-1p.pdf"))
 TEXT_FILE_PATH = osp.abspath(
     osp.join(osp.dirname(__file__), "assets", "book-war-and-peace-1p.txt"))
 
 
+class Test_transformation():
+
+  def __init__(self):
+    pass
+
+  def __call__(self,) -> List:
+    """Applies the transformation.
+    """
+    pass
+
+
 class TestReadyToUsePipelines:
   """Tests for ready to use pipelines."""
 
@@ -77,3 +95,38 @@ def test_pipeline_standard_markdown(self,):
     assert pipeline.transformations[0].__class__.__name__ == 'MarkdownPartition'
     assert pipeline.transformations[1].__class__.__name__ == 'Clean_extra_whitespace'
     assert pipeline.transformations[2].__class__.__name__ == 'Group_broken_paragraphs'
+
+  def test_schema_error(self):
+    # Incorrect type of transformations object
+    with pytest.raises(SchemaError):
+      _ = Pipeline(
+          name='test-1',
+          transformations=(
+              PDFPartition(max_characters=1024, overlap=None),
+              Clean_extra_whitespace(),
+              ExtractDateTimeTz(),
+              ExtractEmailAddress(),
+          ))
+
+    # Incorrect First transformation
+    with pytest.raises(SchemaError):
+      _ = Pipeline(
+          name='test-2',
+          transformations=[
+              Clean_extra_whitespace(),
+              PDFPartition(max_characters=1024, overlap=None),
+              ExtractDateTimeTz(),
+              ExtractEmailAddress(),
+          ])
+
+    # Incorrect Instance of transformation
+    with pytest.raises(SchemaError):
+      _ = Pipeline(
+          name='test-3',
+          transformations=[
+              PDFPartition(max_characters=1024, overlap=None),
+              Clean_extra_whitespace(),
+              ExtractDateTimeTz(),
+              ExtractEmailAddress(),
+              Test_transformation(),
+          ])