Adolfo Reyna
507 lines
20 KiB
Python
507 lines
20 KiB
Python
import sys
|
|
from unittest.mock import MagicMock
|
|
|
|
# MANDATORY: Mock lzma BEFORE any other imports
|
|
try:
|
|
import lzma
|
|
except ImportError:
|
|
mock_lzma = MagicMock()
|
|
mock_lzma.FORMAT_XZ, mock_lzma.FORMAT_ALONE, mock_lzma.FORMAT_RAW = 1, 2, 3
|
|
mock_lzma.CHECK_NONE, mock_lzma.CHECK_CRC32, mock_lzma.CHECK_CRC64, mock_lzma.CHECK_SHA256 = 0, 1, 4, 10
|
|
sys.modules["_lzma"] = MagicMock()
|
|
sys.modules["lzma"] = mock_lzma
|
|
|
|
import time
|
|
import argparse
|
|
import multiprocessing
|
|
import os
|
|
import re
|
|
import json
|
|
import requests
|
|
from collections import deque
|
|
from datetime import datetime, timezone
|
|
|
|
|
|
def build_paragraph_messages(previous_refined_context, previous_source_text, new_source_text):
|
|
system_prompt = """You are a careful live transcript editor working from noisy translated source text.
|
|
|
|
Goal:
|
|
Produce the most faithful readable English update.
|
|
|
|
Decision rules:
|
|
1. NEW SOURCE TEXT is the primary evidence and should dominate the output.
|
|
2. Use PREVIOUS SOURCE TEXT and PREVIOUS REFINED CONTEXT only when they clearly help resolve or continue the new material.
|
|
3. If the new material starts a fresh thought, output only the new material.
|
|
4. Remove exact or near-exact repetition unless the repetition is clearly intentional rhetoric.
|
|
5. Do not add explanations, disclaimers, or meta commentary.
|
|
6. Prefer conservative wording over guessed meaning.
|
|
7. Return only the revised transcript text."""
|
|
|
|
prompt = f"""Edit this transcript update.
|
|
|
|
[PREVIOUS REFINED CONTEXT]
|
|
{previous_refined_context}
|
|
|
|
[PREVIOUS SOURCE TEXT]
|
|
{previous_source_text}
|
|
|
|
[NEW SOURCE TEXT]
|
|
{new_source_text}
|
|
"""
|
|
return [
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": prompt},
|
|
], prompt
|
|
|
|
|
|
def build_line_messages(prev1, prev2, corrected):
|
|
system_prompt = """You are a real-time English caption corrector for live speech.
|
|
|
|
Task:
|
|
Clean only the current caption line.
|
|
|
|
Hard rules:
|
|
1. Preserve meaning exactly. Never replace current content with prior context.
|
|
2. Remove disfluencies and false starts.
|
|
3. Fix punctuation, casing, and obvious STT typos.
|
|
4. If uncertain, return the original line unchanged.
|
|
5. Output only one corrected English line."""
|
|
|
|
prompt = (
|
|
"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:"
|
|
)
|
|
return [
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": prompt},
|
|
], prompt
|
|
|
|
|
|
def append_llm_request_log(log_path, entry):
|
|
if not log_path:
|
|
return
|
|
log_dir = os.path.dirname(log_path)
|
|
if log_dir:
|
|
os.makedirs(log_dir, exist_ok=True)
|
|
with open(log_path, "a", encoding="utf-8") as handle:
|
|
json.dump(entry, handle, ensure_ascii=False)
|
|
handle.write("\n")
|
|
|
|
|
|
def load_llm_request_log_entry(log_path, entry_index):
|
|
if not log_path or not os.path.exists(log_path):
|
|
raise FileNotFoundError(f"LLM request log not found: {log_path}")
|
|
|
|
with open(log_path, "r", encoding="utf-8") as handle:
|
|
entries = [json.loads(line) for line in handle if line.strip()]
|
|
|
|
if not entries:
|
|
raise ValueError(f"LLM request log is empty: {log_path}")
|
|
|
|
if entry_index is None or entry_index == -1:
|
|
return entries[-1]
|
|
|
|
if entry_index < 0:
|
|
entry_index = len(entries) + entry_index
|
|
|
|
if entry_index < 0 or entry_index >= len(entries):
|
|
raise IndexError(f"LLM request log index {entry_index} is out of range for {len(entries)} entries")
|
|
|
|
return entries[entry_index]
|
|
|
|
|
|
def run_llm_prompt_test(args):
|
|
from dotenv import load_dotenv
|
|
load_dotenv()
|
|
|
|
api_key = os.environ.get("OPENAI_API_KEY")
|
|
log_path = getattr(args, "llm_request_log_path", "logs/llm_requests.jsonl")
|
|
|
|
def call_openai(model, messages):
|
|
response = requests.post(
|
|
"https://api.openai.com/v1/chat/completions",
|
|
json={"model": model, "messages": messages},
|
|
headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
|
|
timeout=getattr(args, "post_correct_llm_timeout", 8.0),
|
|
)
|
|
if response.status_code >= 400:
|
|
try:
|
|
payload = response.json()
|
|
detail = json.dumps(payload, ensure_ascii=True)
|
|
except Exception:
|
|
detail = response.text.strip()
|
|
raise RuntimeError(f"OpenAI API error {response.status_code}: {detail}")
|
|
return response.json()["choices"][0]["message"]["content"].strip()
|
|
|
|
def call_ollama(model, prompt, temperature, system=None):
|
|
response = requests.post(
|
|
getattr(args, "post_correct_ollama_url", "http://127.0.0.1:11434/api/generate"),
|
|
json={
|
|
"model": model,
|
|
"prompt": prompt,
|
|
"system": system or "",
|
|
"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 send_test_request(mode, messages, prompt, temperature, metadata, model=None):
|
|
request_model = model or args.post_correct_model
|
|
is_openai = request_model.startswith("gpt-")
|
|
entry = {
|
|
"timestamp": datetime.now(timezone.utc).isoformat(),
|
|
"mode": mode,
|
|
"provider": "openai" if is_openai else "ollama",
|
|
"model": request_model,
|
|
"temperature": temperature,
|
|
"messages": messages,
|
|
"metadata": metadata,
|
|
}
|
|
started_at = time.time()
|
|
if is_openai:
|
|
if not api_key:
|
|
raise RuntimeError("OPENAI_API_KEY is not set")
|
|
try:
|
|
response_text = call_openai(request_model, messages)
|
|
entry["response"] = {"status": "ok", "content": response_text}
|
|
return response_text
|
|
except Exception as exc:
|
|
entry["response"] = {"status": "error", "error": str(exc)}
|
|
raise
|
|
finally:
|
|
entry["duration_ms"] = round((time.time() - started_at) * 1000, 2)
|
|
append_llm_request_log(log_path, entry)
|
|
system = ""
|
|
for message in messages:
|
|
if message.get("role") == "system":
|
|
system = message.get("content", "")
|
|
break
|
|
try:
|
|
response_text = call_ollama(request_model, prompt, temperature, system=system)
|
|
entry["response"] = {"status": "ok", "content": response_text}
|
|
return response_text
|
|
except Exception as exc:
|
|
entry["response"] = {"status": "error", "error": str(exc)}
|
|
raise
|
|
finally:
|
|
entry["duration_ms"] = round((time.time() - started_at) * 1000, 2)
|
|
append_llm_request_log(log_path, entry)
|
|
|
|
if getattr(args, "llm_test_from_log", None) is not None:
|
|
logged_entry = load_llm_request_log_entry(log_path, args.llm_test_from_log)
|
|
logged_messages = logged_entry.get("messages")
|
|
if not logged_messages:
|
|
raise ValueError("Selected log entry does not contain messages")
|
|
logged_prompt = next((m.get("content", "") for m in logged_messages if m.get("role") == "user"), "")
|
|
logged_model = logged_entry.get("model", args.post_correct_model)
|
|
replay_model = logged_model if getattr(args, "llm_test_use_logged_model", False) else args.post_correct_model
|
|
response = send_test_request(
|
|
logged_entry.get("mode", "replay"),
|
|
logged_messages,
|
|
logged_prompt,
|
|
logged_entry.get("temperature", 0.1),
|
|
{
|
|
"replay_source_log_path": log_path,
|
|
"replay_source_index": args.llm_test_from_log,
|
|
"replay_source_timestamp": logged_entry.get("timestamp"),
|
|
"replay_source_model": logged_model,
|
|
},
|
|
model=replay_model,
|
|
)
|
|
print("[LLM TEST] Replayed request from log:")
|
|
print(f"source_index={args.llm_test_from_log} source_model={logged_model} replay_model={replay_model}")
|
|
print(response)
|
|
|
|
if getattr(args, "llm_test_line", None):
|
|
messages, prompt = build_line_messages(
|
|
getattr(args, "llm_test_prev1", ""),
|
|
getattr(args, "llm_test_prev2", ""),
|
|
args.llm_test_line,
|
|
)
|
|
response = send_test_request(
|
|
"line",
|
|
messages,
|
|
prompt,
|
|
0.1,
|
|
{
|
|
"prev1": getattr(args, "llm_test_prev1", ""),
|
|
"prev2": getattr(args, "llm_test_prev2", ""),
|
|
"input": args.llm_test_line,
|
|
},
|
|
)
|
|
print("[LLM TEST] Line response:")
|
|
print(response)
|
|
|
|
if getattr(args, "llm_test_segments", None):
|
|
messages, prompt = build_paragraph_messages(
|
|
getattr(args, "llm_test_context", ""),
|
|
args.llm_test_segments,
|
|
)
|
|
response = send_test_request(
|
|
"paragraph",
|
|
messages,
|
|
prompt,
|
|
0.1,
|
|
{
|
|
"context": getattr(args, "llm_test_context", ""),
|
|
"segments": args.llm_test_segments,
|
|
},
|
|
)
|
|
print("[LLM TEST] Paragraph response:")
|
|
print(response)
|
|
|
|
def run_llm_processor(in_queue, out_queue, args):
|
|
from dotenv import load_dotenv
|
|
load_dotenv()
|
|
|
|
print("[LLM] Starting Rolling Paragraph Engine...")
|
|
|
|
api_key = os.environ.get("OPENAI_API_KEY")
|
|
is_openai = args.post_correct_model.startswith("gpt-")
|
|
llm_state = {"line_warned": False, "paragraph_warned": False}
|
|
log_path = getattr(args, "llm_request_log_path", "logs/llm_requests.jsonl")
|
|
|
|
def check_ollama_health():
|
|
try:
|
|
response = requests.get("http://127.0.0.1:11434/api/tags", timeout=2.5)
|
|
response.raise_for_status()
|
|
data = response.json()
|
|
models = data.get("models", [])
|
|
available = {model.get("name") for model in models if model.get("name")}
|
|
if args.post_correct_model not in available and f"{args.post_correct_model}:latest" not in available:
|
|
print(f"[LLM] Configured Ollama model '{args.post_correct_model}' is not installed.")
|
|
return True
|
|
except Exception as exc:
|
|
print(f"[LLM] Ollama health check failed ({exc}). Local LLM features may fall back.")
|
|
return False
|
|
|
|
def warmup_ollama_model():
|
|
try:
|
|
response = requests.post(
|
|
getattr(args, "post_correct_ollama_url", "http://127.0.0.1:11434/api/generate"),
|
|
json={
|
|
"model": args.post_correct_model,
|
|
"prompt": "Return exactly: ok",
|
|
"system": "You are a concise assistant.",
|
|
"stream": False,
|
|
"options": {"temperature": 0.0},
|
|
"keep_alive": getattr(args, "post_correct_keep_alive", "30m"),
|
|
},
|
|
timeout=getattr(args, "post_correct_warmup_timeout", 20.0),
|
|
)
|
|
response.raise_for_status()
|
|
body = response.json()
|
|
if body.get("response", "").strip():
|
|
total_s = body.get("total_duration", 0) / 1_000_000_000
|
|
load_s = body.get("load_duration", 0) / 1_000_000_000
|
|
print(f"[LLM] Ollama warmup OK for {args.post_correct_model} (load={load_s:.2f}s total={total_s:.2f}s).")
|
|
return True
|
|
except Exception as exc:
|
|
print(f"[LLM] Ollama warmup failed for {args.post_correct_model} ({exc}).")
|
|
return False
|
|
|
|
if not is_openai:
|
|
check_ollama_health()
|
|
if getattr(args, "post_correct_llm", False) or getattr(args, "llm_paragraph", False):
|
|
warmup_ollama_model()
|
|
|
|
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},
|
|
headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
|
|
timeout=getattr(args, "post_correct_llm_timeout", 8.0),
|
|
)
|
|
if response.status_code >= 400:
|
|
try:
|
|
payload = response.json()
|
|
detail = json.dumps(payload, ensure_ascii=True)
|
|
except Exception:
|
|
detail = response.text.strip()
|
|
raise RuntimeError(f"OpenAI API error {response.status_code}: {detail}")
|
|
return response.json()["choices"][0]["message"]["content"].strip()
|
|
|
|
def call_ollama(prompt, temperature, system=None):
|
|
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,
|
|
"system": system or "",
|
|
"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):
|
|
entry = {
|
|
"timestamp": datetime.now(timezone.utc).isoformat(),
|
|
"mode": "paragraph" if warn_key == "paragraph_warned" else "line",
|
|
"provider": "openai" if is_openai else "ollama",
|
|
"model": args.post_correct_model,
|
|
"temperature": temperature,
|
|
"messages": messages,
|
|
}
|
|
started_at = time.time()
|
|
try:
|
|
if is_openai:
|
|
if not api_key:
|
|
raise RuntimeError("OPENAI_API_KEY is not set")
|
|
response_text = call_openai(messages)
|
|
entry["response"] = {"status": "ok", "content": response_text}
|
|
return response_text
|
|
system = ""
|
|
if messages:
|
|
for message in messages:
|
|
if message.get("role") == "system":
|
|
system = message.get("content", "")
|
|
break
|
|
response_text = call_ollama(prompt, temperature, system=system)
|
|
entry["response"] = {"status": "ok", "content": response_text}
|
|
return response_text
|
|
except Exception as exc:
|
|
entry["response"] = {"status": "error", "error": str(exc)}
|
|
if not llm_state[warn_key]:
|
|
label = "Paragraph LLM" if warn_key == "paragraph_warned" else "Line LLM"
|
|
print(f"[LLM] {label} unavailable ({exc}). Falling back to deterministic text.")
|
|
llm_state[warn_key] = True
|
|
return ""
|
|
finally:
|
|
entry["duration_ms"] = round((time.time() - started_at) * 1000, 2)
|
|
append_llm_request_log(log_path, entry)
|
|
|
|
# State: This is the ONLY text we send to the LLM as context
|
|
active_context = ""
|
|
previous_source_window = ""
|
|
recent_lines = deque(maxlen=2)
|
|
|
|
def call_llm_rolling_refine(new_segments_str, context, previous_source_text):
|
|
messages, prompt = build_paragraph_messages(context, previous_source_text, new_segments_str)
|
|
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 ""
|
|
messages, prompt = build_line_messages(prev1, prev2, corrected)
|
|
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()
|
|
|
|
while True:
|
|
try:
|
|
item = in_queue.get()
|
|
if item is None: break
|
|
if "draft" in item:
|
|
out_queue.put(item)
|
|
continue
|
|
|
|
raw_text = item.get("original", "").strip()
|
|
bridge_text = item.get("en_bridge", raw_text).strip()
|
|
if not raw_text: continue
|
|
|
|
corrected_text = post_correct_line(bridge_text)
|
|
|
|
# Hallucination check
|
|
words = corrected_text.lower().split()
|
|
if len(words) > 10 and len(set(words)) < 3: continue
|
|
|
|
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 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
|
|
paragraph_from_llm = False
|
|
if args.llm_paragraph and should_refine:
|
|
new_batch = " ".join(pending_buffer)
|
|
result = call_llm_rolling_refine(new_batch, active_context, previous_source_window)
|
|
|
|
if result:
|
|
# Logic to handle paragraph breaks
|
|
if "\n\n" in result:
|
|
# Split by double newline
|
|
parts = result.split("\n\n")
|
|
# The last part is the new "Active Context"
|
|
active_context = parts[-1].strip()
|
|
# The whole result is sent to the user (contains the break)
|
|
structured_paragraph = result
|
|
else:
|
|
# No break, just update the context
|
|
active_context = result
|
|
structured_paragraph = result
|
|
paragraph_from_llm = True
|
|
else:
|
|
structured_paragraph = new_batch
|
|
active_context = new_batch
|
|
previous_source_window = new_batch
|
|
|
|
pending_buffer = []
|
|
last_llm_call_time = time.time()
|
|
|
|
out_queue.put({
|
|
"raw": raw_text,
|
|
"corrected": corrected_text,
|
|
"paragraph": structured_paragraph,
|
|
"en_bridge": bridge_text,
|
|
"used_llm_line": bool(getattr(args, "post_correct_llm", False)),
|
|
"used_llm_paragraph": paragraph_from_llm,
|
|
"detected_lang": item.get("detected_lang"),
|
|
"speaker": item.get("speaker"),
|
|
"ts": item.get("ts", time.time())
|
|
})
|
|
|
|
except Exception as e:
|
|
print(f"[LLM] Error: {e}")
|
|
|
|
if __name__ == "__main__":
|
|
run_llm_processor(multiprocessing.Queue(), multiprocessing.Queue(), argparse.Namespace())
|