change chunking to overlapping time segments

1 files changed, 52 insertions, 55 deletions

diff --git a/main.py b/main.py
index 283ea9a..045fb2e 100644
--- a/main.py
+++ b/main.py
@@ -9,9 +9,11 @@ from tqdm import tqdm
 
 client = openai.OpenAI(
     #base_url="http://localhost:8080/v1",
-    base_url="https://api.openai.com/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.
 
@@ -23,6 +25,7 @@ You will be provided with segments of a podcast transcript, formatted with times
 - 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:**
 
@@ -162,51 +165,43 @@ def seconds_to_timestamp(total_seconds: float) -> str:
 
 def chunk_transcript(
     segments: list[dict], 
-    trigger_phrase: str = "brought to you by",
-    max_chunk_duration_minutes: int = 4
+    chunk_duration_seconds: int = 240,
+    overlap_seconds: int = 60
 ) -> list[list[dict]]:
-    """
-    Chunks the transcript segments based on occurrences of a trigger phrase,
-    or a maximum time duration, whichever comes first.
-    Each chunk starts at a trigger phrase and ends just before the next,
-    or when max_chunk_duration_minutes is reached, or at the end of transcript.
-    """
-    chunks = []
-    current_chunk_start_idx = 0
-    max_chunk_duration_seconds = max_chunk_duration_minutes * 60
-
     if not segments:
         return []
 
-    for i, segment in enumerate(segments):
-        is_trigger_segment = trigger_phrase.lower() in segment["text"].lower()
-
-        current_chunk_start_time_seconds = timestamp_to_seconds(segments[current_chunk_start_idx]['start_time'])
-        current_segment_end_time_seconds = timestamp_to_seconds(segment['end_time'])
-        current_chunk_duration = current_segment_end_time_seconds - current_chunk_start_time_seconds
-
-        # Condition 1: New trigger phrase found (and it's not the start of the very first chunk)
-        # We also need to ensure the trigger is not the first segment of the current chunk,
-        # otherwise it would create an empty chunk.
-        if is_trigger_segment and (i > current_chunk_start_idx):
-            # End the previous chunk *before* this new trigger segment
-            chunks.append(segments[current_chunk_start_idx:i])
-            current_chunk_start_idx = i # Start a new chunk from this trigger segment
-            continue # Move to the next segment after processing the split
-
-        # Condition 2: Max duration exceeded
-        # This check should happen for any segment after the start of a chunk
-        if current_chunk_duration >= max_chunk_duration_seconds and i > current_chunk_start_idx:
-            # End the current chunk *at* the current segment
-            chunks.append(segments[current_chunk_start_idx:i+1])
-            current_chunk_start_idx = i + 1 # Start a new chunk from the next segment
-            # We don't 'continue' here, as the new chunk needs to start and the loop will increment 'i'
-
-    # Add the last chunk if any segments remain. This covers the case where the transcript
-    # ends without another trigger or exceeding the time limit.
-    if current_chunk_start_idx < len(segments):
-        chunks.append(segments[current_chunk_start_idx:])
+    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)
@@ -220,13 +215,15 @@ def call_openai_api(chunk_text: str, system_prompt: str) -> str:
     response = completions_with_backoff(
         #model="Gemma-3-4B_32K",
         #model="google/gemini-2.0-flash-exp:free",
-        model="gpt-4.1",
+        #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()
@@ -241,7 +238,7 @@ def process_file(input_file: str, output_file: str):
     # 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(parsed_segments, trigger_phrase="brought to you by")
+    transcript_chunks = chunk_transcript(segments=parsed_segments)
 
     all_ad_timestamps = []
 
@@ -250,36 +247,36 @@ def process_file(input_file: str, output_file: str):
         if not chunk:
             continue
 
-        chunk_text_for_api = format_transcript_segment(chunk)
-        
-        # Add the "Your Turn" header before the chunk content
-        user_prompt_content = f"**Your Turn - Input Transcript:**\n{chunk_text_for_api}"
+        user_prompt = format_transcript_segment(chunk)
 
         try:
-            ad_timestamps_str = call_openai_api(user_prompt_content, SYSTEM_PROMPT_TEMPLATE)
+            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:
-            # 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": chunk_text_for_api, "target": "\n".join(found_timestamps)})
+
+        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 = "".join(p.split(".")[:-1]) + ".jsonl"
+        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 = "".join(p.split(".")[:-1]) + ".jsonl"
+        out = "data/" + os.path.splitext(os.path.basename(p))[0] + ".jsonl"
         process_file(p, out)
 
 if __name__ == "__main__":