Fix structure

README.md

@@ -0,0 +1,15 @@
+---
+tags:
+- text-classification
+- generic
+library_name: generic
+widget:
+- text: 'This video is sponsored by squarespace'
+  example_title: Sponsor
+- text: 'Check out the merch at linustechtips.com'
+  example_title: Unpaid/self promotion
+- text: "Don't forget to like, comment and subscribe"
+  example_title: Interaction reminder
+- text: 'pqh4LfPeCYs,824.695,826.267,826.133,829.876,835.933,927.581'
+  example_title: Extract text from video
+---

checkpoint-325000/optimizer.pt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:96765e5aa06e0e6bb3828a8da9c276e30fefada85f8a18852f84b00ff074a1ff
-size 876116189

pipeline.py

@@ -0,0 +1,325 @@
+import json
+from functools import lru_cache
+from youtube_transcript_api import (
+    YouTubeTranscriptApi,
+    TooManyRequests,
+    YouTubeRequestFailed,
+    CouldNotRetrieveTranscript
+)
+import json
+import re
+import requests
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    TextClassificationPipeline,
+)
+from typing import Any, Dict, List
+import os
+import numpy as np
+
+CATEGORIES = [None, 'SPONSOR', 'SELFPROMO', 'INTERACTION']
+
+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 = [
+    '[Music]'
+    '[Applause]'
+    '[Laughter]'
+]
+
+
+@lru_cache(maxsize=16)
+def get_words(video_id, transcript_type='auto', fallback='manual', filter_words_to_remove=True, granularity='word'):
+    """Get parsed video transcript with caching system
+    returns None if not processed yet and process is False
+    """
+
+    raw_transcript_json = None
+    try:
+        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
+        return get_words(video_id, transcript_type, fallback, granularity)
+
+    except CouldNotRetrieveTranscript:  # Retrying won't solve
+        pass  # Mark as empty transcript
+
+    except json.decoder.JSONDecodeError:
+        return get_words(video_id, transcript_type, fallback, granularity)
+
+    if not raw_transcript_json and fallback is not None:
+        return get_words(video_id, 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
+
+
+def word_start(word):
+    return word['start']
+
+
+def word_end(word):
+    return word.get('end', word['start'])
+
+
+def extract_segment(words, start, end, map_function=None):
+    """Extracts all words with time in [start, end]"""
+
+    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)
+
+
+class PreTrainedPipeline():
+    def __init__(self, path: str):
+        self.model2 = AutoModelForSequenceClassification.from_pretrained(path)
+        self.tokenizer2 = AutoTokenizer.from_pretrained(path)
+        self.pipeline2 = SponsorBlockClassificationPipeline(
+            model=self.model2, tokenizer=self.tokenizer2)
+
+    def __call__(self, inputs: str) -> List[Dict[str, Any]]:
+
+        # Automated call (compressed string)
+        if ' ' not in inputs and inputs.count(',') >= 2:
+            split_info = inputs.split(',', 1)
+            times = np.reshape(np.array(split_info[1].split(',')), (-1, 2))
+            data = []
+            for start, end in times:
+                data.append({
+                    'video_id': split_info[0],
+                    'start': float(start),
+                    'end': float(end)
+                })
+        else:
+            data = inputs
+
+        return self.pipeline2(data)
+
+
+class SponsorBlockClassificationPipeline(TextClassificationPipeline):
+    def __init__(self, model, tokenizer):
+        super().__init__(model=model, tokenizer=tokenizer, return_all_scores=True)
+
+    def preprocess(self, data, **tokenizer_kwargs):
+        if isinstance(data, str):  # If string, assume this is what user wants to classify
+            text = data
+        else:  # Otherwise, get data from transcript
+            words = get_words(data['video_id'])
+            segment_words = extract_segment(words, data['start'], data['end'])
+            text = ' '.join(x['text'] for x in segment_words)
+
+        return self.tokenizer(
+            text, return_tensors=self.framework, **tokenizer_kwargs)

requirements.txt

@@ -0,0 +1 @@
+youtube_transcript_api