Add image unit tests (#603)

* Add image unit tests for timm

Signed-off-by: Ryan Wolf <rywolf@nvidia.com>

* Guard imports

Signed-off-by: Ryan Wolf <rywolf@nvidia.com>

* Adjust GPU marker

Signed-off-by: Ryan Wolf <rywolf@nvidia.com>

* Unmock inference in timm tests

Signed-off-by: Ryan Wolf <rywolf@nvidia.com>

* Add test for running inference with a fake model

Signed-off-by: Ryan Wolf <rywolf@nvidia.com>

* Add aesthetic and nsfw tests

Signed-off-by: Ryan Wolf <rywolf@nvidia.com>

* Add more aesthetic and nsfw tests

Signed-off-by: Ryan Wolf <rywolf@nvidia.com>

* Add image text pair dataset tests

Signed-off-by: Ryan Wolf <rywolf@nvidia.com>

* Guard gpu import

Signed-off-by: Ryan Wolf <rywolf@nvidia.com>

* Fix function name

Signed-off-by: Ryan Wolf <rywolf@nvidia.com>

* Fix bug in image text pair dataset

Signed-off-by: Ryan Wolf <rywolf@nvidia.com>

* Address Sarah's review

Signed-off-by: Ryan Wolf <rywolf@nvidia.com>

* Replace broken test

Signed-off-by: Ryan Wolf <rywolf@nvidia.com>

---------

Signed-off-by: Ryan Wolf <rywolf@nvidia.com>

Author: ryantwolf

nemo_curator/datasets/image_text_pair_dataset.py

@@ -149,8 +149,8 @@ def filter_members(member):
     def _get_eligible_samples(self, output_path: str, samples_per_shard: int):
         parquet_glob_str = os.path.join(output_path, "temp_*.parquet")
         tar_glob_str = os.path.join(self.path, "*.tar")
-        parquet_files = open_files(parquet_glob_str)
-        tar_files = open_files(tar_glob_str)
+        parquet_files = sorted(open_files(parquet_glob_str), key=lambda f: f.path)
+        tar_files = sorted(open_files(tar_glob_str), key=lambda f: f.path)
 
         curr_df = None
         total_tar_samples = []
@@ -198,8 +198,9 @@ def _get_eligible_samples(self, output_path: str, samples_per_shard: int):
                     samples_per_shard * entries_per_sample :
                 ]
 
-        # Return the remaining df and samples
-        yield curr_df, total_tar_samples
+        # Return the remaining df and samples if it's not empty
+        if len(curr_df) > 0:
+            yield curr_df, total_tar_samples
 
     @staticmethod
     def _combine_id(shard_id, sample_id, max_shards=5, max_samples_per_shard=4) -> str: