remove columns not in common across datasets

run_speech_recognition_ctc.py

@@ -331,7 +331,7 @@ def create_vocabulary_from_data(
         batched=True,
         batch_size=10000,
         keep_in_memory=False,
-        # remove_columns=datasets["train"].column_names,
+        remove_columns=datasets["train"].column_names,
     )
 
     # take union of all unique characters in each dataset
@@ -418,6 +418,11 @@ def main():
     # 1. First, let's load the dataset
     raw_datasets = DatasetDict()
 
+    def common_cols(dataset_a, dataset_b):
+        col_a = set(dataset_a.column_names)
+        col_b = set(dataset_b.column_names)
+        return [col for col in col_a if col in col_b]
+
     if training_args.do_train:
 
         # Multiple datasets might need to be loaded from HF
@@ -437,18 +442,21 @@ def main():
                         split=train_split_name,
                         use_auth_token=data_args.use_auth_token,
                     )
+                    min_columns_train = raw_datasets["train"].column_names
                 else:
+                    new_dataset = load_dataset(
+                        dataset_name,
+                        dataset_config_name,
+                        split=train_split_name,
+                        use_auth_token=data_args.use_auth_token,
+                    )
                     raw_datasets["train"] = concatenate_datasets(
                         [
                             raw_datasets["train"],
-                            load_dataset(
-                                dataset_name,
-                                dataset_config_name,
-                                split=train_split_name,
-                                use_auth_token=data_args.use_auth_token,
-                            )
+                            new_dataset
                         ]
                     )
+                    min_columns_train = common_cols(min_columns, new_dataset.column_names)
             else:
                 logging.warning(f"{dataset_name} {dataset_config_name} as split is {train_split_name}")
 
@@ -468,6 +476,8 @@ def main():
 
         if data_args.max_train_samples is not None:
             raw_datasets["train"] = raw_datasets["train"].select(range(data_args.max_train_samples))
+        other_columns_train = [col for col in raw_datasets["train"].column_names if col not in min_columns_train]
+        raw_datasets["train"].remove_columns(other_columns_train)
 
     if training_args.do_eval:
         # Multiple datasets might need to be loaded from HF
@@ -486,23 +496,28 @@ def main():
                         split=eval_split_name,
                         use_auth_token=data_args.use_auth_token,
                     )
+                    min_columns_eval = raw_datasets["eval"].column_names
                 else:
+                    new_dataset = load_dataset(
+                        dataset_name,
+                        dataset_config_name,
+                        split=eval_split_name,
+                        use_auth_token=data_args.use_auth_token,
+                    )
                     raw_datasets["eval"] = concatenate_datasets(
                         [
                             raw_datasets["eval"],
-                            load_dataset(
-                                dataset_name,
-                                dataset_config_name,
-                                split=eval_split_name,
-                                use_auth_token=data_args.use_auth_token,
-                            )
+                            new_dataset
                         ]
                     )
+                    min_columns_eval = common_cols(min_columns_eval, new_dataset.column_names)
             else:
                 logging.warning(f"{dataset_name} {dataset_config_name} as split is {eval_split_name}")
 
         if data_args.max_eval_samples is not None:
             raw_datasets["eval"] = raw_datasets["eval"].select(range(data_args.max_eval_samples))
+        other_columns_eval = [col for col in raw_datasets["eval"].column_names if col not in min_columns_eval]
+        raw_datasets["eval"].remove_columns(other_columns_eval)
 
     # 2. We remove some special characters from the datasets
     # that make training complicated and do not help in transcribing the speech