Fix prediction and evaluation arguments

src/evaluate.py

@@ -7,7 +7,7 @@ from transformers import (
 from preprocess import DatasetArguments, ProcessedArguments, get_words
 from model import get_classifier_vectorizer
 from shared import device
-from predict import ClassifierArguments, PredictArguments, predict, filter_predictions
+from predict import ClassifierArguments, predict, filter_predictions, TrainingOutputArguments
 from segment import word_start, word_end, SegmentationArguments, add_labels_to_words
 import pandas as pd
 from dataclasses import dataclass, field
@@ -19,7 +19,7 @@ import random
 
 
 @dataclass
-class EvaluationArguments:
+class EvaluationArguments(TrainingOutputArguments):
     """
     Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
     """
@@ -29,8 +29,7 @@ class EvaluationArguments:
             'help': 'The number of videos to test on'
         }
     )
-    model_path: Optional[str] = PredictArguments.__dataclass_fields__[
-        'model_path']
+
     data_dir: Optional[str] = DatasetArguments.__dataclass_fields__['data_dir']
     dataset: Optional[str] = DatasetArguments.__dataclass_fields__[
         'validation_file']

src/predict.py

@@ -1,3 +1,4 @@
+from transformers.trainer_utils import get_last_checkpoint
 from shared import OutputArguments
 from typing import Optional
 from segment import (
@@ -21,7 +22,6 @@ from dataclasses import dataclass, field
 from transformers import HfArgumentParser
 from shared import device
 import logging
-from transformers.trainer_utils import get_last_checkpoint
 
 
 def seconds_to_time(seconds):
@@ -31,12 +31,7 @@ def seconds_to_time(seconds):
 
 
 @dataclass
-class PredictArguments:
-
-    video_id: str = field(
-        metadata={
-            'help': 'Video to predict sponsorship segments for'}
-    )
+class TrainingOutputArguments:
 
     model_path: str = field(
         default=None,
@@ -59,6 +54,15 @@ class PredictArguments:
                 'Unable to find model, explicitly set `--model_path`')
 
 
+@dataclass
+class PredictArguments(TrainingOutputArguments):
+    video_id: str = field(
+        default=None,
+        metadata={
+            'help': 'Video to predict sponsorship segments for'}
+    )
+
+
 SPONSOR_MATCH_RE = fr'(?<={CustomTokens.START_SPONSOR.value})\s*(.*?)\s*(?={CustomTokens.END_SPONSOR.value}|$)'
 
 MATCH_WINDOW = 25       # Increase for accuracy, but takes longer: O(n^3)
@@ -252,6 +256,10 @@ def main():
     ))
     predict_args, segmentation_args, classifier_args = hf_parser.parse_args_into_dataclasses()
 
+    if predict_args.video_id is None:
+        print('No video ID supplied. Use `--video_id`.')
+        return
+
     model = AutoModelForSeq2SeqLM.from_pretrained(predict_args.model_path)
     model.to(device())
 

src/train.py

@@ -1,10 +1,8 @@
 from preprocess import load_datasets, DatasetArguments
 from predict import ClassifierArguments, SPONSOR_MATCH_RE, DEFAULT_TOKEN_PREFIX
-from shared import device
-from shared import GeneralArguments, OutputArguments
+from shared import device, GeneralArguments, OutputArguments
 from model import ModelArguments
 import transformers
-import logging
 from model import get_model, get_tokenizer
 import logging
 import os