266 lines
10 KiB
Python
Executable File
266 lines
10 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""hymt2_batch.py — batch Hy-MT2 translation.
|
|
|
|
Reads a JSON input listing source strings, runs each through ComfyUI, writes
|
|
one JSON output per item under --out-dir. Already-existing outputs are skipped
|
|
unless --force is set (resumable).
|
|
|
|
Input formats accepted:
|
|
1. A list of strings: ["text 1", "text 2", ...]
|
|
2. A list of objects with "text": [{"id": "a", "text": "..."}, ...]
|
|
3. A dict mapping id → text: {"a": "text 1", "b": "text 2"}
|
|
|
|
Usage:
|
|
/mnt/modelall/ComfyUI/.venv/bin/python hymt2_batch.py vocab.json \\
|
|
--out-dir /mnt/modelall/results/hymt2_vocab \\
|
|
--target Japanese --prompt-lang en --ip 10.8.0.12 --limit 20
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import re
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
from pathlib import Path
|
|
|
|
# Re-use the workflow + client from the single-shot CLI in this same dir.
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
|
from hymt2_cli import ( # noqa: E402
|
|
ComfyClient,
|
|
build_workflow,
|
|
extract_text_from_history,
|
|
)
|
|
|
|
_PREFERRED_TMP = "/mnt/modelall/tmp"
|
|
if os.path.isdir(_PREFERRED_TMP):
|
|
tempfile.tempdir = _PREFERRED_TMP
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# I/O helpers
|
|
# ---------------------------------------------------------------------------
|
|
_BAD_FN_RE = re.compile(r"[^\w\-\.\u4e00-\u9fff]+") # keep word chars + CJK
|
|
|
|
|
|
def sanitize_filename(s: str, *, max_len: int = 80) -> str:
|
|
"""v2.0 §12 套路 C: CJK + special-character-safe filenames."""
|
|
cleaned = _BAD_FN_RE.sub("_", s.strip()).strip("_") or "item"
|
|
return cleaned[:max_len]
|
|
|
|
|
|
def load_items(path: Path) -> list[dict]:
|
|
"""Normalize the input JSON into a list of {"id": ..., "text": ...}."""
|
|
data = json.loads(path.read_text(encoding="utf-8"))
|
|
items: list[dict] = []
|
|
if isinstance(data, list):
|
|
for i, entry in enumerate(data):
|
|
if isinstance(entry, str):
|
|
items.append({"id": f"item_{i:06d}", "text": entry})
|
|
elif isinstance(entry, dict):
|
|
txt = entry.get("text") or entry.get("source") or entry.get("content")
|
|
if not txt:
|
|
print(f"[hymt2/batch] skipping entry {i} — no text field",
|
|
file=sys.stderr)
|
|
continue
|
|
eid = entry.get("id") or sanitize_filename(txt)[:40] or f"item_{i:06d}"
|
|
items.append({"id": str(eid), "text": txt})
|
|
else:
|
|
print(f"[hymt2/batch] skipping entry {i} of type {type(entry)}",
|
|
file=sys.stderr)
|
|
elif isinstance(data, dict):
|
|
for k, v in data.items():
|
|
if isinstance(v, str):
|
|
items.append({"id": str(k), "text": v})
|
|
elif isinstance(v, dict):
|
|
txt = v.get("text") or v.get("source") or v.get("content")
|
|
if txt:
|
|
items.append({"id": str(k), "text": txt})
|
|
else:
|
|
raise ValueError(f"unsupported input shape: {type(data)}")
|
|
return items
|
|
|
|
|
|
def run_one(client: ComfyClient, *, source_text: str, args, force_reload: bool):
|
|
wf = build_workflow(
|
|
model_path=args.model_path,
|
|
device=args.device,
|
|
dtype=args.dtype,
|
|
source_text=source_text,
|
|
target_lang=args.target,
|
|
prompt_lang=args.prompt_lang,
|
|
temperature=args.temperature,
|
|
top_p=args.top_p,
|
|
top_k=args.top_k,
|
|
repetition_penalty=args.repetition_penalty,
|
|
max_new_tokens=args.max_new_tokens,
|
|
seed=args.seed,
|
|
terminology=args.terminology_text,
|
|
target_style=args.style,
|
|
background_text=args.background_text,
|
|
force_reload=force_reload,
|
|
)
|
|
history = client.run(wf, timeout_s=args.timeout)
|
|
text = extract_text_from_history(history, "preview")
|
|
if not text:
|
|
text = extract_text_from_history(history, "translate")
|
|
return text
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(description="Hy-MT2 batch translation.")
|
|
ap.add_argument("input_json", type=Path,
|
|
help="Input JSON (list of strings, list of {id,text}, or dict).")
|
|
|
|
ap.add_argument("--ip", default="127.0.0.1")
|
|
ap.add_argument("--port", type=int, default=8188)
|
|
ap.add_argument("--server", default=None)
|
|
|
|
ap.add_argument("--model-path", default="tencent/Hy-MT2-1.8B")
|
|
ap.add_argument("--gpu", type=int, default=0)
|
|
ap.add_argument("--device", default=None)
|
|
ap.add_argument("--dtype", default="bf16", choices=["bf16", "fp16", "fp32"])
|
|
|
|
ap.add_argument("--target", default="English")
|
|
ap.add_argument("--prompt-lang", default="en", choices=["en", "zh"])
|
|
ap.add_argument("--temperature", type=float, default=0.7)
|
|
ap.add_argument("--top-p", type=float, default=0.6)
|
|
ap.add_argument("--top-k", type=int, default=20)
|
|
ap.add_argument("--repetition-penalty", type=float, default=1.05)
|
|
ap.add_argument("--max-new-tokens", type=int, default=4096)
|
|
ap.add_argument("--seed", type=int, default=-1)
|
|
|
|
ap.add_argument("--terminology", default="",
|
|
help="Path or inline glossary string applied to every item.")
|
|
ap.add_argument("--style", default="")
|
|
ap.add_argument("--background", default="")
|
|
|
|
ap.add_argument("--out-dir", type=Path, required=True,
|
|
help="Output directory. One JSON per item.")
|
|
ap.add_argument("--limit", type=int, default=0,
|
|
help="Process at most N items (0 = no limit).")
|
|
ap.add_argument("--start", type=int, default=0,
|
|
help="Skip the first N items.")
|
|
ap.add_argument("--force", action="store_true",
|
|
help="Overwrite existing output files.")
|
|
ap.add_argument("--retries", type=int, default=2)
|
|
ap.add_argument("--timeout", type=float, default=1800.0)
|
|
ap.add_argument("--ws-recv-timeout", type=float, default=600.0)
|
|
ap.add_argument("--dump-workflow", type=Path, default=None,
|
|
help="Dump the workflow of the FIRST item only (debug).")
|
|
ap.add_argument("-v", "--verbose", action="store_true",
|
|
help="Print every received WebSocket message.")
|
|
|
|
args = ap.parse_args()
|
|
args.server = args.server or f"{args.ip}:{args.port}"
|
|
args.device = args.device or f"cuda:{args.gpu}"
|
|
|
|
# Resolve terminology / background once.
|
|
def _resolve(value: str) -> str:
|
|
if not value:
|
|
return ""
|
|
p = Path(value)
|
|
if p.exists() and p.is_file():
|
|
return p.read_text(encoding="utf-8")
|
|
return value
|
|
|
|
args.terminology_text = _resolve(args.terminology)
|
|
args.background_text = _resolve(args.background)
|
|
|
|
args.out_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
items = load_items(args.input_json)
|
|
print(f"[hymt2/batch] loaded {len(items)} items", file=sys.stderr)
|
|
if args.start:
|
|
items = items[args.start:]
|
|
if args.limit:
|
|
items = items[:args.limit]
|
|
print(f"[hymt2/batch] will process {len(items)} items "
|
|
f"(start={args.start} limit={args.limit})", file=sys.stderr)
|
|
|
|
client = ComfyClient(server=args.server, ws_recv_timeout=args.ws_recv_timeout,
|
|
verbose=args.verbose)
|
|
|
|
stats = {"ok": 0, "skip": 0, "fail": 0}
|
|
t_run_start = time.time()
|
|
first = True
|
|
|
|
for idx, item in enumerate(items):
|
|
eid = sanitize_filename(item["id"])
|
|
out_path = args.out_dir / f"{eid}.json"
|
|
if out_path.exists() and not args.force:
|
|
stats["skip"] += 1
|
|
continue
|
|
|
|
if first and args.dump_workflow:
|
|
wf_preview = build_workflow(
|
|
model_path=args.model_path, device=args.device, dtype=args.dtype,
|
|
source_text=item["text"], target_lang=args.target,
|
|
prompt_lang=args.prompt_lang, temperature=args.temperature,
|
|
top_p=args.top_p, top_k=args.top_k,
|
|
repetition_penalty=args.repetition_penalty,
|
|
max_new_tokens=args.max_new_tokens, seed=args.seed,
|
|
terminology=args.terminology_text, target_style=args.style,
|
|
background_text=args.background_text, force_reload=False,
|
|
)
|
|
args.dump_workflow.parent.mkdir(parents=True, exist_ok=True)
|
|
args.dump_workflow.write_text(
|
|
json.dumps(wf_preview, ensure_ascii=False, indent=2)
|
|
)
|
|
print(f"[hymt2/batch] dumped first workflow to {args.dump_workflow}",
|
|
file=sys.stderr)
|
|
|
|
attempt = 0
|
|
text = None
|
|
last_err = None
|
|
while attempt <= args.retries:
|
|
try:
|
|
text = run_one(client, source_text=item["text"], args=args,
|
|
force_reload=False)
|
|
break
|
|
except Exception as e:
|
|
last_err = e
|
|
attempt += 1
|
|
print(f"[hymt2/batch] item={eid} attempt={attempt} "
|
|
f"failed: {e!r}", file=sys.stderr)
|
|
time.sleep(2.0)
|
|
|
|
first = False
|
|
|
|
if text is None:
|
|
stats["fail"] += 1
|
|
print(f"[hymt2/batch] item={eid} GIVE UP after {args.retries} retries: "
|
|
f"{last_err!r}", file=sys.stderr)
|
|
continue
|
|
|
|
payload = {
|
|
"id": item["id"],
|
|
"source": item["text"],
|
|
"translation": text,
|
|
"target_lang": args.target,
|
|
"prompt_lang": args.prompt_lang,
|
|
"model_path": args.model_path,
|
|
}
|
|
out_path.write_text(json.dumps(payload, ensure_ascii=False, indent=2),
|
|
encoding="utf-8")
|
|
stats["ok"] += 1
|
|
|
|
elapsed = time.time() - t_run_start
|
|
done = stats["ok"] + stats["fail"]
|
|
if done > 0:
|
|
avg = elapsed / done
|
|
eta = avg * (len(items) - idx - 1)
|
|
print(f"[hymt2/batch] [{idx + 1}/{len(items)}] {eid} "
|
|
f"ok={stats['ok']} fail={stats['fail']} skip={stats['skip']} "
|
|
f"avg={avg:.1f}s/item eta={eta / 60:.1f}min",
|
|
file=sys.stderr)
|
|
|
|
print(f"[hymt2/batch] DONE ok={stats['ok']} fail={stats['fail']} "
|
|
f"skip={stats['skip']}", file=sys.stderr)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|