EasyOCR + Traductor (ctranslate2 int8)

Dockerfile

@@ -14,6 +14,15 @@ WORKDIR /app
 COPY requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
 
+# NLLB-200 CTranslate2 formatura bihurtu (INT8 kuantizazioa CPUrako)
+# Build-denboran egiten da: irudia handiagoa baina abiaraztea askoz azkarragoa
+RUN python -c "from transformers import AutoTokenizer; AutoTokenizer.from_pretrained('facebook/nllb-200-distilled-600M')" && \
+    ct2-transformers-converter \
+        --model facebook/nllb-200-distilled-600M \
+        --output_dir /app/nllb-200-distilled-600M-ct2-int8 \
+        --quantization int8 \
+        --force
+
 # Aplikazioaren kodea kopiatu
 COPY app.py .
 

app.py

@@ -3,6 +3,7 @@ app.py  —  FastAPI + EasyOCR + Gemini + NLLB + HiTZ zerbitzaria
 OCR + postzuzenketa (Gemini 2.5 Flash) + itzulpena (NLLB-200 + HiTZ Marian).
 """
 
+import asyncio
 import io
 import logging
 import os
@@ -10,22 +11,18 @@ import re
 import time
 from contextlib import asynccontextmanager
 
+import ctranslate2
 import easyocr
 import httpx
 import numpy as np
-import torch
+import torch  # HiTZ Marian-erako bakarrik (oraindik desaktibatuta)
 from deskew import determine_skew
 from fastapi import FastAPI, File, Form, HTTPException, UploadFile
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
 from PIL import Image
 from skimage.transform import rotate
-from transformers import (
-    AutoModelForSeq2SeqLM,
-    AutoTokenizer,
-    MarianMTModel,
-    MarianTokenizer,
-)
+from transformers import AutoTokenizer
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -41,7 +38,7 @@ SCRIPTS = {
     "korean":     ["en","ko"],
 }
 
-MAX_SIDE = 1500
+MAX_SIDE = 1280
 
 GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")
 GEMINI_MODEL = "gemini-2.5-flash"
@@ -77,6 +74,11 @@ Corrected text:"""
 readers: dict = {}
 
 NLLB_MODEL_NAME = "facebook/nllb-200-distilled-600M"
+NLLB_CT2_DIR = os.environ.get(
+    "NLLB_CT2_DIR", "/app/nllb-200-distilled-600M-ct2-int8"
+)
+CT2_INTRA_THREADS = int(os.environ.get("CT2_INTRA_THREADS", "2"))
+CT2_INTER_THREADS = int(os.environ.get("CT2_INTER_THREADS", "1"))
 
 HITZ_PAIRS = {
     ("en", "eu"): "HiTZ/mt-hitz-en-eu",
@@ -140,7 +142,7 @@ ISO_TO_NLLB = {
     "zu": "zul_Latn",
 }
 
-nllb_model = None
+nllb_translator: ctranslate2.Translator | None = None
 nllb_tokenizer = None
 hitz_models: dict = {}
 
@@ -299,8 +301,29 @@ def _rebuild(blocks_in, translations):
     return "\n".join(out_paragraphs)
 
 
+def _adaptive_max_tokens(sentence: str) -> int:
+    """Token-kopuru maximoa estimatu sarrera-luzeraren arabera."""
+    approx_src_tokens = max(8, len(sentence) // 4)
+    return min(512, max(32, int(approx_src_tokens * 1.8)))
+
+
+_TRANSLATION_CACHE: "dict[tuple[str, str, str], str]" = {}
+_TRANSLATION_CACHE_MAX = 1024
+
+
+def _cache_get(sentence: str, src: str, tgt: str) -> str | None:
+    return _TRANSLATION_CACHE.get((sentence, src, tgt))
+
+
+def _cache_put(sentence: str, src: str, tgt: str, value: str) -> None:
+    if len(_TRANSLATION_CACHE) >= _TRANSLATION_CACHE_MAX:
+        # FIFO sinplea: lehen sartutakoa kendu
+        _TRANSLATION_CACHE.pop(next(iter(_TRANSLATION_CACHE)))
+    _TRANSLATION_CACHE[(sentence, src, tgt)] = value
+
+
 def _nllb_translate(text: str, src_nllb: str, tgt_nllb: str) -> str:
-    """NLLB-200 ereduarekin itzuli, esaldika eta batch-ean."""
+    """NLLB-200 ereduarekin itzuli CTranslate2 motorra erabiliz."""
     if not text.strip():
         return text
     blocks = _flatten_to_sentences(text)
@@ -308,36 +331,51 @@ def _nllb_translate(text: str, src_nllb: str, tgt_nllb: str) -> str:
     if not to_translate:
         return text
 
-    logger.info("[NLLB] %s -> %s | %d esaldi", src_nllb, tgt_nllb, len(to_translate))
+    logger.info("[NLLB-CT2] %s -> %s | %d esaldi", src_nllb, tgt_nllb, len(to_translate))
     t0 = time.time()
+
     nllb_tokenizer.src_lang = src_nllb
-    forced_bos = nllb_tokenizer.convert_tokens_to_ids(tgt_nllb)
-    logger.info("[NLLB] forced_bos_token_id(%s) = %s", tgt_nllb, forced_bos)
+    translations: list[str] = [""] * len(to_translate)
+
+    # 1) Cache-tik bete daitezkeenak bete
+    pending: list[tuple[int, str]] = []
+    cache_hits = 0
+    for i, sentence in enumerate(to_translate):
+        cached = _cache_get(sentence, src_nllb, tgt_nllb)
+        if cached is not None:
+            translations[i] = cached
+            cache_hits += 1
+        else:
+            pending.append((i, sentence))
 
-    translations = []
+    # 2) Falta direnak batch-ean itzuli CTranslate2-rekin
     BATCH = 8
-    for i in range(0, len(to_translate), BATCH):
-        chunk = to_translate[i:i + BATCH]
-        inputs = nllb_tokenizer(
-            chunk, return_tensors="pt", padding=True,
-            truncation=True, max_length=512,
+    for batch_start in range(0, len(pending), BATCH):
+        chunk = pending[batch_start:batch_start + BATCH]
+        chunk_sentences = [s for _, s in chunk]
+        source_tokens_batch = [
+            nllb_tokenizer.convert_ids_to_tokens(nllb_tokenizer.encode(s))
+            for s in chunk_sentences
+        ]
+        max_len = max(_adaptive_max_tokens(s) for s in chunk_sentences)
+        results = nllb_translator.translate_batch(
+            source_tokens_batch,
+            target_prefix=[[tgt_nllb]] * len(chunk_sentences),
+            beam_size=1,
+            max_decoding_length=max_len,
+            no_repeat_ngram_size=3,
         )
-        with torch.no_grad():
-            outputs = nllb_model.generate(
-                **inputs,
-                forced_bos_token_id=forced_bos,
-                max_new_tokens=512,
-                num_beams=2,
-                no_repeat_ngram_size=3,
-                early_stopping=True,
-            )
-        decoded = nllb_tokenizer.batch_decode(outputs, skip_special_tokens=True)
-        for src_s, out_s in zip(chunk, decoded):
-            logger.info("[NLLB]   %r -> %r", src_s[:60], out_s[:60])
-        translations.extend(decoded)
-
-    logger.info("[NLLB] Egina %.1fs-tan", time.time() - t0)
-    return _rebuild(blocks, [t.strip() for t in translations])
+        for (idx, src_s), res in zip(chunk, results):
+            target_ids = nllb_tokenizer.convert_tokens_to_ids(res.hypotheses[0][1:])
+            out_s = nllb_tokenizer.decode(target_ids, skip_special_tokens=True).strip()
+            translations[idx] = out_s
+            _cache_put(src_s, src_nllb, tgt_nllb, out_s)
+            logger.info("[NLLB-CT2]   %r -> %r", src_s[:60], out_s[:60])
+
+    logger.info("[NLLB-CT2] Egina %.1fs-tan (cache hits: %d/%d, tamaina: %d)",
+                time.time() - t0, cache_hits, len(to_translate),
+                len(_TRANSLATION_CACHE))
+    return _rebuild(blocks, translations)
 
 
 def _hitz_translate(text: str, src: str, tgt: str) -> str:
@@ -393,17 +431,22 @@ def translate(text: str, src: str, tgt: str) -> str:
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     for name, langs in SCRIPTS.items():
-        logger.info("Reader kargatzen: %s %s", name, langs)
-        readers[name] = easyocr.Reader(langs, gpu=False)
+        logger.info("Reader kargatzen (quantize=True): %s %s", name, langs)
+        readers[name] = easyocr.Reader(langs, gpu=False, quantize=True)
 
-    global nllb_model, nllb_tokenizer
-    logger.info("[LOAD] NLLB eredua kargatzen: %s", NLLB_MODEL_NAME)
+    global nllb_translator, nllb_tokenizer
+    logger.info("[LOAD] NLLB tokenizer kargatzen: %s", NLLB_MODEL_NAME)
     nllb_tokenizer = AutoTokenizer.from_pretrained(NLLB_MODEL_NAME)
-    nllb_model = AutoModelForSeq2SeqLM.from_pretrained(NLLB_MODEL_NAME)
-    nllb_model.eval()
-    logger.info("[LOAD] NLLB mota: %s | tokenizer: %s",
-                nllb_model.__class__.__name__,
-                nllb_tokenizer.__class__.__name__)
+    logger.info("[LOAD] NLLB CTranslate2 (INT8) kargatzen: %s", NLLB_CT2_DIR)
+    nllb_translator = ctranslate2.Translator(
+        NLLB_CT2_DIR,
+        device="cpu",
+        compute_type="int8",
+        intra_threads=CT2_INTRA_THREADS,
+        inter_threads=CT2_INTER_THREADS,
+    )
+    logger.info("[LOAD] NLLB-CT2 prest | intra=%d inter=%d",
+                CT2_INTRA_THREADS, CT2_INTER_THREADS)
 
     # HiTZ aldi baterako desaktibatuta (transformers bateragarritasun arazoak)
     logger.info("[LOAD] HiTZ karga saltatzen (NLLB soilik modua)")
@@ -417,6 +460,7 @@ async def lifespan(app: FastAPI):
     yield
     readers.clear()
     hitz_models.clear()
+    _TRANSLATION_CACHE.clear()
 
 
 app = FastAPI(title="OCR + Itzulpena API", version="16.0.0", lifespan=lifespan)
@@ -434,8 +478,9 @@ async def health_check():
         "status": "ok",
         "scripts": list(readers.keys()),
         "gemini": bool(GEMINI_API_KEY),
-        "nllb": nllb_model is not None,
-        "nllb_class": nllb_model.__class__.__name__ if nllb_model else None,
+        "nllb": nllb_translator is not None,
+        "nllb_backend": "ctranslate2-int8",
+        "translation_cache_size": len(_TRANSLATION_CACHE),
         "hitz_pairs": [f"{s}-{t}" for (s, t) in hitz_models.keys()],
     }
 
@@ -461,7 +506,10 @@ async def predict(
     img_array = np.array(pil_image)
     img_array = _deskew(img_array)
     reader = readers[script]
-    results = reader.readtext(img_array, detail=1, paragraph=False)
+    loop = asyncio.get_running_loop()
+    results = await loop.run_in_executor(
+        None, lambda: reader.readtext(img_array, detail=1, paragraph=False)
+    )
     raw_text = _group_into_lines(results)
     logger.info("[OCR] Egina %.1fs-tan, %d karaktere", time.time() - t0, len(raw_text))
 

requirements.txt

@@ -10,4 +10,5 @@ httpx==0.27.0
 transformers==4.37.0
 torch==2.2.2
 sentencepiece==0.2.0
-sacremoses==0.1.1
+sacremoses==0.1.1
+ctranslate2==4.3.1