Merge duplicated training dataclasses

src/preprocess.py

@@ -490,54 +490,12 @@ def download_file(url, filename):
 
 @dataclass
 class PreprocessingDatasetArguments(DatasetArguments):
-
-    train_file: Optional[str] = field(
-        default='train.json', metadata={'help': 'The input training data file (a jsonlines file).'}
-    )
-    validation_file: Optional[str] = field(
-        default='valid.json',
-        metadata={
-            'help': 'An optional input evaluation data file to evaluate the metrics on (a jsonlines file).'
-        },
-    )
-    test_file: Optional[str] = field(
-        default='test.json',
-        metadata={
-            'help': 'An optional input test data file to evaluate the metrics on (a jsonlines file).'
-        },
-    )
-
-    c_train_file: Optional[str] = field(
-        default='c_train.json', metadata={'help': 'The input training data file (a jsonlines file).'}
-    )
-    c_validation_file: Optional[str] = field(
-        default='c_valid.json',
-        metadata={
-            'help': 'An optional input evaluation data file to evaluate the metrics on (a jsonlines file).'
-        },
-    )
-    c_test_file: Optional[str] = field(
-        default='c_test.json',
-        metadata={
-            'help': 'An optional input test data file to evaluate the metrics on (a jsonlines file).'
-        },
-    )
-
     # excess_file: Optional[str] = field(
     #     default='excess.json',
     #     metadata={
     #         'help': 'The excess segments left after the split'
     #     },
     # )
-    dataset_cache_dir: Optional[str] = field(
-        default=None,
-        metadata={
-            'help': 'Where to store the cached datasets'
-        },
-    )
-    overwrite_cache: bool = field(
-        default=False, metadata={'help': 'Overwrite the cached training and evaluation sets'}
-    )
 
     positive_file: Optional[str] = field(
         default='sponsor_segments.json', metadata={'help': 'File to output sponsored segments to (a jsonlines file).'}

src/shared.py

@@ -104,6 +104,10 @@ class DatasetArguments:
         },
     )
 
+    overwrite_cache: bool = field(
+        default=False, metadata={'help': 'Overwrite the cached training and evaluation sets'}
+    )
+
     dataset_cache_dir: Optional[str] = field(
         default=None,
         metadata={
@@ -111,6 +115,35 @@ class DatasetArguments:
         },
     )
 
+    train_file: Optional[str] = field(
+        default='train.json', metadata={'help': 'The input training data file (a jsonlines file).'}
+    )
+    validation_file: Optional[str] = field(
+        default='valid.json',
+        metadata={
+            'help': 'An optional input evaluation data file to evaluate the metrics on (a jsonlines file).'
+        },
+    )
+    test_file: Optional[str] = field(
+        default='test.json',
+        metadata={
+            'help': 'An optional input test data file to evaluate the metrics on (a jsonlines file).'
+        },
+    )
+
+    def __post_init__(self):
+        if self.train_file is None or self.validation_file is None:
+            raise ValueError(
+                "Need either a GLUE task, a training/validation file or a dataset name.")
+        else:
+            train_extension = self.train_file.split(".")[-1]
+            assert train_extension in [
+                "csv", "json"], "`train_file` should be a csv or a json file."
+            validation_extension = self.validation_file.split(".")[-1]
+            assert (
+                validation_extension == train_extension
+            ), "`validation_file` should have the same extension (csv or json) as `train_file`."
+
 
 @dataclass
 class OutputArguments:
@@ -178,7 +211,7 @@ def reset():
     print(torch.cuda.memory_summary(device=None, abbreviated=False))
 
 
-def load_datasets(dataset_args):
+def load_datasets(dataset_args: DatasetArguments):
 
     print('Reading datasets')
     data_files = {}
@@ -240,6 +273,39 @@ class CustomTrainingArguments(OutputArguments, TrainingArguments):
     #     * :obj:`"steps"`: Evaluation is done (and logged) every :obj:`eval_steps`.
     #     * :obj:`"epoch"`: Evaluation is done at the end of each epoch.
 
