core/translation.py

"""
Paso 2: Traducción con NLLB-200.
"""

import time
from typing import Dict

import torch

from config import LANG_CODES, TRANSLATION_MODEL_ID
from core.models import ModelManager


def translate_text(
    spanish_text: str,
    manager: ModelManager,
    target_lang: str = "English",
) -> Dict:
    """
    Traduce texto del español al idioma destino.

    Returns:
        dict con 'translated_text', 'processing_time', 'model', 'success'
    """
    start_time = time.time()

    try:
        src_lang = LANG_CODES["Spanish"]
        tgt_lang = LANG_CODES[target_lang]

        manager.translation_tokenizer.src_lang = src_lang

        inputs = manager.translation_tokenizer(
            spanish_text,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=512,
        ).to(manager.device)

        print(f"🌐 Traduciendo a {target_lang}...")
        with torch.no_grad():
            generated_tokens = manager.translation_model.generate(
                **inputs,
                forced_bos_token_id=manager.translation_tokenizer.lang_code_to_id[tgt_lang],
                max_length=512,
                num_beams=5,
                early_stopping=True,
            )

        translated_text = manager.translation_tokenizer.batch_decode(
            generated_tokens, skip_special_tokens=True
        )[0]

        return {
            "translated_text": translated_text.strip(),
            "processing_time": time.time() - start_time,
            "model": TRANSLATION_MODEL_ID,
            "source_lang": "Spanish",
            "target_lang": target_lang,
            "success": True,
        }

    except Exception as e:
        return {
            "translated_text": "",
            "processing_time": time.time() - start_time,
            "model": TRANSLATION_MODEL_ID,
            "success": False,
            "error": str(e),
        }