예제 소스

dataset_omook.py

#!/usr/bin/env python3
"""OMKL JSONL 로그 → VLM SFT 학습 데이터셋.

`vlm_pva_server.py --log-jsonl PATH` 가 누적한 한 수 단위 기록을 읽어,
(보드 이미지, 프롬프트, 정답좌표) 트리플로 변환합니다.

기본 필터:
  - kind == "move" 인 OMKL 레코드만 사용 (서버 predict 응답은 제외 — 모델 자기 응답이라 학습 가치 낮음).
  - 게임이 끝난 (BlackWin/WhiteWin/Draw) 게임만.
  - 그중 *승리한 쪽의 수* 만 학습 타깃 (전문가 시연 가정).
  - 사람/로컬 AI 의 수만 (--all-actors 로 VLM 도 포함 가능).

사용:
  python dataset_omook.py --jsonl logs/omook_vlm.jsonl \\
      --out-dir datasets/omook_sft --save-images
"""

from __future__ import annotations

import argparse
import json
import os
from concurrent.futures import ProcessPoolExecutor
from typing import Iterable, Iterator, List, Optional, Tuple

from PIL import Image, ImageDraw

VALID_SIZES = (13, 15, 19)

ACTOR_HUMAN = "human"
ACTOR_LOCAL_AI = "local_ai"
ACTOR_VLM = "vlm"

STATUS_PLAYING = "Playing"
STATUS_BLACK_WIN = "BlackWin"
STATUS_WHITE_WIN = "WhiteWin"
STATUS_DRAW = "Draw"


def _read_jsonl(path: str) -> Iterator[dict]:
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                yield json.loads(line)
            except json.JSONDecodeError:
                continue


def iter_move_records(path: str) -> Iterator[dict]:
    """OMKL move 레코드만 yield. (서버 predict 로그는 kind 필드 없음 → 자동 제외)"""
    for rec in _read_jsonl(path):
        if rec.get("kind") != "move":
            continue
        yield rec


def group_by_game(records: Iterable[dict]) -> Iterator[List[dict]]:
    """status != Playing 인 레코드를 게임 끝으로 보고 묶음."""
    current: List[dict] = []
    for rec in records:
        current.append(rec)
        if rec.get("status") in (STATUS_BLACK_WIN, STATUS_WHITE_WIN, STATUS_DRAW):
            yield current
            current = []
    # 마지막에 미완 게임이 남으면 버림 (학습엔 완료된 것만 사용).


def winner_stone(game: List[dict]) -> Optional[int]:
    """1=Black, 2=White, None=draw/incomplete."""
    if not game:
        return None
    last = game[-1].get("status")
    if last == STATUS_BLACK_WIN:
        return 1
    if last == STATUS_WHITE_WIN:
        return 2
    return None


# ----- 보드 이미지 렌더링 -----

