예제 소스

dataset_moves_katago.py

#!/usr/bin/env python3
"""대국 국면을 KataGo 최선수로 라벨한 수 학습 데이터 (서버 수 프롬프트 그대로).

dataset_coords_game.py 와 같은 레코드 형식(sys/usr/expect/board)이라 eval_coords_game.py 로 평가한다.
  kmove        일반 수 프롬프트
  kmove_retry  금지 좌표(모델이 그 국면에서 실제로 고른 이미 돌 있는 칸)가 붙은 재시도 프롬프트

  python scripts/dataset_moves_katago.py --source logs/x/matches.sqlite --jsonl logs/x/gemma4.jsonl \
      --out sft_data/gemma_kmove_r1
"""

from __future__ import annotations

import argparse
import json
import random
import sys
import time
from pathlib import Path

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

from scripts.dataset_coords_game import HINTS, board_hash, load_bad_picks, load_positions  # noqa: E402
from vlm_server.dataset_baduk import assistant_text, read_pick_answer, read_pick_answer2  # noqa: E402
from vlm_server.dataset_onpolicy import analyze  # noqa: E402
from vlm_server.server import (  # noqa: E402
    SYSTEM_PROMPT_ALIGNED,
    SYSTEM_PROMPT_ALIGNED_READ,
    SYSTEM_PROMPT_ALIGNED_READ2,
    build_prompt,
)

N = 19


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--source", action="append", required=True)
    ap.add_argument("--jsonl", action="append", default=[])
    ap.add_argument("--out", required=True)
    ap.add_argument("--color", default="W", choices=["B", "W", "BW"])
    ap.add_argument("--retry-frac", type=float, default=0.3)
    ap.add_argument("--max-positions", type=int, default=2400)
    ap.add_argument("--val-n", type=int, default=120)
    ap.add_argument("--visits", type=int, default=64)
    ap.add_argument("--seed", type=int, default=31)
    ap.add_argument("--board-repr", default="aligned", choices=["aligned", "aligned_read", "aligned_read2"])
    ap.add_argument("--holdout-source", action="append", default=[],
                    help="이 대국 국면은 val 로만 쓴다 (지정하면 게임 단위 무작위 val 분할 대신)")
    args = ap.parse_args()

    from vlm_server.dataset_ladder import _KataGo

    rng = random.Random(args.seed)
    sides = {"B": {1}, "W": {2}, "BW": {1, 2}}[args.color]

    def uniq_positions(sources: list[str]) -> list[dict]:
        seen: set[str] = set()
        out = []
        for s in sources:
            for p in load_positions(Path(s)):
                h = board_hash(p["cells"], p["side"])
                if p["side"] in sides and h not in seen:
                    seen.add(h)
                    out.append(p)
        rng.shuffle(out)
        return out

    uniq = uniq_positions(args.source)[: args.max_positions]
    bad = load_bad_picks([Path(p) for p in args.jsonl])
    if args.holdout_source:
        held = uniq_positions(args.holdout_source)[: args.val_n]
        uniq += held
        val_games = {p["src"].rsplit(":", 1)[0] for p in held}
    else:
        games = sorted({p["src"].rsplit(":", 1)[0] for p in uniq})
        rng.shuffle(games)
        val_games = set(games[: max(1, len(games) // 10)])
    system = {
        "aligned_read": SYSTEM_PROMPT_ALIGNED_READ,
        "aligned_read2": SYSTEM_PROMPT_ALIGNED_READ2,
    }.get(args.board_repr, SYSTEM_PROMPT_ALIGNED)

    kg = _KataGo(args.visits, 6.5)
    train, val = [], []
    t0 = time.time()
    for i, pos in enumerate(uniq, 1):
        cells, side = pos["cells"], pos["side"]
        v = analyze(kg, cells, side, 3)
        if not v["top"] or v["top_is_pass"]:
            continue
        best = tuple(v["top"][0])
        retry = rng.random() < args.retry_frac
        forbidden: list[tuple[int, int]] = []
        hint = None
        if retry:
            occ = sorted({q for q in bad.get(board_hash(cells, side), []) if cells[q[0] * N + q[1]]})
            if not occ:
                own = [(k // N, k % N) for k, x in enumerate(cells) if x == side]
                occ = rng.sample(own, min(len(own), rng.randint(1, 3))) if own else []
            forbidden = [q for q in occ if q != best][:6]
            hint = rng.choice(HINTS) if forbidden else None
        user = build_prompt(bytes(cells), side, forbidden=forbidden or None, illegal_hint=hint,
                            board_repr=args.board_repr)
        task = "kmove_retry" if forbidden else "kmove"
        if args.board_repr == "aligned_read2":
            answer = read_pick_answer2(bytes(cells), best[0], best[1])
        elif args.board_repr == "aligned_read":
            answer = read_pick_answer(bytes(cells), best[0], best[1])
        else:
            answer = assistant_text(best)
        rec = {
            "task": task,
            "src": pos["src"],
            "sys": system,
            "usr": user,
            "expect": {"row": best[0], "col": best[1], "side": side, "forbidden": forbidden,
                       "top": v["top"], "lead": v["lead"]},
            "board": "".join(".XO"[x] for x in cells),
            "images": [],
            "messages": [
                {"role": "user", "content": [{"type": "text", "text": f"{system}\n\n{user}"}]},
                {"role": "assistant", "content": [{"type": "text", "text": answer}]},
            ],
        }
        (val if pos["src"].rsplit(":", 1)[0] in val_games else train).append(rec)
        if i % 200 == 0:
            print(f"[kmove] {i}/{len(uniq)} train={len(train)} val={len(val)} {time.time() - t0:.0f}s", flush=True)

    rng.shuffle(val)
    val = val[: args.val_n]
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    for name, rows in (("train", train), ("val", val)):
        with (out / f"{name}.jsonl").open("w", encoding="utf-8") as f:
            for r in rows:
                f.write(json.dumps(r, ensure_ascii=False) + "\n")
    summary = {
        "positions": len(uniq),
        "train": len(train),
        "train_retry": sum(1 for r in train if r["task"] == "kmove_retry"),
        "val": len(val),
        "seconds": round(time.time() - t0),
    }
    (out / "summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(f"[kmove] {summary}", flush=True)
    return 0 if (train or val) else 1


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