Upload 21 files

.github/FUNDING.yml

@@ -0,0 +1,3 @@
+github: xenova
+ko_fi: xenova
+custom: https://www.buymeacoffee.com/xenova

.github/workflows/check_large_file.yml

@@ -0,0 +1,16 @@
+name: Check file size
+on:               # or directly `on: [push]` to run the action on every push on any branch
+  pull_request:
+    branches: [main]
+
+  # to run this workflow manually from the Actions tab
+  workflow_dispatch:
+
+jobs:
+  sync-to-hub:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check large files
+        uses: ActionsDesk/lfs-warning@v2.0
+        with:
+          filesizelimit: 10485760 # this is 10MB so we can sync to HF Spaces

.github/workflows/sync_with_huggingface.yml

@@ -0,0 +1,19 @@
+name: Sync to Hugging Face hub
+on:
+  push:
+    branches: [main]
+
+  # to run this workflow manually from the Actions tab
+  workflow_dispatch:
+
+jobs:
+  sync-to-hub:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+        with:
+          fetch-depth: 0
+      - name: Push to hub
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
+        run: git push https://Xenova:$HF_TOKEN@huggingface.co/spaces/Xenova/sponsorblock-ml main

app.py

@@ -0,0 +1,318 @@
+
+from functools import partial
+from math import ceil, floor
+import streamlit.components.v1 as components
+import streamlit as st
+import sys
+import os
+import json
+from urllib.parse import quote
+
+# Allow direct execution
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), 'src'))  # noqa
+
+from preprocess import get_words
+from predict import PredictArguments, SegmentationArguments, predict as pred
+from shared import GeneralArguments, seconds_to_time, CATGEGORY_OPTIONS
+from utils import regex_search
+from model import get_model_tokenizer_classifier
+from errors import TranscriptError
+
+st.set_page_config(
+    page_title='SponsorBlock ML',
+    page_icon='🤖',
+    layout='wide',
+    #  initial_sidebar_state="expanded",
+    menu_items={
+        # 'Get Help': 'https://github.com/xenova/sponsorblock-ml',
+        # 'Report a bug': 'https://github.com/xenova/sponsorblock-ml/issues/new/choose',
+        #  'About': "# This is a header. This is an *extremely* cool app!"
+    }
+)
+
+
+YT_VIDEO_REGEX = r'''(?x)^
+                (?:
+                    # http(s):// or protocol-independent URL
+                    (?:https?://|//)
+                    (?:(?:(?:(?:\w+\.)?[yY][oO][uU][tT][uU][bB][eE](?:-nocookie|kids)?\.com/|
+                    youtube\.googleapis\.com/)                        # the various hostnames, with wildcard subdomains
+                    (?:.*?\#/)?                                          # handle anchor (#/) redirect urls
+                    (?:                                                  # the various things that can precede the ID:
+                        # v/ or embed/ or e/
+                        (?:(?:v|embed|e)/(?!videoseries))
+                        |(?:                                             # or the v= param in all its forms
+                            # preceding watch(_popup|.php) or nothing (like /?v=xxxx)
+                            (?:(?:watch|movie)(?:_popup)?(?:\.php)?/?)?
+                            (?:\?|\#!?)                                  # the params delimiter ? or # or #!
+                            # any other preceding param (like /?s=tuff&v=xxxx or ?s=tuff&v=V36LpHqtcDY)
+                            (?:.*?[&;])??
+                            v=
+                        )
+                    ))
+                    |(?:
+                    youtu\.be                                        # just youtu.be/xxxx
+                    )/)
+                )?                                                       # all until now is optional -> you can pass the naked ID
+                # here is it! the YouTube video ID
+                (?P<id>[0-9A-Za-z_-]{11})'''
+
+# https://github.com/google-research/text-to-text-transfer-transformer#released-model-checkpoints
+# https://github.com/google-research/text-to-text-transfer-transformer/blob/main/released_checkpoints.md#experimental-t5-pre-trained-model-checkpoints
+
+# https://huggingface.co/docs/transformers/model_doc/t5
+# https://huggingface.co/docs/transformers/model_doc/t5v1.1
+
+
+# Faster caching system for predictions (No need to hash)
+@st.cache(persist=True, allow_output_mutation=True)
+def create_prediction_cache():
+    return {}
+
+
+@st.cache(persist=True, allow_output_mutation=True)
+def create_function_cache():
+    return {}
+
+
+prediction_cache = create_prediction_cache()
+prediction_function_cache = create_function_cache()
+
+MODELS = {
+    'Small (293 MB)': {
+        'pretrained': 'google/t5-v1_1-small',
+        'repo_id': 'Xenova/sponsorblock-small',
+        'num_parameters': '77M'
+    },
+    'Base v1 (850 MB)': {
+        'pretrained': 't5-base',
+        'repo_id': 'Xenova/sponsorblock-base-v1',
+        'num_parameters': '220M'
+    },
+
+    'Base v1.1 (944 MB)': {
+        'pretrained': 'google/t5-v1_1-base',
+        'repo_id': 'Xenova/sponsorblock-base-v1.1',
+        'num_parameters': '250M'
+    }
+}
+
+# Create per-model cache
+for m in MODELS:
+    if m not in prediction_cache:
+        prediction_cache[m] = {}
+
+
+CLASSIFIER_PATH = 'Xenova/sponsorblock-classifier-v2'
+
+
+TRANSCRIPT_TYPES = {
+    'AUTO_MANUAL': {
+        'label': 'Auto-generated (fallback to manual)',
+        'type': 'auto',
+        'fallback': 'manual'
+    },
+    'MANUAL_AUTO': {
+        'label': 'Manual (fallback to auto-generated)',
+        'type': 'manual',
+        'fallback': 'auto'
+    },
+    # 'TRANSLATED': 'Translated to English' # Coming soon
+}
+
+
+def predict_function(model_id, model, tokenizer, segmentation_args, classifier, video_id, words, ts_type_id):
+    cache_id = f'{video_id}_{ts_type_id}'
+
+    if cache_id not in prediction_cache[model_id]:
+        prediction_cache[model_id][cache_id] = pred(
+            video_id, model, tokenizer,
+            segmentation_args=segmentation_args,
+            words=words,
+            classifier=classifier
+        )
+    return prediction_cache[model_id][cache_id]
+
+
+def load_predict(model_id):
+    model_info = MODELS[model_id]
+
+    if model_id not in prediction_function_cache:
+        # Use default segmentation and classification arguments
+        predict_args = PredictArguments(model_name_or_path=model_info['repo_id'])
+        general_args = GeneralArguments()
+        segmentation_args = SegmentationArguments()
+
+        model, tokenizer, classifier = get_model_tokenizer_classifier(predict_args, general_args)
+
+        prediction_function_cache[model_id] = partial(
+            predict_function, model_id, model, tokenizer, segmentation_args, classifier)
+
+
+    return prediction_function_cache[model_id]
+
+
+def create_button(text, url):
+    return f"""<div class="row-widget stButton" style="text-align: center">
+        <a href="{url}" target="_blank" rel="noopener noreferrer" class="btn-link">
+            <button kind="primary" class="btn">{text}</button>
+        </a>
+    </div>"""
+
+
+def main():
+    st.markdown("""<style>
+    .btn {
+        display: inline-flex;
+        -webkit-box-align: center;
+        align-items: center;
+        -webkit-box-pack: center;
+        justify-content: center;
+        font-weight: 600;
+        padding: 0.25rem 0.75rem;
+        border-radius: 0.25rem;
+        margin: 0px;
+        line-height: 1.5;
+        color: inherit;
+        width: auto;
+        user-select: none;
+        background-color: inherit;
+        border: 1px solid rgba(49, 51, 63, 0.2);
+    }
+    .btn-link {
+        color: inherit;
+        text-decoration: none;
+    }
+    </style>""", unsafe_allow_html=True)
+
+    top = st.container()
+    output = st.empty()
+
+    # Display heading and subheading
+    top.markdown('# PromoDetect')
+    top.markdown(
+        '##### Automatically detect in-video YouTube sponsorships, self/unpaid promotions, and interaction reminders.')
+
+    # Add controls
+
+    col1, col2 = top.columns(2)
+
+    with col1:
+        model_id = st.selectbox(
+            'Select model', MODELS.keys(), index=0, on_change=output.empty)
+
+    with col2:
+        ts_type_id = st.selectbox(
+            'Transcript type', TRANSCRIPT_TYPES.keys(), index=0, format_func=lambda x: TRANSCRIPT_TYPES[x]['label'], on_change=output.empty)
+
+    query_params = st.experimental_get_query_params()
+
+    video_id = None
+    
+    if 'v' in query_params:
+        video_id = query_params['v'][0]
+
+    if video_id is None:
+        video_input = top.text_input('Video URL/ID:', on_change=output.empty)
+    else :
+        video_input = top.text_input('Video URL/ID:', on_change=output.empty,value = video_id)
+    categories = top.multiselect('Categories:',
+                                 CATGEGORY_OPTIONS.keys(),
+                                 CATGEGORY_OPTIONS.keys(),
+                                 format_func=CATGEGORY_OPTIONS.get, on_change=output.empty
+                                 )
+
+    # Hide segments with a confidence lower than
+    confidence_threshold = top.slider(
+        'Confidence Threshold (%):', min_value=0, value=50, max_value=100, on_change=output.empty)
+
+    if len(video_input) == 0:  # No input, do not continue
+        return
+
+    # Load prediction function
+    with st.spinner('Loading model...'):
+        predict = load_predict(model_id)
+
+    with output.container():  # Place all content in output container
+        video_id = regex_search(video_input, YT_VIDEO_REGEX)
+        if video_id is None:
+            st.exception(ValueError('Invalid YouTube URL/ID'))
+            return
+
+        try:
+            with st.spinner('Downloading transcript...'):
+                words = get_words(video_id,
+                                  transcript_type=TRANSCRIPT_TYPES[ts_type_id]['type'],
+                                  fallback=TRANSCRIPT_TYPES[ts_type_id]['fallback']
+                                  )
+        except TranscriptError:
+            pass
+
+        if not words:
+            st.error('No transcript found!')
+            return
+
+        with st.spinner('Running model...'):
+            predictions = predict(video_id, words, ts_type_id)
+
+        if len(predictions) == 0:
+            st.success('No segments found!')
+            return
+
+        submit_segments = []
+        for index, prediction in enumerate(predictions, start=1):
+            category_key = prediction['category'].upper()
+            if category_key not in categories:
+                continue  # Skip
+
+            confidence = prediction['probability'] * 100
+
+            if confidence < confidence_threshold:
+                continue
+
+            submit_segments.append({
+                'segment': [prediction['start'], prediction['end']],
+                'category': prediction['category'],
+                'actionType': 'skip'
+            })
+            start_time = seconds_to_time(prediction['start'])
+            end_time = seconds_to_time(prediction['end'])
+            with st.expander(
+                f"[{category_key}] Prediction #{index} ({start_time} \u2192 {end_time})"
+            ):
+
+                url = f"https://www.youtube-nocookie.com/embed/{video_id}?&start={floor(prediction['start'])}&end={ceil(prediction['end'])}"
+                # autoplay=1controls=0&&modestbranding=1&fs=0
+
+                # , width=None, height=None, scrolling=False
+                components.iframe(url, width=670, height=376)
+
+                text = ' '.join(w['text'] for w in prediction['words'])
+                st.write(f"**Times:** {start_time} \u2192 {end_time}")
+                st.write(
+                    f"**Category:** {CATGEGORY_OPTIONS[category_key]}")
+                st.write(f"**Confidence:** {confidence:.2f}%")
+                st.write(f'**Text:** "{text}"')
+
+        if not submit_segments:
+            st.success(
+                f'No segments found! ({len(predictions)} ignored due to filters/settings)')
+            return
+
+        num_hidden = len(predictions) - len(submit_segments)
+        if num_hidden > 0:
+            st.info(
+                f'{num_hidden} predictions hidden (adjust the settings and filters to view them all).')
+
+        json_data = quote(json.dumps(submit_segments))
+        link = f'https://www.youtube.com/watch?v={video_id}#segments={json_data}'
+        st.markdown(create_button('Submit Segments', link),
+                    unsafe_allow_html=True)
+
+        # st.markdown(f"""<div style="text-align: center;font-size: 16px;margin-top: 6px">
+        # <a href="https://wiki.sponsor.ajay.app/w/Automating_Submissions" target="_blank" rel="noopener noreferrer">(Review before submitting!)</a>
+        # </div>""", unsafe_allow_html=True)
+
+
+if __name__ == '__main__':
+    main()

