Upgrade classifier to transformer-based model

src/classify.py

@@ -0,0 +1,41 @@
+from transformers import TextClassificationPipeline
+import preprocess
+import segment
+
+
+class SponsorBlockClassificationPipeline(TextClassificationPipeline):
+    def __init__(self, model, tokenizer):
+        device = next(model.parameters()).device.index
+        super().__init__(model=model, tokenizer=tokenizer,
+                         return_all_scores=True, truncation=True, device=device)
+
+    def preprocess(self, data, **tokenizer_kwargs):
+        # TODO add support for lists
+        texts = []
+
+        if not isinstance(data, list):
+            data = [data]
+
+        for d in data:
+            if isinstance(d, dict):  # Otherwise, get data from transcript
+                words = preprocess.get_words(d['video_id'])
+                segment_words = segment.extract_segment(
+                    words, d['start'], d['end'])
+                text = preprocess.clean_text(
+                    ' '.join(x['text'] for x in segment_words))
+                texts.append(text)
+            elif isinstance(d, str):  # If string, assume this is what user wants to classify
+                texts.append(d)
+            else:
+                raise ValueError(f'Invalid input type: "{type(d)}"')
+
+        return self.tokenizer(
+            texts, return_tensors=self.framework, **tokenizer_kwargs)
+
+
+def main():
+    pass
+
+
+if __name__ == '__main__':
+    main()

src/errors.py

@@ -1,9 +1,10 @@
+
 class SponsorBlockException(Exception):
     """Base class for all sponsor block exceptions"""
     pass
 
 
-class PredictionException(SponsorBlockException):
+class InferenceException(SponsorBlockException):
     """An exception occurred while predicting sponsor segments"""
     pass
 
@@ -21,8 +22,3 @@ class ModelError(SponsorBlockException):
 class ModelLoadError(ModelError):
     """An exception occurred while loading the model"""
     pass
-
-
-class ClassifierLoadError(ModelError):
-    """An exception occurred while loading the classifier"""
-    pass

src/evaluate.py

@@ -1,10 +1,10 @@
 
-from model import get_model_tokenizer
+from model import get_model_tokenizer_classifier, InferenceArguments
 from utils import jaccard
 from transformers import HfArgumentParser
-from preprocess import DatasetArguments, get_words
-from shared import GeneralArguments
-from predict import ClassifierArguments, predict, InferenceArguments
+from preprocess import get_words
+from shared import GeneralArguments, DatasetArguments
+from predict import predict
 from segment import extract_segment, word_start, word_end, SegmentationArguments, add_labels_to_words
 import pandas as pd
 from dataclasses import dataclass, field
@@ -134,11 +134,10 @@ def main():
         EvaluationArguments,
         DatasetArguments,
         SegmentationArguments,
-        ClassifierArguments,
         GeneralArguments
     ))
 
-    evaluation_args, dataset_args, segmentation_args, classifier_args, general_args = hf_parser.parse_args_into_dataclasses()
+    evaluation_args, dataset_args, segmentation_args, general_args = hf_parser.parse_args_into_dataclasses()
 
     # Load labelled data:
     final_path = os.path.join(
@@ -149,8 +148,8 @@ def main():
                      f'Run `python src/preprocess.py --update_database --do_create` to generate "{final_path}".')
         return
 
-    model, tokenizer = get_model_tokenizer(
-        evaluation_args.model_path, evaluation_args.cache_dir, general_args.no_cuda)
+    model, tokenizer, classifier = get_model_tokenizer_classifier(
+        evaluation_args, general_args)
 
     with open(final_path) as fp:
         final_data = json.load(fp)
@@ -187,8 +186,9 @@ def main():
                     continue
 
                 # Make predictions
-                predictions = predict(video_id, model, tokenizer,
-                                      segmentation_args, words, classifier_args)
+                predictions = predict(video_id, model, tokenizer, segmentation_args,
+                                      classifier=classifier,
+                                      min_probability=evaluation_args.min_probability)
 
                 # Get labels
                 sponsor_segments = final_data.get(video_id)

src/model.py

@@ -1,13 +1,68 @@
-from huggingface_hub import hf_hub_download
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
-from shared import CustomTokens
-from errors import ClassifierLoadError, ModelLoadError
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, AutoConfig, AutoModelForSequenceClassification, TrainingArguments
+from shared import CustomTokens, GeneralArguments
 from functools import lru_cache
-import pickle
-import os
 from dataclasses import dataclass, field
-from typing import Optional
+from typing import Optional, Union
 import torch
+import classify
+import base64
+import re
+import requests
+import json
+import logging
+
+logging.basicConfig()
+logger = logging.getLogger(__name__)
+
+# Public innertube key (b64 encoded so that it is not incorrectly flagged)
+INNERTUBE_KEY = base64.b64decode(
+    b'QUl6YVN5QU9fRkoyU2xxVThRNFNURUhMR0NpbHdfWTlfMTFxY1c4').decode()
+
+YT_CONTEXT = {
+    'client': {
+        'userAgent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/96.0.4664.110 Safari/537.36,gzip(gfe)',
+        'clientName': 'WEB',
+        'clientVersion': '2.20211221.00.00',
+    }
+}
+_YT_INITIAL_DATA_RE = r'(?:window\s*\[\s*["\']ytInitialData["\']\s*\]|ytInitialData)\s*=\s*({.+?})\s*;\s*(?:var\s+meta|</script|\n)'
+
+
+def get_all_channel_vids(channel_id):
+    continuation = None
+    while True:
+        if continuation is None:
+            params = {'list': channel_id.replace('UC', 'UU', 1)}
+            response = requests.get(
+                'https://www.youtube.com/playlist', params=params)
+            items = json.loads(re.search(_YT_INITIAL_DATA_RE, response.text).group(1))['contents']['twoColumnBrowseResultsRenderer']['tabs'][0]['tabRenderer']['content'][
+                'sectionListRenderer']['contents'][0]['itemSectionRenderer']['contents'][0]['playlistVideoListRenderer']['contents']
+        else:
+            params = {'key': INNERTUBE_KEY}
+            data = {
+                'context': YT_CONTEXT,
+                'continuation': continuation
+            }
+            response = requests.post(
+                'https://www.youtube.com/youtubei/v1/browse', params=params, json=data)
+            items = response.json()[
+                'onResponseReceivedActions'][0]['appendContinuationItemsAction']['continuationItems']
+
+        new_token = None
+        for vid in items:
+            info = vid.get('playlistVideoRenderer')
+            if info:
+                yield info['videoId']
+                continue
+
+            info = vid.get('continuationItemRenderer')
+            if info:
+                new_token = info['continuationEndpoint']['continuationCommand']['token']
+
+        if new_token is None:
+            break
+        continuation = new_token
+
 
 
 @dataclass
