summaryrefslogtreecommitdiff
path: root/main.py
diff options
context:
space:
mode:
Diffstat (limited to 'main.py')
-rw-r--r--main.py107
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
9 9
10client = openai.OpenAI( 10client = openai.OpenAI(
11 #base_url="http://localhost:8080/v1", 11 #base_url="http://localhost:8080/v1",
12 base_url="https://api.openai.com/v1", 12 #base_url="https://api.openai.com/v1",
13 base_url="https://openrouter.ai/api/v1",
13) 14)
14 15
16MODEL = "deepseek/deepseek-chat-v3-0324"
15SYSTEM_PROMPT_TEMPLATE = """ 17SYSTEM_PROMPT_TEMPLATE = """
16You 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. 18You 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.
17 19
@@ -23,6 +25,7 @@ You will be provided with segments of a podcast transcript, formatted with times
23- 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. 25- 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.
24- Ensure the start and end times correspond precisely to the ad segment in the transcript. 26- Ensure the start and end times correspond precisely to the ad segment in the transcript.
25- Output ONLY the timestamp range(s) in the specified format. No other text, no JSON, no explanations. 27- Output ONLY the timestamp range(s) in the specified format. No other text, no JSON, no explanations.
28- 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.
26 29
27**Examples:** 30**Examples:**
28 31
@@ -162,51 +165,43 @@ def seconds_to_timestamp(total_seconds: float) -> str:
162 165
163def chunk_transcript( 166def chunk_transcript(
164 segments: list[dict], 167 segments: list[dict],
165 trigger_phrase: str = "brought to you by", 168 chunk_duration_seconds: int = 240,
166 max_chunk_duration_minutes: int = 4 169 overlap_seconds: int = 60
167) -> list[list[dict]]: 170) -> list[list[dict]]:
168 """
169 Chunks the transcript segments based on occurrences of a trigger phrase,
170 or a maximum time duration, whichever comes first.
171 Each chunk starts at a trigger phrase and ends just before the next,
172 or when max_chunk_duration_minutes is reached, or at the end of transcript.
173 """
174 chunks = []
175 current_chunk_start_idx = 0
176 max_chunk_duration_seconds = max_chunk_duration_minutes * 60
177
178 if not segments: 171 if not segments:
179 return [] 172 return []
180 173
181 for i, segment in enumerate(segments): 174 chunks = []
182 is_trigger_segment = trigger_phrase.lower() in segment["text"].lower() 175
183 176 # Calculate step size for the window
184 current_chunk_start_time_seconds = timestamp_to_seconds(segments[current_chunk_start_idx]['start_time']) 177 step_seconds = chunk_duration_seconds - overlap_seconds
185 current_segment_end_time_seconds = timestamp_to_seconds(segment['end_time']) 178 if step_seconds <= 0:
186 current_chunk_duration = current_segment_end_time_seconds - current_chunk_start_time_seconds 179 raise ValueError("Overlap must be smaller than chunk duration.")
187 180
188 # Condition 1: New trigger phrase found (and it's not the start of the very first chunk) 181 transcript_end_time = timestamp_to_seconds(segments[-1]['end_time'])
189 # We also need to ensure the trigger is not the first segment of the current chunk, 182
190 # otherwise it would create an empty chunk. 183 window_start_time = 0.0
191 if is_trigger_segment and (i > current_chunk_start_idx): 184 while window_start_time < transcript_end_time:
192 # End the previous chunk *before* this new trigger segment 185 window_end_time = window_start_time + chunk_duration_seconds
193 chunks.append(segments[current_chunk_start_idx:i])
194 current_chunk_start_idx = i # Start a new chunk from this trigger segment
195 continue # Move to the next segment after processing the split
196
197 # Condition 2: Max duration exceeded
198 # This check should happen for any segment after the start of a chunk
199 if current_chunk_duration >= max_chunk_duration_seconds and i > current_chunk_start_idx:
200 # End the current chunk *at* the current segment
201 chunks.append(segments[current_chunk_start_idx:i+1])
202 current_chunk_start_idx = i + 1 # Start a new chunk from the next segment
203 # We don't 'continue' here, as the new chunk needs to start and the loop will increment 'i'
204
205 # Add the last chunk if any segments remain. This covers the case where the transcript
206 # ends without another trigger or exceeding the time limit.
207 if current_chunk_start_idx < len(segments):
208 chunks.append(segments[current_chunk_start_idx:])
209 186
187 current_chunk_segments = []
188 for s in segments:
189 seg_start = timestamp_to_seconds(s['start_time'])
190 seg_end = timestamp_to_seconds(s['end_time'])
191
192 # Add segment if it has any overlap with the current window
193 if seg_start < window_end_time and seg_end > window_start_time:
194 current_chunk_segments.append(s)
195
196 if current_chunk_segments:
197 chunks.append(current_chunk_segments)
198
199 window_start_time += step_seconds
200
201 # Break if the last segment has been fully processed
202 if window_start_time > timestamp_to_seconds(current_chunk_segments[-1]['end_time']):
203 break
204
210 return chunks 205 return chunks
211 206
212@backoff.on_exception(backoff.expo, openai.RateLimitError) 207@backoff.on_exception(backoff.expo, openai.RateLimitError)
@@ -220,13 +215,15 @@ def call_openai_api(chunk_text: str, system_prompt: str) -> str:
220 response = completions_with_backoff( 215 response = completions_with_backoff(
221 #model="Gemma-3-4B_32K", 216 #model="Gemma-3-4B_32K",
222 #model="google/gemini-2.0-flash-exp:free", 217 #model="google/gemini-2.0-flash-exp:free",
223 model="gpt-4.1", 218 #model="gpt-4.1",
219 model=MODEL,
224 messages=[ 220 messages=[
225 {"role": "system", "content": system_prompt}, 221 {"role": "system", "content": system_prompt},
226 {"role": "user", "content": f"Input Transcript:\n{chunk_text}"} 222 {"role": "user", "content": f"Input Transcript:\n{chunk_text}"}
227 ], 223 ],
228 temperature=0.0, 224 temperature=0.0,
229 max_tokens=1024, # realistically it'd be under like 512 225 max_tokens=1024, # realistically it'd be under like 512
226 extra_body={"provider": {"quantizations": ["fp8"], "sort": "price" }},
230 ) 227 )
231 #print(response) 228 #print(response)
232 return response.choices[0].message.content.strip() 229 return response.choices[0].message.content.strip()
@@ -241,7 +238,7 @@ def process_file(input_file: str, output_file: str):
241 # Adjust the trigger phrase if "brought to you by" isn't consistently sufficient 238 # Adjust the trigger phrase if "brought to you by" isn't consistently sufficient
242 # For more complex cases, you might use a list of trigger phrases: 239 # For more complex cases, you might use a list of trigger phrases:
243 # ["brought to you by", "a quick word from our sponsor", "we'll be right back"] 240 # ["brought to you by", "a quick word from our sponsor", "we'll be right back"]
244 transcript_chunks = chunk_transcript(parsed_segments, trigger_phrase="brought to you by") 241 transcript_chunks = chunk_transcript(segments=parsed_segments)
245 242
246 all_ad_timestamps = [] 243 all_ad_timestamps = []
247 244
@@ -250,36 +247,36 @@ def process_file(input_file: str, output_file: str):
250 if not chunk: 247 if not chunk:
251 continue 248 continue
252 249
253 chunk_text_for_api = format_transcript_segment(chunk) 250 user_prompt = format_transcript_segment(chunk)
254
255 # Add the "Your Turn" header before the chunk content
256 user_prompt_content = f"**Your Turn - Input Transcript:**\n{chunk_text_for_api}"
257 251
258 try: 252 try:
259 ad_timestamps_str = call_openai_api(user_prompt_content, SYSTEM_PROMPT_TEMPLATE) 253 ad_timestamps_str = call_openai_api(user_prompt, SYSTEM_PROMPT_TEMPLATE)
260 except Exception as e: 254 except Exception as e:
261 print(f"Error calling OpenAI API: {e}") 255 print(f"Error calling OpenAI API: {e}")
262 continue 256 continue
263 257
264 if ad_timestamps_str: 258 if ad_timestamps_str.startswith("[") or "no output" in ad_timestamps_str.lower() or "no ad" in ad_timestamps_str.lower():
265 # Split by newlines in case the model returns multiple segments 259 ad_timestamps_str = ""
266 found_timestamps = [ts.strip() for ts in ad_timestamps_str.split('\n') if ts.strip()] 260
267 all_ad_timestamps.append({"id": f"{"".join(input_file.split(".")[:-1])}_chunk{i}", f"text": chunk_text_for_api, "target": "\n".join(found_timestamps)}) 261 # Split by newlines in case the model returns multiple segments
262 found_timestamps = [ts.strip() for ts in ad_timestamps_str.split('\n') if ts.strip()]
263 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)})
268 264
269 with open(output_file, "w") as f: 265 with open(output_file, "w") as f:
270 for s in all_ad_timestamps: 266 for s in all_ad_timestamps:
271 f.write(json.dumps(s) + "\n") 267 f.write(json.dumps(s) + "\n")
272 268
273def main(): 269def main():
270 os.makedirs("data/")
274 paths = sys.argv[1:] 271 paths = sys.argv[1:]
275 for p in paths: 272 for p in paths:
276 out = "".join(p.split(".")[:-1]) + ".jsonl" 273 out = "data/" + os.path.splitext(os.path.basename(p))[0] + ".jsonl"
277 if os.path.isfile(out): 274 if os.path.isfile(out):
278 paths.remove(p) 275 paths.remove(p)
279 print(paths) 276 print(paths)
280 277
281 for p in tqdm(paths): 278 for p in tqdm(paths):
282 out = "".join(p.split(".")[:-1]) + ".jsonl" 279 out = "data/" + os.path.splitext(os.path.basename(p))[0] + ".jsonl"
283 process_file(p, out) 280 process_file(p, out)
284 281
285if __name__ == "__main__": 282if __name__ == "__main__":