Fix training arguments dataclasses

src/shared.py

@@ -1,5 +1,5 @@
 from transformers.trainer_utils import get_last_checkpoint as glc
-from transformers import TrainingArguments
+from transformers import Seq2SeqTrainingArguments, TrainingArguments
 import os
 from utils import re_findall
 import logging
@@ -76,14 +76,15 @@ _SEGMENT_END = END_SEGMENT_TEMPLATE.format(r'\w+')
 SEGMENT_MATCH_RE = fr'{_SEGMENT_START}\s*(?P<text>.*?)\s*(?:{_SEGMENT_END}|$)'
 
 
+def extract_sponsor_matches_from_text(text):
+    if CustomTokens.NO_SEGMENT.value in text:
+        return []
+    else:
+        return re_findall(SEGMENT_MATCH_RE, text)
+
+
 def extract_sponsor_matches(texts):
-    to_return = []
-    for text in texts:
-        if CustomTokens.NO_SEGMENT.value in text:
-            to_return.append([])
-        else:
-            to_return.append(re_findall(SEGMENT_MATCH_RE, text))
-    return to_return
+    return list(map(extract_sponsor_matches_from_text, texts))
 
 
 @dataclass
@@ -134,6 +135,22 @@ class DatasetArguments:
         },
     )
 
+    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).'
+        },
+    )
+
     def __post_init__(self):
         if self.train_file is None or self.validation_file is None:
             raise ValueError(
@@ -234,7 +251,7 @@ def load_datasets(dataset_args: DatasetArguments):
 
 
 @dataclass
-class CustomTrainingArguments(OutputArguments, TrainingArguments):
+class AdditionalTrainingArguments:
     seed: Optional[int] = GeneralArguments.__dataclass_fields__['seed']
 
     num_train_epochs: float = field(
@@ -242,7 +259,7 @@ class CustomTrainingArguments(OutputArguments, TrainingArguments):
 
     save_steps: int = field(default=5000, metadata={
                             'help': 'Save checkpoint every X updates steps.'})
-    eval_steps: int = field(default=5000, metadata={
+    eval_steps: int = field(default=25000, metadata={
                             'help': 'Run an evaluation every X steps.'})
     logging_steps: int = field(default=5000, metadata={
                                'help': 'Log every X updates steps.'})
@@ -311,6 +328,11 @@ class CustomTrainingArguments(OutputArguments, TrainingArguments):
     )
 
 
+@dataclass
+class CustomTrainingArguments(OutputArguments, AdditionalTrainingArguments):
+    pass
+
+
 logging.basicConfig()
 logger = logging.getLogger(__name__)
 

src/train.py

@@ -1,6 +1,6 @@
-from preprocess import PreprocessingDatasetArguments
 from shared import (
     CustomTokens,
+    DatasetArguments,
     prepare_datasets,
     load_datasets,
     CustomTrainingArguments,
@@ -17,13 +17,15 @@ from transformers import (
     DataCollatorForSeq2Seq,
     HfArgumentParser,
     Seq2SeqTrainer,
+    Seq2SeqTrainingArguments,
 )
 
 from transformers.utils import check_min_version
 from transformers.utils.versions import require_version
+from dataclasses import dataclass
 
 # Will error if the minimal version of Transformers is not installed. Remove at your own risks.
-check_min_version('4.13.0.dev0')
+check_min_version('4.17.0')
 require_version('datasets>=1.8.0',
                 'To fix: pip install -r requirements.txt')
 
@@ -40,6 +42,11 @@ logging.basicConfig(
 )
 
 
+@dataclass
+class Seq2SeqTrainingArguments(CustomTrainingArguments, Seq2SeqTrainingArguments):
+    pass
+
+
 def main():
 
     # See all possible arguments in src/transformers/training_args.py
@@ -48,8 +55,8 @@ def main():
 
     hf_parser = HfArgumentParser((
         ModelArguments,
-        PreprocessingDatasetArguments,
-        CustomTrainingArguments
+        DatasetArguments,
+        Seq2SeqTrainingArguments
     ))
     model_args, dataset_args, training_args = hf_parser.parse_args_into_dataclasses()
 

src/train_classifier.py

@@ -4,7 +4,7 @@
 import logging
 import os
 import sys
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from typing import Optional
 
 import datasets
@@ -16,11 +16,20 @@ from transformers import (
     EvalPrediction,
     HfArgumentParser,
     Trainer,
+    TrainingArguments,
     set_seed,
 )
 from transformers.utils import check_min_version
 from transformers.utils.versions import require_version
-from shared import CATEGORIES, DatasetArguments, prepare_datasets, load_datasets, CustomTrainingArguments, train_from_checkpoint, get_last_checkpoint
+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.
@@ -32,23 +41,19 @@ os.environ['WANDB_DISABLED'] = 'true'
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class ClassifierTrainingArguments(CustomTrainingArguments, TrainingArguments):
+    pass
+
+
 @dataclass
 class ClassifierDatasetArguments(DatasetArguments):
-    train_file: Optional[str] = field(
-        default='c_train.json', metadata={'help': 'The input training data file (a jsonlines file).'}
-    )
-    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).'
-        },
-    )
-    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).'
-        },
-    )
+    train_file: Optional[str] = DatasetArguments.__dataclass_fields__[
+        'c_train_file']
+    validation_file: Optional[str] = DatasetArguments.__dataclass_fields__[
+        'c_validation_file']
+    test_file: Optional[str] = DatasetArguments.__dataclass_fields__[
+        'c_test_file']
 
 
 def main():
@@ -59,7 +64,7 @@ def main():
     hf_parser = HfArgumentParser((
         ModelArguments,
         ClassifierDatasetArguments,
-        CustomTrainingArguments
+        ClassifierTrainingArguments
     ))
     model_args, dataset_args, training_args = hf_parser.parse_args_into_dataclasses()