clean code. add logs. log audio correctly

run_speech_recognition_ctc.py

@@ -22,7 +22,6 @@ TODO:
 """
 
 import datetime
-import functools
 import json
 import logging
 import os
@@ -34,7 +33,6 @@ from typing import Dict, List, Optional, Union
 
 import datasets
 import numpy as np
-import pandas as pd
 import torch
 import wandb
 from datasets import DatasetDict, concatenate_datasets, load_dataset, load_metric
@@ -382,9 +380,11 @@ def log_to_wandb(training_args):
         wandb.login()
         training_args.report_to = ["wandb"]
         training_args.run_name = run_name
+        wandb.init()
     except Exception as e:
         logger.warning(f"\nFailed logging in to wandb: {e}\nThis experiment will not be logged.\n")
 
+    return training_args
 
 def detect_last_checkpoint(training_args):
 
@@ -417,7 +417,7 @@ def log_small_sumary(training_args):
     logger.info("Training/evaluation parameters %s", training_args)
 
 
-def load_dataset(training_args, data_args):
+def load_datasets(training_args, data_args):
 
     raw_datasets = DatasetDict()
 
@@ -470,7 +470,7 @@ def load_dataset(training_args, data_args):
     return raw_datasets
 
 
-def clean_dataset(raw_datasets, training_args, data_args):
+def clean_datasets(raw_datasets, training_args, data_args):
 
     chars_to_ignore_regex = (
         f'[{"".join(data_args.chars_to_ignore)}]' if data_args.chars_to_ignore is not None else None
@@ -528,7 +528,7 @@ def clean_dataset(raw_datasets, training_args, data_args):
     return raw_datasets
 
 
-def create_tokenizer_kwargs(raw_datasets, training_args, model_args, data_args, config):
+def create_tokenizer_args(cleaned_datasets, training_args, model_args, data_args, config):
 
     tokenizer_name_or_path = model_args.tokenizer_name_or_path
     tokenizer_kwargs = {}
@@ -546,7 +546,7 @@ def create_tokenizer_kwargs(raw_datasets, training_args, model_args, data_args,
             if not os.path.isfile(vocab_file):
                 os.makedirs(tokenizer_name_or_path, exist_ok=True)
                 vocab_dict = create_vocabulary_from_data(
-                    raw_datasets,
+                    cleaned_datasets,
                     word_delimiter_token=data_args.word_delimiter_token,
                     unk_token=data_args.unk_token,
                     pad_token=data_args.pad_token,
@@ -566,17 +566,22 @@ def create_tokenizer_kwargs(raw_datasets, training_args, model_args, data_args,
             "word_delimiter_token": data_args.word_delimiter_token,
         }
 
-    return tokenizer_kwargs
+    return tokenizer_name_or_path, tokenizer_kwargs
 
 
-def vectorize_dataset(raw_datasets, feature_extractor, tokenizer, training_args, data_args):
+def vectorize_dataset(cleaned_datasets, feature_extractor, tokenizer, training_args, data_args):
 
     # make sure that dataset decodes audio with correct sampling rate
-    dataset_sampling_rate = next(iter(raw_datasets.values())).features[data_args.audio_column_name].sampling_rate
+    dataset_sampling_rate = next(iter(cleaned_datasets.values())).features[data_args.audio_column_name].sampling_rate
     if dataset_sampling_rate != feature_extractor.sampling_rate:
-        raw_datasets = raw_datasets.cast_column(
+        cleaned_datasets = cleaned_datasets.cast_column(
             data_args.audio_column_name, datasets.features.Audio(sampling_rate=feature_extractor.sampling_rate)
         )
+    
+    log_metadata_on_wandb(
+        cleaned_datasets=cleaned_datasets,
+        audio_column_name=data_args.audio_column_name
+    )
 
     # derive max & min input length for sample rate & max duration
     max_input_length = data_args.max_duration_in_seconds * feature_extractor.sampling_rate
@@ -606,15 +611,15 @@ def vectorize_dataset(raw_datasets, feature_extractor, tokenizer, training_args,
 
     with training_args.main_process_first(desc="dataset map preprocessing"):
         vectorized_datasets = DatasetDict()
-        vectorized_datasets["train"] = raw_datasets["train"].map(
+        vectorized_datasets["train"] = cleaned_datasets["train"].map(
             prepare_dataset,
-            remove_columns=raw_datasets["train"].column_names,
+            remove_columns=cleaned_datasets["train"].column_names,
             num_proc=data_args.preprocessing_num_workers,
             desc="preprocess datasets",
         )
-        vectorized_datasets["eval"] = raw_datasets["eval"].map(
+        vectorized_datasets["eval"] = cleaned_datasets["eval"].map(
             prepare_dataset,
-            remove_columns=raw_datasets["eval"].column_names,
+            remove_columns=cleaned_datasets["eval"].column_names,
             num_proc=data_args.preprocessing_num_workers,
             desc="preprocess datasets",
         )
@@ -628,30 +633,57 @@ def vectorize_dataset(raw_datasets, feature_extractor, tokenizer, training_args,
             num_proc=data_args.preprocessing_num_workers,
             input_columns=["input_length"],
         )
+    
+    log_audio_on_wandb(
+        vectorized_datasets=vectorized_datasets,
+        audio_column_name="input_values",
+        sampling_rate=feature_extractor.sampling_rate
+    )
+    
+    return vectorized_datasets
 
 
-def log_dataset_sample_on_wandb(vectorized_datasets, audio_column_name):
+def log_metadata_on_wandb(
+    cleaned_datasets,
+    audio_column_name,
+    max_samples=10
+):
+
+    pd_train = cleaned_datasets["train"].select(range(max_samples)).to_pandas()
+    pd_eval = cleaned_datasets["eval"].select(range(max_samples)).to_pandas()
 
-    pd_train = vectorized_datasets["train"].select(range(10)).to_pandas()
-    pd_eval = vectorized_datasets["eval"].select(range(10)).to_pandas()
+    wandb.log({
+        "Training samples": pd_train.drop(labels=audio_column_name, axis=1),
+        "Eval samples": pd_eval.drop(labels=audio_column_name, axis=1),
+    })
+
+
+def log_audio_on_wandb(
+    vectorized_datasets,
+    audio_column_name,
+    sampling_rate,
+    max_samples=10
+):
 
     dict_log = {}
-    for i, audio in pd_train[audio_column_name]:
+    for i, array in enumerate(vectorized_datasets["train"][audio_column_name]):
         dict_log[f"Training sample {i}"] = wandb.Audio(
-            audio["array"],
-            audio_rate=audio["sampling_rate"]
+            array,
+            sample_rate=sampling_rate
         )
-    for i, audio in pd_eval[audio_column_name]:
+        if i+1 == max_samples:
+            break
+    for i, array in enumerate(vectorized_datasets["eval"][audio_column_name]):
         dict_log[f"Eval sample {i}"] = wandb.Audio(
-            audio["array"],
-            audio_rate=audio["sampling_rate"]
+            array,
+            sample_rate=sampling_rate
         )
+        if i+1 == max_samples:
+            break
 
-    wandb.log({
-        "Training samples": pd_train.drop(labels=audio_column_name, axis=1),
-        "Eval samples": pd_eval.drop(labels=audio_column_name, axis=1),
-        "Audio samples": dict_log
-    })
+    print("\nLogging audio to wandb...\n")
+    wandb.log({"Audio samples": dict_log})
+    print("\nLogged audio to wandb...\n")
 
 
 def prepare_training(
@@ -671,11 +703,6 @@ def prepare_training(
     if data_args.dataset_seed is not None:
         vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(seed=data_args.dataset_seed)
     
-    log_dataset_sample_on_wandb(
-        vectorized_datasets=vectorized_datasets,
-        audio_column_name=data_args.audio_column_name
-    )
-    
     # for large datasets it is advised to run the preprocessing on a
     # single machine first with ``args.preprocessing_only`` since there will mostly likely
     # be a timeout when running the script in distributed mode.
@@ -722,7 +749,7 @@ def prepare_training(
     data_collator = DataCollatorCTCWithPadding(processor=processor)
 
     # Initialize Trainer
-    return Trainer(
+    trainer = Trainer(
         model=model,
         data_collator=data_collator,
         args=training_args,
@@ -731,6 +758,7 @@ def prepare_training(
         eval_dataset=vectorized_datasets["eval"] if training_args.do_eval else None,
         tokenizer=feature_extractor,
     )
+    return trainer
 
 
 def do_training(
@@ -786,7 +814,7 @@ def do_eval(
     return trainer
 
 
-def log_results(trainer, training_args, model_args, data_args):
+def log_and_push_results(trainer, training_args, model_args, data_args):
 
     config_name = data_args.dataset_config_name if data_args.dataset_config_name is not None else "na"
     kwargs = {
@@ -806,6 +834,7 @@ def log_results(trainer, training_args, model_args, data_args):
 
 
 def inst_model_tokenizer_feature_extractor(
+    tokenizer_name_or_path,
     tokenizer_kwargs,
     training_args,
     model_args,
@@ -815,7 +844,7 @@ def inst_model_tokenizer_feature_extractor(
 
     # load tokenizer
     tokenizer = AutoTokenizer.from_pretrained(
-        model_args.tokenizer_name_or_path,
+        tokenizer_name_or_path,
         use_auth_token=data_args.use_auth_token,
         **tokenizer_kwargs,
     )
@@ -874,67 +903,78 @@ def main():
     else:
         model_args, data_args, training_args = parser.parse_args_into_dataclasses()
 
-    # 1. Set logging
+    # 1. Set logs
     set_log_config_and_level(local_rank=training_args.local_rank)
     training_args = log_to_wandb(training_args=training_args)
     log_small_sumary(training_args=training_args)
+    logger.info("Logs set\n")
 
     # 2. Set random seed
     set_seed(training_args.seed)
+    logger.info("Seed set\n")
 
-    # 3. First, let's load the dataset
-    raw_datasets = load_dataset(training_args=training_args, data_args=data_args)
+    # 3. First, let's load the datasets
+    raw_datasets = load_datasets(training_args=training_args, data_args=data_args)
+    logger.info("Dataset loaded\n")
     
     # 4. We remove some special characters from the datasets
     # that make training complicated and do not help in transcribing the speech
     # E.g. characters, such as `,` and `.` do not really have an acoustic characteristic
     # that could be easily picked up by the model
-    raw_datasets = clean_dataset(
+    cleaned_datasets = clean_datasets(
         raw_datasets=raw_datasets,
         training_args=training_args,
         data_args=data_args
     )
+    logger.info("Dataset cleaned\n")
 
     # 5. Next, let's load the config as we might need it to create the tokenizer
     config = AutoConfig.from_pretrained(
-        model_args.model_name_or_path, cache_dir=model_args.cache_dir, use_auth_token=data_args.use_auth_token
+        model_args.model_name_or_path,
+        cache_dir=model_args.cache_dir,
+        use_auth_token=data_args.use_auth_token
     )
+    logger.info("Config loaded\n")
 
     # 6. Next, if no tokenizer file is defined,
     # we create the vocabulary of the model by extracting all unique characters from
     # the training and evaluation datasets
     # We need to make sure that only first rank saves vocabulary
     # make sure all processes wait until vocab is created
-    tokenizer_kwargs = create_tokenizer_kwargs(
-        raw_datasets=raw_datasets,
+    tokenizer_name_or_path, tokenizer_kwargs = create_tokenizer_args(
+        cleaned_datasets=cleaned_datasets,
         training_args=training_args,
         model_args=model_args,
         data_args=data_args,
         config=config
     )
+    logger.info("Tokenizer args loaded\n")
 
     # 7. Now we can instantiate the feature extractor, tokenizer and model
     # Note for distributed training, the .from_pretrained methods guarantee that only
     # one local process can concurrently download model & vocab.
     model, tokenizer, feature_extractor, config = inst_model_tokenizer_feature_extractor(
+        tokenizer_name_or_path=tokenizer_name_or_path,
         tokenizer_kwargs=tokenizer_kwargs,
         training_args=training_args,
         model_args=model_args,
         data_args=data_args,
         config=config
     )
+    logger.info("Model, tokenizer, feature_extractor and config loaded\n")
 
     # 8. Now we preprocess the datasets including loading the audio, resampling and normalization
     # Thankfully, `datasets` takes care of automatically loading and resampling the audio,
     # so that we just need to set the correct target sampling rate and normalize the input
     # via the `feature_extractor`
     vectorized_datasets = vectorize_dataset(
-        raw_datasets=raw_datasets,
+        cleaned_datasets=cleaned_datasets,
         feature_extractor=feature_extractor,
         tokenizer=tokenizer,
         training_args=training_args,
         data_args=data_args
     )
+    logger.info("Dataset vectorized\n")
 
     # 9. Next, we can prepare the training.
     # Let's use word error rate (WER) as our evaluation metric,
@@ -948,9 +988,11 @@ def main():
         data_args=data_args,
         config=config
     )
+    logger.info("Trainer instantiated\n")
 
     # 10. Train model
     last_checkpoint = detect_last_checkpoint(training_args=training_args)
+    logger.info("Last checkpoint detected\n")
     if training_args.do_train:
         trainer = do_training(
             trainer=trainer,
@@ -959,6 +1001,7 @@ def main():
             model_args=model_args,
             data_args=data_args
         )
+        logger.info("Training completed\n")
 
     # 11. Eval model
     if training_args.do_eval:
@@ -967,15 +1010,17 @@ def main():
             vectorized_datasets=vectorized_datasets,
             data_args=data_args
         )
+        logger.info("Eval completed\n")
 
     # 12. Push to hub and update model card
-    log_results(
+    log_and_push_results(
         trainer=trainer,
         training_args=training_args,
         model_args=model_args,
         data_args=data_args
     )
-    
+    logger.info("Results logged\n")
+
 
 if __name__ == "__main__":
     main()