@@ -17,34 +72,24 @@ class ModelArguments:
     """
 
     model_name_or_path: str = field(
-        default=None,
-        # default='google/t5-v1_1-small',  # t5-small
         metadata={
             'help': 'Path to pretrained model or model identifier from huggingface.co/models'
         }
     )
 
-    # config_name: Optional[str] = field( # TODO remove?
-    #     default=None, metadata={'help': 'Pretrained config name or path if not the same as model_name'}
-    # )
-    # tokenizer_name: Optional[str] = field(
-    #     default=None, metadata={
-    #         'help': 'Pretrained tokenizer name or path if not the same as model_name'
-    #     }
-    # )
     cache_dir: Optional[str] = field(
         default='models',
         metadata={
             'help': 'Where to store the pretrained models downloaded from huggingface.co'
         },
     )
-    use_fast_tokenizer: bool = field(  # TODO remove?
+    use_fast_tokenizer: bool = field(
         default=True,
         metadata={
             'help': 'Whether to use one of the fast tokenizer (backed by the tokenizers library) or not.'
         },
     )
-    model_revision: str = field(  # TODO remove?
+    model_revision: str = field(
         default='main',
         metadata={
             'help': 'The specific model version to use (can be a branch name, tag name or commit id).'
@@ -57,62 +102,131 @@ class ModelArguments:
             'with private models).'
         },
     )
-    resize_position_embeddings: Optional[bool] = field(
+
+import itertools
+from errors import InferenceException
+
+@dataclass
+class InferenceArguments(ModelArguments):
+
+    model_name_or_path: str = field(
+        default='Xenova/sponsorblock-small',
+        metadata={
+            'help': 'Path to pretrained model used for prediction'
+        }
+    )
+    classifier_model_name_or_path: str = field(
+        default='Xenova/sponsorblock-classifier-v2',
+        metadata={
+            'help': 'Use a pretrained classifier'
+        }
+    )
+
+    max_videos: Optional[int] = field(
         default=None,
         metadata={
-            'help': "Whether to automatically resize the position embeddings if `max_source_length` exceeds the model's position embeddings."
-        },
+            'help': 'The number of videos to test on'
+        }
+    )
+    start_index: int = field(default=None, metadata={
+        'help': 'Video to start the evaluation at.'})
+    channel_id: Optional[str] = field(
+        default=None,
+        metadata={
+            'help': 'Used to evaluate a channel'
+        }
+    )
+    video_ids: str = field(
+        default_factory=lambda: [],
+        metadata={
+            'nargs': '+'
+        }
     )
 
+    output_as_json: bool = field(default=False, metadata={
+                                 'help': 'Output evaluations as JSON'})
 
-@lru_cache(maxsize=None)
-def get_classifier_vectorizer(classifier_args):
-    # Classifier
-    classifier_path = os.path.join(
-        classifier_args.classifier_dir, classifier_args.classifier_file)
-    if not os.path.exists(classifier_path):
-        hf_hub_download(repo_id=classifier_args.classifier_model,
-                        filename=classifier_args.classifier_file,
-                        cache_dir=classifier_args.classifier_dir,
-                        force_filename=classifier_args.classifier_file,
-                        )
-    with open(classifier_path, 'rb') as fp:
-        classifier = pickle.load(fp)
-
-    # Vectorizer
-    vectorizer_path = os.path.join(
-        classifier_args.classifier_dir, classifier_args.vectorizer_file)
-    if not os.path.exists(vectorizer_path):
-        hf_hub_download(repo_id=classifier_args.classifier_model,
-                        filename=classifier_args.vectorizer_file,
-                        cache_dir=classifier_args.classifier_dir,
-                        force_filename=classifier_args.vectorizer_file,
-                        )
-    with open(vectorizer_path, 'rb') as fp:
-        vectorizer = pickle.load(fp)
-
-    return classifier, vectorizer
-
-
-@lru_cache(maxsize=None)
-def get_model_tokenizer(model_name_or_path, cache_dir=None, no_cuda=False):
-    if model_name_or_path is None:
-        raise ModelLoadError('Invalid model_name_or_path.')
-
-    # Load pretrained model and tokenizer
-    model = AutoModelForSeq2SeqLM.from_pretrained(
-        model_name_or_path, cache_dir=cache_dir)
-    if not no_cuda:
-        model.to('cuda' if torch.cuda.is_available() else 'cpu')
+    min_probability: float = field(
+        default=0.5, metadata={'help': 'Remove all predictions whose classification probability is below this threshold.'})
+
+    def __post_init__(self):
+
+        self.video_ids = list(map(str.strip, self.video_ids))
+
+        if any(len(video_id) != 11 for video_id in self.video_ids):
+            raise InferenceException('Invalid video IDs (length not 11)')
+
+        if self.channel_id is not None:
+            start = self.start_index or 0
+            end = None if self.max_videos is None else start + self.max_videos
+
+            channel_video_ids = list(itertools.islice(get_all_channel_vids(
+                self.channel_id), start, end))
+            logger.info(
+                f'Found {len(channel_video_ids)} for channel {self.channel_id}')
+
+            self.video_ids += channel_video_ids
+
+
+
+def get_model_tokenizer_classifier(inference_args: InferenceArguments, general_args: GeneralArguments):
+
+    original_path = inference_args.model_name_or_path
+
+    # Load main model and tokenizer
+    model, tokenizer = get_model_tokenizer(inference_args, general_args)
+
+    # Load classifier
+    inference_args.model_name_or_path = inference_args.classifier_model_name_or_path
+    classifier_model, classifier_tokenizer = get_model_tokenizer(
+        inference_args, general_args, model_type='classifier')
+
+    classifier = classify.SponsorBlockClassificationPipeline(
+        classifier_model, classifier_tokenizer)
+
+    # Reset to original model_name_or_path
+    inference_args.model_name_or_path = original_path
+
+    return model, tokenizer, classifier
+
+
+def get_model_tokenizer(model_args: ModelArguments, general_args: Union[GeneralArguments, TrainingArguments] = None, config_args=None, model_type='seq2seq'):
+    if config_args is None:
+        config_args = {}
+
+    use_auth_token = True if model_args.use_auth_token else None
+
+    config = AutoConfig.from_pretrained(
+        model_args.model_name_or_path,
+        cache_dir=model_args.cache_dir,
+        revision=model_args.model_revision,
+        use_auth_token=use_auth_token,
+        **config_args
+    )
 
     tokenizer = AutoTokenizer.from_pretrained(
-        model_name_or_path, cache_dir=cache_dir)
+        model_args.model_name_or_path,
+        cache_dir=model_args.cache_dir,
+        use_fast=model_args.use_fast_tokenizer,
+        revision=model_args.model_revision,
+        use_auth_token=use_auth_token,
+    )
 
-    # Ensure model and tokenizer contain the custom tokens
+    model_type = AutoModelForSeq2SeqLM if model_type == 'seq2seq' else AutoModelForSequenceClassification
+    model = model_type.from_pretrained(
+        model_args.model_name_or_path,
+        config=config,
+        cache_dir=model_args.cache_dir,
+        revision=model_args.model_revision,
+        use_auth_token=use_auth_token,
+    )
+
+    # Add custom tokens
     CustomTokens.add_custom_tokens(tokenizer)
     model.resize_token_embeddings(len(tokenizer))
 
-    # TODO find a way to adjust based on model's input size
-    # print('tokenizer.model_max_length', tokenizer.model_max_length)
+    # Potentially move model to gpu
+    if general_args is not None and not general_args.no_cuda:
+        model.to('cuda' if torch.cuda.is_available() else 'cpu')
 
     return model, tokenizer

src/moderate.py

@@ -0,0 +1,104 @@
+import torch
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    HfArgumentParser
+)
+
+from train_classifier import ClassifierModelArguments
+from shared import CATEGORIES, DatasetArguments
+from tqdm import tqdm
+
+from preprocess import get_words, clean_text
+from segment import extract_segment
+import os
+import json
+import numpy as np
+
+
+def softmax(_outputs):
+    maxes = np.max(_outputs, axis=-1, keepdims=True)
+    shifted_exp = np.exp(_outputs - maxes)
+    return shifted_exp / shifted_exp.sum(axis=-1, keepdims=True)
+
+
+def main():
+    # See all possible arguments in src/transformers/training_args.py
+    # or by passing the --help flag to this script.
+    # We now keep distinct sets of args, for a cleaner separation of concerns.
+
+    parser = HfArgumentParser((ClassifierModelArguments, DatasetArguments))
+    model_args, dataset_args = parser.parse_args_into_dataclasses()
+
+    model = AutoModelForSequenceClassification.from_pretrained(
+        model_args.model_name_or_path)
+    tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path)
+
+    processed_db_path = os.path.join(
+        dataset_args.data_dir, dataset_args.processed_database)
+    with open(processed_db_path) as fp:
+        data = json.load(fp)
+
+    mapped_categories = {
+        str(v).lower(): k for k, v in enumerate(CATEGORIES)
+    }
+
+    for video_id, segments in tqdm(data.items()):
+
+        words = get_words(video_id)
+
+        if not words:
+            continue  # No/empty transcript for video_id
+
+        valid_segments = []
+        texts = []
+        for segment in segments:
+            segment_words = extract_segment(
+                words, segment['start'], segment['end'])
+            text = clean_text(' '.join(x['text'] for x in segment_words))
+
+            duration = segment['end'] - segment['start']
+            wps = len(segment_words)/duration if duration > 0 else 0
+            if wps < 1.5:
+                continue
+
+            # Do not worry about those that are locked or have enough votes
+            if segment['locked']:  # or segment['votes'] > 5:
+                continue
+
+            texts.append(text)
+            valid_segments.append(segment)
+
+        if not texts:
+            continue  # No valid segments
+
+        model_inputs = tokenizer(
+            texts, return_tensors='pt', padding=True, truncation=True)
+
+        with torch.no_grad():
+            model_outputs = model(**model_inputs)
+            outputs = list(map(lambda x: x.numpy(), model_outputs['logits']))
+
+        scores = softmax(outputs)
+
+        for segment, text, score in zip(valid_segments, texts, scores):
+            predicted_index = score.argmax().item()
+
+            if predicted_index == mapped_categories[segment['category']]:
+                continue  # Ignore correct segments
+
+            a = {k: round(float(score[i]), 3)
+                 for i, k in enumerate(CATEGORIES)}
+
+            del segment['submission_time']
+            segment.update({
+                'predicted': str(CATEGORIES[predicted_index]).lower(),
+                'text': text,
+                'scores': a
+            })
+
+            print(json.dumps(segment))
+
+
+if __name__ == "__main__":
+    main()

src/predict.py

@@ -1,17 +1,8 @@
-import itertools
-import base64
-import re
-import requests
-import json
+
 from transformers import HfArgumentParser
-from transformers.trainer_utils import get_last_checkpoint
 from dataclasses import dataclass, field
 import logging
-import os
-import itertools
-from utils import re_findall
-from shared import CustomTokens, START_SEGMENT_TEMPLATE, END_SEGMENT_TEMPLATE, GeneralArguments, OutputArguments, seconds_to_time
-from typing import Optional
+from shared import CustomTokens, extract_sponsor_matches, GeneralArguments, seconds_to_time
 from segment import (
     generate_segments,
     extract_segment,
@@ -22,129 +13,12 @@ from segment import (
     SegmentationArguments
 )
 import preprocess
-from errors import PredictionException, TranscriptError, ModelLoadError, ClassifierLoadError
-from model import ModelArguments, get_classifier_vectorizer, get_model_tokenizer
+from errors import TranscriptError
+from model import get_model_tokenizer_classifier, InferenceArguments
 
 logging.basicConfig()
 logger = logging.getLogger(__name__)
 
-# Public innertube key (b64 encoded so that it is not incorrectly flagged)
-INNERTUBE_KEY = base64.b64decode(
-    b'QUl6YVN5QU9fRkoyU2xxVThRNFNURUhMR0NpbHdfWTlfMTFxY1c4').decode()
-
-YT_CONTEXT = {
-    'client': {
-        'userAgent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/96.0.4664.110 Safari/537.36,gzip(gfe)',
-        'clientName': 'WEB',
-        'clientVersion': '2.20211221.00.00',
-    }
-}
-_YT_INITIAL_DATA_RE = r'(?:window\s*\[\s*["\']ytInitialData["\']\s*\]|ytInitialData)\s*=\s*({.+?})\s*;\s*(?:var\s+meta|</script|\n)'
-
-
-def get_all_channel_vids(channel_id):
-    continuation = None
-    while True:
-        if continuation is None:
-            params = {'list': channel_id.replace('UC', 'UU', 1)}
-            response = requests.get(
-                'https://www.youtube.com/playlist', params=params)
-            items = json.loads(re.search(_YT_INITIAL_DATA_RE, response.text).group(1))['contents']['twoColumnBrowseResultsRenderer']['tabs'][0]['tabRenderer']['content'][
-                'sectionListRenderer']['contents'][0]['itemSectionRenderer']['contents'][0]['playlistVideoListRenderer']['contents']
-        else:
-            params = {'key': INNERTUBE_KEY}
-            data = {
-                'context': YT_CONTEXT,
-                'continuation': continuation
-            }
-            response = requests.post(
-                'https://www.youtube.com/youtubei/v1/browse', params=params, json=data)
-            items = response.json()[
-                'onResponseReceivedActions'][0]['appendContinuationItemsAction']['continuationItems']
-
-        new_token = None
-        for vid in items:
-            info = vid.get('playlistVideoRenderer')
-            if info:
-                yield info['videoId']
-                continue
-
-            info = vid.get('continuationItemRenderer')
-            if info:
-                new_token = info['continuationEndpoint']['continuationCommand']['token']
-
-        if new_token is None:
-            break
-        continuation = new_token
-
-
-@dataclass
-class InferenceArguments:
-
-    model_path: str = field(
-        default='Xenova/sponsorblock-small',
-        metadata={
-            'help': 'Path to pretrained model used for prediction'
-        }
-    )
-    cache_dir: Optional[str] = ModelArguments.__dataclass_fields__['cache_dir']
-
-    output_dir: Optional[str] = OutputArguments.__dataclass_fields__[
-        'output_dir']
-
-    max_videos: Optional[int] = field(
-        default=None,
-        metadata={
-            'help': 'The number of videos to test on'
-        }
-    )
-    start_index: int = field(default=None, metadata={
-        'help': 'Video to start the evaluation at.'})
-    channel_id: Optional[str] = field(
-        default=None,
-        metadata={
-            'help': 'Used to evaluate a channel'
-        }
-    )
-    video_ids: str = field(
-        default_factory=lambda: [],
-        metadata={
-            'nargs': '+'
-        }
-    )
-
-    output_as_json: bool = field(default=False, metadata={
-                                 'help': 'Output evaluations as JSON'})
-
-    def __post_init__(self):
-        # Try to load model from latest checkpoint
-        if self.model_path is None:
-            if os.path.exists(self.output_dir):
-                last_checkpoint = get_last_checkpoint(self.output_dir)
-                if last_checkpoint is not None:
-                    self.model_path = last_checkpoint
-                else:
-                    raise ModelLoadError(
-                        'Unable to load model from checkpoint, explicitly set `--model_path`')
-            else:
-                raise ModelLoadError(
-                    f'Unable to find model in {self.output_dir}, explicitly set `--model_path`')
-
-        if any(len(video_id) != 11 for video_id in self.video_ids):
-            raise PredictionException('Invalid video IDs (length not 11)')
-
-        if self.channel_id is not None:
-            start = self.start_index or 0
-            end = None if self.max_videos is None else start + self.max_videos
-
-            channel_video_ids = list(itertools.islice(get_all_channel_vids(
-                self.channel_id), start, end))
-            logger.info(
-                f'Found {len(channel_video_ids)} for channel {self.channel_id}')
-
-            self.video_ids += channel_video_ids
-
-
 @dataclass
 class PredictArguments(InferenceArguments):
     video_id: str = field(
@@ -160,10 +34,6 @@ class PredictArguments(InferenceArguments):
         super().__post_init__()
 
 
-_SEGMENT_START = START_SEGMENT_TEMPLATE.format(r'(?P<category>\w+)')
-_SEGMENT_END = END_SEGMENT_TEMPLATE.format(r'\w+')
-SEGMENT_MATCH_RE = fr'{_SEGMENT_START}\s*(?P<text>.*?)\s*(?:{_SEGMENT_END}|$)'
-
 MATCH_WINDOW = 25       # Increase for accuracy, but takes longer: O(n^3)
 MERGE_TIME_WITHIN = 8   # Merge predictions if they are within x seconds
 
@@ -171,70 +41,35 @@ MERGE_TIME_WITHIN = 8   # Merge predictions if they are within x seconds
 START_TIME_ZERO_THRESHOLD = 0.08
 
 
-@dataclass(frozen=True, eq=True)
-class ClassifierArguments:
-    classifier_model: Optional[str] = field(
-        default='Xenova/sponsorblock-classifier',
-        metadata={
-            'help': 'Use a pretrained classifier'
-        }
-    )
-
-    classifier_dir: Optional[str] = field(
-        default='classifiers',
-        metadata={
-            'help': 'The directory that contains the classifier and vectorizer.'
-        }
-    )
-
-    classifier_file: Optional[str] = field(
-        default='classifier.pickle',
-        metadata={
-            'help': 'The name of the classifier'
-        }
-    )
-
-    vectorizer_file: Optional[str] = field(
-        default='vectorizer.pickle',
-        metadata={
-            'help': 'The name of the vectorizer'
-        }
-    )
-
-    min_probability: float = field(
-        default=0.5, metadata={'help': 'Remove all predictions whose classification probability is below this threshold.'})
-
-
-def filter_and_add_probabilities(predictions, classifier_args):
+def filter_and_add_probabilities(predictions, classifier, min_probability):
     """Use classifier to filter predictions"""
     if not predictions:
         return predictions
 
-    classifier, vectorizer = get_classifier_vectorizer(classifier_args)
+    # We update the predicted category from the extractive transformer
+    # if the classifier is confident enough it is another category
 
-    transformed_segments = vectorizer.transform([
+    texts = [
         preprocess.clean_text(' '.join([x['text'] for x in pred['words']]))
         for pred in predictions
-    ])
-    probabilities = classifier.predict_proba(transformed_segments)
+    ]
+    classifications = classifier(texts)
 
-    # Transformer sometimes says segment is of another category, so we
-    # update category and probabilities if classifier is confident it is another category
     filtered_predictions = []
-    for prediction, probabilities in zip(predictions, probabilities):
-        predicted_probabilities = {k: v for k,
-                                   v in zip(CATEGORIES, probabilities)}
+    for prediction, probabilities in zip(predictions, classifications):
+        predicted_probabilities = {
+            p['label'].lower(): p['score'] for p in probabilities}
 
         # Get best category + probability
         classifier_category = max(
             predicted_probabilities, key=predicted_probabilities.get)
         classifier_probability = predicted_probabilities[classifier_category]
 
-        if classifier_category is None and classifier_probability > classifier_args.min_probability:
+        if classifier_category == 'none' and classifier_probability > min_probability:
             continue  # Ignore
 
         if (prediction['category'] not in predicted_probabilities) \
-                or (classifier_category is not None and classifier_probability > 0.5):  # TODO make param
+                or (classifier_category != 'none' and classifier_probability > 0.5):  # TODO make param
             # Unknown category or we are confident enough to overrule,
             # so change category to what was predicted by classifier
             prediction['category'] = classifier_category
@@ -252,7 +87,7 @@ def filter_and_add_probabilities(predictions, classifier_args):
     return filtered_predictions
 
 
-def predict(video_id, model, tokenizer, segmentation_args, words=None, classifier_args=None):
+def predict(video_id, model, tokenizer, segmentation_args, words=None, classifier=None, min_probability=None):
     # Allow words to be passed in so that we don't have to get the words if we already have them
     if words is None:
         words = preprocess.get_words(video_id)
@@ -272,13 +107,9 @@ def predict(video_id, model, tokenizer, segmentation_args, words=None, classifie
         prediction['words'] = extract_segment(
             words, prediction['start'], prediction['end'])
 
-    # TODO add back
-    if classifier_args is not None:
-        try:
-            predictions = filter_and_add_probabilities(
-                predictions, classifier_args)
-        except ClassifierLoadError:
-            print('Unable to load classifer')
+    if classifier is not None:
+        predictions = filter_and_add_probabilities(
+            predictions, classifier, min_probability)
 
     return predictions
 
@@ -300,9 +131,6 @@ def greedy_match(list, sublist):
     return best_i, best_j, best_k
 
 
-CATEGORIES = [None, 'SPONSOR', 'SELFPROMO', 'INTERACTION']
-
-
 def predict_sponsor_text(text, model, tokenizer):
     """Given a body of text, predict the words which are part of the sponsor"""
     model_device = next(model.parameters()).device
@@ -322,11 +150,7 @@ def predict_sponsor_text(text, model, tokenizer):
 
 def predict_sponsor_matches(text, model, tokenizer):
     sponsorship_text = predict_sponsor_text(text, model, tokenizer)
-
-    if CustomTokens.NO_SEGMENT.value in sponsorship_text:
-        return []
-
-    return re_findall(SEGMENT_MATCH_RE, sponsorship_text)
+    return extract_sponsor_matches(sponsorship_text)
 
 
 def segments_to_predictions(segments, model, tokenizer):
@@ -400,24 +224,23 @@ def main():
     hf_parser = HfArgumentParser((
         PredictArguments,
         SegmentationArguments,
-        ClassifierArguments,
         GeneralArguments
     ))
-    predict_args, segmentation_args, classifier_args, general_args = hf_parser.parse_args_into_dataclasses()
+    predict_args, segmentation_args, general_args = hf_parser.parse_args_into_dataclasses()
 
     if not predict_args.video_ids:
         logger.error(
             'No video IDs supplied. Use `--video_id`, `--video_ids`, or `--channel_id`.')
         return
 
-    model, tokenizer = get_model_tokenizer(
-        predict_args.model_path, predict_args.cache_dir, general_args.no_cuda)
+    model, tokenizer, classifier = get_model_tokenizer_classifier(
+        predict_args, general_args)
 
     for video_id in predict_args.video_ids:
-        video_id = video_id.strip()
         try:
-            predictions = predict(video_id, model, tokenizer,
-                                  segmentation_args, classifier_args=classifier_args)
+            predictions = predict(video_id, model, tokenizer, segmentation_args,
+                                  classifier=classifier,
+                                  min_probability=predict_args.min_probability)
         except TranscriptError:
             logger.warning(f'No transcript available for {video_id}')
             continue

src/preprocess.py

@@ -1,14 +1,15 @@
+from shared import DatasetArguments
 from utils import jaccard
 from functools import lru_cache
 from datetime import datetime
 import itertools
-from typing import Optional, List
-from model import ModelArguments
+from typing import Optional
+import model as model_module
 import segment
 from tqdm import tqdm
 from dataclasses import dataclass, field
 from transformers import HfArgumentParser
-from shared import ACTION_OPTIONS, CATGEGORY_OPTIONS, START_SEGMENT_TEMPLATE, END_SEGMENT_TEMPLATE, GeneralArguments, CustomTokens
+from shared import extract_sponsor_matches, ACTION_OPTIONS, CATEGORIES, CATGEGORY_OPTIONS, START_SEGMENT_TEMPLATE, END_SEGMENT_TEMPLATE, GeneralArguments, CustomTokens
 import csv
 import re
 import random
@@ -213,9 +214,10 @@ def get_words(video_id, process=True, transcript_type='auto', fallback='manual',
                 else:
                     ts = transcript_list.find_generated_transcript(
                         LANGUAGE_PREFERENCE_LIST)
-
-                raw_transcript_json = ts._http_client.get(
-                    f'{ts._url}&fmt=json3').json()
+                raw_transcript = ts._http_client.get(
+                    f'{ts._url}&fmt=json3').content
+                if raw_transcript:
+                    raw_transcript_json = json.loads(raw_transcript)
 
     except (TooManyRequests, YouTubeRequestFailed):
         raise  # Cannot recover from these errors and do not mark as empty transcript
@@ -386,9 +388,14 @@ class PreprocessArguments:
 
     max_date: str = field(
         # default='01/01/9999', # Include all
-        default='02/02/2022',
+        default='01/03/2022',
         metadata={'help': 'Only use videos that have some segment from before this date (exclusive). This allows for videos to have segments be corrected, but ignores new videos (posted after this date) to enter the pool.'})
 
+    # max_unseen_date: str = field( # TODO
+    #     default='02/03/2022',
+    #     metadata={'help': 'Generate test and validation data from `max_date` to `max_unseen_date`'})
+    # Specify min/max video id for splitting (seen vs. unseen)
+
     keep_duplicate_segments: bool = field(
         default=False, metadata={'help': 'Keep duplicate segments'}
     )
@@ -482,25 +489,7 @@ def download_file(url, filename):
 
 
 @dataclass
-class DatasetArguments:
-    data_dir: Optional[str] = field(
-        default='data',
-        metadata={
-            'help': 'The directory which stores train, test and/or validation data.'
-        },
-    )
-    processed_file: Optional[str] = field(
-        default='segments.json',
-        metadata={
-            'help': 'Processed data file'
-        },
-    )
-    processed_database: Optional[str] = field(
-        default='processed_database.json',
-        metadata={
-            'help': 'Processed database file'
-        },
-    )
+class PreprocessingDatasetArguments(DatasetArguments):
 
     train_file: Optional[str] = field(
         default='train.json', metadata={'help': 'The input training data file (a jsonlines file).'}
@@ -508,21 +497,38 @@ class DatasetArguments:
     validation_file: Optional[str] = field(
         default='valid.json',
         metadata={
-            'help': 'An optional input evaluation data file to evaluate the metrics (rouge) on (a jsonlines file).'
+            '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 (rouge) on (a jsonlines file).'
+            'help': 'An optional input test data file to evaluate the metrics on (a jsonlines file).'
         },
     )
-    excess_file: Optional[str] = field(
-        default='excess.json',
+
+    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': 'The excess segments left after the split'
+            '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={
@@ -555,9 +561,9 @@ def main():
     # Generate final.json from sponsorTimes.csv
     hf_parser = HfArgumentParser((
         PreprocessArguments,
-        DatasetArguments,
+        PreprocessingDatasetArguments,
         segment.SegmentationArguments,
-        ModelArguments,
+        model_module.ModelArguments,
         GeneralArguments
     ))
     preprocess_args, dataset_args, segmentation_args, model_args, general_args = hf_parser.parse_args_into_dataclasses()
@@ -821,8 +827,7 @@ def main():
         # , max_videos, max_segments
 
         from model import get_model_tokenizer
-        model, tokenizer = get_model_tokenizer(
-            model_args.model_name_or_path, model_args.cache_dir, general_args.no_cuda)
+        model, tokenizer = get_model_tokenizer(model_args, general_args)
 
         # TODO
         # count_videos = 0
@@ -871,8 +876,9 @@ def main():
                         continue
 
                     d = {
-                        'video_index': offset + start_index,
+                        # 'video_index': offset + start_index,
                         'video_id': video_id,
+                        # 'uuid': video_id, # TODO add uuid
                         'text': ' '.join(x['cleaned'] for x in seg),
                         'start': seg_start,
                         'end': seg_end,
@@ -919,7 +925,7 @@ def main():
             z = int(preprocess_args.percentage_positive /
                     percentage_negative * len(non_sponsors))
 
-            excess = sponsors[z:]
+            # excess = sponsors[z:]
             sponsors = sponsors[:z]
 
         else:
@@ -927,7 +933,7 @@ def main():
             z = int(percentage_negative /
                     preprocess_args.percentage_positive * len(sponsors))
 
-            excess = non_sponsors[z:]
+            # excess = non_sponsors[z:]
             non_sponsors = non_sponsors[:z]
 
         logger.info('Join')
@@ -935,6 +941,7 @@ def main():
 
         random.shuffle(all_labelled_segments)
 
+        # TODO split based on video ids
         logger.info('Split')
         ratios = [preprocess_args.train_split,
                   preprocess_args.test_split,
@@ -958,15 +965,52 @@ def main():
             else:
                 logger.info(f'Skipping {name}')
 
+        classifier_splits = {
+            dataset_args.c_train_file: train_data,
+            dataset_args.c_test_file: test_data,
+            dataset_args.c_validation_file: valid_data
+        }
+
+        none_category = CATEGORIES.index(None)
+
+        # Output training, testing and validation data
+        for name, items in classifier_splits.items():
+            outfile = os.path.join(dataset_args.data_dir, name)
+            if not os.path.exists(outfile) or preprocess_args.overwrite:
+                with open(outfile, 'w', encoding='utf-8') as fp:
+                    for i in items:
+                        x = json.loads(i)  # TODO add uuid
+                        labelled_items = []
+
+                        matches = extract_sponsor_matches(x['extracted'])
+
+                        if x['extracted'] == CustomTokens.NO_SEGMENT.value:
+                            labelled_items.append({
+                                'text': x['text'],
+                                'label': none_category
+                            })
+                        else:
+                            for match in matches:
+                                labelled_items.append({
+                                    'text': match['text'],
+                                    'label': CATEGORIES.index(match['category'])
+                                })
+
+                        for labelled_item in labelled_items:
+                            print(json.dumps(labelled_item), file=fp)
+
+            else:
+                logger.info(f'Skipping {name}')
+
         logger.info('Write')
         # Save excess items
-        excess_path = os.path.join(
-            dataset_args.data_dir, dataset_args.excess_file)
-        if not os.path.exists(excess_path) or preprocess_args.overwrite:
-            with open(excess_path, 'w', encoding='utf-8') as fp:
-                fp.writelines(excess)
-        else:
-            logger.info(f'Skipping {dataset_args.excess_file}')
+        # excess_path = os.path.join(
+        #     dataset_args.data_dir, dataset_args.excess_file)
+        # if not os.path.exists(excess_path) or preprocess_args.overwrite:
+        #     with open(excess_path, 'w', encoding='utf-8') as fp:
+        #         fp.writelines(excess)
+        # else:
+        #     logger.info(f'Skipping {dataset_args.excess_file}')
 
         logger.info(
             f'Finished splitting: {len(sponsors)} sponsors, {len(non_sponsors)} non sponsors')

src/segment.py

@@ -121,6 +121,8 @@ def generate_segments(words, tokenizer, segmentation_args):
 
 def extract_segment(words, start, end, map_function=None):
     """Extracts all words with time in [start, end]"""
+    if words is None:
+        words = []
 
     a = max(binary_search_below(words, 0, len(words), start), 0)
     b = min(binary_search_above(words, -1, len(words) - 1, end) + 1, len(words))

src/shared.py

@@ -1,3 +1,10 @@
+from transformers.trainer_utils import get_last_checkpoint as glc
+from transformers import TrainingArguments
+import os
+from utils import re_findall
+import logging
+import sys
+from datasets import load_dataset
 import re
 import gc
 from time import time_ns
@@ -8,6 +15,8 @@ from typing import Optional
 from dataclasses import dataclass, field
 from enum import Enum
 
+CATEGORIES = [None, 'SPONSOR', 'SELFPROMO', 'INTERACTION']
+
 ACTION_OPTIONS = ['skip', 'mute', 'full']
 
 CATGEGORY_OPTIONS = {
@@ -62,6 +71,47 @@ class CustomTokens(Enum):
         tokenizer.add_tokens(cls.custom_tokens())
 
 
+_SEGMENT_START = START_SEGMENT_TEMPLATE.format(r'(?P<category>\w+)')
+_SEGMENT_END = END_SEGMENT_TEMPLATE.format(r'\w+')
+SEGMENT_MATCH_RE = fr'{_SEGMENT_START}\s*(?P<text>.*?)\s*(?:{_SEGMENT_END}|$)'
+
+
+def extract_sponsor_matches(text):
+    if CustomTokens.NO_SEGMENT.value in text:
+        return []
+
+    return re_findall(SEGMENT_MATCH_RE, text)
+
+
+@dataclass
+class DatasetArguments:
+    data_dir: Optional[str] = field(
+        default='data',
+        metadata={
+            'help': 'The directory which stores train, test and/or validation data.'
+        },
+    )
+    processed_file: Optional[str] = field(
+        default='segments.json',
+        metadata={
+            'help': 'Processed data file'
+        },
+    )
+    processed_database: Optional[str] = field(
+        default='processed_database.json',
+        metadata={
+            'help': 'Processed database file'
+        },
+    )
+
+    dataset_cache_dir: Optional[str] = field(
+        default=None,
+        metadata={
+            'help': 'Where to store the cached datasets'
+        },
+    )
+
+
 @dataclass
 class OutputArguments:
 
@@ -126,3 +176,106 @@ def reset():
     torch.cuda.empty_cache()
     gc.collect()
     print(torch.cuda.memory_summary(device=None, abbreviated=False))
+
+
+def load_datasets(dataset_args):
+
+    print('Reading datasets')
+    data_files = {}
+
+    if dataset_args.train_file is not None:
+        data_files['train'] = os.path.join(
+            dataset_args.data_dir, dataset_args.train_file)
+    if dataset_args.validation_file is not None:
+        data_files['validation'] = os.path.join(
+            dataset_args.data_dir, dataset_args.validation_file)
+    if dataset_args.test_file is not None:
+        data_files['test'] = os.path.join(
+            dataset_args.data_dir, dataset_args.test_file)
+
+    return load_dataset('json', data_files=data_files, cache_dir=dataset_args.dataset_cache_dir)
+
+
+@dataclass
+class CustomTrainingArguments(OutputArguments, TrainingArguments):
+    seed: Optional[int] = GeneralArguments.__dataclass_fields__['seed']
+
+    num_train_epochs: float = field(
+        default=1, metadata={'help': 'Total number of training epochs to perform.'})
+
+    save_steps: int = field(default=5000, metadata={
+                            'help': 'Save checkpoint every X updates steps.'})
+    eval_steps: int = field(default=5000, metadata={
+                            'help': 'Run an evaluation every X steps.'})
+    logging_steps: int = field(default=5000, metadata={
+                               'help': 'Log every X updates steps.'})
+
+    # do_eval: bool = field(default=False, metadata={
+    #                       'help': 'Whether to run eval on the dev set.'})
+    # do_predict: bool = field(default=False, metadata={
+    #                          'help': 'Whether to run predictions on the test set.'})
+
+    per_device_train_batch_size: int = field(
+        default=4, metadata={'help': 'Batch size per GPU/TPU core/CPU for training.'}
+    )
+    per_device_eval_batch_size: int = field(
+        default=4, metadata={'help': 'Batch size per GPU/TPU core/CPU for evaluation.'}
+    )
+
+    # report_to: Optional[List[str]] = field(
+    #     default=None, metadata={"help": "The list of integrations to report the results and logs to."}
+    # )
+    evaluation_strategy: str = field(
+        default='steps',
+        metadata={
+            'help': 'The evaluation strategy to use.',
+            'choices': ['no', 'steps', 'epoch']
+        },
+    )
+
+    # evaluation_strategy (:obj:`str` or :class:`~transformers.trainer_utils.IntervalStrategy`, `optional`, defaults to :obj:`"no"`):
+    # The evaluation strategy to adopt during training. Possible values are:
+
+    #     * :obj:`"no"`: No evaluation is done during training.
+    #     * :obj:`"steps"`: Evaluation is done (and logged) every :obj:`eval_steps`.
+    #     * :obj:`"epoch"`: Evaluation is done at the end of each epoch.
+
+
+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:
+        last_checkpoint = glc(training_args.output_dir)
+        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
+            raise ValueError(
+                f'Output directory ({training_args.output_dir}) already exists and is not empty. Use --overwrite_output_dir to overcome.'
+            )
+        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
+            logger.info(
+                f'Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change the `--output_dir` or add `--overwrite_output_dir` to train from scratch.'
+            )
+    return last_checkpoint
+
+
+def train_from_checkpoint(trainer, last_checkpoint, training_args):
+    checkpoint = None
+    if training_args.resume_from_checkpoint is not None:
+        checkpoint = training_args.resume_from_checkpoint
+    elif last_checkpoint is not None:
+        checkpoint = last_checkpoint
+
+    train_result = trainer.train(resume_from_checkpoint=checkpoint)
+
+    trainer.save_model()  # Saves the tokenizer too for easy upload
+
+    return train_result

src/train.py

@@ -1,7 +1,5 @@
-from datasets import load_dataset
-from preprocess import DatasetArguments
-from predict import ClassifierArguments, SEGMENT_MATCH_RE, CATEGORIES
-from shared import CustomTokens, GeneralArguments, OutputArguments
+from preprocess import PreprocessingDatasetArguments
+from shared import CustomTokens, load_datasets, CustomTrainingArguments, get_last_checkpoint, train_from_checkpoint
 from model import ModelArguments
 import transformers
 import logging
@@ -9,21 +7,15 @@ import os
 import sys
 from dataclasses import dataclass, field
 from typing import Optional
-import datasets
-import pickle
+from datasets import utils as d_utils
 from transformers import (
     DataCollatorForSeq2Seq,
     HfArgumentParser,
     Seq2SeqTrainer,
-    Seq2SeqTrainingArguments
 )
 
-from transformers.trainer_utils import get_last_checkpoint
 from transformers.utils import check_min_version
 from transformers.utils.versions import require_version
-from sklearn.linear_model import LogisticRegression
-from sklearn.feature_extraction.text import TfidfVectorizer
-from utils import re_findall
 
 # Will error if the minimal version of Transformers is not installed. Remove at your own risks.
 check_min_version('4.13.0.dev0')
@@ -43,23 +35,6 @@ logging.basicConfig(
 )
 
 
-def load_datasets(dataset_args):
-
-    print('Reading datasets')
-    data_files = {}
-
-    if dataset_args.train_file is not None:
-        data_files['train'] = os.path.join(
-            dataset_args.data_dir, dataset_args.train_file)
-    if dataset_args.validation_file is not None:
-        data_files['validation'] = os.path.join(
-            dataset_args.data_dir, dataset_args.validation_file)
-    if dataset_args.test_file is not None:
-        data_files['test'] = os.path.join(
-            dataset_args.data_dir, dataset_args.test_file)
-
-    return load_dataset('json', data_files=data_files, cache_dir=dataset_args.dataset_cache_dir)
-
 
 @dataclass
 class DataTrainingArguments:
@@ -92,58 +67,7 @@ class DataTrainingArguments:
     )
 
 
-@dataclass
-class SequenceTrainingArguments(OutputArguments, Seq2SeqTrainingArguments):
-    seed: Optional[int] = GeneralArguments.__dataclass_fields__['seed']
-
-    num_train_epochs: float = field(
-        default=1, metadata={'help': 'Total number of training epochs to perform.'})
-
-    save_steps: int = field(default=5000, metadata={
-                            'help': 'Save checkpoint every X updates steps.'})
-    eval_steps: int = field(default=5000, metadata={
-                            'help': 'Run an evaluation every X steps.'})
-    logging_steps: int = field(default=5000, metadata={
-                               'help': 'Log every X updates steps.'})
-
-    skip_train_transformer: bool = field(default=False, metadata={
-        'help': 'Whether to skip training the transformer.'})
-    train_classifier: bool = field(default=False, metadata={
-        'help': 'Whether to run training on the 2nd phase (classifier).'})
-
-    # do_eval: bool = field(default=False, metadata={
-    #                       'help': 'Whether to run eval on the dev set.'})
-    do_predict: bool = field(default=False, metadata={
-                             'help': 'Whether to run predictions on the test set.'})
-
-    per_device_train_batch_size: int = field(
-        default=4, metadata={'help': 'Batch size per GPU/TPU core/CPU for training.'}
-    )
-    per_device_eval_batch_size: int = field(
-        default=4, metadata={'help': 'Batch size per GPU/TPU core/CPU for evaluation.'}
-    )
-
-    # report_to: Optional[List[str]] = field(
-    #     default=None, metadata={"help": "The list of integrations to report the results and logs to."}
-    # )
-    evaluation_strategy: str = field(
-        default='steps',
-        metadata={
-            'help': 'The evaluation strategy to use.',
-            'choices': ['no', 'steps', 'epoch']
-        },
-    )
-
-    # evaluation_strategy (:obj:`str` or :class:`~transformers.trainer_utils.IntervalStrategy`, `optional`, defaults to :obj:`"no"`):
-    # The evaluation strategy to adopt during training. Possible values are:
-
-    #     * :obj:`"no"`: No evaluation is done during training.
-    #     * :obj:`"steps"`: Evaluation is done (and logged) every :obj:`eval_steps`.
-    #     * :obj:`"epoch"`: Evaluation is done at the end of each epoch.
-
-
 def main():
-    # reset()
 
     # See all possible arguments in src/transformers/training_args.py
     # or by passing the --help flag to this script.
@@ -151,16 +75,15 @@ def main():
 
     hf_parser = HfArgumentParser((
         ModelArguments,
-        DatasetArguments,
+        PreprocessingDatasetArguments,
         DataTrainingArguments,
-        SequenceTrainingArguments,
-        ClassifierArguments
+        CustomTrainingArguments
     ))
-    model_args, dataset_args, data_training_args, training_args, classifier_args = hf_parser.parse_args_into_dataclasses()
+    model_args, dataset_args, data_training_args, training_args = hf_parser.parse_args_into_dataclasses()
 
     log_level = training_args.get_process_log_level()
     logger.setLevel(log_level)
-    datasets.utils.logging.set_verbosity(log_level)
+    d_utils.logging.set_verbosity(log_level)
     transformers.utils.logging.set_verbosity(log_level)
     transformers.utils.logging.enable_default_handler()
     transformers.utils.logging.enable_explicit_format()
@@ -199,231 +122,130 @@ def main():
 
     # In distributed training, the load_dataset function guarantees that only one local process can concurrently
     # download the dataset.
-    if training_args.skip_train_transformer and not training_args.train_classifier:
-        print('Nothing to do. Exiting')
-        return
-
     raw_datasets = load_datasets(dataset_args)
     # , cache_dir=model_args.cache_dir
 
     # 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.
 
-    if training_args.train_classifier:
-        print('Train classifier')
-        # 1. Vectorize raw data to pass into classifier
-        # CountVectorizer TfidfVectorizer
-        # TfidfVectorizer - better (comb of CountVectorizer)
-        vectorizer = TfidfVectorizer(  # CountVectorizer
-            # lowercase=False,
-            # stop_words='english',  # TODO optimise stop words?
-            # stop_words=stop_words,
-
-            ngram_range=(1, 2),  # best so far
-            # max_features=8000  # remove for higher accuracy?
-            max_features=20000
-            # max_features=10000
-            # max_features=1000
-        )
 
-        train_test_data = {
-            'train': {
-                'X': [],
-                'y': []
-            },
-            'test': {
-                'X': [],
-                'y': []
-            }
-        }
-
-        print('Splitting')
-        for ds_type in train_test_data:
-            dataset = raw_datasets[ds_type]
-
-            for row in dataset:
-                matches = re_findall(SEGMENT_MATCH_RE, row['extracted'])
-                if matches:
-                    for match in matches:
-                        train_test_data[ds_type]['X'].append(match['text'])
-
-                        class_index = CATEGORIES.index(match['category'])
-                        train_test_data[ds_type]['y'].append(class_index)
-
-                else:
-                    train_test_data[ds_type]['X'].append(row['text'])
-                    train_test_data[ds_type]['y'].append(0)
-
-        print('Fitting')
-        _X_train = vectorizer.fit_transform(train_test_data['train']['X'])
-        _X_test = vectorizer.transform(train_test_data['test']['X'])
-
-        y_train = train_test_data['train']['y']
-        y_test = train_test_data['test']['y']
-
-        # 2. Create classifier
-        classifier = LogisticRegression(max_iter=2000, class_weight='balanced')
-
-        # 3. Fit data
-        print('Fit classifier')
-        classifier.fit(_X_train, y_train)
-
-        # 4. Measure accuracy
-        accuracy = classifier.score(_X_test, y_test)
-
-        print(f'[LogisticRegression] Accuracy percent:',
-              round(accuracy*100, 3))
-
-        # 5. Save classifier and vectorizer
-        with open(os.path.join(classifier_args.classifier_dir, classifier_args.classifier_file), 'wb') as fp:
-            pickle.dump(classifier, fp)
-
-        with open(os.path.join(classifier_args.classifier_dir, classifier_args.vectorizer_file), 'wb') as fp:
-            pickle.dump(vectorizer, fp)
-
-    if not training_args.skip_train_transformer:
-        # Detecting last checkpoint.
-        last_checkpoint = None
-        if os.path.isdir(training_args.output_dir) and not training_args.overwrite_output_dir:
-            last_checkpoint = get_last_checkpoint(training_args.output_dir)
-            if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
-                raise ValueError(
-                    f'Output directory ({training_args.output_dir}) already exists and is not empty. Use --overwrite_output_dir to overcome.'
-                )
-            elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
-                logger.info(
-                    f'Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change the `--output_dir` or add `--overwrite_output_dir` to train from scratch.'
-                )
-
-        from model import get_model_tokenizer
-        model, tokenizer = get_model_tokenizer(
-            model_args.model_name_or_path, model_args.cache_dir, training_args.no_cuda)
-
-        # Preprocessing the datasets.
-        # We need to tokenize inputs and targets.
-        column_names = raw_datasets['train'].column_names
-
-        prefix = CustomTokens.EXTRACT_SEGMENTS_PREFIX.value
-
-        PAD_TOKEN_REPLACE_ID = -100
-
-        # https://github.com/huggingface/transformers/issues/5204
-        def preprocess_function(examples):
-            inputs = examples['text']
-            targets = examples['extracted']
-            inputs = [prefix + inp for inp in inputs]
-            model_inputs = tokenizer(inputs, truncation=True)
-
-            # Setup the tokenizer for targets
-            with tokenizer.as_target_tokenizer():
-                labels = tokenizer(targets, truncation=True)
-
-            # If we are padding here, replace all tokenizer.pad_token_id in the labels by -100
-            # when we want to ignore padding in the loss.
-
-            model_inputs['labels'] = [
-                [(l if l != tokenizer.pad_token_id else PAD_TOKEN_REPLACE_ID)
-                 for l in label]
-                for label in labels['input_ids']
-            ]
-
-            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')
-
-        # Data collator
-        data_collator = DataCollatorForSeq2Seq(
-            tokenizer,
-            model=model,
-            label_pad_token_id=PAD_TOKEN_REPLACE_ID,
-            pad_to_multiple_of=8 if training_args.fp16 else None,
-        )
+    # Detecting last checkpoint.
+    last_checkpoint = get_last_checkpoint(training_args)
+
+    from model import get_model_tokenizer
+    model, tokenizer = get_model_tokenizer(model_args, training_args)
+
+    # Preprocessing the datasets.
+    # We need to tokenize inputs and targets.
+    column_names = raw_datasets['train'].column_names
+
+    prefix = CustomTokens.EXTRACT_SEGMENTS_PREFIX.value
+
+    PAD_TOKEN_REPLACE_ID = -100
+
+    # https://github.com/huggingface/transformers/issues/5204
+    def preprocess_function(examples):
+        inputs = examples['text']
+        targets = examples['extracted']
+        inputs = [prefix + inp for inp in inputs]
+        model_inputs = tokenizer(inputs, truncation=True)
+
+        # Setup the tokenizer for targets
+        with tokenizer.as_target_tokenizer():
+            labels = tokenizer(targets, truncation=True)
+
+        # If we are padding here, replace all tokenizer.pad_token_id in the labels by -100
+        # when we want to ignore padding in the loss.
+
+        model_inputs['labels'] = [
+            [(l if l != tokenizer.pad_token_id else PAD_TOKEN_REPLACE_ID)
+                for l in label]
+            for label in labels['input_ids']
+        ]
 
-        # Done processing datasets
+        return model_inputs
 
-        # Initialize our Trainer
-        trainer = Seq2SeqTrainer(
-            model=model,
-            args=training_args,
-            train_dataset=train_dataset,
-            eval_dataset=eval_dataset,
-            tokenizer=tokenizer,
-            data_collator=data_collator,
+    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')
+
+    # Data collator
+    data_collator = DataCollatorForSeq2Seq(
+        tokenizer,
+        model=model,
+        label_pad_token_id=PAD_TOKEN_REPLACE_ID,
+        pad_to_multiple_of=8 if training_args.fp16 else None,
+    )
+
+    # Done processing datasets
+
+    # Initialize our Trainer
+    trainer = Seq2SeqTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        tokenizer=tokenizer,
+        data_collator=data_collator,
+    )
+
+    # Training
+    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(
+        train_dataset)
+    metrics['train_samples'] = min(max_train_samples, len(train_dataset))
+
+    trainer.log_metrics('train', metrics)
+    trainer.save_metrics('train', metrics)
+    trainer.save_state()
+
+    kwargs = {'finetuned_from': model_args.model_name_or_path,
+                'tasks': 'summarization'}
 
-        # Training
-        checkpoint = None
-        if training_args.resume_from_checkpoint is not None:
-            checkpoint = training_args.resume_from_checkpoint
-        elif last_checkpoint is not None:
-            checkpoint = last_checkpoint
-
-        try:
-            train_result = trainer.train(resume_from_checkpoint=checkpoint)
-            trainer.save_model()  # Saves the tokenizer too for easy upload
-        except KeyboardInterrupt:
-            # TODO add option to save model on interrupt?
-            # print('Saving model')
-            # trainer.save_model(os.path.join(
-            #     training_args.output_dir, 'checkpoint-latest'))  # TODO use dir
-            raise
-
-        metrics = train_result.metrics
-        max_train_samples = data_training_args.max_train_samples or len(
-            train_dataset)
-        metrics['train_samples'] = min(max_train_samples, len(train_dataset))
-
-        trainer.log_metrics('train', metrics)
-        trainer.save_metrics('train', metrics)
-        trainer.save_state()
-
-        kwargs = {'finetuned_from': model_args.model_name_or_path,
-                  'tasks': 'summarization'}
-
-        if training_args.push_to_hub:
-            trainer.push_to_hub(**kwargs)
-        else:
-            trainer.create_model_card(**kwargs)
+    if training_args.push_to_hub:
+        trainer.push_to_hub(**kwargs)
+    else:
+        trainer.create_model_card(**kwargs)
 
 
 if __name__ == '__main__':

src/train_classifier.py

@@ -0,0 +1,287 @@
+
+""" Finetuning the library models for sequence classification."""
+
+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 (
+    DataCollatorWithPadding,
+    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 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")
+
+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."
+        },
+    )
+    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."
+        },
+    )
+
+    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():
+    # See all possible arguments in src/transformers/training_args.py
+    # 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()
+
+    # 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)],
+    )
+
+    log_level = training_args.get_process_log_level()
+    logger.setLevel(log_level)
+    datasets.utils.logging.set_verbosity(log_level)
+    transformers.utils.logging.set_verbosity(log_level)
+    transformers.utils.logging.enable_default_handler()
+    transformers.utils.logging.enable_explicit_format()
+
+    # 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}"
+    )
+    logger.info(f"Training/evaluation parameters {training_args}")
+
+    # Detecting last checkpoint.
+    last_checkpoint = get_last_checkpoint(training_args)
+
+    # Set seed before initializing model.
+    set_seed(training_args.seed)
+
+    # Loading a dataset from your local files.
+    # CSV/JSON training and evaluation files are needed.
+    raw_datasets = load_datasets(data_args)
+
+    # See more about loading any type of standard or custom dataset at
+    # https://huggingface.co/docs/datasets/loading_datasets.html.
+
+    config_args = {
+        'num_labels': len(CATEGORIES),
+        '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')
+
+
+    # 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:
+        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}."
+        )
+    max_seq_length = min(data_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)
+        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")
+
+    # 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.
+    def compute_metrics(p: EvalPrediction):
+        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()}
+
+    # 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:
+        data_collator = DataCollatorWithPadding(
+            tokenizer, pad_to_multiple_of=8)
+    else:
+        data_collator = None
+
+    # Initialize our Trainer
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        compute_metrics=compute_metrics,
+        tokenizer=tokenizer,
+        data_collator=data_collator,
+    )
+
+    # Training
+    train_result = train_from_checkpoint(
+        trainer, last_checkpoint, training_args)
+
+    metrics = train_result.metrics
+    max_train_samples = (
+        data_args.max_train_samples if data_args.max_train_samples is not None else 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.save_state()
+
+    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__":
+    main()

src/utils.py

@@ -1,8 +1,4 @@
 import re
-import logging
-
-logging.basicConfig()
-logger = logging.getLogger(__name__)
 
 
 def re_findall(pattern, string):