Adolfo Reyna
131 lines
4.8 KiB
Python
131 lines
4.8 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 requests
|
|
from collections import deque
|
|
|
|
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-")
|
|
target_lang = "English"
|
|
if args.lang and args.lang != 'en': target_lang = args.lang
|
|
|
|
# State: This is the ONLY text we send to the LLM as context
|
|
active_context = ""
|
|
|
|
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]
|
|
"{context}"
|
|
|
|
[NEW RAW ASR SEGMENTS]
|
|
"{new_segments_str}"
|
|
|
|
[INSTRUCTIONS]
|
|
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}.
|
|
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 ""
|
|
|
|
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
|
|
|
|
# Hallucination check
|
|
words = bridge_text.lower().split()
|
|
if len(words) > 10 and len(set(words)) < 3: continue
|
|
|
|
pending_buffer.append(bridge_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) >= 1 and (time.time() - last_llm_call_time) > 15: should_refine = True
|
|
|
|
structured_paragraph = None
|
|
if args.llm_paragraph and should_refine:
|
|
new_batch = " ".join(pending_buffer)
|
|
result = call_llm_rolling_refine(new_batch, active_context)
|
|
|
|
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
|
|
|
|
pending_buffer = []
|
|
last_llm_call_time = time.time()
|
|
|
|
out_queue.put({
|
|
"raw": raw_text,
|
|
"corrected": bridge_text,
|
|
"paragraph": structured_paragraph,
|
|
"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())
|