[DataComp] Run pipeline at scale

components/filter_image_resolution/Dockerfile

@@ -11,7 +11,7 @@ RUN pip3 install --no-cache-dir -r requirements.txt
 
 # Install Fondant
 # This is split from other requirements to leverage caching
-ARG FONDANT_VERSION=main
+ARG FONDANT_VERSION=f3f3925b8e8f634e2978e5c7fcefa72c53baba7c
 RUN pip3 install fondant[aws,azure,gcp]@git+https://github.com/ml6team/fondant@${FONDANT_VERSION}
 
 # Set the working directory to the component folder

components/filter_image_resolution/fondant_component.yaml

@@ -1,14 +1,14 @@
 name: Filter image resolution
 description: Component that filters images based on minimum size and max aspect ratio
-image: ghcr.io/ml6team/filter_image_resolution:latest
+image: ghcr.io/ml6team/filter_image_resolution:f3f3925b8e8f634e2978e5c7fcefa72c53baba7c
 
 consumes:
   image:
     fields:
       width:
-        type: int16
+        type: int64
       height:
-        type: int16
+        type: int64
 
 args:
   min_image_dim:

components/load_from_hf_hub/Dockerfile

@@ -11,7 +11,7 @@ RUN pip3 install --no-cache-dir -r requirements.txt
 
 # Install Fondant
 # This is split from other requirements to leverage caching
-ARG FONDANT_VERSION=main
+ARG FONDANT_VERSION=edc65f516067937401e325650d5b5409f071bc39
 RUN pip3 install fondant[aws,azure,gcp]@git+https://github.com/ml6team/fondant@${FONDANT_VERSION}
 
 # Set the working directory to the component folder

components/load_from_hf_hub/fondant_component.yaml

@@ -1,6 +1,6 @@
 name: Load from hub
 description: Component that loads a dataset from the hub
-image: ghcr.io/ml6team/load_from_hf_hub:dev
+image: ghcr.io/ml6team/load_from_hf_hub:edc65f516067937401e325650d5b5409f071bc39
 
 produces:
   dummy_variable:  #TODO: fill in here
@@ -23,4 +23,9 @@ args:
   n_rows_to_load:
     description: Optional argument that defines the number of rows to load. Useful for testing pipeline runs on a small scale
     type: int
-    default: None
+    default: None
+  dataset_length:
+    description: Optional argument that defines the length of the dataset. Required in case `n_rows_to_load` is specified.
+    type: int
+    default: None
+

components/load_from_hf_hub/src/main.py

@@ -16,6 +16,7 @@ def __init__(self, *_,
              column_name_mapping: dict,
              image_column_names: t.Optional[list],
              n_rows_to_load: t.Optional[int],
+             dataset_length: int,
     ) -> None:
         """
         Args:
@@ -25,11 +26,14 @@ def __init__(self, *_,
                 format the image from HF hub format to a byte string
             n_rows_to_load: optional argument that defines the number of rows to load. Useful for
               testing pipeline runs on a small scale.
+            dataset_length: optional argument that specifies the length of the entire dataset. Only
+            required in case n_rows_to_load is specified.
         """
         self.dataset_name = dataset_name
         self.column_name_mapping = column_name_mapping
         self.image_column_names = image_column_names
         self.n_rows_to_load = n_rows_to_load
+        self.dataset_length = dataset_length
 
     def load(self) -> dd.DataFrame:
         # 1) Load data, read as Dask dataframe
@@ -44,12 +48,28 @@ def load(self) -> dd.DataFrame:
                 )
 
         # 3) Rename columns
+        logger.info("Renaming columns...")
         dask_df = dask_df.rename(columns=self.column_name_mapping)
 
         # 4) Optional: only return specific amount of rows
-        if self.n_rows_to_load:
-            dask_df = dask_df.head(self.n_rows_to_load)
-            dask_df = dd.from_pandas(dask_df, npartitions=1)
+        if self.n_rows_to_load is not None:
+            if self.dataset_length is None:
+                raise ValueError("""Make sure to also specify the length of the entire
+                                 dataset. This is required as otherwise only the first
+                                 partition can be loaded""")
+            logger.info(f"""Loading approximately {self.n_rows_to_load} rows...
+                        at least one partition""")
+            partition_length = self.dataset_length // dask_df.npartitions
+            npartitions = max(self.n_rows_to_load // partition_length, 1)
+            dask_df = dask_df.head(self.n_rows_to_load, npartitions=npartitions)
+            dask_df = dd.from_pandas(dask_df, npartitions=npartitions)
+            # .reset_index(drop=True) # will reset it from 0 for every partition
+
+        # Set monotonically increasing index
+        logger.info("Setting the index...")
+        dask_df["id"] = 1
+        dask_df["id"] = dask_df.id.cumsum()
+        dask_df = dask_df.set_index("id", sort=True)
 
         return dask_df
 

examples/pipelines/datacomp/components/cluster_image_embeddings/Dockerfile

@@ -11,7 +11,7 @@ RUN pip3 install --no-cache-dir -r requirements.txt
 
 # Install Fondant
 # This is split from other requirements to leverage caching
-ARG FONDANT_VERSION=main
+ARG FONDANT_VERSION=79df895e9d62d2010ccb8d40ee7e4fd4c68f117d
 RUN pip3 install fondant[aws,azure,gcp]@git+https://github.com/ml6team/fondant@${FONDANT_VERSION}
 
 # Set the working directory to the component folder

examples/pipelines/datacomp/components/cluster_image_embeddings/src/main.py

@@ -16,7 +16,7 @@
 class ClusterImageEmbeddingsComponent(DaskTransformComponent):
     """Component that clusters images based on embeddings."""
 
-    def __init__(self, sample_ratio: float, num_clusters: int) -> None:
+    def __init__(self, *_, sample_ratio: float, num_clusters: int) -> None:
         self.sample_ratio = sample_ratio
         self.num_clusters = num_clusters
 

examples/pipelines/datacomp/components/detect_text/Dockerfile

@@ -0,0 +1,24 @@
+FROM --platform=linux/amd64 python:3.8-slim
+
+# System dependencies
+RUN apt-get update && \
+    apt-get upgrade -y && \
+    apt-get install git -y
+
+# Install requirements
+COPY requirements.txt ./
+RUN pip3 install --no-cache-dir -r requirements.txt
+RUN pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cpu
+
+# Install Fondant
+# This is split from other requirements to leverage caching
+ARG FONDANT_VERSION=39c56436e20fb920a50c26a4d0753251993f3251
+RUN pip3 install fondant[aws,azure,gcp]@git+https://github.com/ml6team/fondant@${FONDANT_VERSION}
+
+# Set the working directory to the component folder
+WORKDIR /component/src
+
+# Copy over src-files
+COPY src/ .
+
+ENTRYPOINT ["python", "main.py"]

examples/pipelines/datacomp/components/detect_text/fondant_component.yaml

@@ -0,0 +1,21 @@
+name: Detect text
+description: Component that detects text in images
+image: ghcr.io/ml6team/detect_text:39c56436e20fb920a50c26a4d0753251993f3251
+
+consumes:
+  image:
+    fields:
+      data:
+        type: binary
+
+produces:
+  image:
+    fields:
+      data:
+        type: binary
+      boxes:
+        type: array
+        items:
+          type: array
+          items:
+            type: int64

examples/pipelines/datacomp/components/detect_text/requirements.txt

@@ -0,0 +1,4 @@
+huggingface-hub==0.16.4
+onnxruntime==1.15.1
+opencv-python-headless
+scipy