Merge pull request #34 from Clarifai/DEVX-453-Multi-Processing-Support-for-Data-Ingestion-Pipeline

[DEVX-453]: Added Multi-Processing support for Data Ingestion Pipeline

Author: sanjaychelliah

clarifai_datautils/multimodal/pipeline/base.py

@@ -1,4 +1,5 @@
 import os
+from concurrent.futures import ThreadPoolExecutor
 from typing import List, Type
 
 from tqdm import tqdm
@@ -34,14 +35,16 @@ def run(self,
           files: str = None,
           folder: str = None,
           show_progress: bool = True,
-          loader: bool = True):
+          loader: bool = True,
+          num_workers: int = 4):
     """Runs the Data Ingestion pipeline.
 
     Args:
         files Any[str,List]: List of files to process.
         folder (str): Folder containing the files.
         show_progress (bool): Whether to show progress bar.
         loader (bool): Whether to return a Clarifai Dataloader Object to pass to SDK Dataset Upload Functionality.
+        num_workers (int): Number of workers to use for parallel processing.
 
     Returns:
         List of transformed elements or ClarifaiDataLoader object.
@@ -57,22 +60,31 @@ def run(self,
 
     # Get files
     if files is not None:
-      self.elements = [files] if isinstance(files, str) else files
-      assert isinstance(self.elements, list), ' Files should be a list of strings.'
+      all_files = [files] if isinstance(files, str) else files
+      assert isinstance(all_files, list), 'Files should be a list of strings.'
     elif folder is not None:
-      self.elements = [os.path.join(folder, f) for f in os.listdir(folder)]
+      all_files = [os.path.join(folder, f) for f in os.listdir(folder)]
 
-    # Apply transformations
-    #TODO: num_workers support
-    if show_progress:
-      with tqdm(total=len(self.transformations), desc='Applying Transformations') as progress:
-        for transform in self.transformations:
-          self.elements = transform(self.elements)
-          progress.update()
+    self.elements = []
 
-    else:
+    # Apply transformations
+    def transform_file(file_elements):
       for transform in self.transformations:
-        self.elements = transform(self.elements)
+        file_elements = transform(file_elements)
+      self.elements.extend(file_elements)
+
+    # num_workers support
+    with ThreadPoolExecutor(max_workers=num_workers) as executor:
+      if show_progress:
+        with tqdm(total=len(all_files), desc='Transforming Files') as progress:
+          futures = [executor.submit(transform_file, [file]) for file in all_files]
+          for job in futures:
+            job.result()
+            progress.update()
+      else:
+        futures = [executor.submit(transform_file, [file]) for file in all_files]
+        for job in futures:
+          job.result()
 
     if loader:
       if self.transformations[0].__class__.__name__ == 'PDFPartitionMultimodal':

tests/pipelines/test_pdf_pipelines.py

@@ -32,7 +32,8 @@ def test_pipeline_run(self,):
         ])
     elements = pipeline.run(files=PDF_FILE_PATH)
     assert len(elements) == 3
-    assert elements[0].text[:9] == 'MAIN GAME'
+    assert (elements[0].text[:9] == 'MAIN GAME') and (
+        elements[1].text[1:12] == 'Transcribed') and (elements[2].text[:10] == 'Thankfully')
     assert elements[0].metadata['filename'] == 'DA-1p.pdf'
     assert elements[0].metadata['page_number'] == 1
     assert elements[0].metadata['text_after'] == 'our assault."'

tests/pipelines/test_text_pipelines.py

@@ -36,7 +36,8 @@ def test_pipeline_run(self,):
         ])
     elements = pipeline.run(files=TEXT_FILE_PATH)
     assert len(elements) == 4
-    assert elements[0].text[:9] == 'CHAPTER I'
+    assert (elements[0].text[:9] == 'CHAPTER I') and (elements[1].text[:3] == 'All') and (
+        elements[2].text[1:6] == 'First') and (elements[3].text[1:5] == 'What')
     assert elements[0].metadata['filename'] == 'book-war-and-peace-1p.txt'
     assert elements[0].metadata['text_after'] == 'used only by the elite.'