From d0b5876ae5bfce26e525ee78fe0bcbb41152ed01 Mon Sep 17 00:00:00 2001 From: hbyang Date: Thu, 20 Aug 2026 17:19:41 +0900 Subject: [PATCH] =?UTF-8?q?feat:=20GPU=20=EC=83=9D=EC=84=B1=20=EC=9E=91?= =?UTF-8?q?=EC=97=85=20=EC=83=A4=EB=94=A9=20=EC=A7=80=EC=9B=90?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- scripts/generate_ai_samples.py | 15 ++++++++++++++- tests/test_generate_ai_samples.py | 6 ++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/scripts/generate_ai_samples.py b/scripts/generate_ai_samples.py index ea6828c..679587e 100644 --- a/scripts/generate_ai_samples.py +++ b/scripts/generate_ai_samples.py @@ -218,12 +218,18 @@ def main() -> int: ap.add_argument("--min-hangul-ratio", type=float, default=0.55, help="알파벳 문자 중 최소 한글 비율") ap.add_argument("--seed", type=int, default=20260819) + ap.add_argument("--shard-count", type=int, default=1, + help="소재 행을 서로 겹치지 않는 N개 샤드로 분할") + ap.add_argument("--shard-index", type=int, default=0, + help="이 프로세스가 담당할 0-based 샤드") ap.add_argument("--out", type=Path, default=Path("data/training/ai_samples.jsonl")) ap.add_argument("--dry-run", action="store_true", help="프롬프트만 출력하고 API 를 호출하지 않는다 (비용 0)") args = ap.parse_args() if args.concurrency < 1: ap.error("--concurrency는 1 이상이어야 합니다") + if args.shard_count < 1 or not 0 <= args.shard_index < args.shard_count: + ap.error("--shard-count는 1 이상, --shard-index는 0 이상 shard-count 미만이어야 합니다") columns, rows = load_xlsx_rows(args.xlsx) if args.text_column not in columns: @@ -268,7 +274,14 @@ def main() -> int: rng = random.Random(args.seed) # 소재 행을 결정적으로 섞는다. 같은 seed 면 같은 순서가 나온다. - order = sorted(range(len(rows)), key=lambda i: _stable_int(f"{args.seed}:{i}")) + order = sorted( + ( + i for i in range(len(rows)) + if _stable_int(f"shard:{i}") % args.shard_count == args.shard_index + ), + key=lambda i: _stable_int(f"{args.seed}:{i}"), + ) + logger.info("소재 샤드 %d/%d: %d행", args.shard_index, args.shard_count, len(order)) client = None if not args.dry_run: diff --git a/tests/test_generate_ai_samples.py b/tests/test_generate_ai_samples.py index 5351e0a..84386fa 100644 --- a/tests/test_generate_ai_samples.py +++ b/tests/test_generate_ai_samples.py @@ -108,3 +108,9 @@ class TestDiversity: a = gen._stable_int("abc") % len(gen.STYLE_VARIANTS) b = gen._stable_int("abc") % len(gen.STYLE_VARIANTS) assert a == b + + def test_two_shards_are_disjoint_and_complete(self): + left = {i for i in range(100) if gen._stable_int(f"shard:{i}") % 2 == 0} + right = {i for i in range(100) if gen._stable_int(f"shard:{i}") % 2 == 1} + assert not left & right + assert left | right == set(range(100))