예제 소스

dataset_coords.py

"""좌표 grounding 전용 SFT 샘플 생성 (locate / lookup).

수 예측 LoRA 가 좌표 능력을 붕괴시키지 않도록, 좌표 QA 를 먼저/혼합 학습한다.
"""

from __future__ import annotations

import argparse
import json
import random
from pathlib import Path
from typing import List, Optional, Tuple

from .dataset_baduk import BOARD_SIZE, board_text_aligned

SYSTEM = (
    "당신은 19×19 바둑 보드 텍스트 격자를 읽는 검사기입니다. "
    "좌표는 0-indexed 정수 (row, col) 이며 각각 0..18 입니다. "
    "행 번호는 각 줄 왼쪽, 열 번호는 맨 위 헤더입니다. "
    "각 칸 기호(X/O/.)는 두 글자 폭의 **왼쪽**에 있습니다. "
    "응답은 설명 없이 JSON 한 줄만 출력하세요."
)


def _pack(stones: List[Tuple[int, int, int]]) -> bytes:
    b = bytearray(BOARD_SIZE * BOARD_SIZE)
    for r, c, color in stones:
        b[r * BOARD_SIZE + c] = color
    return bytes(b)


def _sample_points(rng: random.Random, n: int) -> List[Tuple[int, int]]:
    pts = set()
    while len(pts) < n:
        pts.add((rng.randrange(BOARD_SIZE), rng.randrange(BOARD_SIZE)))
    return sorted(pts)


def make_locate(rng: random.Random) -> dict:
    n = 1 if rng.random() < 0.7 else 2
    pts = _sample_points(rng, n)
    color = 1 if rng.random() < 0.5 else 2
    stones = [(r, c, color) for r, c in pts]
    board = _pack(stones)
    grid = board_text_aligned(board)
    color_name = "흑(X)" if color == 1 else "백(O)"
    if n == 1:
        r, c = pts[0]
        user = (
            f"{SYSTEM}\n\n보드 문자 격자:\n{grid}\n\n"
            f"보드 위에 {color_name} 돌이 정확히 1개 있습니다. "
            f'그 돌의 좌표만 JSON 한 줄로 답하세요. 예: {{"row": 3, "col": 3}}'
        )
        assistant = f'{{"row": {r}, "col": {c}}}'
    else:
        pts_json = ",".join(f'{{"row": {r}, "col": {c}}}' for r, c in pts)
        user = (
            f"{SYSTEM}\n\n보드 문자 격자:\n{grid}\n\n"
            f"보드 위에 {color_name} 돌이 정확히 {n}개 있습니다. "
            f'모든 좌표를 JSON 한 줄로: {{"points":[{pts_json}]}}'
        )
        assistant = f'{{"points":[{pts_json}]}}'
    return {
        "task": "locate",
        "images": [],
        "messages": [
            {"role": "user", "content": [{"type": "text", "text": user}]},
            {"role": "assistant", "content": [{"type": "text", "text": assistant}]},
        ],
    }


def make_lookup(rng: random.Random) -> dict:
    r, c = rng.randrange(BOARD_SIZE), rng.randrange(BOARD_SIZE)
    # 70%: 그 칸에 돌, 30%: 빈칸 조회(돌은 다른 곳)
    if rng.random() < 0.7:
        color = 1 if rng.random() < 0.5 else 2
        stones = [(r, c, color)]
        expect = "X" if color == 1 else "O"
    else:
        er, ec = r, c
        while (er, ec) == (r, c):
            er, ec = rng.randrange(BOARD_SIZE), rng.randrange(BOARD_SIZE)
        color = 1 if rng.random() < 0.5 else 2
        stones = [(er, ec, color)]
        expect = "."
    board = _pack(stones)
    grid = board_text_aligned(board)
    user = (
        f"{SYSTEM}\n\n보드 문자 격자:\n{grid}\n\n"
        f"칸 (row={r}, col={c}) 의 내용은 무엇입니까? "
        f'JSON 한 줄만: {{"cell": "."}} 또는 {{"cell": "X"}} 또는 {{"cell": "O"}}'
    )
    assistant = f'{{"cell": "{expect}"}}'
    return {
        "task": "lookup",
        "images": [],
        "messages": [
            {"role": "user", "content": [{"type": "text", "text": user}]},
            {"role": "assistant", "content": [{"type": "text", "text": assistant}]},
        ],
    }


# 게이트에서 틀리기 쉬운 고정 좌표를 오버샘플
HARD_COORDS = [
    (0, 0),
    (0, 18),
    (18, 0),
    (18, 18),
    (9, 9),
    (9, 8),
    (9, 10),
    (15, 15),
    (15, 14),
    (3, 15),
    (15, 3),
    (4, 14),
    (13, 13),
    (18, 9),
]


def make_hard_locate(rng: random.Random) -> dict:
    r, c = HARD_COORDS[rng.randrange(len(HARD_COORDS))]
    color = 1
    board = _pack([(r, c, color)])
    grid = board_text_aligned(board)
    user = (
        f"{SYSTEM}\n\n보드 문자 격자:\n{grid}\n\n"
        f'보드 위에 흑(X) 돌이 정확히 1개 있습니다. '
        f'그 돌의 좌표만 JSON 한 줄로 답하세요. 예: {{"row": 3, "col": 3}}'
    )
    assistant = f'{{"row": {r}, "col": {c}}}'
    return {
        "task": "locate_hard",
        "images": [],
        "messages": [
            {"role": "user", "content": [{"type": "text", "text": user}]},
            {"role": "assistant", "content": [{"type": "text", "text": assistant}]},
        ],
    }


def main(argv: Optional[List[str]] = None) -> int:
    p = argparse.ArgumentParser(description="좌표 grounding SFT jsonl 생성")
    p.add_argument("--out-dir", default="sft_data/coord_v3")
    p.add_argument("--n-samples", type=int, default=4000)
    p.add_argument("--val-split", type=float, default=0.05)
    p.add_argument("--seed", type=int, default=7)
    args = p.parse_args(argv)

    rng = random.Random(args.seed)
    out = Path(args.out_dir)
    out.mkdir(parents=True, exist_ok=True)
    train_path = out / "train.jsonl"
    val_path = out / "val.jsonl"

    n_train = n_val = 0
    with train_path.open("w", encoding="utf-8") as ft, val_path.open("w", encoding="utf-8") as fv:
        for i in range(args.n_samples):
            roll = rng.random()
            if roll < 0.45:
                rec = make_locate(rng)
            elif roll < 0.75:
                rec = make_lookup(rng)
            else:
                rec = make_hard_locate(rng)
            line = json.dumps(rec, ensure_ascii=False) + "\n"
            if rng.random() < args.val_split:
                fv.write(line)
                n_val += 1
            else:
                ft.write(line)
                n_train += 1

    info = {
        "out_dir": str(out.resolve()),
        "n_samples_train": n_train,
        "n_samples_val": n_val,
        "board_format": "aligned_left",
        "tasks": ["locate", "lookup", "locate_hard"],
    }
    (out / "dataset_info.json").write_text(json.dumps(info, indent=2), encoding="utf-8")
    print(json.dumps(info, indent=2, ensure_ascii=False))
    return 0


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