Seed-X translator service client for the /_translate socket
This commit is contained in:
@@ -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/<key>`` —
|
||||
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()
|
||||
Reference in New Issue
Block a user