path: root/main.py

import os
import re
import sys
import backoff
import openai
import json
from datetime import datetime, timedelta
from tqdm import tqdm

client = openai.OpenAI(
    #base_url="http://localhost:8080/v1",
    #base_url="https://api.openai.com/v1",
    base_url="https://openrouter.ai/api/v1",
)

MODEL = "deepseek/deepseek-chat-v3-0324"
SYSTEM_PROMPT_TEMPLATE = """
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.

You will be provided with segments of a podcast transcript, formatted with timestamps. Your output must be ONLY the start and end timestamp of the identified advertising segment, separated by a hyphen, in the format HH:MM:SS.mmm-HH:MM:SS.mmm. If there are multiple ad segments, output each on a new line. If no ad segments are found, output nothing.

**Strict Rules:**
- 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.
- 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.

**Examples:**

---
**Example 1 Input:**
[00:00:00.000 --> 00:00:02.480]   This episode is brought to you by Apple Cash.
[00:00:02.480 --> 00:00:07.040]   Sending payments used to be clunky, unnecessarily difficult, and weirdly invasive at times,
[00:00:07.040 --> 00:00:08.560]   until I discovered Apple Cash.
[00:00:08.560 --> 00:00:11.020]   With Apple Cash, payments are private by design,
[00:00:11.020 --> 00:00:14.900]   so I don't have to deal with public feeds, awkward reactions, or other payment drama.
[00:00:14.900 --> 00:00:18.780]   I can send cash in messages right in the conversations I'm already having,
[00:00:18.780 --> 00:00:20.060]   which is super convenient.
[00:00:20.060 --> 00:00:22.600]   There's also a cool feature called Tap to Cash
[00:00:22.600 --> 00:00:25.900]   that lets you pay somebody nearby by holding your iPhone near theirs.
[00:00:25.900 --> 00:00:28.560]   Switch to Apple Cash and start sending privately.
[00:00:28.620 --> 00:00:31.900]   Apple Cash services are provided by Green Dot Bank, member FDIC.

**Example 1 Output:**
00:00:00.000-00:00:31.900

---
**Example 2 Input:**
[00:38:07.480 --> 00:38:10.820]   This podcast is brought to you by Carvana.
[00:38:10.820 --> 00:38:13.660]   Got a car to sell, but no time to waste?
[00:38:13.660 --> 00:38:17.820]   Hop on to Carvana.com to get a real offer for your car in seconds.
[00:38:17.820 --> 00:38:23.060]   All you have to do is enter your license plate, answer a few quick questions, and if you accept
[00:38:23.060 --> 00:38:26.000]   the offer, Carvana will pay you as soon as you hand the keys over.
[00:38:26.360 --> 00:38:29.060]   They even offer same day pickup in many cities.
[00:38:29.060 --> 00:38:35.740]   Save your time, score some cash, and sell your car the convenient way to Carvana.
[00:38:36.360 --> 00:38:37.000]   Pickup times vary.
[00:38:37.000 --> 00:38:37.620]   Fees may apply.

**Example 2 Output:**
00:38:07.480-00:38:37.620

---
**Example 3 Input:**
[00:14:46.000 --> 00:14:50.860]   This episode is brought to you by Diet Coke.
[00:14:50.860 --> 00:14:53.800]   You know that moment when you just need to hit pause and refresh?
[00:14:54.380 --> 00:14:56.320]   An ice-cold Diet Coke isn't just a break.
[00:14:56.320 --> 00:14:59.800]   It's your chance to catch your breath and savor a moment that's all about you.
[00:15:00.000 --> 00:15:02.540]   Always refreshing, still the same great taste.
[00:15:02.540 --> 00:15:03.500]   Diet Coke.
[00:15:03.500 --> 00:15:04.860]   Make time for you time.

**Example 3 Output:**
00:14:46.000-00:15:04.860

---
**Example 4 Input:**
[00:15:00.000 --> 00:15:20.000] So what are your thoughts on the new policy changes that were just announced?
[00:15:20.000 --> 00:15:45.000] Well, I think it's a step in the right direction, but there are definitely some areas that need more clarification.
[00:15:45.000 --> 00:16:10.000] For instance, the section on renewable energy credits could be more explicit.
[00:16:10.000 --> 00:16:15.000] This next part of our conversation might be a bit technical, but I think it's important to delve into the specifics.

**Example 4 Output:**


---
**Example 5 Input:**
[00:11:27.080 --> 00:11:28.180]   You're going to be using the same thing.
[00:11:28.180 --> 00:11:32.480]   You've got to, like, get outside that bubble or you're just, you're going to bore yourself.
[00:11:32.480 --> 00:11:40.820]   This episode of Cortex is brought to you by Squarespace, the all-in-one website platform designed to help you stand out and succeed online.
[00:11:40.820 --> 00:11:43.920]   Whether you're just getting started or scaling your own business,
[00:11:43.920 --> 00:11:46.840]   Squarespace gives you everything you need to claim your domain,
[00:11:46.840 --> 00:11:49.380]   showcase your offerings of a professional website,
[00:14:50.000 --> 00:14:52.220]   grow your brand, and get paid all in one place.
[00:14:52.220 --> 00:14:54.960]   It's so easy to get started with Squarespace.
[00:14:54.960 --> 00:14:59.380]   In fact, they've made it easier than ever before with their new system, Blueprint AI.
[00:14:59.380 --> 00:15:05.140]   Squarespace's AI-enhanced website builder lets you quickly and easily build a site bespoke to your
business.
[00:15:05.140 --> 00:15:08.140]   Just input some basic information about your industry and goals.
[00:15:08.140 --> 00:15:09.480]   This isn't the only way.
[00:15:09.480 --> 00:15:12.200]   In fact, you can go in and choose from one of their beautiful templates,
[00:15:12.200 --> 00:15:14.620]   their professionally designed beautiful templates,
[00:15:14.940 --> 00:15:16.900]   and customize it to your heart's content.
[00:15:16.900 --> 00:15:20.800]   One of the things that I've always loved about Squarespace is their drag-and-drop tools,
[00:15:20.800 --> 00:15:26.440]   buttons, sliders, selectors, where you can go in and customize your design for your website
[00:15:26.440 --> 00:15:29.420]   by actually looking at the design for your website while you're doing it.
[00:15:29.460 --> 00:15:30.620]   You don't need to know any code.
[00:15:30.620 --> 00:15:31.880]   I love this.
[00:13:31.200 --> 00:13:35.220]   I want to talk about the other devices that are in your working life.
[00:13:35.220 --> 00:13:39.060]   What are the other computers that you use to get things done?
[00:13:39.060 --> 00:13:45.000]   So there are three computers that live in my life.

**Example 5 Output:**
00:11:32.480-00:13:31.200
"""

