add wandb

run_speech_recognition_ctc.py

@@ -15,6 +15,7 @@
 
 """ Fine-tuning a 🤗 Transformers CTC model for automatic speech recognition"""
 
+import datetime
 import functools
 import json
 import logging
@@ -28,6 +29,7 @@ from typing import Dict, List, Optional, Union
 import datasets
 import numpy as np
 import torch
+import wandb
 from datasets import DatasetDict, load_dataset, load_metric
 
 import transformers
@@ -355,6 +357,15 @@ def main():
     else:
         model_args, data_args, training_args = parser.parse_args_into_dataclasses()
 
+    # TODO: Replace with check of wandb env vars
+    try:
+        os.environ["WANDB_PROJECT"] = os.getcwd().split("/")[-1]
+        wandb.login()
+        training_args.report_to = ["wandb"]
+        training_args.run_name = f"{datetime.datetime.utcnow()}".replace(" ", "T")
+    except:
+        pass
+
     # Detecting last checkpoint.
     last_checkpoint = None
     if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir: