75 lines
2.7 KiB
Python
75 lines
2.7 KiB
Python
 
|
#!/usr/bin/env python3
|
||
|
|
"""JSON-lines front end for normalizer.py.
|
||
|
|
|
||
|
|
The normalizer needs `transformers>=5`, which must not be installed into the
|
||
|
|
ComfyUI/Higgs environment, so it lives in its own venv and is driven as a
|
||
|
|
separate process. Loading a 1.1 GB model per subtitle block would be absurd, so
|
||
|
|
this is a long-lived worker: the caller starts it once per run, streams blocks
|
||
|
|
through it, and closes stdin when done.
|
||
|
|
|
||
|
|
Protocol — one JSON object per line, in both directions:
|
||
|
|
|
||
|
|
← {"ready": true, "mode": "model", "device": "cuda:0", "llm": true, …}
|
||
|
|
→ {"id": 1, "text": "Wygrał 1247 złotych."}
|
||
|
|
← {"id": 1, "text": "Wygrał tysiąc dwieście…", "changed": true,
|
||
|
|
"leg": "plain", "guard": "", "error": null}
|
||
|
|
→ {"quit": true} (or just close stdin)
|
||
|
|
|
||
|
|
`leg` is which model answered — "plain" for num2words plus the inflection
|
||
|
|
model, "complex" for the fine-tune, "skipped" when the block needed nothing.
|
||
|
|
`guard` names the check that rejected the model's answer, empty when it passed.
|
||
|
|
|
||
|
|
stdout carries the protocol and nothing else; every log line, warning and
|
||
|
|
traceback goes to stderr. A request that fails is answered with the original
|
||
|
|
text and an `error` field — the worker does not die on one bad block.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import json
|
||
|
|
import logging
|
||
|
|
import sys
|
||
|
|
|
||
|
|
import normalizer
|
||
|
|
|
||
|
|
|
||
|
|
def _send(obj) -> None:
|
||
|
|
sys.stdout.write(json.dumps(obj, ensure_ascii=False) + "\n")
|
||
|
|
sys.stdout.flush()
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> int:
|
||
|
|
logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s",
|
||
|
|
stream=sys.stderr)
|
||
|
|
|
||
|
|
state = normalizer.status() # loads the model, or reports why not
|
||
|
|
state["llm"] = normalizer.load_llm() # the inflection leg is optional
|
||
|
|
state["llm_dir"] = str(normalizer.LLM_DIR)
|
||
|
|
_send({"ready": True, **state})
|
||
|
|
|
||
|
|
for line in sys.stdin:
|
||
|
|
line = line.strip()
|
||
|
|
if not line:
|
||
|
|
continue
|
||
|
|
try:
|
||
|
|
req = json.loads(line)
|
||
|
|
except ValueError as exc:
|
||
|
|
_send({"id": None, "text": None, "error": f"bad request: {exc}"})
|
||
|
|
continue
|
||
|
|
if req.get("quit"):
|
||
|
|
break
|
||
|
|
|
||
|
|
text = req.get("text") or ""
|
||
|
|
leg, guard = "", ""
|
||
|
|
try:
|
||
|
|
result, leg, guard = normalizer.normalize_detail(text)
|
||
|
|
error = None
|
||
|
|
except Exception as exc: # normalize_detail() is not meant to raise
|
||
|
|
logging.exception("normalize_detail() raised")
|
||
|
|
result, error = text, f"{type(exc).__name__}: {exc}"
|
||
|
|
_send({"id": req.get("id"), "text": result, "changed": result != text,
|
||
|
|
"leg": leg, "guard": guard, "error": error})
|
||
|
|
return 0
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
sys.exit(main())
|