From 2f78877f151e74c7ca5d13cc3f701b790635459e Mon Sep 17 00:00:00 2001 From: Leo Vasanko Date: Wed, 2 Sep 2026 03:24:57 +0000 Subject: [PATCH] Seed-X translator service client for the /_translate socket --- scripts/translator.py | 179 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 179 insertions(+) create mode 100644 scripts/translator.py diff --git a/scripts/translator.py b/scripts/translator.py new file mode 100644 index 0000000..f640ef4 --- /dev/null +++ b/scripts/translator.py @@ -0,0 +1,179 @@ +#!/usr/bin/env python3 +# /// script +# requires-python = ">=3.14" +# dependencies = [ +# "accelerate>=1.14.0", +# "msgspec>=0.19.0", +# "torch>=2.13.0", +# "tracerite>=2.6.5", +# "transformers>=5.16.1", +# "websockets>=15.0.1", +# ] +# /// +"""Pagerite translator service: translate site content with Seed-X-PPO-7B. + +Connects to a Pagerite server's translator WebSocket (``/_translate/`` — +deliberately outside /_api, the key is the access control; find it in the +site settings, GET /_api/settings -> ``translate_key``), announces the +languages it handles and translates whatever the server pushes: everything +pending on connect, then deltas as the content changes. Results go back per +job and the server stores them (docs/localization.md). + +Seed-X-PPO-7B (bf16, ~15 GB) is the only supported model; its 28 languages +are the ceiling of what a site can announce through this script. + +Usage: + uv run scripts/translator.py ws://localhost:8080 --key "..." --to finnish + uv run scripts/translator.py wss://example.com --key "..." --to fi german es +""" + +import argparse +import asyncio +import sys +import time + +import msgspec +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer +import tracerite +import websockets + +tracerite.load() + +SEED_X = "ByteDance-Seed/Seed-X-PPO-7B" + +# Seed-X language tags (appended to the prompt; required by its PPO training) +SEED_X_TAGS = { + "arabic": "ar", "chinese": "zh", "czech": "cs", "danish": "da", + "dutch": "nl", "english": "en", "finnish": "fi", "french": "fr", + "german": "de", "greek": "el", "hungarian": "hu", "indonesian": "id", + "italian": "it", "japanese": "ja", "korean": "ko", "malay": "ms", + "norwegian": "no", "persian": "fa", "polish": "pl", "portuguese": "pt", + "romanian": "ro", "russian": "ru", "spanish": "es", "swedish": "sv", + "thai": "th", "turkish": "tr", "ukrainian": "uk", "vietnamese": "vi", +} +SEED_X_NAMES = {v: k for k, v in SEED_X_TAGS.items()} + +#: The fragments are Markdown; Seed-X has no system prompt, so it goes in-line. +NOTE = ", preserving all Markdown formatting, URLs and code exactly unchanged" + + +# The message structs below duplicate pagerite/translate.py 1:1: this script +# runs in its own uv environment and cannot import the server package. The +# "type" tag selects the frame; bytes fields ride as base64. +class Hello(msgspec.Struct, tag="hello"): + """Client greeting on connect: the target languages it handles.""" + + langs: list[str] + + +class TransItem(msgspec.Struct): + """One fragment to translate: original Markdown (or a node title).""" + + key: bytes #: 9-byte chunk hash (base64 in the JSON frame) + text: str + path: str #: article it came from ("" = front page), no leading slash + kind: str #: "chunk" | "title" + + +class Job(msgspec.Struct, tag="job"): + """Server push: pending items for one language.""" + + lang: str + items: list[TransItem] + + +class TransResult(msgspec.Struct): + """One translated fragment.""" + + key: bytes + text: str + + +class Result(msgspec.Struct, tag="result"): + """Client reply: a batch of translations for one language.""" + + lang: str + items: list[TransResult] + + +def load_seed_x(): + t0 = time.monotonic() + tokenizer = AutoTokenizer.from_pretrained(SEED_X) + model = AutoModelForCausalLM.from_pretrained(SEED_X, dtype=torch.bfloat16, device_map="auto") + print(f"[seed-x loaded in {time.monotonic() - t0:.0f}s]", file=sys.stderr) + return tokenizer, model + + +def seed_x_chunk(tokenizer, model, text: str, target_lang: str, tag: str, source_lang: str = "English", + note: str = ""): + """Translate one segment; returns (translation, output_tokens, generation_seconds).""" + # No chat template on this model; the trailing language tag is required (trans/ style prompt). + prompt = f"Translate the following {source_lang} text into {target_lang}{note}:\n{text} <{tag}>" + inputs = tokenizer(prompt, return_tensors="pt").to(model.device) + t0 = time.monotonic() + out = model.generate(**inputs, max_new_tokens=max(1024, 2 * inputs.input_ids.shape[1]), do_sample=False) + dt = time.monotonic() - t0 + n = out.shape[1] - inputs.input_ids.shape[1] + return tokenizer.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True).strip(), n, dt + + +async def do_job(ws, job: Job, tokenizer, model) -> None: + """Translate every item of one job and send the results back as a batch.""" + lang_name = SEED_X_NAMES[job.lang].capitalize() + results = [] + for i, item in enumerate(job.items, 1): + # Deliberately blocking: nothing else needs the loop while a job is + # being answered, and the reconnect loop recovers a dropped connection + # (results already stored stay stored; only the still-pending items + # are re-pushed). + text, tokens, dt = seed_x_chunk(tokenizer, model, item.text, lang_name, job.lang, note=NOTE) + results.append(TransResult(key=item.key, text=text)) + print(f"[{job.lang} {i}/{len(job.items)} {item.kind} {item.path or '/'}: " + f"{tokens} tokens in {dt:.1f}s = {tokens / dt:.1f} tok/s]", file=sys.stderr) + await ws.send(msgspec.json.encode(Result(lang=job.lang, items=results)).decode()) + + +async def serve(url: str, codes: list[str], tokenizer, model) -> None: + """Connect, announce languages, answer jobs; reconnect with backoff.""" + backoff = 1 + while True: + try: + async with websockets.connect(url) as ws: + backoff = 1 + await ws.send(msgspec.json.encode(Hello(langs=codes)).decode()) + print(f"[connected; translating: {', '.join(codes)}]", file=sys.stderr) + async for raw in ws: + await do_job(ws, msgspec.json.decode(raw, type=Job), tokenizer, model) + except websockets.exceptions.InvalidHandshake: + sys.exit("handshake rejected; check --key and the server URL") + except (OSError, websockets.exceptions.ConnectionClosed) as e: + print(f"[connection lost ({e}); reconnecting in {backoff}s]", file=sys.stderr) + await asyncio.sleep(backoff) + backoff = min(backoff * 2, 60) + + +def main(): + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("server", help="server WebSocket base URL, e.g. ws://localhost:8080") + p.add_argument("--key", required=True, help="translator API key (site settings: translate_key)") + p.add_argument("--to", required=True, nargs="+", + help="target language(s): Seed-X names or codes, e.g. finnish de es") + args = p.parse_args() + + codes = [] + for lang in args.to: + low = lang.lower() + code = SEED_X_TAGS.get(low) or (low if low in SEED_X_NAMES else None) + if code is None: + p.error(f"unknown language {lang!r}; supported: {', '.join(sorted(SEED_X_TAGS))}") + if code not in codes: + codes.append(code) + + tokenizer, model = load_seed_x() # once, before the (re)connect loop + url = f"{args.server.rstrip('/')}/_translate/{args.key}" + asyncio.run(serve(url, codes, tokenizer, model)) + + +if __name__ == "__main__": + main()