mejora en el modelo de traduccion

This commit is contained in:
jlimolina 2026-04-04 03:17:37 +02:00
parent be56de7dd4
commit ecd7a3cdf8
21 changed files with 3708 additions and 322 deletions

View file

@ -2,6 +2,7 @@ import os
import time
import logging
import re
import fcntl
from typing import List, Optional
import psycopg2
@ -19,19 +20,21 @@ LOG = logging.getLogger("translator_ct2")
TRANSLATOR_ID = os.environ.get("TRANSLATOR_ID", "")
TRANSLATOR_TOTAL = int(os.environ.get("TRANSLATOR_TOTAL", "1"))
def clean_text(text: str) -> str:
if not text:
return ""
text = re.sub(r'<[^>]+>', '', text)
text = text.replace('<unk>', '')
text = text.replace('&nbsp;', ' ')
text = text.replace('&amp;', '&')
text = text.replace('&lt;', '<')
text = text.replace('&gt;', '>')
text = text.replace('&quot;', '"')
text = re.sub(r'\s+', ' ', text)
text = re.sub(r"<[^>]+>", "", text)
text = text.replace("<unk>", "")
text = text.replace("&nbsp;", " ")
text = text.replace("&amp;", "&")
text = text.replace("&lt;", "<")
text = text.replace("&gt;", ">")
text = text.replace("&quot;", '"')
text = re.sub(r"\s+", " ", text)
return text.strip()
DB_CONFIG = {
"host": os.environ.get("DB_HOST", "localhost"),
"port": int(os.environ.get("DB_PORT", 5432)),
@ -40,12 +43,14 @@ DB_CONFIG = {
"password": os.environ.get("DB_PASS", "x"),
}
def _env_list(name: str, default="es"):
raw = os.environ.get(name)
if raw:
return [s.strip() for s in raw.split(",") if s.strip()]
return [default]
def _env_int(name: str, default: int = 8):
v = os.environ.get(name)
try:
@ -53,14 +58,16 @@ def _env_int(name: str, default: int = 8):
except Exception:
return default
def _env_str(name: str, default=None):
v = os.environ.get(name)
return v if v else default
TARGET_LANGS = _env_list("TARGET_LANGS")
BATCH_SIZE = _env_int("TRANSLATOR_BATCH", 8)
MAX_SRC_TOKENS = _env_int("MAX_SRC_TOKENS", 512)
MAX_NEW_TOKENS = _env_int("MAX_NEW_TOKENS", 512)
BATCH_SIZE = _env_int("TRANSLATOR_BATCH", 128)
MAX_SRC_TOKENS = _env_int("MAX_SRC_TOKENS", 256)
MAX_NEW_TOKENS = _env_int("MAX_NEW_TOKENS", 256)
CT2_MODEL_PATH = _env_str("CT2_MODEL_PATH", "/app/models/nllb-ct2")
CT2_DEVICE = _env_str("CT2_DEVICE", "cpu")
@ -69,87 +76,160 @@ UNIVERSAL_MODEL = _env_str("UNIVERSAL_MODEL", "facebook/nllb-200-distilled-600M"
BODY_CHARS_CHUNK = _env_int("BODY_CHARS_CHUNK", 900)
LANG_CODE_MAP = {
"en": "eng_Latn", "es": "spa_Latn", "fr": "fra_Latn", "de": "deu_Latn",
"it": "ita_Latn", "pt": "por_Latn", "nl": "nld_Latn", "sv": "swe_Latn",
"da": "dan_Latn", "fi": "fin_Latn", "no": "nob_Latn",
"pl": "pol_Latn", "cs": "ces_Latn", "sk": "slk_Latn",
"sl": "slv_Latn", "hu": "hun_Latn", "ro": "ron_Latn",
"el": "ell_Grek", "ru": "rus_Cyrl", "uk": "ukr_Cyrl",
"tr": "tur_Latn", "ar": "arb_Arab", "fa": "pes_Arab",
"he": "heb_Hebr", "zh": "zho_Hans", "ja": "jpn_Jpan",
"ko": "kor_Hang", "vi": "vie_Latn",
"en": "eng_Latn",
"es": "spa_Latn",
"fr": "fra_Latn",
"de": "deu_Latn",
"it": "ita_Latn",
"pt": "por_Latn",
"nl": "nld_Latn",
"sv": "swe_Latn",
"da": "dan_Latn",
"fi": "fin_Latn",
"no": "nob_Latn",
"pl": "pol_Latn",
"cs": "ces_Latn",
"sk": "slk_Latn",
"sl": "slv_Latn",
"hu": "hun_Latn",
"ro": "ron_Latn",
"el": "ell_Grek",
"ru": "rus_Cyrl",
"uk": "ukr_Cyrl",
"tr": "tur_Latn",
"ar": "arb_Arab",
"fa": "pes_Arab",
"he": "heb_Hebr",
"zh": "zho_Hans",
"ja": "jpn_Jpan",
"ko": "kor_Hang",
"vi": "vie_Latn",
}
_tokenizer = None
_translator = None
def ensure_model():
global _tokenizer, _translator
if _translator:
return
model_path = CT2_MODEL_PATH
model_bin = os.path.join(model_path, "model.bin")
if not os.path.exists(model_bin):
LOG.info(f"CTranslate2 model not found at {model_path}, converting from {UNIVERSAL_MODEL}...")
# Check if model exists AND is complete (all required files present and non-empty)
required_files = ["model.bin", "config.json", "shared_vocabulary.json"]
model_exists = os.path.exists(model_bin)
if model_exists:
# Verify all files exist and have reasonable size
all_files_ok = True
for f in required_files:
fpath = os.path.join(model_path, f)
if not os.path.exists(fpath) or os.path.getsize(fpath) < 1000:
all_files_ok = False
break
if not all_files_ok:
LOG.info(f"Model files incomplete or corrupted, re-converting...")
# Clean up corrupted files
for f in required_files:
try:
fpath = os.path.join(model_path, f)
if os.path.exists(fpath):
os.remove(fpath)
except:
pass
model_exists = False
if not model_exists:
LOG.info(
f"CTranslate2 model not found at {model_path}, converting from {UNIVERSAL_MODEL}..."
)
convert_model()
LOG.info(f"Loading CTranslate2 model from {model_path} on {CT2_DEVICE}")
device = os.environ.get("CT2_DEVICE", "cpu")
LOG.info(f"Loading CTranslate2 model from {model_path} on {device}")
_translator = ctranslate2.Translator(
model_path,
device=CT2_DEVICE,
device=device,
compute_type=CT2_COMPUTE_TYPE,
)
_tokenizer = AutoTokenizer.from_pretrained(UNIVERSAL_MODEL)
LOG.info("CTranslate2 model loaded successfully")
def convert_model():
import subprocess
model_path = CT2_MODEL_PATH
lock_file = os.path.join(model_path, ".converting.lock")
# Clean up any corrupted files from previous failed conversions
if os.path.exists(model_path):
for f in os.listdir(model_path):
if f.endswith(".lock") or f.startswith("."):
try:
os.remove(os.path.join(model_path, f))
except:
pass
os.makedirs(model_path, exist_ok=True)
quantization = CT2_COMPUTE_TYPE if CT2_COMPUTE_TYPE != "auto" else "int8"
cmd = [
"ct2-transformers-converter",
"--model", UNIVERSAL_MODEL,
"--output_dir", model_path,
"--quantization", quantization,
"--force"
]
LOG.info(f"Running: {' '.join(cmd)}")
result = subprocess.run(cmd, capture_output=True, text=True, timeout=1800)
if result.returncode != 0:
LOG.error(f"Model conversion failed: {result.stderr}")
raise RuntimeError("Failed to convert model")
LOG.info("Model conversion completed")
# Use lock file to prevent multiple workers from converting simultaneously
lock_fd = os.open(lock_file, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
try:
quantization = CT2_COMPUTE_TYPE if CT2_COMPUTE_TYPE != "auto" else "float16"
cmd = [
"ct2-transformers-converter",
"--model",
UNIVERSAL_MODEL,
"--output_dir",
model_path,
"--quantization",
quantization,
"--force",
]
LOG.info(f"Running: {' '.join(cmd)}")
result = subprocess.run(cmd, capture_output=True, text=True, timeout=3600)
if result.returncode != 0:
LOG.error(f"Model conversion failed: {result.stderr}")
raise RuntimeError("Failed to convert model")
LOG.info("Model conversion completed")
finally:
os.close(lock_fd)
try:
os.remove(lock_file)
except:
pass
def translate_texts(src: str, tgt: str, texts: List[str]) -> List[str]:
if not texts:
return []
ensure_model()
clean = [(t or "").strip() for t in texts]
if all(not t for t in clean):
return ["" for _ in clean]
src_code = LANG_CODE_MAP.get(src, f"{src}_Latn")
tgt_code = LANG_CODE_MAP.get(tgt, "spa_Latn")
try:
_tokenizer.src_lang = src_code
except Exception:
pass
sources = []
for t in clean:
if t:
@ -158,18 +238,18 @@ def translate_texts(src: str, tgt: str, texts: List[str]) -> List[str]:
sources.append(tokens)
else:
sources.append([])
target_prefix = [[tgt_code]] * len(sources)
results = _translator.translate_batch(
sources,
target_prefix=target_prefix,
beam_size=2,
beam_size=1,
max_decoding_length=MAX_NEW_TOKENS,
repetition_penalty=2.0,
no_repeat_ngram_size=3,
repetition_penalty=1.2,
no_repeat_ngram_size=2,
)
translated = []
for result in results:
try:
@ -197,18 +277,19 @@ def translate_texts(src: str, tgt: str, texts: List[str]) -> List[str]:
except Exception as e:
LOG.error(f"Error processing result: {e}")
translated.append("")
return translated
def split_body_into_chunks(text: str) -> List[str]:
text = (text or "").strip()
if len(text) <= BODY_CHARS_CHUNK:
return [text] if text else []
parts = re.split(r'(\n\n+|(?<=[\.\!\?؛؟。])\s+)', text)
parts = re.split(r"(\n\n+|(?<=[\.\!\?؛؟。])\s+)", text)
chunks = []
current = ""
for part in parts:
if not part:
continue
@ -220,31 +301,34 @@ def split_body_into_chunks(text: str) -> List[str]:
current = part
if current.strip():
chunks.append(current.strip())
return chunks if chunks else [text]
def translate_body_long(src: str, tgt: str, body: str) -> str:
body = (body or "").strip()
if not body:
return ""
chunks = split_body_into_chunks(body)
if len(chunks) == 1:
return translate_texts(src, tgt, [body])[0]
translated_chunks = []
for ch in chunks:
tr = translate_texts(src, tgt, [ch])[0]
translated_chunks.append(tr)
return " ".join(translated_chunks)
def normalize_lang(lang: Optional[str], default: str = "es") -> Optional[str]:
if not lang:
return default
lang = lang.strip().lower()[:2]
return lang if lang else default
def detect_lang(text: str) -> str:
if not text or len(text) < 10:
return "en"
@ -253,63 +337,75 @@ def detect_lang(text: str) -> str:
except Exception:
return "en"
def process_batch(conn, rows):
todo = []
for r in rows:
lang_to = normalize_lang(r.get("lang_to"), "es") or "es"
lang_from = normalize_lang(r.get("lang_from")) or detect_lang(r.get("titulo") or "")
lang_from = normalize_lang(r.get("lang_from")) or detect_lang(
r.get("titulo") or ""
)
titulo = (r.get("titulo") or "").strip()
resumen = (r.get("resumen") or "").strip()
if lang_from == lang_to:
# Mark as done and copy original text if languages match
cursor = conn.cursor()
cursor.execute("""
cursor.execute(
"""
UPDATE traducciones
SET titulo_trad = %s, resumen_trad = %s, status = 'done'
WHERE id = %s
""", (titulo, resumen, r.get("tr_id")))
""",
(titulo, resumen, r.get("tr_id")),
)
conn.commit()
cursor.close()
continue
todo.append({
"tr_id": r.get("tr_id"),
"lang_from": lang_from,
"lang_to": lang_to,
"titulo": titulo,
"resumen": resumen,
})
todo.append(
{
"tr_id": r.get("tr_id"),
"lang_from": lang_from,
"lang_to": lang_to,
"titulo": titulo,
"resumen": resumen,
}
)
if not todo:
return
# 1. FAST LOCKING: Commit locked_at immediately to inform other workers
cursor = conn.cursor()
tr_ids = [item["tr_id"] for item in todo]
cursor.execute(f"""
cursor.execute(
f"""
UPDATE traducciones
SET locked_at = NOW()
WHERE id = ANY(ARRAY[{','.join(['%s'] * len(tr_ids))}])
""", tr_ids)
WHERE id = ANY(ARRAY[{",".join(["%s"] * len(tr_ids))}])
""",
tr_ids,
)
conn.commit()
cursor.close()
from collections import defaultdict
groups = defaultdict(list)
for item in todo:
key = (item["lang_from"], item["lang_to"])
groups[key].append(item)
for (lang_from, lang_to), items in groups.items():
LOG.info(f"Translating {lang_from} -> {lang_to} ({len(items)} items)")
try:
titles = [i["titulo"] for i in items]
translated_titles = translate_texts(lang_from, lang_to, titles)
for item, tt in zip(items, translated_titles):
body = (item["resumen"] or "").strip()
tb = ""
@ -319,52 +415,61 @@ def process_batch(conn, rows):
except Exception as e:
LOG.error(f"Body translation error for ID {item['tr_id']}: {e}")
tb = item["resumen"]
tt = clean_text((tt or "").strip())
tb = clean_text((tb or "").strip())
if not tt:
tt = item["titulo"]
if not tb:
tb = item["resumen"]
# 2. INDIVIDUAL COMMIT: Save each item as it's done
try:
cursor = conn.cursor()
cursor.execute("""
cursor.execute(
"""
UPDATE traducciones
SET titulo_trad = %s, resumen_trad = %s, status = 'done', locked_at = NULL
WHERE id = %s
""", (tt, tb, item["tr_id"]))
""",
(tt, tb, item["tr_id"]),
)
conn.commit()
cursor.close()
except Exception as e:
LOG.error(f"Update error for ID {item['tr_id']}: {e}")
conn.rollback()
LOG.info(f"Finished group {lang_from} -> {lang_to}")
except Exception as e:
LOG.error(f"Batch group error {lang_from} -> {lang_to}: {e}")
# Mark these as error to avoid infinite loop if it's a model crash
try:
cursor = conn.cursor()
cursor.execute("""
cursor.execute(
"""
UPDATE traducciones SET status = 'error', locked_at = NULL
WHERE id = ANY(ARRAY[{','.join(['%s'] * len(items))}])
""", [i["tr_id"] for i in items])
""",
[i["tr_id"] for i in items],
)
conn.commit()
cursor.close()
except:
conn.rollback()
def fetch_pending_translations(conn):
cursor = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
worker_id = os.environ.get("HOSTNAME", f"worker-{os.getpid()}")
total_found = 0
for lang in TARGET_LANGS:
cursor.execute("""
cursor.execute(
"""
SELECT t.id as tr_id, t.lang_from, t.lang_to,
n.titulo, n.resumen, n.id as noticia_id
FROM traducciones t
@ -375,31 +480,45 @@ def fetch_pending_translations(conn):
ORDER BY n.fecha DESC
LIMIT %s
FOR UPDATE SKIP LOCKED
""", (lang, BATCH_SIZE))
""",
(lang, BATCH_SIZE),
)
rows = cursor.fetchall()
if rows:
LOG.info(f"Found {len(rows)} pending translations for {lang}")
process_batch(conn, rows)
total_found += len(rows)
cursor.close()
return total_found
def connect_db():
return psycopg2.connect(**DB_CONFIG)
def main():
LOG.info(f"CTranslate2 translator worker started (device={CT2_DEVICE}, instances={TRANSLATOR_TOTAL})")
LOG.info(
f"CTranslate2 translator worker started (device={CT2_DEVICE}, instances={TRANSLATOR_TOTAL})"
)
ensure_model()
while True:
try:
conn = connect_db()
fetch_pending_translations(conn)
total = fetch_pending_translations(conn)
conn.close()
if total == 0:
LOG.info("No pending translations, sleeping...")
else:
LOG.info(f"Processed {total} translations, sleeping...")
except Exception as e:
LOG.error(f"Error: {e}")
time.sleep(30)
if __name__ == "__main__":
main()