+    preprocessing_num_workers: Optional[int] = field(
+        default=None,
+        metadata={'help': 'The number of processes to use for the preprocessing.'},
+    )
+    max_seq_length: int = field(
+        default=512,
+        metadata={
+            "help": "The maximum total input sequence length after tokenization. Sequences longer "
+            "than this will be truncated, sequences shorter will be padded."
+        },
+    )
+    max_train_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "For debugging purposes or quicker training, truncate the number of training examples to this "
+            "value if set."
+        },
+    )
+    max_eval_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
+            "value if set."
+        },
+    )
+    max_predict_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "For debugging purposes or quicker training, truncate the number of prediction examples to this "
+            "value if set."
+        },
+    )
+
 
 logging.basicConfig()
 logger = logging.getLogger(__name__)
@@ -279,3 +345,37 @@ def train_from_checkpoint(trainer, last_checkpoint, training_args):
     trainer.save_model()  # Saves the tokenizer too for easy upload
 
     return train_result
+
+
+def prepare_datasets(raw_datasets, dataset_args: DatasetArguments, training_args: CustomTrainingArguments, preprocess_function):
+
+    with training_args.main_process_first(desc="dataset map pre-processing"):
+        raw_datasets = raw_datasets.map(
+            preprocess_function,
+            batched=True,
+            load_from_cache_file=not dataset_args.overwrite_cache,
+            desc="Running tokenizer on dataset",
+        )
+
+    if 'train' not in raw_datasets:
+        raise ValueError('Train dataset missing')
+    train_dataset = raw_datasets['train']
+    if training_args.max_train_samples is not None:
+        train_dataset = train_dataset.select(
+            range(training_args.max_train_samples))
+
+    if 'validation' not in raw_datasets:
+        raise ValueError('Validation dataset missing')
+    eval_dataset = raw_datasets['validation']
+    if training_args.max_eval_samples is not None:
+        eval_dataset = eval_dataset.select(
+            range(training_args.max_eval_samples))
+
+    if 'test' not in raw_datasets:
+        raise ValueError('Test dataset missing')
+    predict_dataset = raw_datasets['test']
+    if training_args.max_predict_samples is not None:
+        predict_dataset = predict_dataset.select(
+            range(training_args.max_predict_samples))
+
+    return train_dataset, eval_dataset, predict_dataset

src/train.py

@@ -1,12 +1,17 @@
 from preprocess import PreprocessingDatasetArguments
-from shared import CustomTokens, load_datasets, CustomTrainingArguments, get_last_checkpoint, train_from_checkpoint
+from shared import (
+    CustomTokens,
+    prepare_datasets,
+    load_datasets,
+    CustomTrainingArguments,
+    get_last_checkpoint,
+    train_from_checkpoint
+)
 from model import ModelArguments
 import transformers
 import logging
 import os
 import sys