def parse_whisper_transcript(transcript_text: str) -> list[dict]:
    """
    Parses a Whisper-generated transcript into a list of dictionaries,
    each containing 'start_time', 'end_time', and 'text'.
    """
    segments = []
    # Regex to match timestamp format [HH:MM:SS.mmm --> HH:MM:SS.mmm] and the following text
    pattern = re.compile(r"^\[(\d{2}:\d{2}:\d{2}\.\d{3}) --> (\d{2}:\d{2}:\d{2}\.\d{3})\]\s*(.*)$", re.MULTILINE)
    
    for line in transcript_text.strip().split('\n'):
        match = pattern.match(line)
        if match:
            start_time, end_time, text = match.groups()
            segments.append({
                "start_time": start_time,
                "end_time": end_time,
                "text": text.strip()
            })
    return segments

def format_transcript_segment(segments: list[dict]) -> str:
    """Formats a list of segment dictionaries back into the prompt-friendly string."""
    formatted_lines = []
    for s in segments:
        formatted_lines.append(f"[{s['start_time']} --> {s['end_time']}]   {s['text']}")
    return "\n".join(formatted_lines)

def timestamp_to_seconds(ts_str: str) -> float:
    """Converts HH:MM:SS.mmm string to total seconds."""
    parts = ts_str.split(':')
    hours = int(parts[0])
    minutes = int(parts[1])
    seconds_ms_str = parts[2].split('.')
    seconds = int(seconds_ms_str[0])
    milliseconds = int(seconds_ms_str[1])
    return hours * 3600 + minutes * 60 + seconds + milliseconds / 1000.0

def seconds_to_timestamp(total_seconds: float) -> str:
    """Converts total seconds to HH:MM:SS.mmm string."""
    integer_seconds = int(total_seconds)
    milliseconds = int((total_seconds - integer_seconds) * 1000)

    td = timedelta(seconds=integer_seconds)
    hours, remainder = divmod(td.total_seconds(), 3600)
    minutes, seconds = divmod(remainder, 60)
    
    return f"{int(hours):02}:{int(minutes):02}:{int(seconds):02}.{milliseconds:03}"

