Fix image curation on latest RAPIDS

nemo_curator/datasets/image_text_pair_dataset.py

@@ -79,7 +79,11 @@ def from_webdataset(cls, path: str, id_col: str):
             path (str): The path to the WebDataset-like format on disk or cloud storage.
             id_col (str): The column storing the unique identifier for each record.
         """
-        metadata = dask_cudf.read_parquet(path)
+        metadata = dask_cudf.read_parquet(path, split_row_groups=False)
+        # TODO: This is a hack to ensure that the number of partitions is not combined
+        # and remain the same as the number of shards.
+        # DEBUG: Why is this happening?
+        metadata = metadata.repartition(npartitions=metadata.npartitions)
         metadata = metadata.map_partitions(cls._sort_partition, id_col=id_col)
 
         tar_files = cls._get_tar_files(path)

nemo_curator/image/classifiers/base.py

@@ -17,9 +17,9 @@
 import cudf
 import cupy as cp
 import torch
+from crossfit.backend.cudf.series import create_list_series_from_1d_or_2d_ar
 
 from nemo_curator.datasets import ImageTextPairDataset
-from nemo_curator.utils.cudf_utils import create_list_series_from_1d_or_2d_ar
 from nemo_curator.utils.distributed_utils import load_object_on_worker
 
 

nemo_curator/image/embedders/base.py

@@ -16,11 +16,11 @@
 
 import cupy as cp
 import torch
+from crossfit.backend.cudf.series import create_list_series_from_1d_or_2d_ar
 from tqdm import tqdm
 
 from nemo_curator.datasets import ImageTextPairDataset
 from nemo_curator.image.classifiers import ImageClassifier
-from nemo_curator.utils.cudf_utils import create_list_series_from_1d_or_2d_ar
 from nemo_curator.utils.distributed_utils import load_object_on_worker