-from dataclasses import dataclass, field
-from typing import Optional
 from datasets import utils as d_utils
 from transformers import (
     DataCollatorForSeq2Seq,
@@ -35,38 +40,6 @@ logging.basicConfig(
 )
 
 
-
-@dataclass
-class DataTrainingArguments:
-    """
-    Arguments pertaining to what data we are going to input our model for training and eval.
-    """
-
-    preprocessing_num_workers: Optional[int] = field(
-        default=None,
-        metadata={'help': 'The number of processes to use for the preprocessing.'},
-    )
-
-    max_train_samples: Optional[int] = field(
-        default=None,
-        metadata={
-            'help': 'For debugging purposes or quicker training, truncate the number of training examples to this value if set.'
-        },
-    )
-    max_eval_samples: Optional[int] = field(
-        default=None,
-        metadata={
-            'help': 'For debugging purposes or quicker training, truncate the number of evaluation examples to this value if set.'
-        },
-    )
-    max_predict_samples: Optional[int] = field(
-        default=None,
-        metadata={
-            'help': 'For debugging purposes or quicker training, truncate the number of prediction examples to this value if set.'
-        },
-    )
-
-
 def main():
 
     # See all possible arguments in src/transformers/training_args.py
@@ -76,10 +49,9 @@ def main():
     hf_parser = HfArgumentParser((
         ModelArguments,
         PreprocessingDatasetArguments,
-        DataTrainingArguments,
         CustomTrainingArguments
     ))
-    model_args, dataset_args, data_training_args, training_args = hf_parser.parse_args_into_dataclasses()
+    model_args, dataset_args, training_args = hf_parser.parse_args_into_dataclasses()
 
     log_level = training_args.get_process_log_level()
     logger.setLevel(log_level)
@@ -128,7 +100,6 @@ def main():
     # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
     # https://huggingface.co/docs/datasets/loading_datasets.html.
 
-
     # Detecting last checkpoint.
     last_checkpoint = get_last_checkpoint(training_args)
 
@@ -165,47 +136,8 @@ def main():
 
         return model_inputs
 
-    def prepare_dataset(dataset, desc):
-        return dataset.map(
-            preprocess_function,
-            batched=True,
-            num_proc=data_training_args.preprocessing_num_workers,
-            remove_columns=column_names,
-            load_from_cache_file=not dataset_args.overwrite_cache,
-            desc=desc,  # tokenizing train dataset
-        )
-    # train_dataset # TODO shuffle?
-
-    # if training_args.do_train:
-    if 'train' not in raw_datasets:  # TODO do checks above?
-        raise ValueError('Train dataset missing')
-    train_dataset = raw_datasets['train']
-    if data_training_args.max_train_samples is not None:
-        train_dataset = train_dataset.select(
-            range(data_training_args.max_train_samples))
-    with training_args.main_process_first(desc='train dataset map pre-processing'):
-        train_dataset = prepare_dataset(
-            train_dataset, desc='Running tokenizer on train dataset')
-
-    if 'validation' not in raw_datasets:
-        raise ValueError('Validation dataset missing')
-    eval_dataset = raw_datasets['validation']
-    if data_training_args.max_eval_samples is not None:
-        eval_dataset = eval_dataset.select(
-            range(data_training_args.max_eval_samples))
-    with training_args.main_process_first(desc='validation dataset map pre-processing'):
-        eval_dataset = prepare_dataset(
-            eval_dataset, desc='Running tokenizer on validation dataset')
-
-    if 'test' not in raw_datasets:
-        raise ValueError('Test dataset missing')
-    predict_dataset = raw_datasets['test']
-    if data_training_args.max_predict_samples is not None:
-        predict_dataset = predict_dataset.select(
-            range(data_training_args.max_predict_samples))
-    with training_args.main_process_first(desc='prediction dataset map pre-processing'):
-        predict_dataset = prepare_dataset(
-            predict_dataset, desc='Running tokenizer on prediction dataset')
+    train_dataset, eval_dataset, predict_dataset = prepare_datasets(
+        raw_datasets, dataset_args, training_args, preprocess_function)
 
     # Data collator
     data_collator = DataCollatorForSeq2Seq(
@@ -228,10 +160,11 @@ def main():
     )
 
     # Training
-    train_result = train_from_checkpoint(trainer, last_checkpoint, training_args)
+    train_result = train_from_checkpoint(
+        trainer, last_checkpoint, training_args)
 
     metrics = train_result.metrics
-    max_train_samples = data_training_args.max_train_samples or len(
+    max_train_samples = training_args.max_train_samples or len(
         train_dataset)
     metrics['train_samples'] = min(max_train_samples, len(train_dataset))
 
@@ -240,7 +173,7 @@ def main():
     trainer.save_state()
 
     kwargs = {'finetuned_from': model_args.model_name_or_path,
-                'tasks': 'summarization'}
+              'tasks': 'summarization'}
 
     if training_args.push_to_hub:
         trainer.push_to_hub(**kwargs)

src/train_classifier.py

@@ -3,14 +3,12 @@
 
 import logging
 import os
-import random
 import sys
 from dataclasses import dataclass, field
 from typing import Optional
 
 import datasets
 import numpy as np
-from datasets import load_metric
 
 import transformers
 from transformers import (
@@ -18,96 +16,39 @@ from transformers import (
     EvalPrediction,
     HfArgumentParser,
     Trainer,
-    default_data_collator,
     set_seed,
 )
 from transformers.utils import check_min_version
 from transformers.utils.versions import require_version
-from shared import CATEGORIES, load_datasets, CustomTrainingArguments, train_from_checkpoint, get_last_checkpoint
-from preprocess import PreprocessingDatasetArguments
+from shared import CATEGORIES, DatasetArguments, prepare_datasets, load_datasets, CustomTrainingArguments, train_from_checkpoint, get_last_checkpoint
 from model import get_model_tokenizer, ModelArguments
 
 # Will error if the minimal version of Transformers is not installed. Remove at your own risks.
-check_min_version("4.17.0")
-require_version("datasets>=1.8.0", "To fix: pip install -r requirements.txt")
+check_min_version('4.17.0')
+require_version('datasets>=1.8.0', 'To fix: pip install -r requirements.txt')
 
-os.environ["WANDB_DISABLED"] = "true"
+os.environ['WANDB_DISABLED'] = 'true'
 
 logger = logging.getLogger(__name__)
 
 
 @dataclass
-class DataArguments:
-    """
-    Arguments pertaining to what data we are going to input our model for training and eval.
-
-    Using `HfArgumentParser` we can turn this class
-    into argparse arguments to be able to specify them on
-    the command line.
-    """
-
-    max_seq_length: int = field(
-        default=512,
-        metadata={
-            "help": "The maximum total input sequence length after tokenization. Sequences longer "
-            "than this will be truncated, sequences shorter will be padded."
-        },
-    )
-    overwrite_cache: bool = field(
-        default=False, metadata={"help": "Overwrite the cached preprocessed datasets or not."}
-    )
-    pad_to_max_length: bool = field(
-        default=True,
-        metadata={
-            "help": "Whether to pad all samples to `max_seq_length`. "
-            "If False, will pad the samples dynamically when batching to the maximum length in the batch."
-        },
+class ClassifierDatasetArguments(DatasetArguments):
+    train_file: Optional[str] = field(
+        default='c_train.json', metadata={'help': 'The input training data file (a jsonlines file).'}
     )
-    max_train_samples: Optional[int] = field(
-        default=None,
+    validation_file: Optional[str] = field(
+        default='c_valid.json',
         metadata={
-            "help": "For debugging purposes or quicker training, truncate the number of training examples to this "
-            "value if set."
+            'help': 'An optional input evaluation data file to evaluate the metrics on (a jsonlines file).'
         },
     )
-    max_eval_samples: Optional[int] = field(
-        default=None,
+    test_file: Optional[str] = field(
+        default='c_test.json',
         metadata={
-            "help": "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
-            "value if set."
+            'help': 'An optional input test data file to evaluate the metrics on (a jsonlines file).'
         },
     )
-    max_predict_samples: Optional[int] = field(
-        default=None,
-        metadata={
-            "help": "For debugging purposes or quicker training, truncate the number of prediction examples to this "
-            "value if set."
-        },
-    )
-
-    dataset_cache_dir: Optional[str] = PreprocessingDatasetArguments.__dataclass_fields__[
-        'dataset_cache_dir']
-    data_dir: Optional[str] = PreprocessingDatasetArguments.__dataclass_fields__[
-        'data_dir']
-    train_file: Optional[str] = PreprocessingDatasetArguments.__dataclass_fields__[
-        'c_train_file']
-    validation_file: Optional[str] = PreprocessingDatasetArguments.__dataclass_fields__[
-        'c_validation_file']
-    test_file: Optional[str] = PreprocessingDatasetArguments.__dataclass_fields__[
-        'c_test_file']
-
-    def __post_init__(self):
-        if self.train_file is None or self.validation_file is None:
-            raise ValueError(
-                "Need either a GLUE task, a training/validation file or a dataset name.")
-        else:
-            train_extension = self.train_file.split(".")[-1]
-            assert train_extension in [
-                "csv", "json"], "`train_file` should be a csv or a json file."
-            validation_extension = self.validation_file.split(".")[-1]
-            assert (
-                validation_extension == train_extension
-            ), "`validation_file` should have the same extension (csv or json) as `train_file`."
 
 
 def main():
@@ -115,14 +56,17 @@ def main():
     # or by passing the --help flag to this script.
     # We now keep distinct sets of args, for a cleaner separation of concerns.
 
-    parser = HfArgumentParser(
-        (ModelArguments, DataArguments, CustomTrainingArguments))
-    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+    hf_parser = HfArgumentParser((
+        ModelArguments,
+        ClassifierDatasetArguments,
+        CustomTrainingArguments
+    ))
+    model_args, dataset_args, training_args = hf_parser.parse_args_into_dataclasses()
 
     # Setup logging
     logging.basicConfig(
-        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
-        datefmt="%m/%d/%Y %H:%M:%S",
+        format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
+        datefmt='%m/%d/%Y %H:%M:%S',
         handlers=[logging.StreamHandler(sys.stdout)],
     )
 
@@ -135,10 +79,10 @@ def main():
 
     # Log on each process the small summary:
     logger.warning(
-        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
-        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
+        f'Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}'
+        + f'distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}'
     )
-    logger.info(f"Training/evaluation parameters {training_args}")
+    logger.info(f'Training/evaluation parameters {training_args}')
 
     # Detecting last checkpoint.
     last_checkpoint = get_last_checkpoint(training_args)
@@ -148,7 +92,7 @@ def main():
 
     # Loading a dataset from your local files.
     # CSV/JSON training and evaluation files are needed.
-    raw_datasets = load_datasets(data_args)
+    raw_datasets = load_datasets(dataset_args)
 
     # See more about loading any type of standard or custom dataset at
     # https://huggingface.co/docs/datasets/loading_datasets.html.
@@ -158,69 +102,26 @@ def main():
         'id2label': {k: str(v).upper() for k, v in enumerate(CATEGORIES)},
         'label2id': {str(v).upper(): k for k, v in enumerate(CATEGORIES)}
     }
-    model, tokenizer = get_model_tokenizer(model_args, training_args, config_args=config_args, model_type='classifier')
+    model, tokenizer = get_model_tokenizer(
+        model_args, training_args, config_args=config_args, model_type='classifier')
 
-
-    # Padding strategy
-    if data_args.pad_to_max_length:
-        padding = "max_length"
-    else:
-        # We will pad later, dynamically at batch creation, to the max sequence length in each batch
-        padding = False
-
-    if data_args.max_seq_length > tokenizer.model_max_length:
+    if training_args.max_seq_length > tokenizer.model_max_length:
         logger.warning(
-            f"The max_seq_length passed ({data_args.max_seq_length}) is larger than the maximum length for the"
-            f"model ({tokenizer.model_max_length}). Using max_seq_length={tokenizer.model_max_length}."
+            f'The max_seq_length passed ({training_args.max_seq_length}) is larger than the maximum length for the'
+            f'model ({tokenizer.model_max_length}). Using max_seq_length={tokenizer.model_max_length}.'
         )
-    max_seq_length = min(data_args.max_seq_length, tokenizer.model_max_length)
+    max_seq_length = min(training_args.max_seq_length,
+                         tokenizer.model_max_length)
 
     def preprocess_function(examples):
         # Tokenize the texts
         result = tokenizer(
-            examples['text'], padding=padding, max_length=max_seq_length, truncation=True)
+            examples['text'], padding='max_length', max_length=max_seq_length, truncation=True)
         result['label'] = examples['label']
         return result
 
-    with training_args.main_process_first(desc="dataset map pre-processing"):
-        raw_datasets = raw_datasets.map(
-            preprocess_function,
-            batched=True,
-            load_from_cache_file=not data_args.overwrite_cache,
-            desc="Running tokenizer on dataset",
-        )
-    if training_args.do_train:
-        if "train" not in raw_datasets:
-            raise ValueError("--do_train requires a train dataset")
-        train_dataset = raw_datasets["train"]
-        if data_args.max_train_samples is not None:
-            train_dataset = train_dataset.select(
-                range(data_args.max_train_samples))
-
-    if training_args.do_eval:
-        if "validation" not in raw_datasets:
-            raise ValueError("--do_eval requires a validation dataset")
-        eval_dataset = raw_datasets["validation"]
-        if data_args.max_eval_samples is not None:
-            eval_dataset = eval_dataset.select(
-                range(data_args.max_eval_samples))
-
-    if training_args.do_predict or data_args.test_file is not None:
-        if "test" not in raw_datasets:
-            raise ValueError("--do_predict requires a test dataset")
-        predict_dataset = raw_datasets["test"]
-        if data_args.max_predict_samples is not None:
-            predict_dataset = predict_dataset.select(
-                range(data_args.max_predict_samples))
-
-    # Log a few random samples from the training set:
-    if training_args.do_train:
-        for index in random.sample(range(len(train_dataset)), 3):
-            logger.info(
-                f"Sample {index} of the training set: {train_dataset[index]}.")
-
-    # Get the metric function
-    metric = load_metric("accuracy")
+    train_dataset, eval_dataset, predict_dataset = prepare_datasets(
+        raw_datasets, dataset_args, training_args, preprocess_function)
 
     # You can define your custom compute_metrics function. It takes an `EvalPrediction` object (a namedtuple with a
     # predictions and label_ids field) and has to return a dictionary string to float.
@@ -228,20 +129,11 @@ def main():
         preds = p.predictions[0] if isinstance(
             p.predictions, tuple) else p.predictions
         preds = np.argmax(preds, axis=1)
-        if data_args.task_name is not None:
-            result = metric.compute(predictions=preds, references=p.label_ids)
-            if len(result) > 1:
-                result["combined_score"] = np.mean(
-                    list(result.values())).item()
-            return result
-        else:
-            return {"accuracy": (preds == p.label_ids).astype(np.float32).mean().item()}
+        return {'accuracy': (preds == p.label_ids).astype(np.float32).mean().item()}
 
     # Data collator will default to DataCollatorWithPadding when the tokenizer is passed to Trainer, so we change it if
     # we already did the padding.
-    if data_args.pad_to_max_length:
-        data_collator = default_data_collator
-    elif training_args.fp16:
+    if training_args.fp16:
         data_collator = DataCollatorWithPadding(
             tokenizer, pad_to_multiple_of=8)
     else:
@@ -264,24 +156,24 @@ def main():
 
     metrics = train_result.metrics
     max_train_samples = (
-        data_args.max_train_samples if data_args.max_train_samples is not None else len(
+        training_args.max_train_samples if training_args.max_train_samples is not None else len(
             train_dataset)
     )
-    metrics["train_samples"] = min(max_train_samples, len(train_dataset))
+    metrics['train_samples'] = min(max_train_samples, len(train_dataset))
 
     trainer.save_model()  # Saves the tokenizer too for easy upload
 
-    trainer.log_metrics("train", metrics)
-    trainer.save_metrics("train", metrics)
+    trainer.log_metrics('train', metrics)
+    trainer.save_metrics('train', metrics)
     trainer.save_state()
 
-    kwargs = {"finetuned_from": model_args.model_name_or_path,
-              "tasks": "text-classification"}
+    kwargs = {'finetuned_from': model_args.model_name_or_path,
+              'tasks': 'text-classification'}
     if training_args.push_to_hub:
         trainer.push_to_hub(**kwargs)
     else:
         trainer.create_model_card(**kwargs)
 
 
-if __name__ == "__main__":
+if __name__ == '__main__':
     main()