add logs

run_speech_recognition_seq2seq_streaming.py

@@ -460,6 +460,7 @@ def main():
         revision=model_args.model_revision,
         use_auth_token=True if model_args.use_auth_token else None,
     )
+    logger.info("Loaded pretrained model, tokenizer, and feature extractor")
 
     if model.config.decoder_start_token_id is None:
         raise ValueError("Make sure that `config.decoder_start_token_id` is correctly defined")
@@ -498,6 +499,7 @@ def main():
             if data_args.streaming
             else raw_datasets["train"].select(range(data_args.max_train_samples))
         )
+        logger.info("Using %d train samples", data_args.max_train_samples)
 
     if data_args.max_eval_samples is not None:
         raw_datasets["eval"] = (
@@ -505,6 +507,7 @@ def main():
             if data_args.streaming
             else raw_datasets["eval"].select(range(data_args.max_eval_samples))
         )
+        logger.info("Using %d eval samples", data_args.max_eval_samples)
 
     def prepare_dataset(batch):
         # process audio
@@ -526,6 +529,7 @@ def main():
             prepare_dataset,
             remove_columns=raw_datasets_features,
         ).with_format("torch")
+        logger.info("Dataset map pre-processing done")
 
         if training_args.do_train and data_args.streaming:
             # manually shuffle if streaming (done by the trainer for non-streaming)
@@ -533,6 +537,7 @@ def main():
                 buffer_size=data_args.shuffle_buffer_size,
                 seed=training_args.seed,
             )
+            logger.info("Shuffled dataset")
 
     # filter training data that is shorter than min_input_length or longer than
     # max_input_length
@@ -544,10 +549,12 @@ def main():
             is_audio_in_length_range,
             input_columns=["input_length"],
         )
+        logger.info("Filtered training dataset")
 
     # 8. Load Metric
     metric = evaluate.load("wer")
     do_normalize_eval = data_args.do_normalize_eval
+    logger.info("Loaded metric")
 
     def compute_metrics(pred):
         pred_ids = pred.predictions
@@ -577,12 +584,13 @@ def main():
         config.save_pretrained(training_args.output_dir)
 
     processor = AutoProcessor.from_pretrained(training_args.output_dir)
-
+    
     # 10. Define data collator
     data_collator = DataCollatorSpeechSeq2SeqWithPadding(
         processor=processor,
         decoder_start_token_id=model.config.decoder_start_token_id,
     )
+    logger.info("Defined data collator")
 
     # 11. Configure Trainer
     # Trainer callback to reinitialise and reshuffle the streamable datasets at the beginning of each epoch
@@ -594,6 +602,9 @@ def main():
             elif isinstance(train_dataloader.dataset, IterableDataset):
                 train_dataloader.dataset.set_epoch(train_dataloader.dataset._epoch + 1)
 
+    push_to_hub = training_args.push_to_hub
+    training_args.push_to_hub = False
+
     # Initialize Trainer
     trainer = Seq2SeqTrainer(
         model=model,
@@ -605,9 +616,7 @@ def main():
         compute_metrics=compute_metrics if training_args.predict_with_generate else None,
         callbacks=[ShuffleCallback()] if data_args.streaming else None,
     )
-
-    push_to_hub = training_args.push_to_hub
-    training_args.push_to_hub = False
+    logger.info("Initialized Trainer")
 
     # 12. Training
     if training_args.do_train:
@@ -643,6 +652,7 @@ def main():
         trainer.save_metrics("eval", metrics)
 
     # 14. Write Training Stats
+    logger.info("Training completed. Writing training stats")
     kwargs = {
         "finetuned_from": model_args.model_name_or_path,
         "tasks": "automatic-speech-recognition",
@@ -659,11 +669,13 @@ def main():
         if model_args.model_index_name is not None:
             kwargs["model_name"] = model_args.model_index_name
 
-    training_args.push_to_hub = push_to_hub
+    logger.info("Pushing model to the hub") if push_to_hub else logger.info("Not pushing model to the hub - creating model card only")
+    trainer.args.push_to_hub = push_to_hub
     if training_args.push_to_hub:
         trainer.push_to_hub(**kwargs)
     else:
         trainer.create_model_card(**kwargs)
+    logger.info("*** DONE! ***")
 
     return results