lukasandClaude Opus 5.5
21cfaa07ed
The first start tried the default languages once; a DNS failure at that moment left them missing for good. Languages asked for are now kept as pending in state.json and retried in the background (30 s, doubling to 15 min) until they are in, across restarts; an install from the admin panel joins the same queue. A state from before this, with nothing installed, counts as an unfinished first start. The admin panel shows what is waiting. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
472 lines
16 KiB
Python
472 lines
16 KiB
Python
"""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``, with the languages
|
|
still to install. The first start asks for ``TRANSLATOR_LANGUAGES``; after
|
|
that the admin panel adds and removes languages (``POST``/``DELETE
|
|
/languages/{code}``). A language that could not be downloaded (no network
|
|
yet) is tried again in the background, at growing intervals, until it is
|
|
in — across restarts too.
|
|
- 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 _pending() -> list[str]:
|
|
"""Languages asked for and not installed yet (tried again until they are)."""
|
|
return list((_read_state() or {}).get("pending", []))
|
|
|
|
|
|
def _write_state(installed: list[str], pending: list[str]) -> None:
|
|
DATA.mkdir(parents=True, exist_ok=True)
|
|
tmp = STATE.with_suffix(".tmp")
|
|
done = set(installed) - {PIVOT}
|
|
tmp.write_text(
|
|
json.dumps(
|
|
{
|
|
"languages": sorted(done),
|
|
"pending": [c for c in dict.fromkeys(pending) if c not in done and c != PIVOT],
|
|
}
|
|
)
|
|
)
|
|
tmp.replace(STATE)
|
|
|
|
|
|
def _write_installed(codes: list[str]) -> None:
|
|
_write_state(codes, _pending())
|
|
|
|
|
|
# --- 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_state(_installed() + [code], [c for c in _pending() if c != 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 later
|
|
log.warning("installing %s failed: %s", code, exc)
|
|
with _lock:
|
|
_errors[code] = str(exc)[:300]
|
|
finally:
|
|
with _lock:
|
|
_installing.pop(code, None)
|
|
|
|
|
|
_wake = threading.Event() # a new language to install: try now, not after the pause
|
|
_RETRY_FIRST, _RETRY_MAX = 30.0, 900.0
|
|
|
|
|
|
def _installer() -> None:
|
|
"""Install the pending languages, one after another; while some fail
|
|
(no network, Mozilla out of reach), try again after a pause that grows
|
|
from half a minute to a quarter of an hour."""
|
|
pause = _RETRY_FIRST
|
|
while True:
|
|
for code in _pending():
|
|
with _lock:
|
|
if code in _installing or code in _installed():
|
|
continue
|
|
_installing[code] = {"done": 0, "total": 0}
|
|
_install(code)
|
|
if not _pending():
|
|
pause = _RETRY_FIRST
|
|
_wake.wait() # nothing to do until a language is asked for
|
|
else:
|
|
_wake.wait(pause)
|
|
pause = min(pause * 2, _RETRY_MAX)
|
|
_wake.clear()
|
|
|
|
|
|
def _ask_for(code: str) -> None:
|
|
with _lock:
|
|
_write_state(_installed(), _pending() + [code])
|
|
_errors.pop(code, None)
|
|
_wake.set()
|
|
|
|
|
|
def _start() -> None:
|
|
"""The first start asks for the default languages; every start resumes
|
|
what is still pending."""
|
|
state = _read_state()
|
|
# No state yet — or one from before pending languages were kept, with
|
|
# nothing installed (a first start that could not download): ask for the
|
|
# defaults.
|
|
if state is None or ("pending" not in state and not state.get("languages")):
|
|
_write_state([], [c for c in DEFAULT_LANGUAGES if c])
|
|
threading.Thread(target=_installer, daemon=True).start()
|
|
|
|
|
|
@asynccontextmanager
|
|
async def _lifespan(_: FastAPI) -> AsyncIterator[None]:
|
|
_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()},
|
|
"pending": [c for c in _pending() if c not in _installing],
|
|
"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
|
|
except Exception as exc: # noqa: BLE001 — Mozilla out of reach: queue it anyway
|
|
log.info("model catalogue unavailable, %s queued: %s", code, exc)
|
|
_ask_for(code)
|
|
return {"started": True}
|
|
|
|
|
|
@app.delete("/languages/{code}")
|
|
def remove(code: str) -> dict[str, bool]:
|
|
if code == PIVOT:
|
|
raise HTTPException(400, "English is always installed")
|
|
with _lock:
|
|
_write_state([c for c in _installed() if c != code], [c for c in _pending() if c != code])
|
|
_errors.pop(code, None)
|
|
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}
|