Fix main_v2 bridge and pipeline regressions

This commit is contained in:
Adolfo Reyna
2026-03-16 21:15:15 -04:00
parent ffeb996d7e
commit f79c4f7310
7 changed files with 190 additions and 51 deletions
+119 -27
View File
@@ -15,6 +15,7 @@ import time
import argparse
import multiprocessing
import os
import re
import requests
from collections import deque
@@ -26,18 +27,74 @@ def run_llm_processor(in_queue, out_queue, args):
api_key = os.environ.get("OPENAI_API_KEY")
is_openai = args.post_correct_model.startswith("gpt-")
target_lang = "English"
if args.lang and args.lang != 'en': target_lang = args.lang
llm_state = {"line_warned": False, "paragraph_warned": False}
def normalize_english_caption(text):
normalized = " ".join((text or "").split()).strip()
if normalized and normalized[0].isalpha():
normalized = normalized[0].upper() + normalized[1:]
return normalized
def apply_rule_based_post_correction(text):
corrected = normalize_english_caption(text)
corrected = re.sub(r"\s+", " ", corrected).strip()
corrected = re.sub(r"\s+([,.;!?])", r"\1", corrected)
corrected = re.sub(r"([,.;!?]){2,}", r"\1", corrected)
corrected = re.sub(r"\b(uh+|um+|erm+|ah+|hmm+)\b", "", corrected, flags=re.IGNORECASE)
corrected = re.sub(r"\s+", " ", corrected).strip()
corrected = re.sub(r"\b(\w+)\s+\1\b", r"\1", corrected, flags=re.IGNORECASE)
return normalize_english_caption(corrected)
def word_overlap_ratio(source_text, candidate_text):
src_tokens = set(re.findall(r"[a-z0-9']+", source_text.lower()))
cand_tokens = set(re.findall(r"[a-z0-9']+", candidate_text.lower()))
if not src_tokens:
return 1.0
return len(src_tokens & cand_tokens) / len(src_tokens)
def call_openai(messages):
response = requests.post(
"https://api.openai.com/v1/chat/completions",
json={"model": args.post_correct_model, "messages": messages, "temperature": 0.1},
headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
timeout=getattr(args, "post_correct_llm_timeout", 8.0),
)
response.raise_for_status()
return response.json()["choices"][0]["message"]["content"].strip()
def call_ollama(prompt, temperature):
response = requests.post(
getattr(args, "post_correct_ollama_url", "http://127.0.0.1:11434/api/generate"),
json={
"model": args.post_correct_model,
"prompt": prompt,
"stream": False,
"options": {"temperature": temperature},
"keep_alive": getattr(args, "post_correct_keep_alive", "30m"),
},
timeout=getattr(args, "post_correct_llm_timeout", 8.0),
)
response.raise_for_status()
return response.json().get("response", "").strip()
def call_llm(messages, prompt, temperature, warn_key):
try:
if is_openai:
if not api_key:
raise RuntimeError("OPENAI_API_KEY is not set")
return call_openai(messages)
return call_ollama(prompt, temperature)
except Exception as exc:
if not llm_state[warn_key]:
print(f"[LLM] {warn_key.replace('_', ' ').title()} unavailable ({exc}). Falling back to deterministic text.")
llm_state[warn_key] = True
return ""
# State: This is the ONLY text we send to the LLM as context
active_context = ""
recent_lines = deque(maxlen=2)
def call_llm_rolling_refine(new_segments_str, context):
if not is_openai: return ""
url = "https://api.openai.com/v1/chat/completions"
headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
prompt = f"""Task: Refine the following live transcription stream into clean, professional paragraphs.
[PREVIOUS WORKING CONTEXT]
@@ -50,21 +107,50 @@ def run_llm_processor(in_queue, out_queue, args):
1. INTEGRATE: Polished and merge the new segments into the flow of the 'PREVIOUS WORKING CONTEXT'.
2. CONSOLIDATE: Remove redundant repetitions and translator echoes.
3. ORGANIZE: Use a double newline (\\n\\n) to start a new paragraph when a topic changes or the current one is complete.
4. TARGET LANGUAGE: Output ONLY in {target_lang}.
4. TARGET LANGUAGE: Output ONLY in English.
5. OUTPUT: Provide ONLY the refined, consolidated text. Do not explain anything.
"""
payload = {
"model": args.post_correct_model,
"messages": [{"role": "system", "content": "You are a live transcript editor."},
{"role": "user", "content": prompt}],
"temperature": 0.1
}
try:
resp = requests.post(url, json=payload, headers=headers, timeout=20)
if resp.status_code == 200:
return resp.json()["choices"][0]["message"]["content"].strip()
return ""
except: return ""
messages = [
{"role": "system", "content": "You are a live transcript editor."},
{"role": "user", "content": prompt},
]
return call_llm(messages, prompt, 0.1, "paragraph_warned")
def post_correct_line(text):
if not getattr(args, "post_correct", False):
return normalize_english_caption(text)
corrected = apply_rule_based_post_correction(text)
if not getattr(args, "post_correct_llm", False):
return corrected
prev2 = recent_lines[-2] if len(recent_lines) >= 2 else ""
prev1 = recent_lines[-1] if len(recent_lines) >= 1 else ""
prompt = (
"You are a real-time English caption corrector for live speech.\n"
"Task:\n"
"Clean ONLY the current caption line.\n"
"Hard rules:\n"
"1. Preserve meaning exactly. Never replace current content with prior context.\n"
"2. Remove disfluencies and false starts.\n"
"3. Fix punctuation, casing, and obvious STT typos.\n"
"4. If uncertain, return the original line unchanged.\n"
"5. Output ONLY one corrected English line.\n"
"Context (reference only):\n"
f"Previous line 1: {prev1}\n"
f"Previous line 2: {prev2}\n"
"Current line to correct:\n"
f"{corrected}\n"
"Corrected:"
)
messages = [
{"role": "system", "content": "You are a real-time English caption corrector."},
{"role": "user", "content": prompt},
]
candidate = normalize_english_caption(call_llm(messages, prompt, 0.1, "line_warned"))
if candidate and word_overlap_ratio(corrected, candidate) >= getattr(args, "post_correct_min_overlap", 0.45):
return candidate
return corrected
pending_buffer = []
last_llm_call_time = time.time()
@@ -81,15 +167,18 @@ def run_llm_processor(in_queue, out_queue, args):
bridge_text = item.get("en_bridge", raw_text).strip()
if not raw_text: continue
corrected_text = post_correct_line(bridge_text)
# Hallucination check
words = bridge_text.lower().split()
words = corrected_text.lower().split()
if len(words) > 10 and len(set(words)) < 3: continue
pending_buffer.append(bridge_text)
pending_buffer.append(corrected_text)
recent_lines.append(corrected_text)
should_refine = False
if len(pending_buffer) >= 5: should_refine = True
elif len(pending_buffer) >= 2 and bridge_text.endswith((".", "?", "!")): should_refine = True
elif len(pending_buffer) >= 2 and corrected_text.endswith((".", "?", "!")): should_refine = True
elif len(pending_buffer) >= 1 and (time.time() - last_llm_call_time) > 15: should_refine = True
structured_paragraph = None
@@ -110,13 +199,16 @@ def run_llm_processor(in_queue, out_queue, args):
# No break, just update the context
active_context = result
structured_paragraph = result
pending_buffer = []
last_llm_call_time = time.time()
else:
structured_paragraph = new_batch
active_context = new_batch
pending_buffer = []
last_llm_call_time = time.time()
out_queue.put({
"raw": raw_text,
"corrected": bridge_text,
"corrected": corrected_text,
"paragraph": structured_paragraph,
"detected_lang": item.get("detected_lang"),
"speaker": item.get("speaker"),