[Feat][Preprocess] support merged dataset (#752)

Author: Eigensystem

fastvideo/configs/configs.py

@@ -9,6 +9,29 @@
 logger = init_logger(__name__)
 
 
+class DatasetType(str, Enum):
+    """
+    Enumeration for different dataset types.
+    """
+    HF = "hf"
+    MERGED = "merged"
+
+    @classmethod
+    def from_string(cls, value: str) -> "DatasetType":
+        """Convert string to DatasetType enum."""
+        try:
+            return cls(value.lower())
+        except ValueError:
+            raise ValueError(
+                f"Invalid dataset type: {value}. Must be one of: {', '.join([m.value for m in cls])}"
+            ) from None
+
+    @classmethod
+    def choices(cls) -> list[str]:
+        """Get all available choices as strings for argparse."""
+        return [dataset_type.value for dataset_type in cls]
+
+
 class DatasetType(str, Enum):
     """
     Enumeration for different dataset types.

fastvideo/workflow/preprocess/components.py

@@ -14,7 +14,10 @@
 from datasets import Dataset, Video, load_dataset
 
 from fastvideo.configs.configs import DatasetType, PreprocessConfig
+<<<<<<< HEAD
 from fastvideo.distributed.parallel_state import get_world_rank, get_world_size
+=======
+>>>>>>> 15df36ab ([Feat][Preprocess] support merged dataset (#752))
 from fastvideo.logger import init_logger
 from fastvideo.pipelines.pipeline_batch_info import PreprocessBatch
 
@@ -402,13 +405,19 @@ def _default_file_writer_fn(self, args_tuple: tuple) -> int:
         return written_count
 
 
+<<<<<<< HEAD
 def build_dataset(preprocess_config: PreprocessConfig, split: str,
                   validator: Callable[[dict[str, Any]], bool]) -> Dataset:
     if preprocess_config.dataset_type == DatasetType.HF:
         dataset = load_dataset(preprocess_config.dataset_path, split=split)
         dataset = dataset.filter(validator)
         dataset = dataset.shard(num_shards=get_world_size(),
                                 index=get_world_rank())
+=======
+def build_dataset(preprocess_config: PreprocessConfig, split: str) -> Dataset:
+    if preprocess_config.dataset_type == DatasetType.HF:
+        dataset = load_dataset(preprocess_config.dataset_path, split=split)
+>>>>>>> 15df36ab ([Feat][Preprocess] support merged dataset (#752))
     elif preprocess_config.dataset_type == DatasetType.MERGED:
         metadata_json_path = os.path.join(preprocess_config.dataset_path,
                                           "videos2caption.json")
@@ -422,11 +431,14 @@ def build_dataset(preprocess_config: PreprocessConfig, split: str,
             dataset = dataset.rename_column("cap", "caption")
         if "path" in column_names:
             dataset = dataset.rename_column("path", "name")
+<<<<<<< HEAD
 
         dataset = dataset.filter(validator)
         dataset = dataset.shard(num_shards=get_world_size(),
                                 index=get_world_rank())
 
+=======
+>>>>>>> 15df36ab ([Feat][Preprocess] support merged dataset (#752))
         # add video column
         def add_video_column(item: dict[str, Any]) -> dict[str, Any]:
             item["video"] = os.path.join(video_folder, item["name"])

fastvideo/workflow/preprocess/preprocess_workflow.py

@@ -44,9 +44,13 @@ def register_components(self) -> None:
         self.add_component("raw_data_validator", raw_data_validator)
 
         # training dataset
+<<<<<<< HEAD
         training_dataset = build_dataset(preprocess_config,
                                          split="train",
                                          validator=raw_data_validator)
+=======
+        training_dataset = build_dataset(preprocess_config, split="train")
+>>>>>>> 15df36ab ([Feat][Preprocess] support merged dataset (#752))
         # set load_from_cache_file to False to check filter stats
         training_dataset = training_dataset.filter(raw_data_validator)
         # we do not use collate_fn here because we use iterable-style Dataset
@@ -62,8 +66,13 @@ def register_components(self) -> None:
         # try to load validation dataset if it exists
         try:
             validation_dataset = build_dataset(preprocess_config,
+<<<<<<< HEAD
                                                split="validation",
                                                validator=raw_data_validator)
+=======
+                                               split="validation")
+            validation_dataset = validation_dataset.filter(raw_data_validator)
+>>>>>>> 15df36ab ([Feat][Preprocess] support merged dataset (#752))
             validation_dataloader = DataLoader(
                 validation_dataset,
                 batch_size=preprocess_config.preprocess_video_batch_size,