예제 소스

dataset_opening.py

"""오프닝(첫 수) 강제 학습용 합성 데이터셋 생성기.

목표:
- 빈 19x19 바둑판에서 "첫 수는 화점(4-4)"에 두도록 SFT 샘플을 대량 생성.
- 기존 `dataset_baduk.py`가 만드는 JSONL 포맷과 동일한 `messages`/`images` 스키마를 사용.

사용 예:
  python -m vlm_server.dataset_opening --out-dir sft_data/opening_hoshi --samples 20000

그 다음 기존 데이터와 섞기(예):
  cat sft_data/baduk_v2/train.jsonl sft_data/opening_hoshi/train.jsonl > sft_data/mix/train.jsonl
"""

from __future__ import annotations

import argparse
import json
import os
import random
from pathlib import Path

from .dataset_baduk import BOARD_SIZE, IMG_PX, render_board_png
from .server import build_prompt


HOSHI_19 = [
    (3, 3),
    (3, 9),
    (3, 15),
    (9, 3),
    (9, 9),
    (9, 15),
    (15, 3),
    (15, 9),
    (15, 15),
]


def _assistant_text(point: tuple[int, int]) -> str:
    r, c = point
    return json.dumps({"row": int(r), "col": int(c)}, ensure_ascii=False)


def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(description="빈판 오프닝(화점) SFT JSONL 생성")
    p.add_argument("--out-dir", required=True, help="출력 폴더 (train.jsonl + 빈판 PNG)")
    p.add_argument("--samples", type=int, default=20000, help="생성 샘플 수 (기본 2만)")
    p.add_argument(
        "--seed",
        type=int,
        default=0,
        help="랜덤 시드 (0이면 시드 고정 안 함)",
    )
    p.add_argument(
        "--hoshi",
        choices=["corner4", "all9"],
        default="corner4",
        help="정답 후보 화점 집합. corner4=4개 코너(4-4), all9=9개 화점(코너/변/천원)",
    )
    args = p.parse_args(argv)

    out_dir = Path(args.out_dir).resolve()
    out_dir.mkdir(parents=True, exist_ok=True)

    if args.seed:
        random.seed(int(args.seed))

    candidates = HOSHI_19
    if args.hoshi == "corner4":
        candidates = [(3, 3), (3, 15), (15, 3), (15, 15)]

    # 빈판 생성 + PNG 한 장만 저장 (모든 샘플이 동일 이미지 참조)
    empty = bytes([0]) * (BOARD_SIZE * BOARD_SIZE)
    img = render_board_png(empty, BOARD_SIZE, px=IMG_PX)
    img_path = out_dir / "empty.png"
    img.save(img_path, "PNG", optimize=True)

    # user prompt는 서버와 동일한 형태를 사용 (빈판 + 차례=흑 기준)
    user_text = build_prompt(empty, ai_side=1)

    train_path = out_dir / "train.jsonl"
    with open(train_path, "w", encoding="utf-8") as f:
        for i in range(int(args.samples)):
            pt = random.choice(candidates)
            rec = {
                "id": i,
                "size": BOARD_SIZE,
                "ai_side": 1,  # 첫 수 = 흑
                "target_row": pt[0],
                "target_col": pt[1],
                "is_pass": False,
                "is_tactical": False,
                "images": [str(img_path)],
                "messages": [
                    {
                        "role": "user",
                        "content": [
                            {"type": "image", "image": str(img_path)},
                            {"type": "text", "text": user_text},
                        ],
                    },
                    {
                        "role": "assistant",
                        "content": [{"type": "text", "text": _assistant_text(pt)}],
                    },
                ],
            }
            f.write(json.dumps(rec, ensure_ascii=False) + "\n")

    info = {
        "out_dir": str(out_dir),
        "train_jsonl": str(train_path),
        "empty_png": str(img_path),
        "samples": int(args.samples),
        "hoshi_mode": str(args.hoshi),
        "note": "빈판 오프닝(화점) 합성 데이터셋",
    }
    with open(out_dir / "dataset_info.json", "w", encoding="utf-8") as f:
        json.dump(info, f, ensure_ascii=False, indent=2)

    print(f"OK: wrote {train_path} ({info['samples']} samples)  img={img_path}")
    return 0


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