data/.gitignore

@@ -0,0 +1,4 @@
+# Ignore everything in this directory
+*
+# Except this file
+!.gitignore

models/.gitignore

@@ -0,0 +1,4 @@
+# Ignore everything in this directory
+*
+# Except this file
+!.gitignore

raw/.gitignore

@@ -0,0 +1,4 @@
+# Ignore everything in this directory
+*
+# Except this file
+!.gitignore

requirements.txt

@@ -0,0 +1,7 @@
+transformers
+datasets
+youtube_transcript_api
+torch
+pandas
+numpy
+sentencepiece

src/classify.py

@@ -0,0 +1,43 @@
+from transformers import TextClassificationPipeline
+import preprocess
+import segment
+
+
+class SponsorBlockClassificationPipeline(TextClassificationPipeline):
+    def __init__(self, model, tokenizer):
+        device = next(model.parameters()).device.index
+        if device is None:
+            device = -1
+        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

@@ -0,0 +1,24 @@
+
+class SponsorBlockException(Exception):
+    """Base class for all sponsor block exceptions"""
+    pass
+
+
+class InferenceException(SponsorBlockException):
+    """An exception occurred while predicting sponsor segments"""
+    pass
+
+
+class TranscriptError(SponsorBlockException):
+    """An exception occurred while retrieving the video transcript"""
+    pass
+
+
+class ModelError(SponsorBlockException):
+    """Base class for model-related errors"""
+    pass
+
+
+class ModelLoadError(ModelError):
+    """An exception occurred while loading the model"""
+    pass

src/evaluate.py

@@ -0,0 +1,408 @@
+
+from model import get_model_tokenizer_classifier, InferenceArguments
+from utils import jaccard, safe_print
+from transformers import HfArgumentParser
+from preprocess import get_words, clean_text
+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
+from typing import Optional
+from tqdm import tqdm
+import json
+import os
+import random
+from shared import seconds_to_time
+from urllib.parse import quote
+import logging
+
+logging.basicConfig()
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class EvaluationArguments(InferenceArguments):
+    """Arguments pertaining to how evaluation will occur."""
+    output_file: Optional[str] = field(
+        default='metrics.csv',
+        metadata={
+            'help': 'Save metrics to output file'
+        }
+    )
+
+    skip_missing: bool = field(
+        default=False,
+        metadata={
+            'help': 'Whether to skip checking for missing segments. If False, predictions will be made.'
+        }
+    )
+    skip_incorrect: bool = field(
+        default=False,
+        metadata={
+            'help': 'Whether to skip checking for incorrect segments. If False, classifications will be made on existing segments.'
+        }
+    )
+
+
+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:
+            j = jaccard(prediction['start'], prediction['end'],
+                        sponsor_segment['start'], sponsor_segment['end'])
+            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
+
+        duration = word_end(word) - word_start(word)
+
+        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.get('category') is not None:  # Is actual sponsor
+                metrics['true_positive'] += duration
+            else:
+                metrics['false_positive'] += duration
+        else:
+            # total_negative_time += duration
+            if word.get('category') is not None:  # 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():
+    logger.setLevel(logging.DEBUG)
+
+    hf_parser = HfArgumentParser((
+        EvaluationArguments,
+        DatasetArguments,
+        SegmentationArguments,
+        GeneralArguments
+    ))
+
+    evaluation_args, dataset_args, segmentation_args, general_args = hf_parser.parse_args_into_dataclasses()
+
+    if evaluation_args.skip_missing and evaluation_args.skip_incorrect:
+        logger.error('ERROR: Nothing to do')
+        return
+
+    # Load labelled data:
+    final_path = os.path.join(
+        dataset_args.data_dir, dataset_args.processed_file)
+
+    if not os.path.exists(final_path):
+        logger.error('ERROR: Processed database not found.\n'
+                     f'Run `python src/preprocess.py --update_database --do_create` to generate "{final_path}".')
+        return
+
+    model, tokenizer, classifier = get_model_tokenizer_classifier(
+        evaluation_args, general_args)
+
+    with open(final_path) as fp:
+        final_data = json.load(fp)
+
+    if evaluation_args.video_ids:  # Use specified
+        video_ids = evaluation_args.video_ids
+
+    else:  # Use items found in preprocessed database
+        video_ids = list(final_data.keys())
+        random.shuffle(video_ids)
+
+        if evaluation_args.start_index is not None:
+            video_ids = video_ids[evaluation_args.start_index:]
+
+        if evaluation_args.max_videos is not None:
+            video_ids = video_ids[:evaluation_args.max_videos]
+
+    out_metrics = []
+
+    all_metrics = {}
+    if not evaluation_args.skip_missing:
+        all_metrics['total_prediction_accuracy'] = 0
+        all_metrics['total_prediction_precision'] = 0
+        all_metrics['total_prediction_recall'] = 0
+        all_metrics['total_prediction_fscore'] = 0
+
+    if not evaluation_args.skip_incorrect:
+        all_metrics['classifier_segment_correct'] = 0
+        all_metrics['classifier_segment_count'] = 0
+
+    metric_count = 0
+
+    postfix_info = {}
+
+    try:
+        with tqdm(video_ids) as progress:
+            for video_index, video_id in enumerate(progress):
+                progress.set_description(f'Processing {video_id}')
+
+                words = get_words(video_id)
+                if not words:
+                    continue
+
+                # Get labels
+                sponsor_segments = final_data.get(video_id)
+
+                # Reset previous
+                missed_segments = []
+                incorrect_segments = []
+
+                current_metrics = {
+                    'video_id': video_id
+                }
+                metric_count += 1
+
+                if not evaluation_args.skip_missing:  # Make predictions
+                    predictions = predict(video_id, model, tokenizer, segmentation_args,
+                                          classifier=classifier,
+                                          min_probability=evaluation_args.min_probability)
+
+                    if sponsor_segments:
+                        labelled_words = add_labels_to_words(
+                            words, sponsor_segments)
+
+                        current_metrics.update(
+                            calculate_metrics(labelled_words, predictions))
+
+                        all_metrics['total_prediction_accuracy'] += current_metrics['accuracy']
+                        all_metrics['total_prediction_precision'] += current_metrics['precision']
+                        all_metrics['total_prediction_recall'] += current_metrics['recall']
+                        all_metrics['total_prediction_fscore'] += current_metrics['f-score']
+
+                        # Just for display purposes
+                        postfix_info.update({
+                            'accuracy': all_metrics['total_prediction_accuracy']/metric_count,
+                            'precision':  all_metrics['total_prediction_precision']/metric_count,
+                            'recall':  all_metrics['total_prediction_recall']/metric_count,
+                            'f-score': all_metrics['total_prediction_fscore']/metric_count,
+                        })
+
+                        sponsor_segments = attach_predictions_to_sponsor_segments(
+                            predictions, sponsor_segments)
+
+                        # Identify possible issues:
+                        for prediction in predictions:
+                            if prediction['best_sponsorship'] is not None:
+                                continue
+
+                            prediction_words = prediction.pop('words', [])
+
+                            # Attach original text to missed segments
+                            prediction['text'] = ' '.join(
+                                x['text'] for x in prediction_words)
+                            missed_segments.append(prediction)
+
+                    else:
+                        # Not in database (all segments missed)
+                        missed_segments = predictions
+
+                if not evaluation_args.skip_incorrect and sponsor_segments:
+                    # Check for incorrect segments using the classifier
+
+                    segments_to_check = []
+                    cleaned_texts = []  # Texts to send through tokenizer
+                    for sponsor_segment in sponsor_segments:
+                        segment_words = extract_segment(
+                            words,  sponsor_segment['start'],  sponsor_segment['end'])
+                        sponsor_segment['text'] = ' '.join(
+                            x['text'] for x in segment_words)
+
+                        duration = sponsor_segment['end'] - \
+                            sponsor_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
+                        # or segment['votes'] > 5:
+                        if sponsor_segment['locked']:
+                            continue
+
+                        cleaned_texts.append(
+                            clean_text(sponsor_segment['text']))
+                        segments_to_check.append(sponsor_segment)
+
+                    if segments_to_check:  # Some segments to check
+
+                        segments_scores = classifier(cleaned_texts)
+
+                        num_correct = 0
+                        for segment, scores in zip(segments_to_check, segments_scores):
+
+                            fixed_scores = {
+                                score['label']: score['score']
+                                for score in scores
+                            }
+
+                            all_metrics['classifier_segment_count'] += 1
+
+                            prediction = max(scores, key=lambda x: x['score'])
+                            predicted_category = prediction['label'].lower()
+
+                            if predicted_category == segment['category']:
+                                num_correct += 1
+                                continue  # Ignore correct segments
+
+                            segment.update({
+                                'predicted': predicted_category,
+                                'scores': fixed_scores
+                            })
+
+                            incorrect_segments.append(segment)
+
+                        current_metrics['num_segments'] = len(
+                            segments_to_check)
+                        current_metrics['classified_correct'] = num_correct
+
+                        all_metrics['classifier_segment_correct'] += num_correct
+
+                    if all_metrics['classifier_segment_count'] > 0:
+                        postfix_info['classifier_accuracy'] = all_metrics['classifier_segment_correct'] / \
+                            all_metrics['classifier_segment_count']
+
+                out_metrics.append(current_metrics)
+                progress.set_postfix(postfix_info)
+
+                if missed_segments or incorrect_segments:
+
+                    if evaluation_args.output_as_json:
+                        to_print = {'video_id': video_id}
+
+                        if missed_segments:
+                            to_print['missed'] = missed_segments
+
+                        if incorrect_segments:
+                            to_print['incorrect'] = incorrect_segments
+
+                        safe_print(json.dumps(to_print))
+
+                    else:
+                        safe_print(
+                            f'Issues identified for {video_id} (#{video_index})')
+                        # Potentially missed segments (model predicted, but not in database)
+                        if missed_segments:
+                            safe_print(' - Missed segments:')
+                            segments_to_submit = []
+                            for i, missed_segment in enumerate(missed_segments, start=1):
+                                safe_print(f'\t#{i}:', seconds_to_time(
+                                    missed_segment['start']), '-->', seconds_to_time(missed_segment['end']))
+                                safe_print('\t\tText: "',
+                                           missed_segment['text'], '"', sep='')
+                                safe_print('\t\tCategory:',
+                                           missed_segment.get('category'))
+                                if 'probability' in missed_segment:
+                                    safe_print('\t\tProbability:',
+                                               missed_segment['probability'])
+
+                                segments_to_submit.append({
+                                    'segment': [missed_segment['start'], missed_segment['end']],
+                                    'category': missed_segment['category'].lower(),
+                                    'actionType': 'skip'
+                                })
+
+                            json_data = quote(json.dumps(segments_to_submit))
+                            safe_print(
+                                f'\tSubmit: https://www.youtube.com/watch?v={video_id}#segments={json_data}')
+
+                        # Incorrect segments (in database, but incorrectly classified)
+                        if incorrect_segments:
+                            safe_print(' - Incorrect segments:')
+                            for i, incorrect_segment in enumerate(incorrect_segments, start=1):
+                                safe_print(f'\t#{i}:', seconds_to_time(
+                                    incorrect_segment['start']), '-->', seconds_to_time(incorrect_segment['end']))
+
+                                safe_print(
+                                    '\t\tText: "', incorrect_segment['text'], '"', sep='')
+                                safe_print(
+                                    '\t\tUUID:', incorrect_segment['uuid'])
+                                safe_print(
+                                    '\t\tVotes:', incorrect_segment['votes'])
+                                safe_print(
+                                    '\t\tViews:', incorrect_segment['views'])
+                                safe_print('\t\tLocked:',
+                                           incorrect_segment['locked'])
+
+                                safe_print('\t\tCurrent Category:',
+                                           incorrect_segment['category'])
+                                safe_print('\t\tPredicted Category:',
+                                           incorrect_segment['predicted'])
+                                safe_print('\t\tProbabilities:')
+                                for label, score in incorrect_segment['scores'].items():
+                                    safe_print(
+                                        f"\t\t\t{label}: {score}")
+
+                        safe_print()
+
+    except KeyboardInterrupt:
+        pass
+
+    df = pd.DataFrame(out_metrics)
+
+    df.to_csv(evaluation_args.output_file)
+    logger.info(df.mean())
+
+
+if __name__ == '__main__':
+    main()

