Add source code

src/errors.py

@@ -0,0 +1,13 @@
+class SponsorBlockException(Exception):
+    """Base class for all sponsor block exceptions"""
+    pass
+
+
+class PredictionException(SponsorBlockException):
+    """An exception was occurred while predicting sponsor segments"""
+    pass
+
+
+class TranscriptError(SponsorBlockException):
+    """An exception was occurred while retrieving the video transcript"""
+    pass

src/evaluate.py

@@ -0,0 +1,244 @@
+from datasets import load_dataset
+from transformers import (
+    AutoModelForSeq2SeqLM,
+    AutoTokenizer,
+    HfArgumentParser
+)
+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 segment import word_start, word_end, SegmentationArguments, add_labels_to_words
+import pandas as pd
+from dataclasses import dataclass, field
+from typing import Optional
+from tqdm import tqdm
+import json
+import os
+import random
+
+
+@dataclass
+class EvaluationArguments:
+    """
+    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
+    """
+    max_videos: Optional[int] = field(
+        default=None,
+        metadata={
+            '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']
+
+    output_file: Optional[str] = field(
+        default='metrics.csv',
+        metadata={
+            'help': 'Save metrics to output file'
+        }
+    )
+
+
+def jaccard(x1, x2, y1, y2):
+    # Calculate jaccard index
+    intersection = max(0, min(x2, y2)-max(x1, y1))
+    filled_union = max(x2, y2) - min(x1, y1)
+    return intersection/filled_union
+
+
+def attach_predictions_to_sponsor_segments(predictions, sponsor_segments):
+    """Attach sponsor segments to closest prediction"""
+    for prediction in predictions:
+        prediction['best_overlap'] = 0
+        prediction['best_sponsorship'] = None
+
+    # Assign predictions to actual (labelled) sponsored segments
+    for sponsor_segment in sponsor_segments:
+        sponsor_segment['best_overlap'] = 0
+        sponsor_segment['best_prediction'] = None
+
+        for prediction in predictions:
+
+            j = jaccard(prediction['start'], prediction['end'],
+                        sponsor_segment['start'], sponsor_segment['end'])
+            if sponsor_segment['best_overlap'] < j:
+                sponsor_segment['best_overlap'] = j
+                sponsor_segment['best_prediction'] = prediction
+
+            if prediction['best_overlap'] < j:
+                prediction['best_overlap'] = j
+                prediction['best_sponsorship'] = sponsor_segment
+
+    return sponsor_segments
+
+
+def calculate_metrics(labelled_words, predictions):
+
+    metrics = {
+        'true_positive': 0,  # Is sponsor, predicted sponsor
+        # Is sponsor, predicted not sponsor (i.e., missed it - bad)
+        'false_negative': 0,
+        # Is not sponsor, predicted sponsor (classified incorectly, not that bad since we do manual checking afterwards)
+        'false_positive': 0,
+        'true_negative': 0,  # Is not sponsor, predicted not sponsor
+    }
+
+    metrics['video_duration'] = word_end(
+        labelled_words[-1])-word_start(labelled_words[0])
+
+    for index, word in enumerate(labelled_words):
+        if index >= len(labelled_words) - 1:
+            continue
+
+        # TODO make sure words with manual transcripts
+        duration = labelled_words[index+1]['start'] - word['start']
+
+        predicted_sponsor = False
+        for p in predictions:
+            # Is in some prediction
+            if p['start'] <= word['start'] <= p['end']:
+                predicted_sponsor = True
+                break
+
+        if predicted_sponsor:
+            # total_positive_time += duration
+            if word['sponsor']:  # Is actual sponsor
+                metrics['true_positive'] += duration
+            else:
+                metrics['false_positive'] += duration
+        else:
+            # total_negative_time += duration
+            if word['sponsor']:  # Is actual sponsor
+                metrics['false_negative'] += duration
+            else:
+                metrics['true_negative'] += duration
+
+    # NOTE In cases where we encounter division by 0, we say that the value is 1
+    # https://stats.stackexchange.com/a/1775
+    # (Precision) TP+FP=0: means that all instances were predicted as negative
+    # (Recall)    TP+FN=0: means that there were no positive cases in the input data
+
+    # The fraction of predictions our model got right
+    # Can simplify, but use full formula
+    z = metrics['true_positive'] + metrics['true_negative'] + \
+        metrics['false_positive'] + metrics['false_negative']
+    metrics['accuracy'] = (
+        (metrics['true_positive'] + metrics['true_negative']) / z) if z > 0 else 1
+
+    # What proportion of positive identifications was actually correct?
+    z = metrics['true_positive'] + metrics['false_positive']
+    metrics['precision'] = (metrics['true_positive'] / z) if z > 0 else 1
+
+    # What proportion of actual positives was identified correctly?
+    z = metrics['true_positive'] + metrics['false_negative']
+    metrics['recall'] = (metrics['true_positive'] / z) if z > 0 else 1
+
+    # https://deepai.org/machine-learning-glossary-and-terms/f-score
+
+    s = metrics['precision'] + metrics['recall']
+    metrics['f-score'] = (2 * (metrics['precision'] *
+                               metrics['recall']) / s) if s > 0 else 0
+
+    return metrics
+
+
+def main():
+    hf_parser = HfArgumentParser((
+        EvaluationArguments,
+        ProcessedArguments,
+        SegmentationArguments,
+        ClassifierArguments
+    ))
+
+    evaluation_args, processed_args, segmentation_args, classifier_args = hf_parser.parse_args_into_dataclasses()
+
+    model = AutoModelForSeq2SeqLM.from_pretrained(evaluation_args.model_path)
+    model.to(device())
+
+    tokenizer = AutoTokenizer.from_pretrained(evaluation_args.model_path)
+
+    dataset = load_dataset('json', data_files=os.path.join(
+        evaluation_args.data_dir, evaluation_args.dataset))['train']
+
+    video_ids = [row['video_id'] for row in dataset]
+    random.shuffle(video_ids)  # TODO Make param
+
+    if evaluation_args.max_videos is not None:
+        video_ids = video_ids[:evaluation_args.max_videos]
+
+    # Load labelled data:
+    final_path = os.path.join(
+        processed_args.processed_dir, processed_args.processed_file)
+
+    with open(final_path) as fp:
+        final_data = json.load(fp)
+
+    classifier, vectorizer = get_classifier_vectorizer(classifier_args)
+
+    total_accuracy = 0
+    total_precision = 0
+    total_recall = 0
+    total_fscore = 0
+
+    count = 0
+    out_metrics = []
+
+    try:
+        with tqdm(video_ids) as progress:
+            for video_id in progress:
+                progress.set_description(f'Processing {video_id}')
+                sponsor_segments = final_data.get(video_id, [])
+
+                words = get_words(video_id)
+                if not words:
+                    continue
+
+                count += 1
+
+                # Make predictions
+                predictions = predict(video_id, model, tokenizer,
+                                      segmentation_args, words)
+
+                # Filter predictions
+                predictions = filter_predictions(
+                    predictions, classifier, vectorizer, classifier_args)
+
+                labelled_words = add_labels_to_words(words, sponsor_segments)
+                met = calculate_metrics(labelled_words, predictions)
+                met['video_id'] = video_id
+
+                out_metrics.append(met)
+
+                total_accuracy += met['accuracy']
+                total_precision += met['precision']
+                total_recall += met['recall']
+                total_fscore += met['f-score']
+
+                progress.set_postfix({
+                    'accuracy': total_accuracy/count,
+                    'precision':  total_precision/count,
+                    'recall':  total_recall/count,
+                    'f-score': total_fscore/count
+                })
+
+                labelled_predicted_segments = attach_predictions_to_sponsor_segments(
+                    predictions, sponsor_segments)
+                for seg in labelled_predicted_segments:
+                    if seg['best_prediction'] is None:
+                        print('\nNo match found for', seg)
+
+    except KeyboardInterrupt:
+        pass
+
+    df = pd.DataFrame(out_metrics)
+
+    df.to_csv(evaluation_args.output_file)
+    print(df.mean())
+
+
+if __name__ == '__main__':
+    main()

src/model.py

@@ -0,0 +1,111 @@
+import pickle
+import os
+from shared import CustomTokens
+from transformers import AutoTokenizer, AutoConfig, AutoModelForSeq2SeqLM
+from dataclasses import dataclass, field
+from typing import Optional
+
+
+@dataclass
+class ModelArguments:
+    """
+    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
+    """
+
+    model_name_or_path: str = field(
+        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=None,
+        metadata={
+            'help': 'Where to store the pretrained models downloaded from huggingface.co'},
+    )
+    use_fast_tokenizer: bool = field(  # TODO remove?
+        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?
+        default='main',
+        metadata={
+            'help': 'The specific model version to use (can be a branch name, tag name or commit id).'},
+    )
+    use_auth_token: bool = field(
+        default=False,
+        metadata={
+            'help': 'Will use the token generated when running `transformers-cli login` (necessary to use this script '
+            'with private models).'
+        },
+    )
+    resize_position_embeddings: Optional[bool] = field(
+        default=None,
+        metadata={
+            'help': "Whether to automatically resize the position embeddings if `max_source_length` exceeds the model's position embeddings."
+        },
+    )
+
+
+def get_model(model_args, use_cache=True):
+    name = model_args.model_name_or_path
+    cached_path = f'models/{name}'
+
+    # Model created after tokenizer:
+    if use_cache and os.path.exists(os.path.join(cached_path, 'pytorch_model.bin')):
+        name = cached_path
+
+    config = AutoConfig.from_pretrained(
+        name,
+        cache_dir=model_args.cache_dir,
+        revision=model_args.model_revision,
+        use_auth_token=True if model_args.use_auth_token else None,
+    )
+
+    model = AutoModelForSeq2SeqLM.from_pretrained(
+        name,
+        from_tf='.ckpt' in name,
+        config=config,
+        cache_dir=model_args.cache_dir,
+        revision=model_args.model_revision,
+        use_auth_token=True if model_args.use_auth_token else None,
+    )
+
+    return model
+
+
+def get_tokenizer(model_args, use_cache=True):
+    name = model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path
+
+    cached_path = f'models/{name}'
+
+    if use_cache and os.path.exists(os.path.join(cached_path, 'tokenizer.json')):
+        name = cached_path
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        name,
+        cache_dir=model_args.cache_dir,
+        use_fast=model_args.use_fast_tokenizer,
+        revision=model_args.model_revision,
+        use_auth_token=True if model_args.use_auth_token else None,
+    )
+
+    CustomTokens.add_custom_tokens(tokenizer)
+
+    return tokenizer
+
+
+def get_classifier_vectorizer(classifier_args):
+    with open(os.path.join(classifier_args.classifier_dir, classifier_args.classifier_file), 'rb') as fp:
+        classifier = pickle.load(fp)
+
+    with open(os.path.join(classifier_args.classifier_dir, classifier_args.vectorizer_file), 'rb') as fp:
+        vectorizer = pickle.load(fp)
+
+    return classifier, vectorizer

