webmail-client/translator/translator.py

"""Ultimo's translation service: Mozilla's Firefox Translations models (the
Bergamot engine) and fastText language identification, behind a small HTTP
API that only the webmail API calls (backend/app/translation.py), on the
internal network. A message's text never leaves the server.

- Firefox's models translate to and from English; any other pair goes
  through English (Polish → English → German). So a language is installed as
  its two models, to and from English, and English is always there.
- What is installed is kept in ``$DATA_DIR/state.json``. The first start
  installs ``TRANSLATOR_LANGUAGES``; after that the admin panel adds and
  removes languages (``POST``/``DELETE /languages/{code}``).
- Models come from Mozilla's Remote Settings (where Firefox gets them), each
  file checked against its SHA-256; only release models (no test builds).
- Language identification: fastText's lid.176 (176 languages), baked into
  the image.

API (JSON):
  GET    /health
  GET    /languages            installed / installing / available, errors
  POST   /languages/{code}     install (in the background)
  DELETE /languages/{code}     remove
  POST   /detect     {"q": str | [str]}                     → language(s)
  POST   /translate  {"q": str | [str], "source", "target"} → translation(s)
"""

from __future__ import annotations

import hashlib
import json
import logging
import os
import re
import shutil
import threading
import time
import urllib.request
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Any

import fasttext
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel


def _load_engine() -> Any:
    """The Bergamot engine: only its compiled module. Importing the bergamot
    package itself fetches model catalogues from translatelocally.com and
    OPUS on every start (and fails without them); we load our own models."""
    import importlib.machinery
    import importlib.util

    spec = importlib.util.find_spec("bergamot")
    if spec is None or not spec.submodule_search_locations:
        raise ImportError("the bergamot package is not installed")
    folder = Path(next(iter(spec.submodule_search_locations)))
    path = next(folder.glob("_bergamot*.so"))
    loader = importlib.machinery.ExtensionFileLoader("_bergamot", str(path))
    module_spec = importlib.util.spec_from_loader("_bergamot", loader)
    assert module_spec is not None
    module = importlib.util.module_from_spec(module_spec)
    loader.exec_module(module)
    return module


bergamot = _load_engine()

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("translator")

DATA = Path(os.environ.get("DATA_DIR", "/data"))
MODELS = DATA / "models"
STATE = DATA / "state.json"
LID_MODEL = os.environ.get("LID_MODEL", "/opt/lid.176.ftz")
DEFAULT_LANGUAGES = [
    c.strip() for c in os.environ.get("TRANSLATOR_LANGUAGES", "pl,de,fr,es,it").split(",")
]
THREADS = int(os.environ.get("TRANSLATOR_THREADS", "2"))
RECORDS_URL = os.environ.get(
    "MODELS_RECORDS_URL",
    "https://firefox.settings.services.mozilla.com/v1/buckets/main/collections/"
    "translations-models/records",
)
ATTACHMENTS_URL = os.environ.get(
    "MODELS_ATTACHMENTS_URL", "https://firefox-settings-attachments.cdn.mozilla.net/"
)
PIVOT = "en"
_RECORDS_TTL = 3600
_CODE = re.compile(r"^[a-z]{2,3}(-[A-Za-z]{2,4})?$")

_service = bergamot.Service(bergamot.ServiceConfig(numWorkers=THREADS, logLevel="off"))
_lid = fasttext.load_model(LID_MODEL)
_lock = threading.Lock()  # state, loaded models, installs
_loaded: dict[str, Any] = {}  # "en-pl" → bergamot model
_installing: dict[str, dict[str, int]] = {}  # "pl" → {"done": bytes, "total": bytes}
_errors: dict[str, str] = {}  # "pl" → why the last install failed
_records: tuple[float, list[dict[str, Any]]] = (0.0, [])


# --- what is installed --------------------------------------------------------------


def _read_state() -> dict[str, Any] | None:
    try:
        return json.loads(STATE.read_text())
    except (OSError, ValueError):
        return None


def _installed() -> list[str]:
    state = _read_state() or {}
    return [PIVOT] + [c for c in state.get("languages", []) if c != PIVOT]


def _write_installed(codes: list[str]) -> None:
    DATA.mkdir(parents=True, exist_ok=True)
    tmp = STATE.with_suffix(".tmp")
    tmp.write_text(json.dumps({"languages": sorted(set(codes) - {PIVOT})}))
    tmp.replace(STATE)


# --- Mozilla's model catalogue -----------------------------------------------------------


def _fetch(url: str, timeout: float = 60) -> bytes:
    req = urllib.request.Request(url, headers={"User-Agent": "ultimo-translator"})
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        return resp.read()


def _catalogue() -> list[dict[str, Any]]:
    """Mozilla's model records (cached for an hour)."""
    global _records
    stamp, cached = _records
    if cached and time.monotonic() - stamp < _RECORDS_TTL:
        return cached
    data = json.loads(_fetch(RECORDS_URL))["data"]
    _records = (time.monotonic(), data)
    return data


