add highlight/visualizer script and update system prompt

2 files changed, 155 insertions, 30 deletions

diff --git a/highlight.py b/highlight.py
new file mode 100644
index 0000000..361a0d2
--- /dev/null
+++ b/highlight.py
@@ -0,0 +1,118 @@
+import sys
+import json
+import re
+from datetime import timedelta
+
+# ANSI escape codes for terminal colors
+class Colors:
+    GREEN = '\033[92m'  # Highlight for valid ad text
+    RED = '\033[91m'    # Highlight for invalid target ranges
+    BLUE = '\033[94m'   # Header color
+    RESET = '\033[0m'   # Reset to default color
+
+def timestamp_to_seconds(ts_str: str) -> float:
+    """Converts HH:MM:SS.mmm string to total seconds. Returns -1.0 on failure."""
+    try:
+        parts = ts_str.split(':')
+        h, m = int(parts[0]), int(parts[1])
+        s_ms = parts[2].split('.')
+        s, ms = int(s_ms[0]), int(s_ms[1])
+        return h * 3600 + m * 60 + s + ms / 1000.0
+    except (ValueError, IndexError):
+        return -1.0
+
+def visualize_entry(entry: dict):
+    """Parses and prints a single entry from the jsonl file with highlighting."""
+    entry_id = entry.get("id", "N/A")
+    full_text = entry.get("text", "")
+    target_str = entry.get("target", "")
+
+    print(f"{Colors.BLUE}{'='*80}\nID: {entry_id}{Colors.RESET}")
+
+    # --- 1. Parse the transcript from the 'text' field ---
+    transcript_lines = []
+    # Regex to find the start of the actual transcript text
+    transcript_match = re.search(r"Transcript:\n(.*)", full_text, re.DOTALL)
+    if not transcript_match:
+        print(f"{Colors.RED}Could not find transcript text in entry.{Colors.RESET}")
+        return
+
+    raw_transcript = transcript_match.group(1).strip()
+    line_pattern = re.compile(r"\[(\d{2}:\d{2}:\d{2}\.\d{3}) --> (\d{2}:\d{2}:\d{2}\.\d{3})\](.*)")
+    
+    min_transcript_time, max_transcript_time = float('inf'), float('-inf')
+
+    for line in raw_transcript.split('\n'):
+        match = line_pattern.match(line)
+        if match:
+            start_str, end_str, text = match.groups()
+            start_sec = timestamp_to_seconds(start_str)
+            end_sec = timestamp_to_seconds(end_str)
+            transcript_lines.append({"start": start_sec, "end": end_sec, "line": line})
+            min_transcript_time = min(min_transcript_time, start_sec)
+            max_transcript_time = max(max_transcript_time, end_sec)
+
+    # --- 2. Parse and validate the target ranges ---
+    valid_target_ranges = []
+    has_invalid_target = False
+    
+    if target_str:
+        for target_line in target_str.strip().split('\n'):
+            if '-' not in target_line:
+                has_invalid_target = True
+                continue
+            
+            start_target_str, end_target_str = target_line.split('-', 1)
+            start_target_sec = timestamp_to_seconds(start_target_str.strip())
+            end_target_sec = timestamp_to_seconds(end_target_str.strip())
+
+            # Validation checks
+            is_valid_format = start_target_sec != -1.0 and end_target_sec != -1.0
+            is_valid_range = is_valid_format and start_target_sec < end_target_sec
+            # Check if the target overlaps with the transcript's time frame at all
+            is_overlapping = is_valid_range and start_target_sec < max_transcript_time and end_target_sec > min_transcript_time
+
+            if is_overlapping:
+                valid_target_ranges.append({"start": start_target_sec, "end": end_target_sec})
+            else:
+                has_invalid_target = True
+    
+    # --- 3. Print results with highlighting ---
+    print(f"TARGET: ", end="")
+    if has_invalid_target:
+        print(f"{Colors.RED}{target_str or 'EMPTY'}{Colors.RESET}")
+    else:
+        print(f"{Colors.GREEN}{target_str or 'EMPTY'}{Colors.RESET}")
+    print("-" * 30)
+
+    for line_info in transcript_lines:
+        is_ad_line = False
+        for target_range in valid_target_ranges:
+            # Check for overlap between the line's time and the target's time
+            if line_info['start'] < target_range['end'] and line_info['end'] > target_range['start']:
+                is_ad_line = True
+                break
+        
+        if is_ad_line:
+            print(f"{Colors.GREEN}{line_info['line']}{Colors.RESET}")
+        else:
+            print(line_info['line'])
+
+def main(filepath: str):
+    try:
+        with open(filepath, 'r', encoding='utf-8') as f:
+            for line in f:
+                try:
+                    entry = json.loads(line)
+                    visualize_entry(entry)
+                except json.JSONDecodeError:
+                    print(f"{Colors.RED}Error: Could not decode JSON line: {line.strip()}{Colors.RESET}")
+    except FileNotFoundError:
+        print(f"{Colors.RED}Error: File not found at {filepath}{Colors.RESET}")
+
+if __name__ == "__main__":
+    if len(sys.argv) < 2:
+        print("Usage: python highlight.py <path_to_labels.jsonl>")
+        sys.exit(1)
+    
+    main(sys.argv[1])
diff --git a/main.py b/main.py
index c7a1ff5..a58e02f 100644
--- a/main.py
+++ b/main.py
@@ -13,12 +13,8 @@ client = openai.OpenAI(
     api_key=os.getenv("OPENROUTER_API_KEY")
 )
 
