feat: implement accumulative LLM paragraph engine and harden 4-process pipeline
This commit is contained in:
Adolfo Reyna
+71
-104
@@ -1,79 +1,57 @@
|
||||
import sys
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
# Comprehensive workaround for missing _lzma in some Python builds
|
||||
try:
|
||||
import lzma
|
||||
except ImportError:
|
||||
mock_lzma = MagicMock()
|
||||
mock_lzma.FORMAT_XZ = 1
|
||||
mock_lzma.FORMAT_ALONE = 2
|
||||
mock_lzma.FORMAT_RAW = 3
|
||||
mock_lzma.CHECK_NONE = 0
|
||||
mock_lzma.CHECK_CRC32 = 1
|
||||
mock_lzma.CHECK_CRC64 = 4
|
||||
mock_lzma.CHECK_SHA256 = 10
|
||||
sys.modules["_lzma"] = MagicMock()
|
||||
sys.modules["lzma"] = mock_lzma
|
||||
|
||||
import torch
|
||||
from transformers import MarianMTModel, MarianTokenizer
|
||||
import re
|
||||
import argparse
|
||||
import multiprocessing
|
||||
import time
|
||||
import requests
|
||||
|
||||
TARGET_LANGS = {
|
||||
"es": "Helsinki-NLP/opus-mt-en-es",
|
||||
"fr": "Helsinki-NLP/opus-mt-en-fr",
|
||||
"ar": "Helsinki-NLP/opus-mt-en-ar"
|
||||
}
|
||||
|
||||
def split_text_for_translation(text, max_chars=250):
|
||||
normalized = " ".join(text.split()).strip()
|
||||
if not normalized or len(normalized) <= max_chars:
|
||||
return [normalized] if normalized else []
|
||||
chunks = []
|
||||
current = ""
|
||||
sentences = [s for s in re.split(r"(?<=[.!?])\s+", normalized) if s]
|
||||
for sentence in sentences:
|
||||
if len(sentence) > max_chars:
|
||||
words = sentence.split()
|
||||
word_chunk = ""
|
||||
for word in words:
|
||||
candidate = f"{word_chunk} {word}".strip()
|
||||
if len(candidate) <= max_chars: word_chunk = candidate
|
||||
else:
|
||||
if word_chunk: chunks.append(word_chunk)
|
||||
word_chunk = word
|
||||
if word_chunk: chunks.append(word_chunk)
|
||||
continue
|
||||
candidate = f"{current} {sentence}".strip()
|
||||
if candidate and len(candidate) <= max_chars: current = candidate
|
||||
else:
|
||||
if current: chunks.append(current)
|
||||
current = sentence
|
||||
if current: chunks.append(current)
|
||||
return chunks
|
||||
|
||||
def call_ollama_generate(args, prompt):
|
||||
payload = {"model": args.post_correct_model, "prompt": prompt, "stream": False}
|
||||
try:
|
||||
resp = requests.post("http://127.0.0.1:11434/api/generate", json=payload, timeout=10)
|
||||
return resp.json().get("response", "").strip()
|
||||
except: return ""
|
||||
|
||||
def structure_paragraph_with_llm(text, args):
|
||||
prompt = f"Format the following transcribed segments into a coherent paragraph. Output ONLY the paragraph:\n{text}"
|
||||
return call_ollama_generate(args, prompt) or text
|
||||
|
||||
def post_correct_with_llm(text, args):
|
||||
prompt = f"Correct grammar and typos in this caption. Output ONLY the corrected text:\n{text}"
|
||||
return call_ollama_generate(args, prompt) or text
|
||||
import re
|
||||
|
||||
def run_translation(in_queue, out_queue, args):
|
||||
import sys
|
||||
from unittest.mock import MagicMock
|
||||
# Comprehensive workaround for missing _lzma in some Python builds
|
||||
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 torch
|
||||
from transformers import MarianMTModel, MarianTokenizer
|
||||
|
||||
TARGET_LANGS = {
|
||||
"es": "Helsinki-NLP/opus-mt-en-es",
|
||||
"fr": "Helsinki-NLP/opus-mt-en-fr",
|
||||
"ar": "Helsinki-NLP/opus-mt-en-ar"
|
||||
}
|
||||
|
||||
def split_text_for_translation(text, max_chars=250):
|
||||
normalized = " ".join(text.split()).strip()
|
||||
if not normalized or len(normalized) <= max_chars: return [normalized] if normalized else []
|
||||
chunks, current = [], ""
|
||||
sentences = [s for s in re.split(r"(?<=[.!?])\s+", normalized) if s]
|
||||
for sentence in sentences:
|
||||
if len(sentence) > max_chars:
|
||||
words, word_chunk = sentence.split(), ""
|
||||
for word in words:
|
||||
candidate = f"{word_chunk} {word}".strip()
|
||||
if len(candidate) <= max_chars: word_chunk = candidate
|
||||
else:
|
||||
if word_chunk: chunks.append(word_chunk)
|
||||
word_chunk = word
|
||||
if word_chunk: chunks.append(word_chunk)
|
||||
continue
|
||||
candidate = f"{current} {sentence}".strip()
|
||||
if candidate and len(candidate) <= max_chars: current = candidate
|
||||
else:
|
||||
if current: chunks.append(current)
|
||||
current = sentence
|
||||
if current: chunks.append(current)
|
||||
return chunks
|
||||
|
||||
device = "mps" if torch.backends.mps.is_available() else "cpu"
|
||||
print(f"[Translate] Loading translation models on {device}...")
|
||||
|
||||
translation_engines = {}
|
||||
for lang_key, model_id in TARGET_LANGS.items():
|
||||
if getattr(args, lang_key, False):
|
||||
@@ -83,51 +61,40 @@ def run_translation(in_queue, out_queue, args):
|
||||
MarianTokenizer.from_pretrained(model_id)
|
||||
)
|
||||
|
||||
paragraph_buffer = []
|
||||
|
||||
while True:
|
||||
try:
|
||||
item = in_queue.get()
|
||||
if item is None: break
|
||||
|
||||
# Draft support
|
||||
if "draft" in item:
|
||||
out_queue.put({"draft": item["draft"]})
|
||||
out_queue.put(item)
|
||||
continue
|
||||
|
||||
text = item.get("original", "").strip()
|
||||
detected_lang = item.get("detected_lang", "en")
|
||||
|
||||
if args.post_correct and args.post_correct_llm:
|
||||
text = post_correct_with_llm(text, args)
|
||||
raw_text, corrected_text, paragraph_text = item.get("raw"), item.get("corrected"), item.get("paragraph")
|
||||
payload = {
|
||||
"original": raw_text,
|
||||
"corrected": corrected_text,
|
||||
"paragraph": paragraph_text,
|
||||
"speaker": item.get("speaker"),
|
||||
"ts": item.get("ts", time.time())
|
||||
}
|
||||
if args.en: payload["en"] = corrected_text or raw_text
|
||||
|
||||
paragraph_buffer.append(text)
|
||||
|
||||
# Paragraph logic
|
||||
if not args.llm_paragraph or text.endswith((".", "?", "!")):
|
||||
full_text = " ".join(paragraph_buffer)
|
||||
text_to_translate = paragraph_text or corrected_text or raw_text
|
||||
if text_to_translate:
|
||||
if translation_engines:
|
||||
for lang_key, (model, tokenizer) in translation_engines.items():
|
||||
chunks = split_text_for_translation(text_to_translate, args.mt_max_chars)
|
||||
translated_parts = []
|
||||
for chunk in chunks:
|
||||
inputs = tokenizer(chunk, return_tensors="pt", padding=True).to(device)
|
||||
with torch.no_grad():
|
||||
translated_tokens = model.generate(**inputs, max_new_tokens=args.mt_max_new_tokens, num_beams=args.mt_num_beams)
|
||||
translated_parts.append(tokenizer.decode(translated_tokens[0], skip_special_tokens=True).strip())
|
||||
payload[lang_key] = " ".join(translated_parts)
|
||||
|
||||
if args.llm_paragraph:
|
||||
full_text = structure_paragraph_with_llm(full_text, args)
|
||||
|
||||
payload = {"original": full_text, "ts": item.get("ts", time.time())}
|
||||
if args.en: payload["en"] = full_text
|
||||
|
||||
for lang_key, (model, tokenizer) in translation_engines.items():
|
||||
chunks = split_text_for_translation(full_text, args.mt_max_chars)
|
||||
translated_parts = []
|
||||
for chunk in chunks:
|
||||
inputs = tokenizer(chunk, return_tensors="pt", padding=True).to(device)
|
||||
with torch.no_grad():
|
||||
translated_tokens = model.generate(**inputs, max_new_tokens=args.mt_max_new_tokens, num_beams=args.mt_num_beams)
|
||||
translated_parts.append(tokenizer.decode(translated_tokens[0], skip_special_tokens=True).strip())
|
||||
payload[lang_key] = " ".join(translated_parts)
|
||||
|
||||
# ALWAYS put in queue, even if no translations were done
|
||||
out_queue.put(payload)
|
||||
paragraph_buffer = []
|
||||
|
||||
except Exception as e: print(f"[Translate] Error: {e}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
import multiprocessing
|
||||
run_translation(multiprocessing.Queue(), multiprocessing.Queue(), argparse.Namespace())
|
||||
|
||||
Reference in New Issue
Block a user