custom_pipeline.py

import youtube_transcript_api2
import json
import re
import requests
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    TextClassificationPipeline,
)
from typing import Any, Dict, List

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 youtube_transcript_api2.YouTubeTranscriptApi.list_transcripts(video_id)
    except json.decoder.JSONDecodeError:
        return None


WORDS_TO_REMOVE = [
    '[Music]'
    '[Applause]'
    '[Laughter]'
]


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 (youtube_transcript_api2.TooManyRequests, youtube_transcript_api2.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 youtube_transcript_api2.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 SponsorBlockClassificationPipeline(TextClassificationPipeline):
    def __init__(self, model, tokenizer):
        super().__init__(model=model, tokenizer=tokenizer, return_all_scores=True)

    def preprocess(self, video, **tokenizer_kwargs):

        words = get_words(video['video_id'])
        segment_words = extract_segment(words, video['start'], video['end'])
        text = ' '.join(x['text'] for x in segment_words)

        model_inputs = self.tokenizer(
            text, return_tensors=self.framework, **tokenizer_kwargs)
        return {'video': video, 'model_inputs': model_inputs}

    def _forward(self, data):
        model_outputs = self.model(**data['model_inputs'])
        return {'video': data['video'], 'model_outputs': model_outputs}

    def postprocess(self, data, function_to_apply=None, return_all_scores=False):
        model_outputs = data['model_outputs']

        results = super().postprocess(model_outputs, function_to_apply, return_all_scores)

        for result in results:
            result['label_text'] = CATEGORIES[result['label']]

        return results # {**data['video'], 'result': results}


# model_id = "Xenova/sponsorblock-classifier-v2"
# model = AutoModelForSequenceClassification.from_pretrained(model_id)
# tokenizer = AutoTokenizer.from_pretrained(model_id)

# pl = SponsorBlockClassificationPipeline(model=model, tokenizer=tokenizer)
data = [{
    'video_id': 'pqh4LfPeCYs',
    'start': 835.933,
    'end': 927.581,
    'category': 'sponsor'
}]
# print(pl(data))

# MODEL_ID = "Xenova/sponsorblock-classifier-v2"
class PreTrainedPipeline():
    def __init__(self, path: str):
        # load the model
        self.model = AutoModelForSequenceClassification.from_pretrained(path)
        self.tokenizer = AutoTokenizer.from_pretrained(path)
        self.pipeline = SponsorBlockClassificationPipeline(
            model=self.model, tokenizer=self.tokenizer)

    def __call__(self, inputs: str) -> List[Dict[str, Any]]:
        json_data = json.loads(inputs)
        return self.pipeline(json_data)

# a = PreTrainedPipeline('Xenova/sponsorblock-classifier-v2')(json.dumps(data))
# print(a)