Adolfo Reyna
131 lines
4.7 KiB
Python
131 lines
4.7 KiB
Python
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)
|