Spaces:
Sleeping
Sleeping
File size: 4,717 Bytes
37b0787 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | from __future__ import annotations
import logging
import random
import time
from typing import Any
from deep_translator import GoogleTranslator
logger = logging.getLogger(__name__)
# ISO 639-1 codes supported by Google Translate
SUPPORTED_LANGUAGES = {
"hi": "hindi", "ta": "tamil", "te": "telugu", "mr": "marathi",
"bn": "bengali", "kn": "kannada", "ml": "malayalam", "gu": "gujarati",
"pa": "punjabi", "or": "odia", "as": "assamese",
"ur": "urdu", "sd": "sindhi", "ks": "kashmiri", "ne": "nepali",
}
class TranslationPipeline:
def __init__(self) -> None:
self._translator: GoogleTranslator | None = None
def _get_translator(self) -> GoogleTranslator:
if self._translator is None:
self._translator = GoogleTranslator(source="auto", target="en")
return self._translator
def translate_to_english(
self, text: str, source_lang: str, is_batch_call: bool = False,
) -> dict[str, str | float | bool]: # type: ignore[return-value]
if source_lang == "en":
return {
"original": text,
"translated": text,
"confidence": 1.0,
"model_used": "none",
"translation_fallback": False,
}
if not is_batch_call:
time.sleep(random.uniform(0.1, 0.5))
try:
translator = self._get_translator()
translated_text = translator.translate(text)
return {
"original": text,
"translated": translated_text or text,
"confidence": 0.85,
"model_used": "GoogleTranslate",
"translation_fallback": False,
}
except Exception as e:
logger.warning(
"Google Translation failed for %s, trying MBART: %s",
source_lang, e
)
try:
from src.core.config import get_app_config
cfg = get_app_config()
mbart_model_name = (
cfg.get("translation", {})
.get("fallback", "facebook/mbart-large-50-many-to-many-mmt")
)
except Exception:
mbart_model_name = "facebook/mbart-large-50-many-to-many-mmt"
try:
from transformers import MBart50TokenizerFast, MBartForConditionalGeneration
if not hasattr(self, "_mbart_tokenizer") or self._mbart_tokenizer is None:
logger.info(f"Loading MBART model: {mbart_model_name}")
self._mbart_tokenizer = MBart50TokenizerFast.from_pretrained(mbart_model_name)
self._mbart_model = (
MBartForConditionalGeneration.from_pretrained(
mbart_model_name
)
)
mbart_lang_map = {
"hi": "hi_IN", "ta": "ta_IN", "te": "te_IN", "mr": "mr_IN",
"bn": "bn_IN", "kn": "kn_IN", "ml": "ml_IN", "gu": "gu_IN",
"pa": "pa_IN", "ur": "ur_PK",
}
mbart_code = mbart_lang_map.get(source_lang, "hi_IN")
self._mbart_tokenizer.src_lang = mbart_code
encoded = self._mbart_tokenizer(text, return_tensors="pt")
generated_tokens = self._mbart_model.generate(
**encoded,
forced_bos_token_id=self._mbart_tokenizer.lang_code_to_id["en_XX"]
)
translated_text = self._mbart_tokenizer.batch_decode(
generated_tokens, skip_special_tokens=True
)[0]
return {
"original": text,
"translated": translated_text,
"confidence": 0.70,
"model_used": "MBart-50",
"translation_fallback": True,
}
except Exception as mbart_err:
logger.warning(f"MBART translation fallback failed: {mbart_err}")
return {
"original": text,
"translated": text,
"confidence": 0.0,
"model_used": "none",
"translation_fallback": True,
}
def translate_batch(
self, texts: list[tuple[str, str]]
) -> list[dict[str, Any]]:
results: list[dict] = []
for i, (text, lang) in enumerate(texts):
if i > 0 and i % 10 == 0:
time.sleep(0.3)
results.append(self.translate_to_english(text, lang, is_batch_call=True))
return results
|