def render_board_png(board: List[int], size: int, cell_px: int = 32, pad: int = 16) -> Image.Image:
    """클라이언트 BoardWidget 의 모양(나무바닥+격자+화점+돌) 을 모사한 PIL 렌더."""
    side_px = pad * 2 + cell_px * (size - 1)
    img = Image.new("RGB", (side_px, side_px), color=(222, 191, 130))
    draw = ImageDraw.Draw(img)
    grid_color = (120, 80, 30)

    for i in range(size):
        y = pad + i * cell_px
        draw.line([(pad, y), (pad + cell_px * (size - 1), y)], fill=grid_color, width=1)
    for j in range(size):
        x = pad + j * cell_px
        draw.line([(x, pad), (x, pad + cell_px * (size - 1))], fill=grid_color, width=1)

    # 화점 (중앙 + 네 모서리, BoardWidget 와 동일 규칙).
    mid = size // 2
    corner = 3
    far = size - 1 - corner
    hoshi = [(mid, mid)]
    if corner < far:
        hoshi += [(corner, corner), (corner, far), (far, corner), (far, far)]
    hoshi_r = max(2, cell_px // 12)
    for hr, hc in hoshi:
        cx = pad + hc * cell_px
        cy = pad + hr * cell_px
        draw.ellipse([cx - hoshi_r, cy - hoshi_r, cx + hoshi_r, cy + hoshi_r], fill=(60, 40, 18))

    # 돌
    rr = int(cell_px * 0.42)
    for r in range(size):
        for c in range(size):
            v = board[r * size + c]
            if v == 0:
                continue
            cx = pad + c * cell_px
            cy = pad + r * cell_px
            if v == 1:
                draw.ellipse(
                    [cx - rr, cy - rr, cx + rr, cy + rr],
                    fill=(20, 20, 20),
                    outline=(0, 0, 0),
                )
            else:
                draw.ellipse(
                    [cx - rr, cy - rr, cx + rr, cy + rr],
                    fill=(240, 240, 240),
                    outline=(90, 90, 90),
                )
    return img


# ----- 프롬프트 (서버 predict 와 같은 포맷) -----

def board_to_text(board: List[int], size: int) -> str:
    last_idx = size - 1
    lines = []
    for r in range(size):
        cells = []
        for c in range(size):
            v = board[r * size + c]
            cells.append("." if v == 0 else ("X" if v == 1 else "O"))
        lines.append(" ".join(cells))
    occ_b = [(r, c) for r in range(size) for c in range(size) if board[r * size + c] == 1]
    occ_w = [(r, c) for r in range(size) for c in range(size) if board[r * size + c] == 2]
    return (
        f"Board (X=Black, O=White, .=empty), rows top→bottom 0..{last_idx}, "
        f"cols left→right 0..{last_idx}:\n"
        + "\n".join(lines)
        + f"\nBlack stones: {occ_b}\nWhite stones: {occ_w}"
    )


def build_prompt(pre_board: List[int], size: int, ai_side: int) -> str:
    you_label = "Black (X)" if ai_side == 1 else "White (O)"
    opp_label = "White" if ai_side == 1 else "Black"
    return (
        f"You are a {size}x{size} Gomoku (오목) player. {opp_label} has just moved; "
        f"you play {you_label} next.\n"
        f"{board_to_text(pre_board, size)}\n"
        f"Choose your next {you_label} move on an EMPTY cell (a '.'). "
        "Reply with EXACTLY two integers separated by a single space: \"row col\". "
        f"row and col are both in 0..{size - 1}. No other text."
    )


# ----- 학습 예제 생성 -----

def iter_meta_examples(
    path: str,
    only_winners: bool = True,
    actors: Tuple[str, ...] = (ACTOR_HUMAN, ACTOR_LOCAL_AI),
    sizes: Tuple[int, ...] = VALID_SIZES,
) -> Iterator[dict]:
    """Yield {pre_board, size, stone, actor, winner, move_row, move_col} — 가벼운 메타.

    실제 이미지/프롬프트 생성은 _build_one_example 에서 별도로 (선택적으로 워커풀에서) 수행한다.
    """
    for game in group_by_game(iter_move_records(path)):
        winner = winner_stone(game)
        if only_winners and winner is None:
            continue
        for rec in game:
            size = rec.get("size")
            if size not in sizes:
                continue
            actor = rec.get("actor")
            if actors and actor not in actors:
                continue
            stone = rec.get("stone")
            if only_winners and winner is not None and stone != winner:
                continue
            board_post = rec.get("board")
            if not board_post or len(board_post) != size * size:
                continue
            mr = rec.get("move_row")
            mc = rec.get("move_col")
            if mr is None or mc is None:
                continue
            if not (0 <= mr < size and 0 <= mc < size):
                continue
            pre = list(board_post)
            idx = mr * size + mc
            if pre[idx] != stone:
                continue
            pre[idx] = 0
            yield {
                "pre_board": pre,
                "size": size,
                "stone": stone,
                "actor": actor,
                "winner": winner if winner is not None else 0,
                "move_row": mr,
                "move_col": mc,
            }


def _build_and_emit(task: tuple) -> dict:
    """워커 프로세스에서 호출 — 이미지를 디스크에 저장까지 끝내고 *작은* 메타만 반환.

    PIL.Image 를 마스터로 pickle 해 돌려보내는 비용이 PNG 저장보다 훨씬 커서, 워커
    내부에서 저장까지 마무리해야 병렬화 효과가 난다.
    """
    idx, meta, out_dir, save_images = task
    pre = meta["pre_board"]
    size = meta["size"]
    stone = meta["stone"]
    rec: dict = {
        "id": idx,
        "size": size,
        "stone": stone,
        "actor": meta["actor"],
        "winner": meta["winner"],
        "prompt": build_prompt(pre, size, ai_side=stone),
        "answer": f"{meta['move_row']} {meta['move_col']}",
    }
    if save_images:
        img_path = os.path.join(out_dir, f"img_{idx:08d}.png")
        render_board_png(pre, size).save(img_path, optimize=False, compress_level=1)
        rec["image_path"] = img_path
    return rec


def _build_one_example(meta: dict) -> dict:
    """단일 프로세스 호환 헬퍼 (이미지를 메모리로 들고 있는 옛 시그니처)."""
    pre = meta["pre_board"]
    size = meta["size"]
    stone = meta["stone"]
    return {
        "image": render_board_png(pre, size),
        "prompt": build_prompt(pre, size, ai_side=stone),
        "answer": f"{meta['move_row']} {meta['move_col']}",
        "size": size,
        "actor": meta["actor"],
        "winner": meta["winner"],
        "stone": stone,
    }


def make_examples(
    path: str,
    only_winners: bool = True,
    actors: Tuple[str, ...] = (ACTOR_HUMAN, ACTOR_LOCAL_AI),
    sizes: Tuple[int, ...] = VALID_SIZES,
) -> Iterator[dict]:
    """단일-프로세스 generator. 호환성용 — 새 코드는 iter_meta_examples 를 직접 쓰세요."""
    for meta in iter_meta_examples(path, only_winners=only_winners, actors=actors, sizes=sizes):
        yield _build_one_example(meta)


def main() -> None:
    parser = argparse.ArgumentParser(description="Build SFT dataset from omook JSONL log")
    parser.add_argument("--jsonl", required=True, help="OMKL/VLM JSONL 로그 경로")
    parser.add_argument(
        "--out-dir",
        default="sft_data/omook_sft",
        help=(
            "데이터셋 출력 디렉토리. 'datasets' 라는 이름은 PyPI 의 huggingface "
            "`datasets` 패키지와 namespace 충돌이 나므로 피하세요."
        ),
    )
    parser.add_argument(
        "--all-actors",
        action="store_true",
        help="VLM 의 수도 포함 (기본: human + local_ai 만)",
    )
    parser.add_argument(
        "--include-losers",
        action="store_true",
        help="패배한 쪽의 수도 포함 (기본: 승리한 쪽만)",
    )
    parser.add_argument(
        "--save-images",
        action="store_true",
        help="이미지를 PNG 파일로 dump (학습 시 디스크에서 lazy 로드 가능). "
        "끄면 train.jsonl 만 만들고 학습 단계에서 직접 렌더해야 함.",
    )
    parser.add_argument(
        "--workers",
        type=int,
        default=max(2, (os.cpu_count() or 4) // 2),
        help="이미지 렌더+저장을 병렬화할 프로세스 수 (기본: cpu/2). 1=직렬.",
    )
    args = parser.parse_args()

    os.makedirs(args.out_dir, exist_ok=True)
    actors = (
        (ACTOR_HUMAN, ACTOR_LOCAL_AI, ACTOR_VLM) if args.all_actors else (ACTOR_HUMAN, ACTOR_LOCAL_AI)
    )
    only_winners = not args.include_losers

    metas = list(iter_meta_examples(args.jsonl, only_winners=only_winners, actors=actors))
    out_jsonl = os.path.join(args.out_dir, "train.jsonl")

    n = len(metas)
    tasks = [(i, m, args.out_dir, args.save_images) for i, m in enumerate(metas)]
    with open(out_jsonl, "w", encoding="utf-8") as f:
        if args.workers <= 1 or n < 64:
            for task in tasks:
                rec = _build_and_emit(task)
                f.write(json.dumps(rec, ensure_ascii=False) + "\n")
        else:
            chunk = max(8, n // (args.workers * 4) or 1)
            with ProcessPoolExecutor(max_workers=args.workers) as pool:
                for rec in pool.map(_build_and_emit, tasks, chunksize=chunk):
                    f.write(json.dumps(rec, ensure_ascii=False) + "\n")

    summary = {
        "examples": n,
        "out_jsonl": out_jsonl,
        "out_dir": args.out_dir,
        "filters": {
            "only_winners": only_winners,
            "actors": list(actors),
            "save_images": args.save_images,
        },
    }
    with open(os.path.join(args.out_dir, "dataset_info.json"), "w", encoding="utf-8") as f:
        json.dump(summary, f, ensure_ascii=False, indent=2)
    print(f"wrote {n} examples → {out_jsonl}")


if __name__ == "__main__":
    main()