예제 소스

dataset_r5_moves.py

#!/usr/bin/env python3
"""R5가 sqlite에 둔 수를 학습 레코드로 만들고, 좌표 혼합 데이터에 넣는다.

휴리스틱 대국에서 R5는 백이다. gemma4 대국은 server_rec.model 이 r5 인 수만 쓴다.
패스·불법·이미 돌이 있는 칸은 넣지 않는다.

  python scripts/dataset_r5_moves.py
"""

from __future__ import annotations

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 scripts.dataset_coords_game import mon, record  # noqa: E402
from vlm_server.dataset_baduk import assistant_text  # noqa: E402
from vlm_server.server import SYSTEM_PROMPT_ALIGNED, build_prompt  # noqa: E402

N = 19
MIX = ROOT / "sft_data" / "mix_gemma_coord_r2" / "train.jsonl"
# 돌아가는 학습 루프가 앞에서부터 2600개만 읽는다.
TRAIN_N = 2600
R5_N = 600
COORD = {"lookup", "row", "col", "locate", "locate_hard"}
SOURCES = (
    (ROOT / "logs/heuristic_r5_eye_20261008/matches.sqlite", "W", None),
    (ROOT / "logs/heuristic_r5_sqlite_nolimit_20261007/matches.sqlite", "W", None),
    (ROOT / "logs/gemma4_r5_20261009/matches.sqlite", "W", "r5"),
)


def r5_record(cells: bytes, row: int, col: int, src: str) -> dict | None:
    if not (0 <= row < N and 0 <= col < N) or len(cells) != N * N:
        return None
    if cells[row * N + col] != 0:
        return None
    ok, *_ = mon.try_play(list(cells), row, col, 2, None)
    if not ok:
        return None
    user = build_prompt(cells, 2, board_repr="aligned")
    return record(
        "r5move",
        SYSTEM_PROMPT_ALIGNED,
        user,
        assistant_text((row, col)),
        src,
        {"row": row, "col": col, "side": 2},
    )


def load_r5_moves(sources=SOURCES) -> list[dict]:
    seen: set[tuple[str, int, int]] = set()
    out: list[dict] = []
    for path, color, model in sources:
        con = sqlite3.connect(path)
        rows = con.execute(
            "select game_id, ply, color, kind, row, col, board, server_rec "
            "from moves where kind='move' and board is not null order by game_id, ply"
        )
        for gid, ply, move_color, _kind, row, col, board, server_rec in rows:
            if move_color != color or not board or len(board) != N * N:
                continue
            if model is not None:
                try:
                    got = json.loads(server_rec or "{}").get("model")
                except json.JSONDecodeError:
                    continue
                if got != model:
                    continue
            if row is None or col is None:
                continue
            key = (board, int(row), int(col))
            if key in seen:
                continue
            seen.add(key)
            try:
                cells = bytes(".XO".index(ch) for ch in board)
            except ValueError:
                continue
            rec = r5_record(cells, int(row), int(col), f"{path.parent.name}:{gid}:{ply}")
            if rec:
                rec["board"] = board
                out.append(rec)
        con.close()
    return out


def mix_with_r5(existing: list[dict], r5: list[dict], train_n: int = TRAIN_N, r5_n: int = R5_N, seed: int = 21) -> list[dict]:
    coord = [r for r in existing if r.get("task") in COORD]
    moves = [r for r in existing if r.get("task") not in COORD and r.get("task") != "r5move"]
    rng = random.Random(seed)
    rng.shuffle(r5)
    chosen_r5 = r5[:r5_n]
    room = train_n - len(coord) - len(chosen_r5)
    if room < 0:
        raise SystemExit(f"coord {len(coord)} + r5 {len(chosen_r5)} exceeds {train_n}")
    rng.shuffle(moves)
    chosen = coord + moves[:room] + chosen_r5
    rng.shuffle(chosen)
    return chosen


def main() -> int:
    r5 = load_r5_moves()
    existing = [json.loads(x) for x in MIX.read_text(encoding="utf-8").splitlines() if x.strip()]
    mixed = mix_with_r5(existing, r5)
    MIX.write_text("".join(json.dumps(r, ensure_ascii=False) + "\n" for r in mixed), encoding="utf-8")
    from collections import Counter
    counts = Counter(r.get("task") or "none" for r in mixed)
    print(f"r5_legal {len(r5)} train {len(mixed)} {dict(counts)}", flush=True)
    return 0


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