예제 소스

dataset_heuristic_teacher.py

#!/usr/bin/env python3
"""계가 휴리스틱이 실제로 둔 수를 Gemma 학습 정답으로 만든다.

기본 원본은 모델 제안이 아니라 서버가 계가로 고른 수다.
"""

from __future__ import annotations

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

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

from vlm_server.dataset_baduk import assistant_text, user_prompt_text  # noqa: E402

N = 19


def edge_line(row: int, col: int) -> int:
    return min(row, col, N - 1 - row, N - 1 - col) + 1


def cells_of(text: str) -> bytes:
    return bytes(".XO".index(ch) for ch in text)


_mon = None


def monitor():
    global _mon
    if _mon is None:
        import importlib.util

        path = ROOT / "scripts" / "monitor_mariago_web_play.py"
        spec = importlib.util.spec_from_file_location("monitor_mariago_web_play", path)
        if spec is None or spec.loader is None:
            raise RuntimeError("heuristic import failed")
        mod = importlib.util.module_from_spec(spec)
        sys.modules["monitor_mariago_web_play"] = mod
        spec.loader.exec_module(mod)
        _mon = mod
    return _mon


def rows_from(db: Path, colors: str, relabel: bool = False) -> list[dict]:
    con = sqlite3.connect(db)
    found = con.execute(
        """select game_id, ply, color, kind, row, col, board
           from moves
           where kind in ('move', 'pass') and color in ({})
           order by game_id, ply""".format(",".join("?" * len(colors))),
        tuple(colors),
    ).fetchall()
    con.close()
    rows = []
    for gid, ply, color, kind, row, col, board in found:
        if not board or len(board) != N * N:
            continue
        try:
            cells = cells_of(board)
        except ValueError:
            continue
        side = 1 if color == "B" else 2
        point = None if kind == "pass" else (int(row), int(col))
        if relabel:
            board = list(cells)
            hr, hc, is_pass = monitor().heuristic_move(
                board, side, None, random.Random(int(gid) * 1000 + int(ply))
            )
            point = None if is_pass or hr is None or hc is None else (int(hr), int(hc))
        rows.append(
            {
                "id": 0,
                "size": N,
                "ai_side": side,
                "target_row": None if point is None else point[0],
                "target_col": None if point is None else point[1],
                "is_pass": point is None,
                "is_tactical": False,
                "board_format": "aligned",
                "text_only": True,
                "images": [],
                "onpolicy": {
                    "kind": "heuristic_pass" if point is None else "heuristic_move",
                    "game_id": gid,
                    "ply": ply,
                    "source": str(db),
                },
                "messages": [
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "text",
                                "text": user_prompt_text(cells, side, board_format="aligned", text_only=True),
                            }
                        ],
                    },
                    {"role": "assistant", "content": [{"type": "text", "text": assistant_text(point)}]},
                ],
            }
        )
    return rows


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--source", action="append", required=True, help="sqlite경로:색깔  예: logs/a.sqlite:BW")
    ap.add_argument("--out", default="sft_data/heuristic_gemma4")
    ap.add_argument("--val-frac", type=float, default=0.05)
    ap.add_argument("--pass-repeat", type=int, default=4)
    ap.add_argument("--seed", type=int, default=13)
    ap.add_argument("--relabel", action="store_true", help="기록된 수 대신 휴리스틱이 둘 수를 정답으로 한다")
    args = ap.parse_args()

    raw: list[dict] = []
    for item in args.source:
        path_s, colors = item.rsplit(":", 1)
        raw.extend(rows_from(Path(path_s), colors, relabel=args.relabel))
    seen = set()
    uniq = []
    for row in raw:
        key = (
            row["ai_side"],
            row["is_pass"],
            row["target_row"],
            row["target_col"],
            row["messages"][0]["content"][0]["text"],
        )
        if key in seen:
            continue
        seen.add(key)
        row["id"] = len(uniq)
        uniq.append(row)

    rng = random.Random(args.seed)
    order = list(range(len(uniq)))
    rng.shuffle(order)
    n_val = max(1, int(len(order) * args.val_frac)) if len(order) > 1 else 0
    val_ids = {uniq[i]["id"] for i in order[:n_val]}
    train_rows = [r for r in uniq if r["id"] not in val_ids]
    val_rows = [r for r in uniq if r["id"] in val_ids]
    passes = [r for r in train_rows if r["is_pass"]]
    train_rows = train_rows + passes * max(0, args.pass_repeat - 1)
    rng.shuffle(train_rows)

    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    with (out / "train.jsonl").open("w", encoding="utf-8") as ft:
        for row in train_rows:
            ft.write(json.dumps(row, ensure_ascii=False) + "\n")
    with (out / "val.jsonl").open("w", encoding="utf-8") as fv:
        for row in val_rows:
            fv.write(json.dumps(row, ensure_ascii=False) + "\n")
    train_n = len(train_rows)
    val_n = len(val_rows)
    label_lines: dict[str, int] = {}
    for row in uniq:
        if row["is_pass"]:
            key = "pass"
        else:
            key = str(edge_line(int(row["target_row"]), int(row["target_col"])))
        label_lines[key] = label_lines.get(key, 0) + 1
    summary = {
        "raw": len(raw),
        "unique": len(uniq),
        "passes": sum(1 for r in uniq if r["is_pass"]),
        "train": train_n,
        "val": val_n,
        "label_lines": label_lines,
    }
    (out / "summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(f"[heuristic-data] {summary}", flush=True)
    return 0 if uniq else 1


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