Update pipeline.py

pipeline.py

@@ -1,78 +1,326 @@
+import youtube_transcript_api2
 import json
+import re
+import requests
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    TextClassificationPipeline,
+)
 from typing import Any, Dict, List
 
-import tensorflow as tf
-from tensorflow import keras
-import base64
-import io
-import os
-import numpy as np
 from PIL import Image
 
+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}
 
 class PreTrainedPipeline():
     def __init__(self, path: str):
         # load the model
-        self.model = keras.models.load_model(os.path.join(path, "tf_model.h5"))
+        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: "Image.Image")-> List[Dict[str, Any]]:
+    def __call__(self, inputs: str) -> List[Dict[str, Any]]:
+        json_data = json.loads(inputs)
+        return self.pipeline(json_data)
 
-        # convert img to numpy array, resize and normalize to make the prediction
-        img = np.array(inputs)
-
-        im = tf.image.resize(img, (128, 128))
-        im = tf.cast(im, tf.float32) / 255.0
-        pred_mask = self.model.predict(im[tf.newaxis, ...])
-        
-        # take the best performing class for each pixel
-        # the output of argmax looks like this [[1, 2, 0], ...]
-        pred_mask_arg = tf.argmax(pred_mask, axis=-1)
-
-        labels = []
-        
-        # convert the prediction mask into binary masks for each class
-        binary_masks = {}
-        mask_codes = {}
-        
-        # when we take tf.argmax() over pred_mask, it becomes a tensor object
-        # the shape becomes TensorShape object, looking like this TensorShape([128]) 
-        # we need to take get shape, convert to list and take the best one
-        
-        rows = pred_mask_arg[0][1].get_shape().as_list()[0]
-        cols = pred_mask_arg[0][2].get_shape().as_list()[0]
-        
-        for cls in range(pred_mask.shape[-1]):
-
-            binary_masks[f"mask_{cls}"] = np.zeros(shape = (pred_mask.shape[1], pred_mask.shape[2])) #create masks for each class
-            
-            for row in range(rows):
-
-                for col in range(cols):
-
-                    if pred_mask_arg[0][row][col] == cls:
-                        
-                        binary_masks[f"mask_{cls}"][row][col] = 1
-                    else:
-                        binary_masks[f"mask_{cls}"][row][col] = 0
-
-            mask = binary_masks[f"mask_{cls}"]
-            mask *= 255
-            img = Image.fromarray(mask.astype(np.int8), mode="L")
-               
-            # we need to make it readable for the widget
-            with io.BytesIO() as out:
-                img.save(out, format="PNG")
-                png_string = out.getvalue()
-                mask = base64.b64encode(png_string).decode("utf-8")
-
-            mask_codes[f"mask_{cls}"] = mask
-    
-
-            # widget needs the below format, for each class we return label and mask string
-            labels.append({
-                "label": f"LABEL_{cls}",
-                "mask": mask_codes[f"mask_{cls}"],
-                "score": 1.0,
-            })
-        return labels
+    def __call__(self, inputs: "Image.Image")-> List[Dict[str, Any]]:
+        json_data = [{
+            'video_id': 'pqh4LfPeCYs',
+            'start': 835.933,
+            'end': 927.581,
+            'category': 'sponsor'
+        }]
+        return self.pipeline(json_data)