def _version(text: str) -> tuple[int, ...] | None:
    """ "2.1" → (2, 1); test builds ("2.1a1") → None."""
    return tuple(int(x) for x in text.split(".")) if re.fullmatch(r"\d+(\.\d+)*", text) else None


def _release_files(source: str, target: str) -> list[dict[str, Any]]:
    """The newest release model of one direction: its files' records."""
    by_version: dict[tuple[int, ...], list[dict[str, Any]]] = {}
    for r in _catalogue():
        if r.get("fromLang") != source or r.get("toLang") != target:
            continue
        if r.get("filter_expression"):
            continue  # a test build (Nightly only)
        v = _version(str(r.get("version", "")))
        if v is not None:
            by_version.setdefault(v, []).append(r)
    if not by_version:
        raise LookupError(f"no model for {source} → {target}")
    return by_version[max(by_version)]


def _available() -> list[str]:
    """Languages Mozilla has release models for, both ways with English."""
    pairs: set[tuple[str, str]] = set()
    for r in _catalogue():
        if not r.get("filter_expression") and _version(str(r.get("version", ""))):
            pairs.add((r["fromLang"], r["toLang"]))
    return sorted({a for a, b in pairs if b == PIVOT and (PIVOT, a) in pairs} | {PIVOT})


# --- installing ---------------------------------------------------------------------------


def _model_dir(source: str, target: str) -> Path:
    return MODELS / f"{source}-{target}"


def _install_direction(code: str, source: str, target: str) -> None:
    files = _release_files(source, target)
    final = _model_dir(source, target)
    tmp = final.with_name(final.name + ".partial")
    shutil.rmtree(tmp, ignore_errors=True)
    tmp.mkdir(parents=True)
    for r in files:
        att = r["attachment"]
        blob = _fetch(ATTACHMENTS_URL + att["location"], timeout=600)
        if hashlib.sha256(blob).hexdigest() != att["hash"]:
            raise ValueError(f"{att['filename']}: checksum mismatch")
        (tmp / att["filename"]).write_bytes(blob)
        with _lock:
            _installing[code]["done"] += len(blob)
    (tmp / "files.json").write_text(
        json.dumps(
            {r["fileType"]: r["attachment"]["filename"] for r in files}
            | {"version": files[0]["version"]}
        )
    )
    shutil.rmtree(final, ignore_errors=True)
    tmp.replace(final)


def _install(code: str) -> None:
    """Both directions of ``code`` with English, then mark it installed."""
    try:
        directions = [(code, PIVOT), (PIVOT, code)]
        total = sum(r["attachment"]["size"] for s, t in directions for r in _release_files(s, t))
        with _lock:
            _installing[code]["total"] = total
        for source, target in directions:
            _install_direction(code, source, target)
        with _lock:
            _write_installed(_installed() + [code])
            for pair in (f"{code}-{PIVOT}", f"{PIVOT}-{code}"):
                _loaded.pop(pair, None)  # a reinstall: load the new files
            _errors.pop(code, None)
        log.info("installed %s", code)
    except Exception as exc:  # noqa: BLE001 — reported in /languages, tried again by hand
        log.warning("installing %s failed: %s", code, exc)
        with _lock:
            _errors[code] = str(exc)[:300]
    finally:
        with _lock:
            _installing.pop(code, None)


def _start_install(code: str) -> bool:
    with _lock:
        if code in _installing:
            return False
        _installing[code] = {"done": 0, "total": 0}
    threading.Thread(target=_install, args=(code,), daemon=True).start()
    return True


def _first_start() -> None:
    """The first start installs the default languages, one after another."""
    if _read_state() is not None:
        return
    _write_installed([])

    def run() -> None:
        for code in DEFAULT_LANGUAGES:
            if code and code != PIVOT:
                with _lock:
                    _installing[code] = {"done": 0, "total": 0}
                _install(code)

    threading.Thread(target=run, daemon=True).start()


@asynccontextmanager
async def _lifespan(_: FastAPI) -> AsyncIterator[None]:
    _first_start()
    yield


app = FastAPI(title="Ultimo translator", docs_url=None, redoc_url=None, lifespan=_lifespan)


# --- translating -------------------------------------------------------------------------


def _config(directory: Path) -> str:
    files = json.loads((directory / "files.json").read_text())
    model = files["model"]
    vocabs = (
        [files["vocab"], files["vocab"]]
        if "vocab" in files
        else [files["srcvocab"], files["trgvocab"]]
    )
    # The models are 8-bit; their file name says which kind (as Firefox decides).
    precision = "int8shiftAlphaAll" if model.endswith("alphas.bin") else "int8shiftAll"
    lines = [
        "relative-paths: false",
        "models:",
        f"  - {directory / model}",
        "vocabs:",
        *[f"  - {directory / v}" for v in vocabs],
        "beam-size: 1",
        "normalize: 1.0",
        "word-penalty: 0",
        "mini-batch: 64",
        "maxi-batch: 1000",
        "maxi-batch-sort: src",
        "workspace: 2000",
        "max-length-factor: 2.5",
        "max-length-break: 128",
        "mini-batch-words: 1024",
        "skip-cost: true",
        "quiet: true",
        f"gemm-precision: {precision}",
    ]
    if "lex" in files:
        lines += ["shortlist:", f"  - {directory / files['lex']}", "  - false"]
    return "\n".join(lines) + "\n"


