fix casting single dataset

run_speech_recognition_seq2seq_streaming.py

@@ -363,6 +363,41 @@ def notify_me(recipient, message=None):
         smtp_obj.quit()
 
 
+def rename_col_and_resample(dataset, dataset_name, text_column_names, text_col_name_ref, audio_column_name, sampling_rate):
+    raw_datasets_features = list(dataset.features.keys())
+    logger.info(f"Dataset {dataset_name} - Features: {raw_datasets_features}")
+                
+    if text_col_name_ref not in raw_datasets_features:
+        if len(text_column_names) == 1:
+            raise ValueError("None of the text column names provided found in dataset."
+                                f"Text columns: {text_column_names}"
+                                f"Dataset columns: {raw_datasets_features}")
+        flag = False
+        for text_column_name in text_column_names:
+            if text_column_name in raw_datasets_features:
+                logger.info(f"Renaming text column {text_column_name} to {text_col_name_ref}")
+                dataset = dataset.rename_column(text_column_name, text_col_name_ref)
+                flag = True
+                break
+        if flag is False:
+            raise ValueError("None of the text column names provided found in dataset."
+                                f"Text columns: {text_column_names}"
+                                f"Dataset columns: {raw_datasets_features}")
+    if audio_column_name is not None and sampling_rate is not None:
+        ds_sr = int(dataset.features[audio_column_name].sampling_rate)
+        if ds_sr != sampling_rate:
+            dataset = dataset.cast_column(
+                audio_column_name, datasets.features.Audio(sampling_rate=sampling_rate)
+            )
+    
+    raw_datasets_features = list(dataset.features.keys())
+    raw_datasets_features.remove(audio_column_name)
+    raw_datasets_features.remove(text_col_name_ref)
+    # Keep only audio and sentence
+    dataset = dataset.remove_columns(column_names=raw_datasets_features)
+    return dataset
+
+
 def load_maybe_streaming_dataset(
     dataset_names,
     dataset_config_names,
@@ -393,34 +428,16 @@ def load_maybe_streaming_dataset(
                     dataset = load_dataset(dataset_name, dataset_config_name, split=split_name, streaming=streaming, **kwargs)
                 else:
                     dataset = load_dataset(dataset_name, split=split_name, streaming=streaming, **kwargs)              
-                raw_datasets_features = list(dataset.features.keys())
-                logger.info(f"Dataset {dataset_name} - Features: {raw_datasets_features}")
-                if text_col_name_ref not in raw_datasets_features:
-                    if len(text_column_names) == 1:
-                        raise ValueError("None of the text column names provided found in dataset."
-                                         f"Text columns: {text_column_names}"
-                                         f"Dataset columns: {raw_datasets_features}")
-                    flag = False
-                    for text_column_name in text_column_names:
-                        if text_column_name in raw_datasets_features:
-                            logger.info(f"Renaming text column {text_column_name} to {text_col_name_ref}")
-                            dataset = dataset.rename_column(text_column_name, text_col_name_ref)
-                            flag = True
-                            break
-                    if flag is False:
-                        raise ValueError("None of the text column names provided found in dataset."
-                                         f"Text columns: {text_column_names}"
-                                         f"Dataset columns: {raw_datasets_features}")
-                if audio_column_name is not None and sampling_rate is not None:
-                    ds_sr = int(dataset.features[audio_column_name].sampling_rate)
-                    if ds_sr != sampling_rate:
-                        dataset = dataset.cast_column(
-                            audio_column_name, datasets.features.Audio(sampling_rate=sampling_rate)
-                        )
-                raw_datasets_features = list(dataset.features.keys())
-                raw_datasets_features.remove(audio_column_name)
-                raw_datasets_features.remove(text_col_name_ref)
-                dataset = dataset.remove_columns(column_names=raw_datasets_features)
+                
+                dataset = rename_col_and_resample(
+                    dataset,
+                    dataset_name,
+                    text_column_names,
+                    text_col_name_ref,
+                    audio_column_name,
+                    sampling_rate
+                )
+
                 dataset_splits.append(dataset)
 
         # interleave multiple splits to form one dataset
@@ -428,7 +445,16 @@ def load_maybe_streaming_dataset(
         return interleaved_dataset
     else:
         # load a single split *with* streaming mode
+        
         dataset = load_dataset(dataset_names, dataset_config_names, split=split, streaming=streaming, **kwargs)
+        dataset = rename_col_and_resample(
+            dataset,
+            dataset_names,
+            text_column_names,
+            text_col_name_ref,
+            audio_column_name,
+            sampling_rate
+        )
         return dataset