예제 소스

replay_pass_cases.py

"""pass_cases.jsonl 의 R5 PASS 국면을 현재 서버 코드로 다시 판정한다.

모델 출력은 PASS로 고정(가장 불리한 경우)하고, 서버의 계가·안전 필터만으로
착수가 나오는지 본다. --katago 를 주면 그 착수와 PASS를 KataGo로 비교한다
(gain = 착수 후 백 기준 집 차이 - PASS 후 백 기준 집 차이).

  .venv/bin/python scripts/replay_pass_cases.py \
      logs/heuristic_r5_nolimit_recheck_20261005/pass_cases.jsonl --katago
"""

from __future__ import annotations

import argparse
import json
import sys
import threading
import types
from collections import Counter
from pathlib import Path
from typing import Optional

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from vlm_server.server import ServerConfig, _GoTcpHandler, _Logger  # noqa: E402

N = 19
WHITE = 2


class _PassModel:
    def predict_move(self, *args, **kwargs):
        return None, None, True


class _Sink:
    def sendall(self, data: bytes) -> None:
        pass


def server_move(cells: bytes, side: int = WHITE) -> tuple[bool, Optional[tuple[int, int]]]:
    cfg = ServerConfig(
        host="127.0.0.1",
        port=0,
        model_id="replay",
        max_new_tokens=1,
        temperature=0.0,
        load_model=True,
        log_jsonl_path=None,
    )
    handler = _GoTcpHandler.__new__(_GoTcpHandler)
    handler.server = types.SimpleNamespace(
        cfg=cfg,
        model=_PassModel(),
        policy=None,
        policy_lock=threading.Lock(),
        logger=_Logger(None),
    )
    return handler._handle_one(_Sink(), None, cells, N, side, [], None)


def katago_gain(kg, cells: bytes, move: tuple[int, int], side: int = WHITE) -> Optional[float]:
    from vlm_server.dataset_ladder import _gtp

    player = "B" if side == 1 else "W"
    stones = [["B" if v == 1 else "W", _gtp(divmod(i, N))] for i, v in enumerate(cells) if v]
    flip = -1.0 if side == 2 else 1.0
    lead: dict[str, float] = {}
    for mv in (_gtp(move), "pass"):
        res = kg.j._query(
            {
                "initialStones": stones,
                "initialPlayer": player,
                "moves": [],
                "rules": "korean",
                "komi": kg.komi,
                "boardXSize": N,
                "boardYSize": N,
                "analyzeTurns": [0],
                "maxVisits": kg.visits,
                "allowMoves": [{"player": player, "moves": [mv], "untilDepth": 1}],
            }
        )
        infos = res.get("moveInfos") or []
        if not infos:
            return None
        lead[mv] = flip * float(infos[0].get("scoreLead", 0.0))
    return round(lead[_gtp(move)] - lead["pass"], 2)


def main(argv: Optional[list[str]] = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("cases", nargs="+")
    ap.add_argument("--katago", action="store_true")
    ap.add_argument("--katago-visits", type=int, default=64)
    ap.add_argument("--out", default="", help="결과 jsonl (기본: 첫 입력 옆 replay_fixed.jsonl)")
    args = ap.parse_args(argv)

    kg = None
    if args.katago:
        from vlm_server.dataset_ladder import _KataGo

        kg = _KataGo(args.katago_visits, 6.5)
    out_path = Path(args.out) if args.out else Path(args.cases[0]).with_name("replay_fixed.jsonl")
    counts: Counter = Counter()
    gains: dict[str, list[float]] = {"midgame": [], "endgame": []}
    with out_path.open("w", encoding="utf-8") as fo:
        for path in args.cases:
            for line in Path(path).open(encoding="utf-8"):
                if not line.strip():
                    continue
                d = json.loads(line)
                cells = bytes(int(v) for v in d["board"])
                phase = d.get("phase", "midgame")
                is_pass, move = server_move(cells)
                rec = {
                    "game_id": d.get("game_id"),
                    "ply": d.get("ply"),
                    "empty": d.get("empty"),
                    "phase": phase,
                    "old": d.get("result"),
                    "fixed_pass": bool(is_pass),
                    "fixed_move": list(move) if move else None,
                }
                counts[f"{phase}:{'pass' if is_pass else 'move'}"] += 1
                if kg is not None and move is not None:
                    g = katago_gain(kg, cells, move)
                    rec["katago_gain"] = g
                    if g is not None:
                        gains[phase].append(g)
                fo.write(json.dumps(rec, ensure_ascii=False) + "\n")
    summary: dict = {"counts": dict(counts), "out": str(out_path)}
    for phase, gs in gains.items():
        if gs:
            summary[f"{phase}_gain_avg"] = round(sum(gs) / len(gs), 2)
            summary[f"{phase}_gain_pos_rate"] = round(sum(1 for g in gs if g > 0) / len(gs), 3)
            summary[f"{phase}_n"] = len(gs)
    print(json.dumps(summary, ensure_ascii=False, indent=2))
    return 0


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