-#MODEL = "deepseek/deepseek-r1-0528:free"
-#MODEL = "deepseek/deepseek-chat-v3-0324" # good with some minor editing of output ranges
-MODEL = "deepseek/deepseek-chat-v3-0324:free" # twice as slow as paid
-#MODEL = "google/gemini-2.5-flash-lite-preview-06-17" # too many false-positives
-#MODEL = "google/gemini-2.0-flash-exp:free" # fine
-#MODEL = "qwen/qwen3-30b-a3b" # sort of works, but slow
+# smaller models than this (besides dense llama 3.1 405b) really struggle
+MODEL = "deepseek/deepseek-chat-v3.1"
 
 SYSTEM_PROMPT = """
 You are an expert podcast content analyzer. Your task is to identify and extract pre-recorded dynamic advertising segments from podcast transcripts. These ads typically have a distinct tone shift, often using more direct, persuasive language, and frequently include calls to action, product mentions, or specific brand names.
@@ -29,6 +25,7 @@ You will be provided with segments of a podcast transcript, formatted with times
 - Only identify segments that are clearly pre-recorded dynamic ads with a noticeable tone shift.
 - Do NOT include host-read sponsorships that blend naturally with the content. Focus only on the "distracting" pre-recorded elements.
 - The end of an ad segment is marked by the return to the original podcast discussion, a clear transition to new content, or the end of the ad's persuasive language and calls to action.
+- A segment should only be classified as an ad if it contains an abrupt topic change AND specific marketing language (e.g., product names, calls to action, special offers). A topic change alone is NOT an ad.
 - Ensure the start and end times correspond precisely to the ad segment in the transcript.
 - Output ONLY the timestamp range(s) in the specified format. No other text, no JSON, no explanations.
 - If no ad segments are found, your output MUST be completely empty. Generating any text, including explanations like '[No output...]', is a failure to follow instructions. The only valid output in this case is a blank string.
@@ -107,8 +104,7 @@ You will be provided with segments of a podcast transcript, formatted with times
 **Example 4 Input:**
 [01:46:23.940 --> 01:46:32.220]   It has so rapidly become one of the most critical apps on my whole system, I cannot even believe it.
 [01:46:32.220 --> 01:46:36.320]   So if you don't use shortcuts, this is completely worthless to you.
-[01:46:36.320 --> 01:46:54.380]   But if you make even mildly complicated shortcuts, you have to be using Logger and it will just completely unlock your ability to be able to make much more complicated shortcuts because you can always know exac
-tly what's happening and exactly where things are going wrong.
+[01:46:36.320 --> 01:46:54.380]   But if you make even mildly complicated shortcuts, you have to be using Logger and it will just completely unlock your ability to be able to make much more complicated shortcuts because you can always know exactly what's happening and exactly where things are going wrong.
 [01:46:54.380 --> 01:46:57.460]   So Logger is my app of the year.
 [01:46:59.320 --> 01:47:02.180]   This episode of Cortex is brought to you by FitBod.
 [01:47:02.180 --> 01:47:06.040]   When you're looking to change your fitness level, it's hard to know where to get started.
@@ -120,26 +116,25 @@ tly what's happening and exactly where things are going wrong.
 [01:47:40.940 --> 01:47:44.500]   And it builds your best possible workout with the use of exercise science.
 [01:47:44.500 --> 01:47:50.260]   They have analyzed billions of data points that have been fine-tuned by FitBod certified personal trainers.
 [01:47:50.260 --> 01:47:57.180]   And you can make sure you're going to be learning new movements the right way thanks to the more than 1,000 demonstration videos that are all in the FitBod app.
-[01:47:57.180 --> 01:48:07.380]   This is my favorite feature of FitBod that when I get shown a new exercise, I have a perfect way to be able to understand how to do it because of all the videos that they can show me from different angles along
- with the instructions.
- [01:48:07.380 --> 01:48:10.100]   And it makes me feel confident about the work that I'm going to be doing.
- [01:48:10.100 --> 01:48:14.340]   Muscles improve when they're working in concert with the entire musculoskeletal system.
- [01:48:14.340 --> 01:48:19.360]   So overworking some muscles while underworking others can negatively impact your overall results.
- [01:48:19.360 --> 01:48:24.280]   This is why FitBod tracks muscle fatigue and recovery to design a well-balanced workout routine,
- [01:48:24.280 --> 01:48:30.960]   which also means you'll never get bored as the app mixes up your workouts with new exercises, rep schemes, supersets, and circuits.
- [01:48:30.960 --> 01:48:37.760]   The FitBod app is super easy to use and it integrates with your Apple Watch, WearOS smartwatch, and apps like Strava, Fitbit, and Apple Health.
- [01:48:37.760 --> 01:48:44.020]   You've already heard in this episode why Gray loves FitBod, and Gray Loving FitBod is the reason why they became a sponsor of the show.
- [01:48:44.020 --> 01:48:45.960]   So you should go and check it out for yourself.
- [01:48:45.960 --> 01:48:51.740]   You'll be able to benefit from all of the incredible work from FitBod to help you in your fitness journey.
- [01:48:51.740 --> 01:48:58.440]   Personalized training of this quality can be expensive, but FitBod is just $12.99 a month or $79.99 a year.
- [01:48:58.440 --> 01:49:04.380]   But you can get 25% off your membership by signing up today at FitBod.me slash Cortex.
- [01:49:04.380 --> 01:49:09.840]   So go now and get your customized fitness plan at FitBod.me slash Cortex.
- [01:49:09.980 --> 01:49:14.180]   Once again, that is FitBod.me slash Cortex, and you'll get 25% off your membership.
- [01:49:14.180 --> 01:49:17.980]   A thanks to FitBod for their continued support of this show and Relay.
- [01:49:17.980 --> 01:49:20.120]   We've made it to home screens.
- [01:49:20.120 --> 01:49:22.360]   Oh, home screens. Right, right.
- [01:49:22.360 --> 01:49:23.820]   How complicated can it be?
- [01:49:23.920 --> 01:49:29.820]   So as like last time, my home screens are broken down into a selection of focus modes too.
+[01:47:57.180 --> 01:48:07.380]   This is my favorite feature of FitBod that when I get shown a new exercise, I have a perfect way to be able to understand how to do it because of all the videos that they can show me from different angles along with the instructions.
+[01:48:07.380 --> 01:48:10.100]   And it makes me feel confident about the work that I'm going to be doing.
+[01:48:10.100 --> 01:48:14.340]   Muscles improve when they're working in concert with the entire musculoskeletal system.
+[01:48:14.340 --> 01:48:19.360]   So overworking some muscles while underworking others can negatively impact your overall results.
+[01:48:19.360 --> 01:48:24.280]   This is why FitBod tracks muscle fatigue and recovery to design a well-balanced workout routine,
+[01:48:24.280 --> 01:48:30.960]   which also means you'll never get bored as the app mixes up your workouts with new exercises, rep schemes, supersets, and circuits.
+[01:48:30.960 --> 01:48:37.760]   The FitBod app is super easy to use and it integrates with your Apple Watch, WearOS smartwatch, and apps like Strava, Fitbit, and Apple Health.
+[01:48:37.760 --> 01:48:44.020]   You've already heard in this episode why Gray loves FitBod, and Gray Loving FitBod is the reason why they became a sponsor of the show.
+[01:48:44.020 --> 01:48:45.960]   So you should go and check it out for yourself.
+[01:48:45.960 --> 01:48:51.740]   You'll be able to benefit from all of the incredible work from FitBod to help you in your fitness journey.
+[01:48:51.740 --> 01:48:58.440]   Personalized training of this quality can be expensive, but FitBod is just $12.99 a month or $79.99 a year.
+[01:48:58.440 --> 01:49:04.380]   But you can get 25% off your membership by signing up today at FitBod.me slash Cortex.
+[01:49:04.380 --> 01:49:09.840]   So go now and get your customized fitness plan at FitBod.me slash Cortex.
+[01:49:09.980 --> 01:49:14.180]   Once again, that is FitBod.me slash Cortex, and you'll get 25% off your membership.
+[01:49:14.180 --> 01:49:17.980]   A thanks to FitBod for their continued support of this show and Relay.
+[01:49:17.980 --> 01:49:20.120]   We've made it to home screens.
+[01:49:20.120 --> 01:49:22.360]   Oh, home screens. Right, right.
+[01:49:22.360 --> 01:49:23.820]   How complicated can it be?
+[01:49:23.920 --> 01:49:29.820]   So as like last time, my home screens are broken down into a selection of focus modes too.
 
 **Example 4 Output:**
 01:46:59.320-01:49:17.980
@@ -253,6 +248,17 @@ tly what's happening and exactly where things are going wrong.
 
 **Example 7 Output:**
 00:14:17.500-00:17:05.040
+
+---
+**Example 8 Input:**
+[00:27:11.039 --> 00:27:14.480] Why aren't stadium seats just toilets?
+[00:27:14.720 --> 00:27:19.839] Like the Romans of old, like the ancient Greek bathhouse.
+[00:27:20.480 --> 00:27:29.920] Everything was cool and you could chit sit next to your bro while you talked out some philosophy or watched a little sporting event, maybe a little wrestling.
+[00:27:30.160 --> 00:27:37.359] Why are we because in the stadium, one of the biggest problems is you oh like, oh, it's between innings, oh, is it's halftime.
+[00:27:37.440 --> 00:27:38.480] Oh, I gotta go to the bathroom.
+
+**Example 8 Output:**
+
 """
 
 # --- Helper Functions ---
@@ -293,7 +299,7 @@ def segment_by_punctuation(asr_results: list, min_words_per_segment: int = 3) ->
         for i, token in enumerate(tokens):
             is_end_of_sentence = any(p in token.strip() for p in punctuations)
             
-            if is_end_of_sentence or (i == len(tokens) - 1):
+            if (is_end_of_sentence or (i == len(tokens) - 1)) and len(tokens) > current_segment_start_idx:
                 segment_tokens = tokens[current_segment_start_idx : i + 1]
 
                 if len(segment_tokens) >= min_words_per_segment:
@@ -306,7 +312,8 @@ def segment_by_punctuation(asr_results: list, min_words_per_segment: int = 3) ->
                         "end_time": end_time,
                         "text": text
                     })
-                current_segment_tokens = []
+
+                current_segment_start_idx = i + 1
     return segments
 
 def chunk_segments(segments: list[dict], chunk_duration_seconds: int = 240, overlap_seconds: int = 60) -> list[list[dict]]: