feat: refactor transcription, translation, and distribution into multiple processes

engine_translate.py

@@ -0,0 +1,130 @@
+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
+
+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 run_translation(in_queue, out_queue, args):
+    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):
+            print(f"[Translate] Loading {lang_key} model...")
+            tokenizer = MarianTokenizer.from_pretrained(model_id)
+            model = MarianMTModel.from_pretrained(model_id).to(device)
+            translation_engines[lang_key] = (model, tokenizer)
+
+    paragraph_buffer = []
+    
+    print("[Translate] Ready.")
+    
+    while True:
+        try:
+            item = in_queue.get()
+            if item is None: break
+            
+            text = item.get("original", "").strip()
+            detected_lang = item.get("detected_lang", "en")
+            
+            # Simple paragraph logic: 
+            # If the segment ends with sentence-terminal punctuation, flush the paragraph.
+            paragraph_buffer.append(text)
+            
+            if text.endswith((".", "?", "!")):
+                full_text = " ".join(paragraph_buffer)
+                payload = {"original": full_text, "en": full_text if detected_lang == "en" else None}
+                
+                # If detected language is not English, we'd normally bridge to English first.
+                # For now, let's assume direct translation for simplicity or bridge if needed.
+                # (Refining bridge logic can come later)
+                
+                for lang_key, (model, tokenizer) in translation_engines.items():
+                    chunks = split_text_for_translation(full_text)
+                    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=150)
+                        translated_parts.append(tokenizer.decode(translated_tokens[0], skip_special_tokens=True).strip())
+                    
+                    translated_text = " ".join(part for part in translated_parts if part).strip()
+                    payload[lang_key] = translated_text
+                    print(f"[Translate] {lang_key.upper()}: {translated_text}")
+
+                out_queue.put(payload)
+                paragraph_buffer = []
+                
+        except Exception as e:
+            print(f"[Translate] Error: {e}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-es", action="store_true")
+    parser.add_argument("-fr", action="store_true")
+    parser.add_argument("-ar", action="store_true")
+    args = parser.parse_args()
+    
+    # Dummy queues for testing
+    in_q = multiprocessing.Queue()
+    out_q = multiprocessing.Queue()
+    run_translation(in_q, out_q, args)