def chunk_transcript(
    segments: list[dict], 
    chunk_duration_seconds: int = 240,
    overlap_seconds: int = 60
) -> list[list[dict]]:
    if not segments:
        return []

    chunks = []
    
    # Calculate step size for the window
    step_seconds = chunk_duration_seconds - overlap_seconds
    if step_seconds <= 0:
        raise ValueError("Overlap must be smaller than chunk duration.")

    transcript_end_time = timestamp_to_seconds(segments[-1]['end_time'])
    
    window_start_time = 0.0
    while window_start_time < transcript_end_time:
        window_end_time = window_start_time + chunk_duration_seconds
        
        current_chunk_segments = []
        for s in segments:
            seg_start = timestamp_to_seconds(s['start_time'])
            seg_end = timestamp_to_seconds(s['end_time'])
            
            # Add segment if it has any overlap with the current window
            if seg_start < window_end_time and seg_end > window_start_time:
                current_chunk_segments.append(s)

        if current_chunk_segments:
            chunks.append(current_chunk_segments)
        
        window_start_time += step_seconds
        
        # Break if the last segment has been fully processed
        if window_start_time > timestamp_to_seconds(current_chunk_segments[-1]['end_time']):
            break

    return chunks

@backoff.on_exception(backoff.expo, openai.RateLimitError)
def completions_with_backoff(**kwargs):
    return client.chat.completions.create(**kwargs)

def call_openai_api(chunk_text: str, system_prompt: str) -> str:
    """
    Makes an API call to OpenAI with the given system prompt and transcript chunk.
    """
    response = completions_with_backoff(
        #model="Gemma-3-4B_32K",
        #model="google/gemini-2.0-flash-exp:free",
        #model="gpt-4.1",
        model=MODEL,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f"Input Transcript:\n{chunk_text}"}
        ],
        temperature=0.0,
        max_tokens=1024,    # realistically it'd be under like 512
        extra_body={"provider": {"quantizations": ["fp8"], "sort": "price" }},
    )
    #print(response)
    return response.choices[0].message.content.strip()

def process_file(input_file: str, output_file: str):
    with open(input_file, "r", encoding="utf-8") as f:
        full_transcript_content = f.read()

    # --- Processing ---
    parsed_segments = parse_whisper_transcript(full_transcript_content)

    # Adjust the trigger phrase if "brought to you by" isn't consistently sufficient
    # For more complex cases, you might use a list of trigger phrases:
    # ["brought to you by", "a quick word from our sponsor", "we'll be right back"]
    transcript_chunks = chunk_transcript(segments=parsed_segments)

    all_ad_timestamps = []

    # Process each chunk with the OpenAI API
    for i, chunk in enumerate(transcript_chunks):
        if not chunk:
            continue

        user_prompt = format_transcript_segment(chunk)

        try:
            ad_timestamps_str = call_openai_api(user_prompt, SYSTEM_PROMPT_TEMPLATE)
        except Exception as e:
            print(f"Error calling OpenAI API: {e}")
            continue

        if ad_timestamps_str.startswith("[") or "no output" in ad_timestamps_str.lower() or "no ad" in ad_timestamps_str.lower():
            ad_timestamps_str = ""

        # Split by newlines in case the model returns multiple segments
        found_timestamps = [ts.strip() for ts in ad_timestamps_str.split('\n') if ts.strip()]
        all_ad_timestamps.append({"id": f"{"".join(input_file.split(".")[:-1])}_chunk{i}", f"text": "Instruction: Identify the timestamp range of the pre-recorded ad in the following transcript. Output only the HH:MM:SS.mmm-HH:MM:SS.mmm range.\n\nTranscript:\n" + user_prompt, "target": "\n".join(found_timestamps)})

    with open(output_file, "w") as f:
        for s in all_ad_timestamps:
            f.write(json.dumps(s) + "\n")

def main():
    os.makedirs("data/")
    paths = sys.argv[1:]
    for p in paths:
        out = "data/" + os.path.splitext(os.path.basename(p))[0] + ".jsonl"
        if os.path.isfile(out):
            paths.remove(p)
    print(paths)

    for p in tqdm(paths):
        out = "data/" + os.path.splitext(os.path.basename(p))[0] + ".jsonl"
        process_file(p, out)

if __name__ == "__main__":
    main()