twlat/cli.py

"""twlat 命令列介面。

    twlat "这个程序有bug"              # 單段轉換
    cat in.txt | twlat                  # 從 stdin 逐行轉換
    twlat -i in.txt -o out.txt          # 檔案轉檔案
    twlat --preset taiwanize "视频"     # 選操作點
    twlat --explain "他在那里"           # 顯示逐項決策
"""
from __future__ import annotations

import argparse
import json
import sys


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(
        prog="twlat",
        description="中國大陸中文 → 臺灣正體中文（確定性、可解釋、離線）")
    ap.add_argument("text", nargs="*", help="要轉換的文字；省略則讀 stdin")
    ap.add_argument("-i", "--input", help="輸入檔（每行一段）")
    ap.add_argument("-o", "--output", help="輸出檔")
    ap.add_argument("--preset", default="balanced",
                    choices=["accuracy", "balanced", "taiwanize", "aggressive"],
                    help="操作點（預設 balanced）")
    ap.add_argument("--device", default="cpu",
                    help="cpu/mps/cuda（預設 cpu；單執行緒吞吐最佳）")
    ap.add_argument("--threads", type=int, default=1,
                    help="torch 執行緒數（預設 1，實測最快）")
    ap.add_argument("--batch-size", type=int, default=8)
    ap.add_argument("--explain", action="store_true", help="輸出逐項決策 JSON")
    ap.add_argument("--ckpt", default=None)
    ap.add_argument("--version", action="store_true")
    a = ap.parse_args(argv)

    import twlat
    if a.version:
        print(twlat.__version__)
        return 0

    import torch
    torch.set_num_threads(max(1, a.threads))

    if a.text:
        lines = [" ".join(a.text)]
    elif a.input:
        with open(a.input, encoding="utf-8") as fh:
            lines = [ln.rstrip("\n") for ln in fh]
    else:
        lines = [ln.rstrip("\n") for ln in sys.stdin]
    if not lines:
        return 0

    conv = twlat.Converter(ckpt=a.ckpt, device=a.device, preset=a.preset)

    if a.explain:
        out = []
        for r in conv.explain_batch(lines, batch_size=a.batch_size):
            out.append({"text": r.text,
                        "decisions": [{"span": [d.start, d.end],
                                       "from": d.source, "to": d.target,
                                       "utility": round(d.utility, 3),
                                       "rule_type": d.rule_type}
                                      for d in r.decisions]})
        payload = json.dumps(out, ensure_ascii=False, indent=1)
    else:
        payload = "\n".join(conv.convert_batch(lines, batch_size=a.batch_size))

    if a.output:
        with open(a.output, "w", encoding="utf-8") as fh:
            fh.write(payload + "\n")
    else:
        print(payload)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())