def _model(source: str, target: str) -> Any:
    pair = f"{source}-{target}"
    with _lock:
        if pair not in _loaded:
            directory = _model_dir(source, target)
            if not (directory / "files.json").exists():
                raise HTTPException(400, f"{source} → {target} is not installed")
            _loaded[pair] = _service.modelFromConfig(_config(directory))
        return _loaded[pair]


_OPTIONS = bergamot.ResponseOptions(alignment=False, qualityScores=False, HTML=False)


def _translate(texts: list[str], source: str, target: str) -> list[str]:
    todo = [i for i, t in enumerate(texts) if t.strip()]
    out = list(texts)
    if not todo or source == target:
        return out
    batch = [texts[i] for i in todo]
    steps = [(source, target)] if PIVOT in (source, target) else [(source, PIVOT), (PIVOT, target)]
    for s, t in steps:
        model = _model(s, t)
        batch = [
            r.target.text for r in _service.translate(model, bergamot.VectorString(batch), _OPTIONS)
        ]
    for i, text in zip(todo, batch, strict=True):
        out[i] = text
    return out


# --- identifying languages -----------------------------------------------------------------

# fastText's labels where Firefox's codes differ.
_LID_CODES = {"zh": "zh-Hans", "no": "nb", "iw": "he", "fil": "tl"}
# Signs only traditional Chinese uses (a few common ones): tell zh-Hant apart.
_TRADITIONAL = set("們這個來說會對時國過還後點學東車長門開關問電話員")


def _detect(text: str) -> dict[str, Any]:
    sample = " ".join(text.split())[:2000]
    if not sample:
        return {"language": "", "confidence": 0.0}
    labels, scores = _lid.predict(sample, k=1)
    code = labels[0].replace("__label__", "") if labels else ""
    code = _LID_CODES.get(code, code)
    if code == "zh-Hans" and sum(c in _TRADITIONAL for c in sample) >= 2:
        code = "zh-Hant"
    return {"language": code, "confidence": round(float(scores[0]) * 100, 1) if scores else 0.0}


# --- the API ---------------------------------------------------------------------------------


class DetectIn(BaseModel):
    q: str | list[str]


class TranslateIn(BaseModel):
    q: str | list[str]
    source: str
    target: str


@app.get("/health")
def health() -> dict[str, bool]:
    return {"ok": True}


@app.get("/languages")
def languages() -> dict[str, Any]:
    try:
        available = _available()
    except Exception as exc:  # noqa: BLE001 — offline: what is installed still works
        log.info("model catalogue unavailable: %s", exc)
        available = []
    with _lock:
        return {
            "installed": _installed(),
            "installing": {k: dict(v) for k, v in _installing.items()},
            "available": available,
            "errors": dict(_errors),
        }


@app.post("/languages/{code}", status_code=202)
def install(code: str) -> dict[str, Any]:
    if not _CODE.fullmatch(code) or code == PIVOT:
        raise HTTPException(400, "not a language code")
    try:
        _release_files(code, PIVOT)
        _release_files(PIVOT, code)
    except LookupError as exc:
        raise HTTPException(404, str(exc)) from exc
    return {"started": _start_install(code)}


@app.delete("/languages/{code}")
def remove(code: str) -> dict[str, bool]:
    if code == PIVOT:
        raise HTTPException(400, "English is always installed")
    with _lock:
        _write_installed([c for c in _installed() if c != code])
        for pair in (f"{code}-{PIVOT}", f"{PIVOT}-{code}"):
            _loaded.pop(pair, None)
            shutil.rmtree(MODELS / pair, ignore_errors=True)
    return {"removed": True}


@app.post("/detect")
def detect(body: DetectIn) -> Any:
    if isinstance(body.q, str):
        return [_detect(body.q)]
    return [_detect(q) for q in body.q]


@app.post("/translate")
def translate(body: TranslateIn) -> dict[str, Any]:
    texts = [body.q] if isinstance(body.q, str) else body.q
    source = body.source
    if source == "auto":
        source = _detect(" ".join(texts))["language"]
    installed = _installed()
    for code in (source, body.target):
        if code not in installed:
            raise HTTPException(400, f"{code or 'this language'} is not installed")
    out = _translate(texts, source, body.target)
    return {"translatedText": out[0] if isinstance(body.q, str) else out, "source": source}