Use logger instead of printing when loading datasets

src/shared.py

@@ -1,5 +1,4 @@
 from transformers.trainer_utils import get_last_checkpoint as glc
-from transformers import Seq2SeqTrainingArguments, TrainingArguments
 import os
 from utils import re_findall
 import logging
@@ -15,6 +14,17 @@ from typing import Optional
 from dataclasses import dataclass, field
 from enum import Enum
 
+
+logging.basicConfig()
+logger = logging.getLogger(__name__)
+
+# Setup logging
+logging.basicConfig(
+    format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
+    datefmt='%m/%d/%Y %H:%M:%S',
+    handlers=[logging.StreamHandler(sys.stdout)],
+)
+
 CATEGORIES = [None, 'SPONSOR', 'SELFPROMO', 'INTERACTION']
 
 ACTION_OPTIONS = ['skip', 'mute', 'full']
@@ -234,7 +244,7 @@ def reset():
 
 def load_datasets(dataset_args: DatasetArguments):
 
-    print('Reading datasets')
+    logger.info('Reading datasets')
     data_files = {}
 
     if dataset_args.train_file is not None:
@@ -333,17 +343,6 @@ class CustomTrainingArguments(OutputArguments, AdditionalTrainingArguments):
     pass
 
 
-logging.basicConfig()
-logger = logging.getLogger(__name__)
-
-# Setup logging
-logging.basicConfig(
-    format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
-    datefmt='%m/%d/%Y %H:%M:%S',
-    handlers=[logging.StreamHandler(sys.stdout)],
-)
-
-
 def get_last_checkpoint(training_args):
     last_checkpoint = None
     if os.path.isdir(training_args.output_dir) and not training_args.overwrite_output_dir: