Add time conversions from outputs

app.py

@@ -3,6 +3,7 @@ import re
 import gradio as gr
 from functions.punctuation import punctuate
 from functions.model_infer import predict_from_document
+from functions.convert_time import match_mask_and_transcript
 
 
 title = "sponsoredBye - never listen to sponsors again"
@@ -12,16 +13,18 @@ article = "Check out [the original Rick and Morty Bot](https://huggingface.co/sp
 
 def pipeline(video_url):
     video_id = video_url.split("?v=")[-1]
-    punctuated_text = punctuate(video_id)
+    punctuated_text, transcript = punctuate(video_id)
     sentences = re.split(r"[\.\!\?]\s", punctuated_text)
     classification, probs = predict_from_document(sentences)
     #    return punctuated_text
+    times = match_mask_and_transcript(sentences, transcript, classification)
     return [
         {
             "start": "12:05",
             "end": "12:52",
             "classification": str(classification),
             "probabilities": probs,
+            "times": times,
         }
     ]
 

functions/convert_time.py

@@ -0,0 +1,50 @@
+import re
+from thefuzz import fuzz
+import numpy as np
+
+
+def match_mask_and_transcript(split_punct, transcript, classification):
+    """
+    Input:
+        split_punct: the punctuated text, split on ?/!/.\s,
+        transcript: original transcript with timestamps
+        classification: classification object (list of numbers 0,1)
+    Output: times
+    """
+
+    # Get the sponsored part
+    sponsored_segment = []
+    for i, val in enumerate(classification):
+        if val == 1:
+            sponsored_segment.append(split_punct[i])
+
+    segment = " ".join(sponsored_segment)
+    sim_scores = list()
+
+    # Check the similarity scores between the sponsored part and the transcript parts
+    for elem in transcript:
+        sim_scores.append(fuzz.partial_ratio(segment, elem["text"]))
+
+    # Get the scores and check if they are above mean + 2*stdev
+    scores = np.array(sim_scores)
+    timestamp_mask = (scores > np.mean(scores) + np.std(scores) * 2).astype(int)
+    timestamps = [
+        (transcript[i]["start"], transcript[i]["duration"])
+        for i, elem in enumerate(timestamp_mask)
+        if elem == 1
+    ]
+
+    # Get the timestamp segments
+    times = []
+    current = -1
+    current_time = 0
+    for elem in timestamps:
+        # Threshold of 5 to see if it is a jump to another segment (also to make sure smaller segments are added together
+        if elem[0] > (current_time + 5):
+            current += 1
+            times.append((elem[0], elem[0] + elem[1]))
+            current_time = elem[0] + elem[1]
+        else:
+            times[current] = (times[current][0], elem[0] + elem[1])
+            current_time = elem[0] + elem[1]
+    return times

functions/model_infer.py

@@ -41,6 +41,6 @@ def predict_from_document(sentences):
     # Set the prediction threshold to 0.8 instead of 0.5, now use mean
     output = (
         prediction.flatten()[: len(sentences)]
-        >= np.mean(prediction) + np.var(prediction) * 2
+        >= np.mean(prediction) + np.std(prediction) * 2
     ).astype(int)
     return output, prediction.flatten()[: len(sentences)]

functions/punctuation.py

@@ -55,4 +55,4 @@ def punctuate(video_id):
     )  # Get the transcript from the YoutubeTranscriptApi
     resp = query_punctuation(splits)  # Get the response from the Inference API
     punctuated_transcript = parse_output(resp, splits)
-    return punctuated_transcript
+    return punctuated_transcript, transcript

requirements.txt

@@ -1,4 +1,6 @@
 youtube_transcript_api
+thefuzz
+numpy
 tensorflow==2.15
 keras
 keras-nlp