src/model.py

@@ -0,0 +1,235 @@
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, AutoConfig, AutoModelForSequenceClassification, TrainingArguments
+from shared import CustomTokens, GeneralArguments
+from dataclasses import dataclass, field
+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
+class ModelArguments:
+    """
+    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
+    """
+
+    model_name_or_path: str = field(
+        default=None,
+        metadata={
+            'help': 'Path to pretrained model or model identifier from huggingface.co/models'
+        }
+    )
+
+    cache_dir: Optional[str] = field(
+        default='models',
+        metadata={
+            'help': 'Where to store the pretrained models downloaded from huggingface.co'
+        },
+    )
+    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(
+        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).'
+        },
+    )
+
+import itertools
+from errors import InferenceException, ModelLoadError
+
+@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='EColi/SB_Classifier',
+        metadata={
+            'help': 'Use a pretrained classifier'
+        }
+    )
+
+    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'})
+
+    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 model_args.model_name_or_path is None:
+        raise ModelLoadError('Must specify --model_name_or_path')
+
+    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_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,
+    )
+
+    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))
+
+    # 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/predict.py

@@ -0,0 +1,285 @@
+
+from transformers import HfArgumentParser
+from dataclasses import dataclass, field
+import logging
+from shared import CustomTokens, extract_sponsor_matches, GeneralArguments, seconds_to_time
+from segment import (
+    generate_segments,
+    extract_segment,
+    MIN_SAFETY_TOKENS,
+    SAFETY_TOKENS_PERCENTAGE,
+    word_start,
+    word_end,
+    SegmentationArguments
+)
+import preprocess
+from errors import TranscriptError
+from model import get_model_tokenizer_classifier, InferenceArguments
+
+logging.basicConfig()
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class PredictArguments(InferenceArguments):
+    video_id: str = field(
+        default=None,
+        metadata={
+            'help': 'Video to predict segments for'}
+    )
+
+    def __post_init__(self):
+        if self.video_id is not None:
+            self.video_ids.append(self.video_id)
+
+        super().__post_init__()
+
+
+MATCH_WINDOW = 25       # Increase for accuracy, but takes longer: O(n^3)
+MERGE_TIME_WITHIN = 8   # Merge predictions if they are within x seconds
+
+# Any prediction whose start time is <= this will be set to start at 0
+START_TIME_ZERO_THRESHOLD = 0.08
+
+
+def filter_and_add_probabilities(predictions, classifier, min_probability):
+    """Use classifier to filter predictions"""
+    if not predictions:
+        return predictions
+
+    # We update the predicted category from the extractive transformer
+    # if the classifier is confident enough it is another category
+
+    texts = [
+        preprocess.clean_text(' '.join([x['text'] for x in pred['words']]))
+        for pred in predictions
+    ]
+    classifications = classifier(texts)
+
+    filtered_predictions = []
+    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 (prediction['category'] not in predicted_probabilities) \
+                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
+
+        if prediction['category'] == 'none':
+            continue  # Ignore if categorised as nothing
+
+        prediction['probability'] = predicted_probabilities[prediction['category']]
+
+        if min_probability is not None and prediction['probability'] < min_probability:
+            continue  # Ignore if below threshold
+
+        # TODO add probabilities, but remove None and normalise rest
+        prediction['probabilities'] = predicted_probabilities
+
+        # if prediction['probability'] < classifier_args.min_probability:
+        #     continue
+
+        filtered_predictions.append(prediction)
+
+    return filtered_predictions
+
+
+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)
+        if not words:
+            raise TranscriptError('Unable to retrieve transcript')
+
+    segments = generate_segments(
+        words,
+        tokenizer,
+        segmentation_args
+    )
+
+    predictions = segments_to_predictions(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 is not None:
+        predictions = filter_and_add_probabilities(
+            predictions, classifier, min_probability)
+
+    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
+
+
+def predict_sponsor_from_texts(texts, model, tokenizer):
+    clean_texts = list(map(preprocess.clean_text, texts))
+    return predict_sponsor_from_cleaned_texts(clean_texts, model, tokenizer)
+
+
+def predict_sponsor_from_cleaned_texts(cleaned_texts, model, tokenizer):
+    """Given a body of text, predict the words which are part of the sponsor"""
+    model_device = next(model.parameters()).device
+
+    decoded_outputs = []
+    # Do individually, to avoid running out of memory for long videos
+    for cleaned_words in cleaned_texts:
+        text = CustomTokens.EXTRACT_SEGMENTS_PREFIX.value + \
+            ' '.join(cleaned_words)
+        input_ids = tokenizer(text, return_tensors='pt',
+                              truncation=True).input_ids.to(model_device)
+
+        # Optimise output length so that we do not generate unnecessarily long texts
+        max_out_len = round(min(
+            max(
+                len(input_ids[0])/SAFETY_TOKENS_PERCENTAGE,
+                len(input_ids[0]) + MIN_SAFETY_TOKENS
+            ),
+            model.model_dim)
+        )
+
+        outputs = model.generate(input_ids, max_length=max_out_len)
+        decoded_outputs.append(tokenizer.decode(
+            outputs[0], skip_special_tokens=True))
+
+    return decoded_outputs
+
+
+def segments_to_predictions(segments, model, tokenizer):
+    predicted_time_ranges = []
+
+    cleaned_texts = [
+        [x['cleaned'] for x in cleaned_segment]
+        for cleaned_segment in segments
+    ]
+
+    sponsorship_texts = predict_sponsor_from_cleaned_texts(
+        cleaned_texts, model, tokenizer)
+
+    matches = extract_sponsor_matches(sponsorship_texts)
+
+    for segment_matches, cleaned_batch, segment in zip(matches, cleaned_texts, segments):
+
+        for match in segment_matches:  # one segment might contain multiple sponsors/ir/selfpromos
+
+            matched_text = match['text'].split()
+
+            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]),
+                'category': match['category']
+            })
+
+    # 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
+
+    prev_prediction = None
+
+    final_predicted_time_ranges = []
+    for range in predicted_time_ranges:
+        start_time = range['start'] if range['start'] > START_TIME_ZERO_THRESHOLD else 0
+        end_time = range['end']
+
+        if prev_prediction is not None and \
+                (start_time <= prev_prediction['end'] <= end_time or    # Merge overlapping segments
+                    (range['category'] == prev_prediction['category']   # Merge disconnected segments if same category and within threshold
+                        and start_time - prev_prediction['end'] <= MERGE_TIME_WITHIN)):
+            # 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,
+                'category': range['category']
+            })
+
+        prev_prediction = range
+
+    return final_predicted_time_ranges
+
+
+def main():
+    # Test on unseen data
+    logger.setLevel(logging.DEBUG)
+
+    hf_parser = HfArgumentParser((
+        PredictArguments,
+        SegmentationArguments,
+        GeneralArguments
+    ))
+    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, classifier = get_model_tokenizer_classifier(
+        predict_args, general_args)
+
+    for video_id in predict_args.video_ids:
+        try:
+            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
+        video_url = f'https://www.youtube.com/watch?v={video_id}'
+        if not predictions:
+            logger.info(f'No predictions found for {video_url}')
+            continue
+
+        # TODO use predict_args.output_as_json
+        print(len(predictions), 'predictions found for', video_url)
+        for index, prediction in enumerate(predictions, start=1):
+            print(f'Prediction #{index}:')
+            print('Text: "',
+                  ' '.join([w['text'] for w in prediction['words']]), '"', sep='')
+            print('Time:', seconds_to_time(
+                prediction['start']), '\u2192', seconds_to_time(prediction['end']))
+            print('Category:', prediction.get('category'))
+            if 'probability' in prediction:
+                print('Probability:', prediction['probability'])
+            print()
+        print()
+
+
+if __name__ == '__main__':
+    main()

