142 lines
4.8 KiB
Python
142 lines
4.8 KiB
Python
#!/usr/bin/env python3
|
|
"""JSON-lines bridge for Argos Translate and LibreTranslate."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import re
|
|
import sys
|
|
import urllib.error
|
|
import urllib.request
|
|
from dataclasses import dataclass
|
|
from typing import Callable
|
|
|
|
|
|
TOKEN_RE = re.compile(r"%(?:\d+\$?)?[A-Za-z]|\\[nrt]|<[^>]+>|&(?:#\d+|#x[0-9A-Fa-f]+|\w+);")
|
|
|
|
|
|
def protect(text: str) -> tuple[str, dict[str, str]]:
|
|
tokens: dict[str, str] = {}
|
|
|
|
def replace(match: re.Match[str]) -> str:
|
|
key = f"ZXQPH{len(tokens)}QXZ"
|
|
tokens[key] = match.group(0)
|
|
return key
|
|
|
|
return TOKEN_RE.sub(replace, text), tokens
|
|
|
|
|
|
def restore(text: str, tokens: dict[str, str]) -> str:
|
|
for key, value in tokens.items():
|
|
text = text.replace(key, value)
|
|
missing = [key for key in tokens if key in text]
|
|
if missing:
|
|
raise RuntimeError(f"translator damaged protected tokens {missing}")
|
|
return text
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Request:
|
|
texts: list[str]
|
|
source: str
|
|
target: str
|
|
backend: str
|
|
url: str | None
|
|
api_key: str | None
|
|
|
|
|
|
def libre_translator(request: Request) -> Callable[[str], str]:
|
|
if not request.url:
|
|
raise RuntimeError("LibreTranslate requires a URL")
|
|
endpoint = request.url.rstrip("/") + "/translate"
|
|
|
|
def translate(text: str) -> str:
|
|
payload: dict[str, object] = {
|
|
"q": text,
|
|
"source": request.source,
|
|
"target": request.target,
|
|
"format": "text",
|
|
}
|
|
if request.api_key:
|
|
payload["api_key"] = request.api_key
|
|
http_request = urllib.request.Request(
|
|
endpoint,
|
|
data=json.dumps(payload).encode("utf-8"),
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST",
|
|
)
|
|
try:
|
|
with urllib.request.urlopen(http_request, timeout=60) as response:
|
|
body = json.load(response)
|
|
except urllib.error.URLError as error:
|
|
raise RuntimeError(f"LibreTranslate request failed at {endpoint}: {error}") from error
|
|
translated = body.get("translatedText")
|
|
if not isinstance(translated, str):
|
|
raise RuntimeError("LibreTranslate returned no translatedText")
|
|
return translated
|
|
|
|
return translate
|
|
|
|
|
|
def argos_translator(request: Request) -> Callable[[str], str]:
|
|
os.environ.setdefault("ARGOS_CHUNK_TYPE", "ARGOSTRANSLATE")
|
|
os.environ.setdefault("ARGOS_DEVICE_TYPE", "cpu")
|
|
try:
|
|
import argostranslate.translate # type: ignore[import-not-found]
|
|
except ImportError as error:
|
|
raise RuntimeError("Argos Translate is unavailable. Install requirements-translator.txt") from error
|
|
|
|
installed = argostranslate.translate.get_installed_languages()
|
|
source = next((language for language in installed if language.code == request.source), None)
|
|
target = next((language for language in installed if language.code == request.target), None)
|
|
if source is None or target is None:
|
|
raise RuntimeError(f"Argos model {request.source} to {request.target} is not installed")
|
|
translation = source.get_translation(target)
|
|
# Android resources are individual UI strings. Avoid network-backed sentence
|
|
# model downloads and pass each protected resource to the installed model.
|
|
class ResourceStringSentencizer:
|
|
@staticmethod
|
|
def split_sentences(text: str) -> list[str]:
|
|
return [text]
|
|
|
|
underlying = getattr(translation, "underlying", translation)
|
|
if hasattr(underlying, "sentencizer"):
|
|
underlying.sentencizer = ResourceStringSentencizer()
|
|
return translation.translate
|
|
|
|
|
|
def process(raw: dict[str, object]) -> dict[str, object]:
|
|
request = Request(
|
|
texts=[str(value) for value in raw.get("texts", [])],
|
|
source=str(raw.get("source", "zh")),
|
|
target=str(raw.get("target", "en")),
|
|
backend=str(raw.get("backend", "argos")),
|
|
url=str(raw["url"]) if raw.get("url") else None,
|
|
api_key=str(raw["apiKey"]) if raw.get("apiKey") else None,
|
|
)
|
|
translator = libre_translator(request) if request.backend == "libre" else argos_translator(request)
|
|
output: list[str] = []
|
|
for text in request.texts:
|
|
protected, tokens = protect(text)
|
|
output.append(restore(translator(protected), tokens))
|
|
return {"translations": output}
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--request", help="single JSON request, otherwise stdin is used")
|
|
args = parser.parse_args()
|
|
try:
|
|
raw = json.loads(args.request) if args.request else json.load(sys.stdin)
|
|
json.dump(process(raw), sys.stdout, ensure_ascii=False)
|
|
sys.stdout.write("\n")
|
|
return 0
|
|
except Exception as error:
|
|
print(str(error), file=sys.stderr)
|
|
return 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|