src/predict.py

@@ -0,0 +1,278 @@
+from shared import OutputArguments
+from typing import Optional
+from segment import (
+    generate_segments,
+    extract_segment,
+    SAFETY_TOKENS,
+    CustomTokens,
+    word_start,
+    word_end,
+    SegmentationArguments
+)
+import preprocess
+import re
+from errors import TranscriptError
+from model import get_classifier_vectorizer
+from transformers import (
+    AutoModelForSeq2SeqLM,
+    AutoTokenizer
+)
+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):
+    h, remainder = divmod(abs(int(seconds)), 3600)
+    m, s = divmod(remainder, 60)
+    return f"{'-' if seconds < 0 else ''}{h:02}:{m:02}:{s:02}"
+
+
+@dataclass
+class PredictArguments:
+
+    video_id: str = field(
+        metadata={
+            'help': 'Video to predict sponsorship segments for'}
+    )
+
+    model_path: str = field(
+        default=None,
+        metadata={
+            'help': 'Path to pretrained model used for prediction'}
+    )
+
+    output_dir: Optional[str] = OutputArguments.__dataclass_fields__[
+        'output_dir']
+
+    def __post_init__(self):
+        if self.model_path is not None:
+            return
+
+        last_checkpoint = get_last_checkpoint(self.output_dir)
+        if last_checkpoint is not None:
+            self.model_path = last_checkpoint
+        else:
+            raise Exception(
+                'Unable to find model, explicitly set `--model_path`')
+
+
+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)
+MERGE_TIME_WITHIN = 8   # Merge predictions if they are within x seconds
+
+
+@dataclass
+class ClassifierArguments:
+    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_predictions(predictions, classifier, vectorizer, classifier_args):
+    """Use classifier to filter predictions"""
+    if not predictions:
+        return predictions
+
+    transformed_segments = vectorizer.transform([
+        preprocess.clean_text(' '.join([x['text'] for x in pred['words']]))
+        for pred in predictions
+    ])
+    probabilities = classifier.predict_proba(transformed_segments)
+
+    filtered_predictions = []
+    for prediction, probability in zip(predictions, probabilities):
+        prediction['probability'] = probability[1]
+
+        if prediction['probability'] >= classifier_args.min_probability:
+            filtered_predictions.append(prediction)
+        # else:
+            # print('removing segment', prediction)
+
+    return filtered_predictions
+
+
+def predict(video_id, model, tokenizer, segmentation_args, words=None, classifier_args=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)
+        if not words:
+            raise TranscriptError('Unable to retrieve transcript')
+
+    segments = generate_segments(
+        words,
+        tokenizer,
+        segmentation_args
+    )
+
+    predictions = segments_to_prediction_times(segments, model, tokenizer)
+
+    # Add words back to time_ranges
+    for prediction in predictions:
+        # Stores words in the range
+        prediction['words'] = extract_segment(
+            words, prediction['start'], prediction['end'])
+
+    if classifier_args is not None:
+        classifier, vectorizer = get_classifier_vectorizer(classifier_args)
+        predictions = filter_predictions(
+            predictions, classifier, vectorizer, classifier_args)
+
+    return predictions
+
+
+def greedy_match(list, sublist):
+    # Return index and length of longest matching sublist
+
+    best_i = -1
+    best_j = -1
+    best_k = 0
+
+    for i in range(len(list)):  # Start position in main list
+        for j in range(len(sublist)):  # Start position in sublist
+            for k in range(len(sublist)-j, 0, -1):  # Width of sublist window
+                if k > best_k and list[i:i+k] == sublist[j:j+k]:
+                    best_i, best_j, best_k = i, j, k
+                    break  # Since window size decreases
+
+    return best_i, best_j, best_k
+
+
+DEFAULT_TOKEN_PREFIX = 'summarize: '
+
+
+def predict_sponsor_text(text, model, tokenizer):
+    """Given a body of text, predict the words which are part of the sponsor"""
+    input_ids = tokenizer(
+        f'{DEFAULT_TOKEN_PREFIX}{text}', return_tensors='pt', truncation=True).input_ids
+
+    # Can't be longer than input length + SAFETY_TOKENS or model input dim
+    max_out_len = min(len(input_ids[0]) + SAFETY_TOKENS, model.model_dim)
+    outputs = model.generate(input_ids, max_length=max_out_len)
+
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+
+def predict_sponsor_matches(text, model, tokenizer):
+    sponsorship_text = predict_sponsor_text(text, model, tokenizer)
+    if CustomTokens.NO_SPONSOR.value in sponsorship_text:
+        return []
+
+    return re.findall(SPONSOR_MATCH_RE, sponsorship_text)
+
+
+def segments_to_prediction_times(segments, model, tokenizer):
+    predicted_time_ranges = []
+
+    # TODO pass to model simultaneously, not in for loop
+    # use 2d array for input ids
+    for segment in segments:
+        cleaned_batch = [preprocess.clean_text(
+            word['text']) for word in segment]
+        batch_text = ' '.join(cleaned_batch)
+
+        matches = predict_sponsor_matches(batch_text, model, tokenizer)
+
+        for match in matches:
+            matched_text = match.split()
+            # TODO skip if too short
+
+            i1, j1, k1 = greedy_match(
+                cleaned_batch, matched_text[:MATCH_WINDOW])
+            i2, j2, k2 = greedy_match(
+                cleaned_batch, matched_text[-MATCH_WINDOW:])
+
+            extracted_words = segment[i1:i2+k2]
+
+            if not extracted_words:
+                continue
+
+            predicted_time_ranges.append({
+                'start': word_start(extracted_words[0]),
+                'end': word_end(extracted_words[-1])
+            })
+
+    # Necessary to sort matches by start time
+    predicted_time_ranges.sort(key=word_start)
+
+    # Merge overlapping predictions and sponsorships that are close together
+    # Caused by model having max input size
+    last_end_time = -1
+    final_predicted_time_ranges = []
+    for range in predicted_time_ranges:
+        start_time = range['start']
+        end_time = range['end']
+
+        if (start_time <= last_end_time <= end_time) or (last_end_time != -1 and start_time - last_end_time <= MERGE_TIME_WITHIN):
+            # Ending time of last segment is in this segment, so we extend last prediction range
+            final_predicted_time_ranges[-1]['end'] = end_time
+
+        else:  # No overlap, is a new prediction
+            final_predicted_time_ranges.append({
+                'start': start_time,
+                'end': end_time,
+            })
+
+        last_end_time = end_time
+
+    return final_predicted_time_ranges
+
+
+def main():
+    # Test on unseen data
+    logging.getLogger().setLevel(logging.DEBUG)
+
+    hf_parser = HfArgumentParser((
+        PredictArguments,
+        SegmentationArguments,
+        ClassifierArguments
+    ))
+    predict_args, segmentation_args, classifier_args = hf_parser.parse_args_into_dataclasses()
+
+    model = AutoModelForSeq2SeqLM.from_pretrained(predict_args.model_path)
+    model.to(device())
+
+    tokenizer = AutoTokenizer.from_pretrained(predict_args.model_path)
+
+    predict_args.video_id = predict_args.video_id.strip()
+    print(
+        f'Predicting for https://www.youtube.com/watch?v={predict_args.video_id}')
+    predictions = predict(predict_args.video_id, model, tokenizer,
+                          segmentation_args, classifier_args=classifier_args)
+
+    for prediction in predictions:
+        print(' '.join([w['text'] for w in prediction['words']]))
+        print(seconds_to_time(prediction['start']),
+              '-->', seconds_to_time(prediction['end']))
+        print(prediction['start'], '-->', prediction['end'])
+        print(prediction['probability'])
+        print()
+
+    print()
+
+
+if __name__ == '__main__':
+    main()