src/preprocess.py

@@ -0,0 +1,979 @@
+from shared import DatasetArguments
+from utils import jaccard
+from functools import lru_cache
+from datetime import datetime
+import itertools
+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 extract_sponsor_matches_from_text, ACTION_OPTIONS, CATEGORIES, CATGEGORY_OPTIONS, START_SEGMENT_TEMPLATE, END_SEGMENT_TEMPLATE, GeneralArguments, CustomTokens
+import csv
+import re
+import random
+import logging
+from youtube_transcript_api import YouTubeTranscriptApi, CouldNotRetrieveTranscript, YouTubeRequestFailed, TooManyRequests
+import os
+import json
+import time
+import requests
+
+
+logging.basicConfig()
+logger = logging.getLogger(__name__)
+
+
+PROFANITY_RAW = '[ __ ]'  # How YouTube transcribes profanity
+PROFANITY_CONVERTED = '*****'  # Safer version for tokenizing
+
+
+NUM_DECIMALS = 3
+
+# https://www.fincher.org/Utilities/CountryLanguageList.shtml
+# https://lingohub.com/developers/supported-locales/language-designators-with-regions
+LANGUAGE_PREFERENCE_LIST = ['en-GB', 'en-US', 'en-CA', 'en-AU', 'en-NZ', 'en-ZA',
+                            'en-IE', 'en-IN', 'en-JM', 'en-BZ', 'en-TT', 'en-PH', 'en-ZW',
+                            'en']
+
+
+def parse_transcript_json(json_data, granularity):
+    assert json_data['wireMagic'] == 'pb3'
+
+    assert granularity in ('word', 'chunk')
+
+    # TODO remove bracketed words?
+    # (kiss smacks)
+    # (upbeat music)
+    # [text goes here]
+
+    # Some manual transcripts aren't that well formatted... but do have punctuation
+    # https://www.youtube.com/watch?v=LR9FtWVjk2c
+
+    parsed_transcript = []
+
+    events = json_data['events']
+
+    for event_index, event in enumerate(events):
+        segments = event.get('segs')
+        if not segments:
+            continue
+
+        # This value is known (when phrase appears on screen)
+        start_ms = event['tStartMs']
+        total_characters = 0
+
+        new_segments = []
+        for seg in segments:
+            # Replace \n, \t, etc. with space
+            text = ' '.join(seg['utf8'].split())
+
+            # Remove zero-width spaces and strip trailing and leading whitespace
+            text = text.replace('\u200b', '').replace('\u200c', '').replace(
+                '\u200d', '').replace('\ufeff', '').strip()
+
+            # Alternatively,
+            # text = text.encode('ascii', 'ignore').decode()
+
+            # Needed for auto-generated transcripts
+            text = text.replace(PROFANITY_RAW, PROFANITY_CONVERTED)
+
+            if not text:
+                continue
+
+            offset_ms = seg.get('tOffsetMs', 0)
+
+            new_segments.append({
+                'text': text,
+                'start': round((start_ms + offset_ms)/1000, NUM_DECIMALS)
+            })
+
+            total_characters += len(text)
+
+        if not new_segments:
+            continue
+
+        if event_index < len(events) - 1:
+            next_start_ms = events[event_index + 1]['tStartMs']
+            total_event_duration_ms = min(
+                event.get('dDurationMs', float('inf')), next_start_ms - start_ms)
+        else:
+            total_event_duration_ms = event.get('dDurationMs', 0)
+
+        # Ensure duration is non-negative
+        total_event_duration_ms = max(total_event_duration_ms, 0)
+
+        avg_seconds_per_character = (
+            total_event_duration_ms/total_characters)/1000
+
+        num_char_count = 0
+        for seg_index, seg in enumerate(new_segments):
+            num_char_count += len(seg['text'])
+
+            # Estimate segment end
+            seg_end = seg['start'] + \
+                (num_char_count * avg_seconds_per_character)
+
+            if seg_index < len(new_segments) - 1:
+                # Do not allow longer than next
+                seg_end = min(seg_end, new_segments[seg_index+1]['start'])
+
+            seg['end'] = round(seg_end, NUM_DECIMALS)
+            parsed_transcript.append(seg)
+
+    final_parsed_transcript = []
+    for i in range(len(parsed_transcript)):
+
+        word_level = granularity == 'word'
+        if word_level:
+            split_text = parsed_transcript[i]['text'].split()
+        elif granularity == 'chunk':
+            # Split on space after punctuation
+            split_text = re.split(
+                r'(?<=[.!?,-;])\s+', parsed_transcript[i]['text'])
+            if len(split_text) == 1:
+                split_on_whitespace = parsed_transcript[i]['text'].split()
+
+                if len(split_on_whitespace) >= 8:  # Too many words
+                    # Rather split on whitespace instead of punctuation
+                    split_text = split_on_whitespace
+                else:
+                    word_level = True
+        else:
+            raise ValueError('Unknown granularity')
+
+        segment_end = parsed_transcript[i]['end']
+        if i < len(parsed_transcript) - 1:
+            segment_end = min(segment_end, parsed_transcript[i+1]['start'])
+
+        segment_duration = segment_end - parsed_transcript[i]['start']
+
+        num_chars_in_text = sum(map(len, split_text))
+
+        num_char_count = 0
+        current_offset = 0
+        for s in split_text:
+            num_char_count += len(s)
+
+            next_offset = (num_char_count/num_chars_in_text) * segment_duration
+
+            word_start = round(
+                parsed_transcript[i]['start'] + current_offset, NUM_DECIMALS)
+            word_end = round(
+                parsed_transcript[i]['start'] + next_offset, NUM_DECIMALS)
+
+            # Make the reasonable assumption that min wps is 1.5
+            final_parsed_transcript.append({
+                'text': s,
+                'start': word_start,
+                'end': min(word_end, word_start + 1.5) if word_level else word_end
+            })
+            current_offset = next_offset
+
+    return final_parsed_transcript
+
+
+def list_transcripts(video_id):
+    try:
+        return YouTubeTranscriptApi.list_transcripts(video_id)
+    except json.decoder.JSONDecodeError:
+        return None
+
+
+WORDS_TO_REMOVE = [
+    CustomTokens.MUSIC.value,
+    CustomTokens.APPLAUSE.value,
+    CustomTokens.LAUGHTER.value
+]
+
+
+@lru_cache(maxsize=16)
+def get_words(video_id, process=True, transcript_type='auto', fallback='manual', filter_words_to_remove=True, download=False, granularity='word'):
+    """Get parsed video transcript with caching system
+    returns None if not processed yet and process is False
+    """
+    # NOTE: granularity='chunk' should only be used for generating training data... nowhere else
+
+    transcript_path = os.path.join(  # TODO use relative path to this
+        'transcripts', transcript_type, f'{video_id}.json')
+
+    raw_transcript_json = None
+    try:
+        if not download and os.path.exists(transcript_path):  # Load from file
+            with open(transcript_path) as fp:
+                raw_transcript_json = json.load(fp)  # May be empty
+
+        elif process:
+            transcript_list = list_transcripts(video_id)
+
+            if transcript_list is not None:
+                if transcript_type == 'manual':
+                    ts = transcript_list.find_manually_created_transcript(
+                        LANGUAGE_PREFERENCE_LIST)
+                else:
+                    ts = transcript_list.find_generated_transcript(
+                        LANGUAGE_PREFERENCE_LIST)
+                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
+
+    except requests.exceptions.RequestException:  # Can recover
+        time.sleep(10)  # Timeout
+        return get_words(video_id, process, transcript_type, fallback, granularity)
+
+    except CouldNotRetrieveTranscript:  # Retrying won't solve
+        pass  # Mark as empty transcript
+
+    except json.decoder.JSONDecodeError:
+        logger.warning(f'JSONDecodeError for {video_id}')
+        if os.path.exists(transcript_path):
+            os.remove(transcript_path)  # Remove file and try again
+        return get_words(video_id, process, transcript_type, fallback, granularity)
+
+    # Tried to process it, but it was empty...
+    if download or (process and not os.path.exists(transcript_path)):
+        with open(transcript_path, 'w') as fp:
+            json.dump(raw_transcript_json, fp)
+
+    if not raw_transcript_json and fallback is not None:
+        return get_words(video_id, process, transcript_type=fallback, fallback=None, granularity=granularity)
+
+    if raw_transcript_json:
+        processed_transcript = parse_transcript_json(
+            raw_transcript_json, granularity)
+        if filter_words_to_remove:
+            processed_transcript = list(
+                filter(lambda x: x['text'] not in WORDS_TO_REMOVE, processed_transcript))
+    else:
+        processed_transcript = raw_transcript_json  # Either None or []
+
+    return processed_transcript
+
+
+# TODO make min_sponsor_segment_length param
+# TODO rename to extract_segments
+def extract_sponsors(words, min_sponsor_segment_length=3):
+    if not words:
+        return []
+
+    paragraphs = []
+    current = []
+    prev_category = None
+
+    for i in range(len(words) + 1):
+        unimportant = i == len(words) or words[i].get('category') is None
+
+        if unimportant or words[i].get('category') != prev_category:
+            if current:  # Save the current batch
+                paragraphs.append({
+                    'words': current,
+                    'category': current[-1].get('category'),
+                })
+
+                current = []
+
+        if not unimportant:  # Some useful information to save
+            current.append(words[i])
+            prev_category = words[i].get('category')
+
+    # Remove all too short:
+    return list(filter(lambda x: len(x['words']) >= min_sponsor_segment_length, 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(PROFANITY_RAW, CustomTokens.PROFANITY.value)
+    text = text.replace(PROFANITY_CONVERTED, CustomTokens.PROFANITY.value)
+
+    return text.strip()
+
+
+def remove_duplicate_segments(segments):
+    # Algorithm based on SponsorBlock algorithm
+    # https://blog.ajay.app/voting-and-pseudo-randomness-or-sponsorblock-or-youtube-sponsorship-segment-blocker
+    # Find sponsors that are overlapping
+
+    best = []
+    for i in segments:
+        similar_segments = []
+        for j in segments:
+            if jaccard(i['start'], i['end'], j['start'], j['end']) > 0.1:  # Some overlap
+                similar_segments.append(j)
+
+        if similar_segments:
+            best_similar_seg = max(similar_segments, key=lambda item: (
+                item['locked'],
+                item['votes'],
+                item['views'],
+                item['reputation']
+            ))
+            if best_similar_seg not in best:
+                best.append(best_similar_seg)
+
+    if len(segments) != len(best):  # Saw some reduction... try again
+        return remove_duplicate_segments(best)
+
+    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
+
+    max_segment_duration: float = field(
+        default=180,  # 3 minutes
+        # >180 => 2.8%
+        # >200 => 2.1%
+        # >250 => 1.1%
+        # >300 => 0.06%
+        metadata={'help': 'Ignore all segments whose duration in seconds is longer than this value (negative means no limit)'})
+
+    min_views: int = field(
+        default=5, metadata={'help': 'Minimum number of views a segment must have to be considered. 0 = show all'})
+
+    # min_reputation: int = field(
+    #     default=0, metadata={'help': 'Minimum reputation a user must have for the segment to be included'})
+
+    min_date: str = field(
+        # default='08/06/2020', # release of v2.0 (https://github.com/ajayyy/SponsorBlock/releases/tag/2.0)
+        # release of v3.0 (https://github.com/ajayyy/SponsorBlock/releases/tag/3.0)
+        default='20/08/2021',
+        # default='01/10/2020', # No more autovote
+        metadata={'help': 'Only use submissions from after this date (inclusive)'})
+
+    max_date: str = field(
+        # default='01/01/9999', # Include all
+        default='15/04/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'}
+    )
+
+    do_process_database: bool = field(
+        default=False, metadata={'help': 'Process the raw database'}
+    )
+    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.'}
+    )
+
+    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).'}
+    )
+
+    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.'})
+
+    start_index: int = field(default=None, metadata={
+        'help': 'Video to start at.'})
+
+    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=1.5, 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)
+
+
+def main():
+    # Responsible for getting transcrips using youtube_transcript_api,
+    # then labelling it according to SponsorBlock's API
+    logger.setLevel(logging.DEBUG)
+
+    # Generate final.json from sponsorTimes.csv
+    hf_parser = HfArgumentParser((
+        PreprocessArguments,
+        DatasetArguments,
+        segment.SegmentationArguments,
+        model_module.ModelArguments,
+        GeneralArguments
+    ))
+    preprocess_args, dataset_args, segmentation_args, model_args, general_args = hf_parser.parse_args_into_dataclasses()
+
+    raw_dataset_path = os.path.join(
+        preprocess_args.raw_data_dir, preprocess_args.raw_data_file)
+
+    if preprocess_args.update_database:
+        logger.info('Updating database')
+        for mirror in MIRRORS:
+            logger.info(f'Downloading from {mirror}')
+            if download_file(mirror, raw_dataset_path):
+                break
+            logger.warning('Failed, trying next')
+
+    os.makedirs(dataset_args.data_dir, exist_ok=True)
+    processed_db_path = os.path.join(
+        dataset_args.data_dir, dataset_args.processed_database)
+
+    # TODO process all valid possible items and then do filtering only later
+    @lru_cache(maxsize=1)
+    def read_db():
+        # if not preprocess_args.overwrite and os.path.exists(processed_db_path):
+        #     logger.info(
+        #         'Using cached processed database (use `--overwrite` to avoid this behaviour).')
+        #     with open(processed_db_path) as fp:
+        #         return json.load(fp)
+        logger.info('Processing raw database')
+        db = {}
+
+        allowed_categories = list(map(str.lower, CATGEGORY_OPTIONS))
+        with open(raw_dataset_path, newline='') as csvfile:
+            reader = csv.DictReader(csvfile)
+
+            for line in reader:
+
+                # Never show:
+                if line['service'] != 'YouTube':
+                    continue
+                if len(line['videoID']) != 11:
+                    continue  # Invalid youtube video ID
+
+                if line['category'] not in allowed_categories:
+                    continue
+                if line['actionType'] not in ACTION_OPTIONS:
+                    continue
+
+                # Ignore hidden items
+                if line['hidden'] == '1' or line['shadowHidden'] == '1':
+                    continue
+
+                # Skip those that aren't highly voted
+                votes = int(line['votes'])
+                if votes < preprocess_args.min_votes:
+                    continue
+
+                locked = line['locked'] == '1'
+
+                reputation = float(line['reputation'])
+                # if reputation < preprocess_args.min_reputation:
+                #     continue # TODO add back?
+                # Problems like mGVn1wCkBrE
+
+                # TODO ignore if over max_duration
+
+                if line['videoID'] not in db:
+                    db[line['videoID']] = []
+
+                db[line['videoID']].append({
+                    'uuid': line['UUID'],
+                    'start': float(line['startTime']),
+                    'end': float(line['endTime']),
+                    'votes': votes,
+                    'locked': locked,
+                    'views': int(line['views']),
+                    'submission_time': float(line['timeSubmitted'])/1e3,
+                    'reputation': reputation,
+                    'category': line['category'],
+                    'action': line['actionType'],
+                })
+
+        # First, remove videos that contain a full-video label
+        # (may confuse model since disclaimers and such aren't labelled)
+        # Must do it here before removing duplicate segments
+        for key in list(db):
+            if any(x['action'] == 'full' for x in db[key]):
+                del db[key]
+
+        # Remove duplicate sponsor segments by choosing best (most votes)
+        if not preprocess_args.keep_duplicate_segments:
+            logger.info('Remove duplicate segments')
+            for key in db:
+                db[key] = remove_duplicate_segments(db[key])
+
+        # We now remove whole videos from the list
+        # Helps with obtaining "fully-labelled" videos
+        min_date = datetime.strptime(preprocess_args.min_date, '%d/%m/%Y')
+        max_date = datetime.strptime(preprocess_args.max_date, '%d/%m/%Y')
+        for key in list(db):
+            if preprocess_args.max_segment_duration >= 0 and any(x['end'] - x['start'] > preprocess_args.max_segment_duration for x in db[key]):
+                # Remove videos that have at least one segment that is longer than
+                # the maximum allowed segment duration. This avoids introducing
+                # segments into training that might contain ignored context (since
+                # they are too long, so the middle might be normal content)
+                del db[key]
+            elif any(datetime.fromtimestamp(x['submission_time']) < min_date for x in db[key]):
+                # Remove videos where any of its segments were submitted before min_date
+                # (essentially removes videos uploaded before min_date)
+                # Prevents issues where some segments of a video are excluded
+                del db[key]
+            elif all(datetime.fromtimestamp(x['submission_time']) > max_date for x in db[key]):
+                # Remove videos where all of its segments were submitted after max_date
+                # (essentially removes videos uploaded after max_date)
+                # Allows for segments to be corrected for past videos
+                del db[key]
+            elif any(not x['locked'] and x['views'] < preprocess_args.min_views for x in db[key]):
+                # Remove videos where any of its non-locked segments do not have enough views
+                # (essentially skips videos that have not been fully watched/reviewed)
+                # Always include segments locked by VIPs, regardless of view count
+                del db[key]
+
+        logger.info(f'Saved {len(db)} videos')
+
+        with open(processed_db_path, 'w') as fp:
+            json.dump(db, fp)
+
+        return db
+
+    if preprocess_args.do_process_database:
+        read_db()
+
+    # 'videoID', 'startTime', 'endTime', 'votes', 'locked', 'incorrectVotes', 'UUID',
+    # 'userID', 'timeSubmitted', 'views', 'category', 'actionType', 'service', 'videoDuration',
+    # 'hidden', 'reputation', 'shadowHidden', 'hashedVideoID', 'userAgent', 'description'
+    if preprocess_args.do_transcribe:
+        logger.info('Collecting videos')
+        parsed_database = read_db()
+
+        # Remove transcripts already processed
+        finished = set(x.split('.')[0] for x in os.listdir(
+            'transcripts/auto/') + os.listdir('transcripts/manual/'))
+
+        video_ids = list(parsed_database.keys() - finished)
+
+        # https://stackoverflow.com/a/63495323
+        import concurrent
+        POLL_INTERVAL = 0.1
+
+        # Wrap get words function to return video_id after completion
+        def get_words_wrapper(video_id):
+            get_words(video_id)
+            return video_id
+
+        logger.info('Setting up ThreadPoolExecutor')
+        with concurrent.futures.ThreadPoolExecutor(max_workers=preprocess_args.num_jobs) as pool, \
+                tqdm(total=len(video_ids)) as progress:
+
+            all_futures = (pool.submit(get_words_wrapper, video_id)
+                           for video_id in video_ids)
+            to_process = set(itertools.islice(
+                all_futures, preprocess_args.num_jobs))
+            try:
+                while to_process:
+                    just_finished, to_process = concurrent.futures.wait(
+                        to_process, timeout=POLL_INTERVAL)
+                    to_process |= set(itertools.islice(
+                        all_futures, len(just_finished)))
+
+                    for d in just_finished:
+                        progress.set_description(f'Processed {d.result()}')
+                        progress.update()
+
+            except KeyboardInterrupt:
+                logger.info(
+                    'Gracefully shutting down: Cancelling unscheduled tasks')
+
+                # only futures that are not done will prevent exiting
+                for future in to_process:
+                    future.cancel()
+
+                logger.info('Waiting for in-progress tasks to complete')
+                concurrent.futures.wait(to_process, timeout=None)
+                logger.info('Cancellation successful')
+
+    final_path = os.path.join(
+        dataset_args.data_dir, dataset_args.processed_file)
+
+    if preprocess_args.do_create:
+        logger.info('Create final data')
+
+        final_data = {}
+
+        parsed_database = read_db()
+
+        transcribed = set(x.split('.')[0] for x in os.listdir(
+            'transcripts/auto/') + os.listdir('transcripts/manual/'))
+
+        # Only consider videos that have been transcribed already
+        video_ids = parsed_database.keys() & transcribed
+
+        with tqdm(total=len(video_ids)) as progress:
+            for index, video_id in enumerate(video_ids):
+                if preprocess_args.max_videos is not None and index >= preprocess_args.max_videos:
+                    break
+                progress.set_description(f'Processing {video_id}')
+                progress.update()
+
+                video_words = get_words(video_id, process=False)
+                if not video_words:
+                    continue
+
+                final_vid_segs = []
+                # Only add segments with high enough wps
+                for seg in parsed_database[video_id]:
+                    segment_words = segment.extract_segment(
+                        video_words, seg['start'], seg['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 = seg['end'] - seg['start']
+                    wps = len(segment_words)/duration if duration > 0 else 0
+
+                    # print(video_id, wps)
+                    if wps < preprocess_args.min_wps:
+                        # Skip sponsor segments without many words
+                        # e.g. music ads with some words on each side
+                        # progress.set_description(f'Skipping bad segment in {video_id} (wps={wps})')
+                        continue
+                    final_vid_segs.append(seg)
+
+                if final_vid_segs:
+                    final_data[video_id] = final_vid_segs
+
+        # 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?
+
+    elif os.path.exists(final_path):
+        # Already exists
+        logging.info(f'{final_path} exists, opening file')
+        with open(final_path) as fp:
+            final_data = json.load(fp)
+        logging.info(f'Found {len(final_data)} videos')
+    else:
+        return  # Do not continue
+
+    # TODO shuffle final_data
+    # if not os.path.exists(excess_path) or preprocess_args.overwrite
+    # TODO use overwrite param
+
+    positive_file = os.path.join(
+        dataset_args.data_dir, preprocess_args.positive_file)
+    negative_file = os.path.join(
+        dataset_args.data_dir, preprocess_args.negative_file)
+
+    if preprocess_args.do_generate:
+        logger.info('Generating')
+        # max_videos=preprocess_args.max_videos,
+        # max_segments=preprocess_args.max_segments,
+        # , max_videos, max_segments
+
+        from model import get_model_tokenizer
+        model, tokenizer = get_model_tokenizer(model_args, general_args)
+
+        # TODO
+        # count_videos = 0
+        # count_segments = 0
+
+        data = final_data.items()
+
+        start_index = preprocess_args.start_index or 0
+        end_index = (preprocess_args.max_videos or len(data)) + start_index
+
+        data = list(itertools.islice(data, start_index, end_index))
+
+        write_mode = 'w'  # if preprocess_args.overwrite else 'a'
+        with open(positive_file, write_mode, encoding='utf-8') as positive, \
+                open(negative_file, write_mode, encoding='utf-8') as negative, \
+                tqdm(data) as progress:
+
+            for offset, (video_id, sponsor_segments) in enumerate(data):
+
+                progress.set_description(f'Processing {video_id}')
+                progress.update()
+
+                # Use chunk granularity to improve manual transcripts
+                words = get_words(video_id, process=False, granularity='chunk')
+                if not words:
+                    continue
+
+                if len(words) <= 1:
+                    continue
+
+                segments = segment.generate_labelled_segments(
+                    words, tokenizer, segmentation_args, sponsor_segments)
+
+                if not segments:
+                    continue
+
+                for seg in segments:
+                    seg_start = segment.word_start(seg[0])
+                    seg_end = segment.word_end(seg[-1])
+                    duration = seg_end - seg_start
+                    wps = len(seg)/duration if duration > 0 else 0
+
+                    # Ignore segments with "not enough words" in the transcript
+                    # Must do here since this includes non-sponsor segments
+                    if wps < preprocess_args.min_wps:
+                        continue
+
+                    d = {
+                        # '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,
+                    }
+
+                    extracted_segments = extract_sponsors(seg)
+                    if extracted_segments:
+                        extracted_texts = []
+                        for s in extracted_segments:
+                            w = ' '.join(q['cleaned'] for q in s['words'])
+                            category = s['category'].upper()
+                            extracted_texts.append(
+                                f'{START_SEGMENT_TEMPLATE.format(category)} {w} {END_SEGMENT_TEMPLATE.format(category)}'
+                            )
+
+                        d['extracted'] = f' {CustomTokens.BETWEEN_SEGMENTS.value} '.join(
+                            extracted_texts)
+                        print(json.dumps(d), file=positive)
+
+                    else:
+                        d['extracted'] = CustomTokens.NO_SEGMENT.value
+                        print(json.dumps(d), file=negative)
+
+    if preprocess_args.do_split:
+        logger.info('Splitting')
+        logger.info('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()
+
+        logger.info('Shuffle')
+        random.shuffle(sponsors)
+        random.shuffle(non_sponsors)
+
+        logger.info('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]
+
+        logger.info('Join')
+        all_labelled_segments = sponsors + non_sponsors
+
+        random.shuffle(all_labelled_segments)
+
+        # TODO split based on video ids
+        logger.info('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)
+            with open(outfile, 'w', encoding='utf-8') as fp:
+                fp.writelines(items)
+
+        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)
+            with open(outfile, 'w', encoding='utf-8') as fp:
+                for item in items:
+                    parsed_item = json.loads(item)  # TODO add uuid
+
+                    matches = extract_sponsor_matches_from_text(
+                        parsed_item['extracted'])
+
+                    if matches:
+                        for match in matches:
+                            print(json.dumps({
+                                'text': match['text'],
+                                'label': CATEGORIES.index(match['category'])
+                            }), file=fp)
+                    else:
+                        print(json.dumps({
+                            'text': parsed_item['text'],
+                            'label': none_category
+                        }), file=fp)
+
+        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}')
+
+        logger.info(
+            f'Finished splitting: {len(sponsors)} sponsors, {len(non_sponsors)} non sponsors')
+
+
+def split(arr, ratios):
+    """Split array according to ratios. Sum of ratios should be <= 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,166 @@
+import preprocess
+from dataclasses import dataclass, field
+
+
+@dataclass
+class SegmentationArguments:
+    pause_threshold: int = field(default=2.5, metadata={
+        'help': 'When the time between words is greater than pause threshold, force into a new segment'})
+
+
+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 SAFETY_TOKENS_PERCENTAGE*max_tokens tokens
+MIN_SAFETY_TOKENS = 8
+SAFETY_TOKENS_PERCENTAGE = 0.9765625
+# e.g. 512 -> 500, 768 -> 750
+
+
+# TODO play around with this?
+OVERLAP_TOKEN_PERCENTAGE = 0.5  # 0.25
+
+
+def add_labels_to_words(words, sponsor_segments):
+
+    for sponsor_segment in sponsor_segments:
+        for w in extract_segment(words, sponsor_segment['start'], sponsor_segment['end']):
+            w['category'] = sponsor_segment['category']
+
+    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):
+
+    cleaned_words_list = []
+    for w in words:
+        w['cleaned'] = preprocess.clean_text(w['text'])
+        cleaned_words_list.append(w['cleaned'])
+
+    # Get lengths of tokenized words
+    num_tokens_list = tokenizer(cleaned_words_list, add_special_tokens=False,
+                                truncation=True, return_attention_mask=False, return_length=True).length
+
+    first_pass_segments = []
+    for index, (word, num_tokens) in enumerate(zip(words, num_tokens_list)):
+        word['num_tokens'] = num_tokens
+
+        # Add new segment
+        if index == 0 or word_start(words[index]) - word_end(words[index-1]) >= segmentation_args.pause_threshold:
+            first_pass_segments.append([word])
+
+        else:  # Add to current segment
+            first_pass_segments[-1].append(word)
+
+    max_q_size = round(SAFETY_TOKENS_PERCENTAGE * tokenizer.model_max_length)
+
+    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 = []
+        after_split_segments = []
+        for word in segment:
+            new_seg = current_segment_num_tokens + \
+                word['num_tokens'] >= max_q_size
+            if new_seg:
+                # Adding this token would make it have too many tokens
+                # We save this batch and create new
+                after_split_segments.append(current_segment)
+
+            # Add tokens to current segment
+            current_segment.append(word)
+            current_segment_num_tokens += word['num_tokens']
+
+            if not new_seg:
+                continue
+
+            # Just created a new segment, so we remove until we only have buffer_size tokens
+            last_index = 0
+            while current_segment_num_tokens > buffer_size and current_segment:
+                current_segment_num_tokens -= current_segment[last_index]['num_tokens']
+                last_index += 1
+
+            current_segment = current_segment[last_index:]
+
+        if current_segment:  # Add remaining segment
+            after_split_segments.append(current_segment)
+
+        # TODO if len(after_split_segments) > 1, a split occurred
+
+        second_pass_segments.extend(after_split_segments)
+
+    # Cleaning up, delete 'num_tokens' from each word
+    for word in words:
+        word.pop('num_tokens', None)
+
+    return second_pass_segments
+
+
+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))
+
+    to_transform = map_function is not None and callable(map_function)
+
+    return [
+        map_function(words[i]) if to_transform else words[i] for i in range(a, b)
+    ]
+
+
+def avg(*items):
+    return sum(items)/len(items)
+
+
+def binary_search_below(transcript, start_index, end_index, time):
+    if start_index >= end_index:
+        return end_index
+
+    middle_index = (start_index + end_index) // 2
+    middle = transcript[middle_index]
+    middle_time = avg(word_start(middle), word_end(middle))
+
+    if time <= middle_time:
+        return binary_search_below(transcript, start_index, middle_index, time)
+    else:
+        return binary_search_below(transcript, middle_index + 1, end_index, time)
+
+
+def binary_search_above(transcript, start_index, end_index, time):
+    if start_index >= end_index:
+        return end_index
+
+    middle_index = (start_index + end_index + 1) // 2
+    middle = transcript[middle_index]
+    middle_time = avg(word_start(middle), word_end(middle))
+
+    if time >= middle_time:
+        return binary_search_above(transcript, middle_index, end_index, time)
+    else:
+        return binary_search_above(transcript, start_index, middle_index - 1, time)

src/shared.py

@@ -0,0 +1,406 @@
+from transformers.trainer_utils import get_last_checkpoint as glc
+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
+import random
+import numpy as np
+import torch
+from typing import Optional
+from dataclasses import dataclass, field
+from enum import Enum
+
+
+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)],
+)
+
+CATEGORIES = [None, 'SPONSOR', 'SELFPROMO', 'INTERACTION']
+
+ACTION_OPTIONS = ['skip', 'mute', 'full']
+
+CATGEGORY_OPTIONS = {
+    'SPONSOR': 'Sponsor',
+    'SELFPROMO': 'Self/unpaid promo',
+    'INTERACTION': 'Interaction reminder',
+}
+
+START_SEGMENT_TEMPLATE = 'START_{}_TOKEN'
+END_SEGMENT_TEMPLATE = 'END_{}_TOKEN'
+
+
+class CustomTokens(Enum):
+    EXTRACT_SEGMENTS_PREFIX = 'EXTRACT_SEGMENTS: '
+
+    # Preprocessing tokens
+    URL = 'URL_TOKEN'
+    HYPHENATED_URL = 'HYPHENATED_URL_TOKEN'
+    NUMBER_PERCENTAGE = 'NUMBER_PERCENTAGE_TOKEN'
+    NUMBER = 'NUMBER_TOKEN'
+
+    SHORT_HYPHENATED = 'SHORT_HYPHENATED_TOKEN'
+    LONG_WORD = 'LONG_WORD_TOKEN'
+
+    # Custom YouTube tokens
+    MUSIC = '[Music]'
+    APPLAUSE = '[Applause]'
+    LAUGHTER = '[Laughter]'
+
+    PROFANITY = 'PROFANITY_TOKEN'
+
+    # Segment tokens
+    NO_SEGMENT = 'NO_SEGMENT_TOKEN'
+
+    START_SPONSOR = START_SEGMENT_TEMPLATE.format('SPONSOR')
+    END_SPONSOR = END_SEGMENT_TEMPLATE.format('SPONSOR')
+
+    START_SELFPROMO = START_SEGMENT_TEMPLATE.format('SELFPROMO')
+    END_SELFPROMO = END_SEGMENT_TEMPLATE.format('SELFPROMO')
+
+    START_INTERACTION = START_SEGMENT_TEMPLATE.format('INTERACTION')
+    END_INTERACTION = END_SEGMENT_TEMPLATE.format('INTERACTION')
+
+    BETWEEN_SEGMENTS = 'BETWEEN_SEGMENTS_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())
+
+
+_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_from_text(text):
+    if CustomTokens.NO_SEGMENT.value in text:
+        return []
+    else:
+        return re_findall(SEGMENT_MATCH_RE, text)
+
+
+def extract_sponsor_matches(texts):
+    return list(map(extract_sponsor_matches_from_text, texts))
+
+
+@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'
+        },
+    )
+
+    overwrite_cache: bool = field(
+        default=False, metadata={'help': 'Overwrite the cached training and evaluation sets'}
+    )
+
+    dataset_cache_dir: Optional[str] = field(
+        default=None,
+        metadata={
+            'help': 'Where to store the cached datasets'
+        },
+    )
+
+    train_file: Optional[str] = field(
+        default='train.json', metadata={'help': 'The input training data file (a jsonlines file).'}
+    )
+    validation_file: Optional[str] = field(
+        default='valid.json',
+        metadata={
+            'help': 'An optional input evaluation data file to evaluate the metrics on (a jsonlines file).'
+        },
+    )
+    test_file: Optional[str] = field(
+        default='test.json',
+        metadata={
+            'help': 'An optional input test data file to evaluate the metrics on (a jsonlines file).'
+        },
+    )
+
+    c_train_file: Optional[str] = field(
+        default='c_train.json', metadata={'help': 'The input training data file (a jsonlines file).'}
+    )
+    c_validation_file: Optional[str] = field(
+        default='c_valid.json',
+        metadata={
+            'help': 'An optional input evaluation data file to evaluate the metrics on (a jsonlines file).'
+        },
+    )
+    c_test_file: Optional[str] = field(
+        default='c_test.json',
+        metadata={
+            'help': 'An optional input test data file to evaluate the metrics on (a jsonlines file).'
+        },
+    )
+
+    def __post_init__(self):
+        if self.train_file is None or self.validation_file is None:
+            raise ValueError(
+                'Need either a dataset name or a training/validation file.')
+
+        else:
+            train_extension = self.train_file.split(".")[-1]
+            assert train_extension in [
+                "csv", "json"], "`train_file` should be a csv or a json file."
+            validation_extension = self.validation_file.split(".")[-1]
+            assert (
+                validation_extension == train_extension
+            ), "`validation_file` should have the same extension (csv or json) as `train_file`."
+
+
+@dataclass
+class OutputArguments:
+
+    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).'
+    })
+    no_cuda: bool = field(default=False, metadata={
+                          'help': 'Do not use CUDA even when it is available'})
+
+    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 seconds_to_time(seconds, remove_leading_zeroes=False):
+    fractional = round(seconds % 1, 3)
+    fractional = '' if fractional == 0 else str(fractional)[1:]
+    h, remainder = divmod(abs(int(seconds)), 3600)
+    m, s = divmod(remainder, 60)
+    hms = f'{h:02}:{m:02}:{s:02}'
+    if remove_leading_zeroes:
+        hms = re.sub(r'^0(?:0:0?)?', '', hms)
+    return f"{'-' if seconds < 0 else ''}{hms}{fractional}"
+
+
+def reset():
+    torch.clear_autocast_cache()
+    torch.cuda.empty_cache()
+    gc.collect()
+    print(torch.cuda.memory_summary(device=None, abbreviated=False))
+
+
+def load_datasets(dataset_args: DatasetArguments):
+
+    logger.info('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 AdditionalTrainingArguments:
+    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=25000, 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.
+
+    preprocessing_num_workers: Optional[int] = field(
+        default=None,
+        metadata={'help': 'The number of processes to use for the preprocessing.'},
+    )
+    max_seq_length: int = field(
+        default=512,
+        metadata={
+            "help": "The maximum total input sequence length after tokenization. Sequences longer "
+            "than this will be truncated, sequences shorter will be padded."
+        },
+    )
+    max_train_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "For debugging purposes or quicker training, truncate the number of training examples to this "
+            "value if set."
+        },
+    )
+    max_eval_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
+            "value if set."
+        },
+    )
+    max_predict_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "For debugging purposes or quicker training, truncate the number of prediction examples to this "
+            "value if set."
+        },
+    )
+
+
+@dataclass
+class CustomTrainingArguments(OutputArguments, AdditionalTrainingArguments):
+    pass
+
+
+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
+
+
+def prepare_datasets(raw_datasets, dataset_args: DatasetArguments, training_args: CustomTrainingArguments, preprocess_function):
+
+    with training_args.main_process_first(desc="dataset map pre-processing"):
+        raw_datasets = raw_datasets.map(
+            preprocess_function,
+            batched=True,
+            load_from_cache_file=not dataset_args.overwrite_cache,
+            desc="Running tokenizer on dataset",
+        )
+
+    if 'train' not in raw_datasets:
+        raise ValueError('Train dataset missing')
+    train_dataset = raw_datasets['train']
+    if training_args.max_train_samples is not None:
+        train_dataset = train_dataset.select(
+            range(training_args.max_train_samples))
+
+    if 'validation' not in raw_datasets:
+        raise ValueError('Validation dataset missing')
+    eval_dataset = raw_datasets['validation']
+    if training_args.max_eval_samples is not None:
+        eval_dataset = eval_dataset.select(
+            range(training_args.max_eval_samples))
+
+    if 'test' not in raw_datasets:
+        raise ValueError('Test dataset missing')
+    predict_dataset = raw_datasets['test']
+    if training_args.max_predict_samples is not None:
+        predict_dataset = predict_dataset.select(
+            range(training_args.max_predict_samples))
+
+    return train_dataset, eval_dataset, predict_dataset

src/train.py

@@ -0,0 +1,191 @@
+from shared import (
+    CustomTokens,
+    DatasetArguments,
+    prepare_datasets,
+    load_datasets,
+    CustomTrainingArguments,
+    get_last_checkpoint,
+    train_from_checkpoint
+)
+from model import ModelArguments
+import transformers
+import logging
+import os
+import sys
+from datasets import utils as d_utils
+from transformers import (
+    DataCollatorForSeq2Seq,
+    HfArgumentParser,
+    Seq2SeqTrainer,
+    Seq2SeqTrainingArguments,
+)
+
+from transformers.utils import check_min_version
+from transformers.utils.versions import require_version
+from dataclasses import dataclass
+
+# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
+check_min_version('4.17.0')
+require_version('datasets>=1.8.0',
+                'To fix: pip install -r requirements.txt')
+
+os.environ['WANDB_DISABLED'] = 'true'
+
+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)],
+)
+
+
+@dataclass
+class Seq2SeqTrainingArguments(CustomTrainingArguments, Seq2SeqTrainingArguments):
+    pass
+
+
+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.
+
+    hf_parser = HfArgumentParser((
+        ModelArguments,
+        DatasetArguments,
+        Seq2SeqTrainingArguments
+    ))
+    model_args, dataset_args, training_args = hf_parser.parse_args_into_dataclasses()
+
+    log_level = training_args.get_process_log_level()
+    logger.setLevel(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()
+
+    # 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.
+    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.
+
+    # 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.
+
+    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
+
+    train_dataset, eval_dataset, predict_dataset = prepare_datasets(
+        raw_datasets, dataset_args, training_args, preprocess_function)
+
+    # 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 = 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/train_classifier.py

@@ -0,0 +1,182 @@
+
+""" Finetuning the library models for sequence classification."""
+
+import logging
+import os
+import sys
+from dataclasses import dataclass
+from typing import Optional
+
+import datasets
+import numpy as np
+
+import transformers
+from transformers import (
+    DataCollatorWithPadding,
+    EvalPrediction,
+    HfArgumentParser,
+    Trainer,
+    TrainingArguments,
+    set_seed,
+)
+from transformers.utils import check_min_version
+from transformers.utils.versions import require_version
+from shared import (
+    CATEGORIES,
+    DatasetArguments,
+    prepare_datasets,
+    load_datasets,
+    CustomTrainingArguments,
+    train_from_checkpoint,
+    get_last_checkpoint
+)
+from 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 ClassifierTrainingArguments(CustomTrainingArguments, TrainingArguments):
+    pass
+
+
+@dataclass
+class ClassifierDatasetArguments(DatasetArguments):
+    def __post_init__(self):
+        self.train_file = self.c_train_file
+        self.validation_file = self.c_validation_file
+        self.test_file = self.c_test_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.
+
+    hf_parser = HfArgumentParser((
+        ModelArguments,
+        ClassifierDatasetArguments,
+        ClassifierTrainingArguments
+    ))
+    model_args, dataset_args, training_args = hf_parser.parse_args_into_dataclasses()
+
+    # Setup logging
+    logging.basicConfig(
+        format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
+        datefmt='%m/%d/%Y %H:%M:%S',
+        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(dataset_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')
+
+    if training_args.max_seq_length > tokenizer.model_max_length:
+        logger.warning(
+            f'The max_seq_length passed ({training_args.max_seq_length}) is larger than the maximum length for the'
+            f'model ({tokenizer.model_max_length}). Using max_seq_length={tokenizer.model_max_length}.'
+        )
+    max_seq_length = min(training_args.max_seq_length,
+                         tokenizer.model_max_length)
+
+    def preprocess_function(examples):
+        # Tokenize the texts
+        result = tokenizer(
+            examples['text'], padding='max_length', max_length=max_seq_length, truncation=True)
+        result['label'] = examples['label']
+        return result
+
+    train_dataset, eval_dataset, predict_dataset = prepare_datasets(
+        raw_datasets, dataset_args, training_args, preprocess_function)
+
+    # You can define your custom compute_metrics function. It takes an `EvalPrediction` object (a namedtuple with a
+    # predictions and label_ids field) and has to return a dictionary string to float.
+    def compute_metrics(p: EvalPrediction):
+        preds = p.predictions[0] if isinstance(
+            p.predictions, tuple) else p.predictions
+        preds = np.argmax(preds, axis=1)
+        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 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 = (
+        training_args.max_train_samples if training_args.max_train_samples is not None else len(
+            train_dataset)
+    )
+    metrics['train_samples'] = min(max_train_samples, len(train_dataset))
+
+    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

@@ -0,0 +1,135 @@
+import re
+import sys
+import locale
+import io
+
+def re_findall(pattern, string):
+    return [m.groupdict() for m in re.finditer(pattern, string)]
+
+
+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 if filled_union > 0 else 0
+
+
+def regex_search(text, pattern, group=1, default=None):
+    match = re.search(pattern, text)
+    return match.group(group) if match else default
+
+
+def _windows_write_string(s, out, skip_errors=True):
+    """ Returns True if the string was written using special methods,
+    False if it has yet to be written out."""
+    # Adapted from http://stackoverflow.com/a/3259271/35070
+
+    import ctypes
+    import ctypes.wintypes
+
+    WIN_OUTPUT_IDS = {
+        1: -11,
+        2: -12,
+    }
+
+    try:
+        fileno = out.fileno()
+    except AttributeError:
+        # If the output stream doesn't have a fileno, it's virtual
+        return False
+    except io.UnsupportedOperation:
+        # Some strange Windows pseudo files?
+        return False
+    if fileno not in WIN_OUTPUT_IDS:
+        return False
+
+    GetStdHandle = ctypes.WINFUNCTYPE(
+        ctypes.wintypes.HANDLE, ctypes.wintypes.DWORD)(
+        ('GetStdHandle', ctypes.windll.kernel32))
+    h = GetStdHandle(WIN_OUTPUT_IDS[fileno])
+
+    WriteConsoleW = ctypes.WINFUNCTYPE(
+        ctypes.wintypes.BOOL, ctypes.wintypes.HANDLE, ctypes.wintypes.LPWSTR,
+        ctypes.wintypes.DWORD, ctypes.POINTER(ctypes.wintypes.DWORD),
+        ctypes.wintypes.LPVOID)(('WriteConsoleW', ctypes.windll.kernel32))
+    written = ctypes.wintypes.DWORD(0)
+
+    GetFileType = ctypes.WINFUNCTYPE(ctypes.wintypes.DWORD, ctypes.wintypes.DWORD)(
+        ('GetFileType', ctypes.windll.kernel32))
+    FILE_TYPE_CHAR = 0x0002
+    FILE_TYPE_REMOTE = 0x8000
+    GetConsoleMode = ctypes.WINFUNCTYPE(
+        ctypes.wintypes.BOOL, ctypes.wintypes.HANDLE,
+        ctypes.POINTER(ctypes.wintypes.DWORD))(
+        ('GetConsoleMode', ctypes.windll.kernel32))
+    INVALID_HANDLE_VALUE = ctypes.wintypes.DWORD(-1).value
+
+    def not_a_console(handle):
+        if handle == INVALID_HANDLE_VALUE or handle is None:
+            return True
+        return ((GetFileType(handle) & ~FILE_TYPE_REMOTE) != FILE_TYPE_CHAR or GetConsoleMode(handle, ctypes.byref(ctypes.wintypes.DWORD())) == 0)
+
+    if not_a_console(h):
+        return False
+
+    def next_nonbmp_pos(s):
+        try:
+            return next(i for i, c in enumerate(s) if ord(c) > 0xffff)
+        except StopIteration:
+            return len(s)
+
+    while s:
+        count = min(next_nonbmp_pos(s), 1024)
+
+        ret = WriteConsoleW(
+            h, s, count if count else 2, ctypes.byref(written), None)
+        if ret == 0:
+            if skip_errors:
+                continue
+            else:
+                raise OSError('Failed to write string')
+        if not count:  # We just wrote a non-BMP character
+            assert written.value == 2
+            s = s[1:]
+        else:
+            assert written.value > 0
+            s = s[written.value:]
+    return True
+
+def preferredencoding():
+    """Get preferred encoding.
+    Returns the best encoding scheme for the system, based on
+    locale.getpreferredencoding() and some further tweaks.
+    """
+    try:
+        pref = locale.getpreferredencoding()
+        'TEST'.encode(pref)
+    except Exception:
+        pref = 'utf-8'
+
+    return pref
+
+def safe_print(*objects, sep=' ', end='\n', out=None, encoding=None, flush=False):
+    """
+    Ensure printing to standard output can be done safely (especially on Windows).
+    There are usually issues with printing emojis and non utf-8 characters.
+    """
+
+    output_string = sep.join(map(lambda x: str(x), objects)) + end
+
+    if out is None:
+        out = sys.stdout
+
+    if sys.platform == 'win32' and encoding is None and hasattr(out, 'fileno'):
+        if _windows_write_string(output_string, out):
+            return
+
+    if 'b' in getattr(out, 'mode', '') or not hasattr(out, 'buffer'):
+        out.write(output_string)
+    else:
+        enc = encoding or getattr(out, 'encoding', None) or preferredencoding()
+        byt = output_string.encode(enc, 'ignore')
+        out.buffer.write(byt)
+
+    if flush and hasattr(out, 'flush'):
+        out.flush()

transcripts/auto/.gitignore

@@ -0,0 +1,5 @@
+# Cached transcripts
+# Ignore everything in this directory
+*
+# Except this file
+!.gitignore

transcripts/manual/.gitignore

@@ -0,0 +1,5 @@
+# Cached transcripts
+# Ignore everything in this directory
+*
+# Except this file
+!.gitignore