src/preprocess.py

@@ -0,0 +1,786 @@
+import itertools
+from typing import Optional
+from datasets import load_dataset
+from model import ModelArguments
+import segment
+from tqdm import tqdm
+from dataclasses import dataclass, field
+from transformers import HfArgumentParser
+from shared import GeneralArguments, CustomTokens
+import csv
+import re
+import random
+import logging
+from youtube_transcript_api import YouTubeTranscriptApi
+from youtube_transcript_api._errors import CouldNotRetrieveTranscript, YouTubeRequestFailed
+import os
+import json
+import time
+import requests
+from utils import InterruptibleThreadPool, Job
+
+
+def find(s, ch):
+    return [i for i, ltr in enumerate(s) if ltr == ch]
+
+
+def wordify(transcript):
+    """Try to replicate format for automatically generated transcripts"""
+    words = []
+
+    for line_index, line in enumerate(transcript):
+        text = line['text'].replace('\n', ' ').strip()
+        if not text:
+            continue
+
+        start = line['start']
+        next_start = transcript[line_index +
+                                1]['start'] if line_index < len(transcript) - 1 else float('inf')
+        end = min(start + line['duration'], next_start)
+        duration = end - start
+
+        indices = find(text, ' ') + [len(text)]
+        start_index = 0
+        for i in range(len(indices)):
+            word = text[start_index:indices[i]].strip()
+            if not word:
+                continue  # Skip empty words (e.g., \n)
+            percentage = start_index/indices[-1]
+
+            w_duration = len(word)/indices[-1] * duration
+
+            w_start = start + percentage * duration
+
+            words.append({
+                'start': round(w_start, 5),
+                'duration': round(w_duration, 5),
+                'end': round(w_start + w_duration, 5),
+                'text': word,
+            })
+
+            start_index = indices[i] + 1
+
+    return words
+
+
+def get_manual_words(transcript_list):
+    transcript = transcript_list.find_manually_created_transcript(
+        ['en-GB', 'en-US', 'en']).fetch()
+    return wordify(transcript)
+
+
+def get_auto_words(transcript_list):
+    words = []
+    transcript = transcript_list.find_generated_transcript(['en'])
+    url = transcript._url + '&fmt=json3'
+    info = transcript._http_client.get(url)
+
+    for event in info.json()['events']:
+        start_ms = event.get('tStartMs', 0)
+
+        for word in event.get('segs') or []:
+            offset_ms = word.get('tOffsetMs', 0)
+
+            texts = word['utf8'].replace(
+                CustomTokens.PROFANITY_RAW.value, CustomTokens.PROFANITY_CONVERTED.value
+            ).strip().split()
+
+            for text in texts:
+                words.append({
+                    'start': (start_ms + offset_ms)/1000,
+                    'text': text
+                })
+
+    return words
+
+
+def get_words(video_id, process=True, fallback=False, transcript_type='auto'):
+    """Get parsed video transcript with caching system
+    returns None if not processed yet and process is False
+    """
+    get_manual_if_fail = fallback and transcript_type == 'auto'
+    transcript_path = os.path.join(
+        'transcripts', transcript_type, f'{video_id}.json')
+    words = []
+    try:
+        if os.path.exists(transcript_path):
+            with open(transcript_path) as fp:
+                wds = json.load(fp)
+
+            if not wds and get_manual_if_fail:
+                return get_words(video_id, process, fallback, 'manual')
+            return wds
+
+        elif not process:
+            return None
+
+        transcript_list = YouTubeTranscriptApi.list_transcripts(video_id)
+
+        if transcript_type == 'manual':
+            words = get_manual_words(transcript_list)
+        else:
+            words = get_auto_words(transcript_list)
+
+    except YouTubeRequestFailed as e:
+        print(e)
+        time.sleep(30)  # Timeout
+        return get_words(video_id, process, fallback, transcript_type)
+
+    except CouldNotRetrieveTranscript:
+        if get_manual_if_fail:
+            print('fallback')
+            return get_words(video_id, process, fallback, 'manual')
+
+    except json.decoder.JSONDecodeError:
+        # Warning, unable to parse JSON
+        pass
+
+    with open(transcript_path, 'w') as fp:
+        json.dump(words, fp)
+
+    return words
+
+
+# TODO make min_sponsor_segment_length param
+def extract_sponsors(words, min_sponsor_segment_length=5):
+    if len(words) < min_sponsor_segment_length:
+        return []  # Force short phrases to not be sponsors
+
+    paragraphs = []
+    current = []
+    for word in words:
+        if not word.get('sponsor') and not current:
+            continue
+
+        if word['sponsor']:
+            current.append(word['text'])
+        else:
+            paragraphs.append(current)
+            current = []
+    if current:
+        paragraphs.append(current)
+
+    # Remove all too short:
+    paragraphs = list(filter(lambda x: len(
+        x) >= min_sponsor_segment_length, paragraphs))
+
+    return paragraphs
+
+
+def clean_text(text):
+
+    # Replace impossibly long words with a special token
+    # Usually the result of incorrect labelling
+    text = re.sub(r'\w{64,}', CustomTokens.LONG_WORD.value, text)
+
+    SHORT_HYPHENATED_REGEX = r'\w{1,2}(?:-\w{1,2}){3,}(?:-?\w*)'
+
+    # Replace hyphenated URLs with special token
+    # For some reason, youtube sometimes transcribes urls in this form:
+    # 'b-a-b-b-e-l-dot-com', 'g-e-t-r-o-m-a-n-com'
+    # not 'e-commerce'
+    text = re.sub(f'{SHORT_HYPHENATED_REGEX}(?:com|org|net)',
+                  CustomTokens.HYPHENATED_URL.value, text)
+
+    # Replace short+hyphenated text with a special token. Of the form:
+    # 'i-i-i-i-i-i-i-i-i-i-i-i', 'b-u-m-f-u-z-z-l-e', 'v-e-r-i-t-a-s-i-u-m', 'do-do-do-do-do'
+    text = re.sub(SHORT_HYPHENATED_REGEX,
+                  CustomTokens.SHORT_HYPHENATED.value, text)
+
+    # Replace URLs with URL_TOKEN
+    URL_REGEX = r'(?:(?:http|https)\:\/\/)?[a-zA-Z0-9\.\/\?\:@\-_=#]+\.(?:[a-zA-Z]){2,6}(?:[a-zA-Z0-9\.\&\/\?\:@\-_=#%])*'
+    text = re.sub(URL_REGEX, CustomTokens.URL.value, text)
+
+    NUM_REGEX = r'(?:\d+,)*(?:\d*[.])?\d+'
+
+    # Encode specific numeric words
+    # Of the form: 12%, 12.34%
+    # Usually included in sponsorships
+    text = re.sub(f'{NUM_REGEX}%',
+                  CustomTokens.NUMBER_PERCENTAGE.value, text)
+
+    # Normal numbers, should not have an effect on sponsorship
+    text = re.sub(NUM_REGEX, CustomTokens.NUMBER.value, text)
+
+    # Replace profanity with special token
+    text = text.replace(CustomTokens.PROFANITY_RAW.value,
+                        CustomTokens.PROFANITY.value)
+    text = text.replace(CustomTokens.PROFANITY_CONVERTED.value,
+                        CustomTokens.PROFANITY.value)
+
+    return text.strip()
+
+
+def remove_duplicate_sponsor_segments(sponsor_segments):
+    """Choose the best sponsor segment if overlapping with others"""
+
+    # Algorithm based on SponsorBlock algorithm
+    # Find sponsors that are overlapping
+    similar = []
+    for i in sponsor_segments:
+        for j in sponsor_segments:
+            # Since we do pairwise, we only check one direction
+            if (j['start'] >= i['start'] and j['start'] <= i['end']):
+                similar.append([i, j])
+
+    # Within each group, choose the segment with the most votes.
+    processed = []
+    best = []
+    for i in similar:
+        if i in processed:
+            continue
+        group = i
+        for j in similar:
+            if j[0] in group or j[1] in group:  # If either in, append both
+                group.append(j[0])
+                group.append(j[1])
+                processed.append(j)
+
+        best.append(max(group, key=lambda item: (
+            item['votes'], item['reputation'], item['views'])))
+
+    return best
+
+
+@dataclass
+class PreprocessArguments:
+    """
+    Arguments pertaining to what data we are going to preprocess.
+    """
+    update_database: bool = field(
+        default=False, metadata={'help': 'Download the raw database.'}
+    )
+
+    do_create: bool = field(
+        default=False, metadata={'help': 'Merge sponsor segments into single file'}
+    )
+    min_votes: int = field(
+        default=0, metadata={'help': 'Minimum number of votes'})
+    # Downvotes will make this negative.
+    # 1 = At least one positive vote
+
+    do_transcribe: bool = field(
+        default=False, metadata={'help': 'Get transcripts for videos'}
+    )
+    num_jobs: int = field(
+        default=4, metadata={'help': 'Number of transcripts to download in parallel'})
+
+    overwrite: bool = field(
+        default=False, metadata={'help': 'Overwrite training, testing and validation data, if present.'}
+    )
+
+    do_generate: bool = field(
+        default=False, metadata={'help': 'Generate labelled data.'}
+    )
+
+    do_split: bool = field(
+        default=False, metadata={'help': 'Generate training, testing and validation data.'}
+    )
+    percentage_positive: float = field(
+        default=0.5, metadata={'help': 'Ratio of positive (sponsor) segments to include in final output'})
+
+    train_split: float = field(
+        default=0.9, metadata={'help': 'Ratio of training data. Value between 0 and 1.'})
+
+    # TODO play around with ratios? lower test/validation split?
+    test_split: float = field(
+        default=0.05, metadata={'help': 'Ratio of testing data. Value between 0 and 1.'})
+    valid_split: float = field(
+        default=0.05, metadata={'help': 'Ratio of validation data. Value between 0 and 1.'})
+
+    skip_videos: int = field(default=None, metadata={
+        'help': 'Number of videos to skip. Set this to the latest video index to append to the current file'})
+
+    max_videos: int = field(default=None, metadata={
+        'help': 'Maximum number of videos to preprocess.'})
+
+    max_segments: int = field(default=None, metadata={
+        'help': 'Maximum number of segments to produce to preprocess.'})
+
+    raw_data_dir: Optional[str] = field(
+        default='raw',
+        metadata={
+            'help': 'Raw data directory'
+        },
+    )
+    raw_data_file: Optional[str] = field(
+        default='sponsorTimes.csv',
+        metadata={
+            'help': 'Raw data file'
+        },
+    )
+
+    min_wps: float = field(
+        default=0.4, metadata={'help': 'Ignore videos with not enough words spoken per second. This is usually indicitive of video whose captions aren\'t English.'})
+    # 0.1 ~ 1%
+    # 0.4 ~ 2.5%
+    # 0.9 ~ 5%
+
+
+# Mirrors for database
+MIRRORS = [
+    'https://sponsor.ajay.app/database/sponsorTimes.csv',  # Latest
+    'https://sb-mirror.mchang.xyz/sponsorTimes.csv',  # 5 minute delay
+    'https://sb.ltn.fi/database/sponsorTimes.csv',  # 5 minute delay
+]
+# TODO only download latest (updates/changes)
+
+
+def download_file(url, filename):
+    """
+    Helper method handling downloading large files from `url` to `filename`.
+
+    Adapted from https://stackoverflow.com/a/42071418
+    """
+    chunk_size = 1024
+    r = requests.get(url, stream=True)
+    total_bytes = int(r.headers['Content-Length'])
+    with open(filename, 'wb') as f, tqdm(unit='B', total=total_bytes) as progress:
+        for chunk in r.iter_content(chunk_size=chunk_size):
+            if chunk:  # filter out keep-alive new chunks
+                progress.update(len(chunk))
+                f.write(chunk)
+
+    return total_bytes == os.path.getsize(filename)
+
+
+@dataclass
+class ProcessedArguments:
+    processed_dir: Optional[str] = field(
+        default='processed',
+        metadata={
+            'help': 'Processed data directory'
+        },
+    )
+    processed_file: Optional[str] = field(
+        default='final.json',
+        metadata={
+            'help': 'Processed data file'
+        },
+    )
+
+
+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)
+
+
+@dataclass
+class DatasetArguments:
+    data_dir: Optional[str] = field(
+        default='data',
+        metadata={
+            'help': 'The directory which stores train, test and/or validation data.'
+        },
+    )
+
+    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 (rouge) 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).'
+        },
+    )
+    excess_file: Optional[str] = field(
+        default='excess.json',
+        metadata={
+            'help': 'The excess segments left after the split'
+        },
+    )
+
+    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).'}
+    )
+    negative_file: Optional[str] = field(
+        default='normal_segments.json', metadata={'help': 'File to output normal segments to (a jsonlines file).'}
+    )
+
+    def __post_init__(self):
+        # TODO check if train/validation datasets exist
+        if self.train_file is None and self.validation_file is None:
+            raise ValueError(
+                'Need either a dataset name or a training/validation file.')
+
+
+def main():
+    # Responsible for getting transcrips using youtube_transcript_api,
+    # then labelling it according to SponsorBlock's API
+
+    logging.getLogger().setLevel(logging.INFO)  # TODO make param
+
+    # Generate final.json from sponsorTimes.csv
+    hf_parser = HfArgumentParser((
+        PreprocessArguments,
+        ProcessedArguments,
+        DatasetArguments,
+        segment.SegmentationArguments,
+        ModelArguments,
+        GeneralArguments
+    ))
+    preprocess_args, processed_args, dataset_args, segmentation_args, model_args, _ = hf_parser.parse_args_into_dataclasses()
+
+    raw_dataset_path = os.path.join(
+        preprocess_args.raw_data_dir, preprocess_args.raw_data_file)
+
+    def get_rows():
+        with open(raw_dataset_path, newline='') as csvfile:
+            reader = csv.DictReader(csvfile)
+            for line in reader:
+                if line['service'] != 'YouTube':
+                    continue
+
+                # TODO add support for other categories and action types?
+                if line['category'] != 'sponsor':
+                    continue
+                if line['actionType'] != 'skip':
+                    continue
+
+                # Ignore hidden items
+                if line['hidden'] == '1' or line['shadowHidden'] == '1':
+                    continue
+
+                if len(line['videoID']) != 11:
+                    continue  # Invalid youtube video ID
+
+                # Skip those that aren't highly voted
+                line['votes'] = int(line['votes'])
+                # incorrect_votes = int(line['incorrectVotes'])
+
+                if line['votes'] < preprocess_args.min_votes:
+                    continue
+
+                yield line
+
+    if preprocess_args.update_database:
+        print('Updating database')
+        for mirror in MIRRORS:
+            print('Downloading from', mirror)
+            if download_file(mirror, raw_dataset_path):
+                break
+            print('Failed, trying next')
+
+    # 'videoID', 'startTime', 'endTime', 'votes', 'locked', 'incorrectVotes', 'UUID',
+    # 'userID', 'timeSubmitted', 'views', 'category', 'actionType', 'service', 'videoDuration',
+    # 'hidden', 'reputation', 'shadowHidden', 'hashedVideoID', 'userAgent', 'description'
+    data_rows = None
+    if preprocess_args.do_transcribe:
+        print('Collecting videos')
+        video_ids = set()
+        data_rows = get_rows()
+        for row in data_rows:
+            video_ids.add(row['videoID'])
+
+        print('Start transcribing')
+        with tqdm(total=len(video_ids)) as progress:
+            def on_job_complete(job):
+                progress.set_description(f'Processed {job.video_id}')
+                progress.update()
+
+            pool = InterruptibleThreadPool(
+                preprocess_args.num_jobs, on_job_complete=on_job_complete)
+
+            print('Adding jobs to pool')
+            for video_id in video_ids:
+                job = Job(get_words, video_id)
+                job.video_id = video_id
+                pool.add_job(job)
+
+            print('Start processing')
+            pool.run()
+
+        print('Finished transcribing')
+
+    final_path = os.path.join(
+        processed_args.processed_dir, processed_args.processed_file)
+
+    if os.path.exists(final_path) and not preprocess_args.do_create:
+        logging.info(f'{final_path} exists, opening file')
+        with open(final_path) as fp:
+            final_data = json.load(fp)
+    else:
+        print('Create final data')
+
+        final_data = {}
+
+        if data_rows is None:
+            data_rows = get_rows()
+
+        # TODO add progress bar
+        # TODO parallelise?
+        for line in data_rows:
+            video_id = line['videoID']
+
+            if video_id not in final_data:
+                final_data[video_id] = []
+
+            segment_start = float(line['startTime'])
+            segment_end = float(line['endTime'])
+
+            video_words = get_words(video_id, process=True)
+            segment_words = segment.extract_segment(
+                video_words, segment_start, segment_end)
+
+            if len(segment_words) <= 1:
+                continue  # Useless to add segment since no words
+
+            # duration = segment.word_end(segment_words[-1]) - segment.word_start(segment_words[0])
+            duration = segment_end - segment_start
+            wps = len(segment_words)/duration if duration > 0 else 0
+
+            if wps < preprocess_args.min_wps:
+                print('bad segment in', video_id, '| wps =', wps)
+                continue
+
+            final_data[video_id].append({
+                'start': segment_start,
+                'end': segment_end,
+                'votes': line['votes'],
+                'locked': line['locked'] == '1',
+                'views': line['views'],
+                'reputation': line['reputation'],
+                'category': line['category'],
+                'action': line['actionType'],
+                'uuid': line['UUID'],
+            })
+
+        # Remove duplicate sponsor segments by choosing best (most votes)
+        for key in final_data:
+            final_data[key] = remove_duplicate_sponsor_segments(
+                final_data[key])
+
+        # Save data
+        with open(final_path, 'w') as fp:
+            json.dump(final_data, fp)
+
+        # final_data = preprocess(
+        #     raw_dataset_path, final_path, preprocess_args.min_votes)
+        # # TODO save metadata in final.json?
+
+    logging.info(f'Found {len(final_data)} videos')
+
+    # TODO shuffle final_data
+
+    # if not os.path.exists(excess_path) or preprocess_args.overwrite
+    # TODO use overwrite param
+
+    os.makedirs(dataset_args.data_dir, exist_ok=True)
+
+    positive_file = os.path.join(
+        dataset_args.data_dir, dataset_args.positive_file)
+    negative_file = os.path.join(
+        dataset_args.data_dir, dataset_args.negative_file)
+
+    if preprocess_args.do_generate:
+        print('Generating')
+        from model import get_tokenizer
+
+        # max_videos=preprocess_args.max_videos,
+        # max_segments=preprocess_args.max_segments,
+        # , max_videos, max_segments
+
+        tokenizer = get_tokenizer(model_args)
+
+        count_videos = 0
+        count_segments = 0  # TODO
+
+        write_mode = 'w' if preprocess_args.overwrite else 'a'
+
+        get_all = preprocess_args.max_videos is None
+        if get_all:
+            total = len(final_data)
+        else:
+            total = preprocess_args.max_videos
+
+        index = 0
+        data = final_data.items()
+        if preprocess_args.skip_videos is not None:
+            print('Skipping first', preprocess_args.skip_videos, 'videos')
+            data = itertools.islice(data, preprocess_args.skip_videos, None)
+            index = preprocess_args.skip_videos
+
+            if get_all:
+                total = max(0, total - preprocess_args.skip_videos)
+            else:
+                total = min(len(final_data) -
+                            preprocess_args.skip_videos, total)
+
+        with open(positive_file, write_mode, encoding='utf-8') as positive, \
+                open(negative_file, write_mode, encoding='utf-8') as negative, \
+                tqdm(total=total) as progress:
+
+            for video_id, sponsor_segments in data:
+                index += 1  # TODO FIX index + incrementing
+                progress.set_description(f'Processing {video_id}')
+
+                if get_all:
+                    progress.update()
+                elif count_videos >= preprocess_args.max_videos:
+                    break
+
+                words = get_words(video_id, False)
+                if not words:
+                    continue
+
+                num_words = len(words)
+                if num_words <= 1:
+                    continue
+
+                # TODO only count words that aren't [Music], [Applause], etc.
+
+                segments = segment.generate_labelled_segments(
+                    words, tokenizer, segmentation_args, sponsor_segments)
+
+                if not segments:
+                    continue
+
+                count_videos += 1
+                if not get_all:
+                    progress.update()
+
+                for seg in segments:
+
+                    segment_text = ' '.join((x['text'] for x in seg))
+
+                    extracted_text = ''
+                    for p in extract_sponsors(seg):
+                        p_text = ' '.join(p)
+                        extracted_text += f'{CustomTokens.START_SPONSOR.value} {p_text} {CustomTokens.END_SPONSOR.value}. '
+
+                    duration = segment.word_end(
+                        seg[-1]) - segment.word_start(seg[0])
+                    wps = len(seg)/duration if duration > 0 else 0
+                    # Ignore segments with "not enough words" in the transcript
+                    if wps < preprocess_args.min_wps:
+                        continue
+
+                    d = {
+                        'video_index': index,
+                        'video_id': video_id,
+                        'text': clean_text(segment_text),
+                        'words_per_second': wps,
+                    }
+
+                    d['sponsor'] = bool(extracted_text)
+                    d['extracted'] = clean_text(
+                        extracted_text) if d['sponsor'] else CustomTokens.NO_SPONSOR.value
+
+                    print(json.dumps(d), file=(
+                        positive if d['sponsor'] else negative))
+
+    if preprocess_args.do_split:
+        print('Splitting')
+        print('Read files')
+
+        with open(positive_file, encoding='utf-8') as positive:
+            sponsors = positive.readlines()
+
+        with open(negative_file, encoding='utf-8') as negative:
+            non_sponsors = negative.readlines()
+
+        print('Shuffle')
+        random.shuffle(sponsors)
+        random.shuffle(non_sponsors)
+
+        print('Calculate ratios')
+        # Ensure correct ratio of positive to negative segments
+        percentage_negative = 1 - preprocess_args.percentage_positive
+
+        if preprocess_args.percentage_positive * len(sponsors) > len(non_sponsors):
+            # Negative is limiting
+            z = int(preprocess_args.percentage_positive /
+                    percentage_negative * len(non_sponsors))
+
+            excess = sponsors[z:]
+            sponsors = sponsors[:z]
+
+        else:
+            # Positive is limiting
+            z = int(percentage_negative /
+                    preprocess_args.percentage_positive * len(sponsors))
+
+            excess = non_sponsors[z:]
+            non_sponsors = non_sponsors[:z]
+
+        print('Join')
+        all_labelled_segments = sponsors + non_sponsors
+
+        random.shuffle(all_labelled_segments)
+
+        print('Split')
+        ratios = [preprocess_args.train_split,
+                  preprocess_args.test_split,
+                  preprocess_args.valid_split]
+
+        train_data, test_data, valid_data = split(
+            all_labelled_segments, ratios)
+
+        splits = {
+            dataset_args.train_file: train_data,
+            dataset_args.test_file: test_data,
+            dataset_args.validation_file: valid_data
+        }
+
+        # Output training, testing and validation data
+        for name, items in 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:
+                    fp.writelines(items)
+            else:
+                print('Skipping', name)
+
+        print('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:
+            print('Skipping', dataset_args.excess_file)
+
+        print('Finished splitting:', len(sponsors),
+              'sponsors,', len(non_sponsors), 'non sponsors')
+
+
+def split(arr, ratios):
+    """Split array according to ratios. Sum of ratios should be less than 1"""
+
+    to_return = []
+
+    cumulative_sum = 0
+    for r in ratios:
+        current = cumulative_sum
+
+        cumulative_sum += r * len(arr)
+        to_return.append(arr[int(current):int(cumulative_sum)])
+
+    return to_return
+
+
+if __name__ == '__main__':
+    main()

src/segment.py

@@ -0,0 +1,142 @@
+import preprocess
+from shared import CustomTokens
+from dataclasses import dataclass, field
+
+
+@dataclass
+class SegmentationArguments:
+    pause_threshold: int = field(default=2, metadata={
+        'help': 'When the time between words is greater than pause threshold, force into a new segment'})
+
+
+# WORDS TO ALWAYS HAVE ON THEIR OWN
+# always_split_re = re.compile(r'\[\w+\]')
+# e.g., [Laughter], [Applause], [Music]
+always_split = [
+    CustomTokens.MUSIC.value,
+    CustomTokens.APPLAUSE.value,
+    CustomTokens.LAUGHTER.value
+]
+
+
+def get_overlapping_chunks_of_tokens(tokens, size, overlap):
+    for i in range(0, len(tokens), size-overlap+1):
+        yield tokens[i:i+size]
+
+
+# Generate up to max_tokens - SAFETY_TOKENS
+SAFETY_TOKENS = 8
+
+
+# TODO play around with this?
+OVERLAP_TOKEN_PERCENTAGE = 0.5  # 0.25
+
+
+def add_labels_to_words(words, sponsor_segments):
+
+    # TODO binary search
+    for word in words:
+        word['sponsor'] = False
+        for sponsor_segment in sponsor_segments:
+            if sponsor_segment['start'] <= word['start'] <= sponsor_segment['end']:
+                word['sponsor'] = True
+
+    # TODO use extract_segment with mapping function?
+    # TODO remove sponsor segments that contain mostly empty space?
+
+    return words
+
+
+def generate_labelled_segments(words, tokenizer, segmentation_args, sponsor_segments):
+    segments = generate_segments(words, tokenizer, segmentation_args)
+
+    labelled_segments = list(
+        map(lambda x: add_labels_to_words(x, sponsor_segments), segments))
+
+    return labelled_segments
+
+
+def word_start(word):
+    return word['start']
+
+
+def word_end(word):
+    return word.get('end', word['start'])
+
+
+def generate_segments(words, tokenizer, segmentation_args):
+    first_pass_segments = []
+
+    for index, word in enumerate(words):
+        # Get length of tokenized word
+        cleaned = preprocess.clean_text(word['text'])
+        word['num_tokens'] = len(
+            tokenizer(cleaned, add_special_tokens=False, truncation=True).input_ids)
+
+        add_new_segment = index == 0
+        if not add_new_segment:
+
+            if word['text'] in always_split or words[index-1]['text'] in always_split:
+                add_new_segment = True
+
+            # Pause too small, do not split
+            elif word_start(words[index]) - word_end(words[index-1]) >= segmentation_args.pause_threshold:
+                add_new_segment = True
+
+        if add_new_segment:  # New segment
+            first_pass_segments.append([word])
+
+        else:  # Add to current segment
+            first_pass_segments[-1].append(word)
+
+    max_q_size = tokenizer.model_max_length - SAFETY_TOKENS
+
+    buffer_size = OVERLAP_TOKEN_PERCENTAGE*max_q_size  # tokenizer.model_max_length
+
+    # In second pass, we split those segments if too big
+    second_pass_segments = []
+    for segment in first_pass_segments:
+        current_segment_num_tokens = 0
+        current_segment = []
+        for word in segment:
+            if current_segment_num_tokens + word['num_tokens'] < max_q_size:
+                # Can add tokens to current segment
+                current_segment.append(word)
+                current_segment_num_tokens += word['num_tokens']
+            else:
+                # Adding this token would make it have too many tokens
+                # We save this batch and create new
+                second_pass_segments.append(current_segment.copy())
+
+                current_segment.append(word)
+                current_segment_num_tokens += word['num_tokens']
+
+                while current_segment_num_tokens > buffer_size and current_segment:
+                    first_word = current_segment.pop(0)
+                    current_segment_num_tokens -= first_word['num_tokens']
+
+        if current_segment:
+            second_pass_segments.append(current_segment.copy())
+
+    return second_pass_segments
+
+
+def extract_segment(words, start, end, map_function=None):
+    """Extract a segment of words that are between (inclusive) the start and end points"""
+    segment_words = []
+
+    if start > end:
+        return segment_words
+
+    # TODO change to binary search
+    for w in words:  # Assumes words are sorted
+        if word_end(w) < start:
+            continue  # Ignore
+        if word_start(w) > end:
+            break  # Done with range
+        if map_function is not None and callable(map_function):
+            w = map_function(w)
+
+        segment_words.append(w)
+
+    return segment_words

src/shared.py

@@ -0,0 +1,96 @@
+import gc
+from time import time_ns
+import random
+import numpy as np
+import torch
+from typing import Optional
+from dataclasses import dataclass, field
+from enum import Enum
+
+
+class CustomTokens(Enum):
+    URL = 'URL_TOKEN'
+    HYPHENATED_URL = 'HYPHENATED_URL_TOKEN'
+    NUMBER_PERCENTAGE = 'NUMBER_PERCENTAGE_TOKEN'
+    NUMBER = 'NUMBER_TOKEN'
+
+    START_SPONSOR = 'START_SPONSOR'
+    END_SPONSOR = 'END_SPONSOR'
+    NO_SPONSOR = 'NO_SPONSOR_FOUND'
+
+    SHORT_HYPHENATED = 'SHORT_HYPHENATED_TOKEN'
+    LONG_WORD = 'LONG_WORD_TOKEN'
+
+    # Custom YouTube tokens
+    MUSIC = '[Music]'
+    APPLAUSE = '[Applause]'
+    LAUGHTER = '[Laughter]'
+
+    PROFANITY_RAW = '[ __ ]'  # How YouTube transcribes profanity
+    PROFANITY_CONVERTED = '*****'  # Safer version for tokenizing
+    PROFANITY = 'PROFANITY_TOKEN'
+
+    @classmethod
+    def custom_tokens(cls):
+        return [e.value for e in cls]
+
+    @classmethod
+    def add_custom_tokens(cls, tokenizer):
+        tokenizer.add_tokens(cls.custom_tokens())
+
+
+@dataclass
+class OutputArguments:
+
+    output_dir: str = field(
+        default='out',
+        metadata={
+            'help': 'The output directory where the model predictions and checkpoints will be written to and read from.'
+        },
+    )
+    checkpoint: Optional[str] = field(
+        default=None,
+        metadata={
+            'help': 'Choose the checkpoint/model to train from or test with. Defaults to the latest checkpoint found in `output_dir`.'
+        },
+    )
+    models_dir: str = field(
+        default='models',
+        metadata={
+            'help': 'The output directory where the model predictions and checkpoints will be written to and read from.'
+        },
+    )
+    # classifier_dir: str = field(
+    #     default='out',
+    #     metadata={
+    #         'help': 'The output directory where the model predictions and checkpoints will be written to and read from.'
+    #     },
+    # )
+
+
+def seed_factory():
+    return time_ns() % (2**32 - 1)
+
+
+@dataclass
+class GeneralArguments:
+    seed: Optional[int] = field(default_factory=seed_factory, metadata={
+        'help': 'Set seed for deterministic training and testing. By default, it uses the current time (results in essentially random results).'
+    })
+
+    def __post_init__(self):
+        random.seed(self.seed)
+        np.random.seed(self.seed)
+        torch.manual_seed(self.seed)
+        torch.cuda.manual_seed_all(self.seed)
+
+
+def device():
+    return torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
+def reset():
+    torch.clear_autocast_cache()
+    torch.cuda.empty_cache()
+    gc.collect()
+    print(torch.cuda.memory_summary(device=None, abbreviated=False))

src/train.py

@@ -0,0 +1,508 @@
+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 model import ModelArguments
+import transformers
+import logging
+from model import get_model, get_tokenizer
+import logging
+import os
+import sys
+from dataclasses import dataclass, field
+from typing import Optional
+import datasets
+import pickle
+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
+
+import re
+
+# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
+check_min_version('4.13.0.dev0')
+require_version('datasets>=1.8.0',
+                'To fix: pip install -r requirements.txt')
+
+os.environ['WANDB_DISABLED'] = 'true'
+
+
+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)],
+)
+
+
+@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.'},
+    )
+    # https://github.com/huggingface/transformers/issues/5204
+    max_source_length: Optional[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_target_length: Optional[int] = field(
+        default=512,
+        metadata={
+            'help': 'The maximum total sequence length for target text after tokenization. Sequences longer '
+            'than this will be truncated, sequences shorter will be padded.'
+        },
+    )
+    val_max_target_length: Optional[int] = field(
+        default=None,
+        metadata={
+            'help': 'The maximum total sequence length for validation target text after tokenization. Sequences longer '
+            'than this will be truncated, sequences shorter will be padded. Will default to `max_target_length`.'
+            'This argument is also used to override the ``max_length`` param of ``model.generate``, which is used '
+            'during ``evaluate`` and ``predict``.'
+        },
+    )
+    pad_to_max_length: bool = field(
+        default=False,
+        metadata={
+            'help': 'Whether to pad all samples to model maximum sentence length. '
+            'If False, will pad the samples dynamically when batching to the maximum length in the batch. More '
+            'efficient on GPU but very bad for TPU.'
+        },
+    )
+    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.'
+        },
+    )
+    num_beams: Optional[int] = field(
+        default=None,
+        metadata={
+            'help': 'Number of beams to use for evaluation. This argument will be passed to ``model.generate``, '
+            'which is used during ``evaluate`` and ``predict``.'
+        },
+    )
+    ignore_pad_token_for_loss: bool = field(
+        default=True,
+        metadata={
+            'help': 'Whether to ignore the tokens corresponding to padded labels in the loss computation or not.'
+        },
+    )
+    source_prefix: Optional[str] = field(
+        default=DEFAULT_TOKEN_PREFIX, metadata={
+            'help': 'A prefix to add before every source text (useful for T5 models).'}
+    )
+
+    # TODO add vectorizer params
+
+    def __post_init__(self):
+        if self.val_max_target_length is None:
+            self.val_max_target_length = self.max_target_length
+
+
+@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=2500, metadata={
+                            'help': 'Save checkpoint every X updates steps.'})
+    eval_steps: int = field(default=2500, metadata={
+                            'help': 'Run an evaluation every X steps.'})
+    logging_steps: int = field(default=2500, 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.
+    # We now keep distinct sets of args, for a cleaner separation of concerns.
+
+    hf_parser = HfArgumentParser((
+        ModelArguments,
+        DatasetArguments,
+        DataTrainingArguments,
+        SequenceTrainingArguments,
+        ClassifierArguments
+    ))
+    model_args, dataset_args, data_training_args, training_args, classifier_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)
+    transformers.utils.logging.set_verbosity(log_level)
+    transformers.utils.logging.enable_default_handler()
+    transformers.utils.logging.enable_explicit_format()
+
+    # Set seed before initializing model.
+    # set_seed(training_args.seed)
+
+    # 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}')
+
+    # FP16 https://github.com/huggingface/transformers/issues/9295
+
+    # Works:
+    # https://huggingface.co/docs/transformers/model_doc/t5v1.1
+    # google/t5-v1_1-small
+    # google/t5-v1_1-base
+    # google/t5-v1_1-large
+    # google/t5-v1_1-xl
+    # google/t5-v1_1-xxl
+
+    # https://huggingface.co/docs/transformers/model_doc/t5
+    # t5-small
+    # t5-base
+    # t5-large
+    # t5-3b
+    # t5-11b
+
+    # allenai/led-base-16384 - https://github.com/huggingface/transformers/issues/9810
+
+    # Further work:
+    # Multilingual- https://huggingface.co/docs/transformers/model_doc/mt5
+
+    # 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=50000
+            # max_features=10000
+        )
+
+        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:
+
+                # Get matches:
+                if row['sponsor']:
+                    matches = re.findall(SPONSOR_MATCH_RE, row['extracted'])
+                else:
+                    matches = [row['text']]
+
+                for match in matches:
+                    train_test_data[ds_type]['X'].append(match)
+                    train_test_data[ds_type]['y'].append(row['sponsor'])
+
+        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=500)
+
+        # 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:
+
+        if data_training_args.source_prefix is None and 't5-' in model_args.model_name_or_path:
+            logger.warning(
+                "You're running a t5 model but didn't provide a source prefix, which is the expected, e.g. with `--source_prefix 'summarize: ' `"
+            )
+
+        # 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.'
+                )
+
+        # Load pretrained model and tokenizer
+        tokenizer = get_tokenizer(model_args)
+        model = get_model(model_args)
+        model.to(device())
+        model.resize_token_embeddings(len(tokenizer))
+
+        if model.config.decoder_start_token_id is None:
+            raise ValueError(
+                'Make sure that `config.decoder_start_token_id` is correctly defined')
+
+        if hasattr(model.config, 'max_position_embeddings') and model.config.max_position_embeddings < data_training_args.max_source_length:
+            if model_args.resize_position_embeddings is None:
+                logger.warning(
+                    f"Increasing the model's number of position embedding vectors from {model.config.max_position_embeddings} to {data_training_args.max_source_length}."
+                )
+                model.resize_position_embeddings(
+                    data_training_args.max_source_length)
+
+            elif model_args.resize_position_embeddings:
+                model.resize_position_embeddings(
+                    data_training_args.max_source_length)
+
+            else:
+                raise ValueError(
+                    f'`--max_source_length` is set to {data_training_args.max_source_length}, but the model only has {model.config.max_position_embeddings}'
+                    f' position encodings. Consider either reducing `--max_source_length` to {model.config.max_position_embeddings} or to automatically '
+                    "resize the model's position encodings by passing `--resize_position_embeddings`."
+                )
+
+        # Preprocessing the datasets.
+        # We need to tokenize inputs and targets.
+        column_names = raw_datasets['train'].column_names
+
+        # Temporarily set max_target_length for training.
+        max_target_length = data_training_args.max_target_length
+        padding = 'max_length' if data_training_args.pad_to_max_length else False
+
+        if training_args.label_smoothing_factor > 0 and not hasattr(model, 'prepare_decoder_input_ids_from_labels'):
+            logger.warning(
+                'label_smoothing is enabled but the `prepare_decoder_input_ids_from_labels` method is not defined for'
+                f'`{model.__class__.__name__}`. This will lead to loss being calculated twice and will take up more memory'
+            )
+
+        prefix = data_training_args.source_prefix if data_training_args.source_prefix is not None else ''
+
+        # 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, max_length=data_training_args.max_source_length, padding=padding, truncation=True)
+
+            # Setup the tokenizer for targets
+            with tokenizer.as_target_tokenizer():
+                labels = tokenizer(
+                    targets, max_length=max_target_length, padding=padding, 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.
+            if padding == 'max_length' and data_training_args.ignore_pad_token_for_loss:
+                labels['input_ids'] = [
+                    [(l if l != tokenizer.pad_token_id else -100) for l in label] for label in labels['input_ids']
+                ]
+            model_inputs['labels'] = 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')
+
+        max_target_length = data_training_args.val_max_target_length
+        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
+        label_pad_token_id = - \
+            100 if data_training_args.ignore_pad_token_for_loss else tokenizer.pad_token_id
+        data_collator = DataCollatorForSeq2Seq(
+            tokenizer,
+            model=model,
+            label_pad_token_id=label_pad_token_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
+        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:
+            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 __name__ == '__main__':
+    main()

src/utils.py

@@ -0,0 +1,86 @@
+import asyncio
+import os
+
+class Job:
+    def __init__(self, function, *args, **kwargs) -> None:
+        self.function = function
+        self.args = args
+        self.kwargs = kwargs
+
+        self.result = None
+
+
+class InterruptibleThreadPool:
+    def __init__(self,
+                 num_workers=None,
+                 loop=None,
+                 shutdown_message='\nAttempting graceful shutdown, press Ctrl+C again to exit...',
+                 on_job_complete=None,  # Useful for monitoring progress
+                 raise_after_interrupt=False,
+                 ) -> None:
+        self.num_workers = os.cpu_count() if num_workers is None else num_workers
+        self.loop = asyncio.get_event_loop() if loop is None else loop
+        self.shutdown_message = shutdown_message
+
+        self.sem = asyncio.Semaphore(num_workers)
+
+        self.jobs = []
+
+        self.on_job_complete = on_job_complete
+        self.raise_after_interrupt = raise_after_interrupt
+
+    async def _sync_to_async(self, job):
+        async with self.sem:  # Limit number of parallel tasks
+            job.result = await self.loop.run_in_executor(None, job.function, *job.args, **job.kwargs)
+
+            if callable(self.on_job_complete):
+                self.on_job_complete(job)
+
+            return job
+
+    def add_job(self, job):
+        self.jobs.append(job)
+
+    def run(self):
+        try:
+            tasks = [
+                # creating task starts coroutine
+                asyncio.ensure_future(self._sync_to_async(job))
+                for job in self.jobs
+            ]
+
+            # https://stackoverflow.com/a/42097478
+            self.loop.run_until_complete(
+                asyncio.gather(*tasks, return_exceptions=True)
+            )
+
+        except KeyboardInterrupt:
+            # Optionally show a message if the shutdown may take a while
+            print(self.shutdown_message, flush=True)
+
+            # Do not show `asyncio.CancelledError` exceptions during shutdown
+            # (a lot of these may be generated, skip this if you prefer to see them)
+            def shutdown_exception_handler(loop, context):
+                if "exception" not in context \
+                        or not isinstance(context["exception"], asyncio.CancelledError):
+                    loop.default_exception_handler(context)
+            self.loop.set_exception_handler(shutdown_exception_handler)
+
+            # Handle shutdown gracefully by waiting for all tasks to be cancelled
+            cancelled_tasks = asyncio.gather(
+                *asyncio.all_tasks(loop=self.loop), loop=self.loop, return_exceptions=True)
+            cancelled_tasks.add_done_callback(lambda t: self.loop.stop())
+            cancelled_tasks.cancel()
+
+            # Keep the event loop running until it is either destroyed or all
+            # tasks have really terminated
+            while not cancelled_tasks.done() and not self.loop.is_closed():
+                self.loop.run_forever()
+
+            if self.raise_after_interrupt:
+                raise
+        finally:
+            self.loop.run_until_complete(self.loop.shutdown_asyncgens())
+            self.loop.close()
+
+        return self.jobs