diff --git a/scripts/generate_ai_samples.py b/scripts/generate_ai_samples.py index 7165036..a68c861 100644 --- a/scripts/generate_ai_samples.py +++ b/scripts/generate_ai_samples.py @@ -127,10 +127,34 @@ def _build_prompt(meta: dict[str, str], target_chars: int, style: str) -> str: ) +def _column_values(frame, column: str): + """pandas DataFrame(테스트/하위호환) 또는 dict 행 목록에서 컬럼을 읽는다.""" + if isinstance(frame, list): + return [row.get(column) for row in frame] + return frame[column].tolist() + + +def load_xlsx_rows(path: Path) -> tuple[list[str], list[dict[str, object]]]: + from openpyxl import load_workbook + + workbook = load_workbook(path, read_only=True, data_only=True) + worksheet = workbook.worksheets[0] + iterator = worksheet.iter_rows(values_only=True) + try: + headers = [str(v).strip() if v is not None else "" for v in next(iterator)] + except StopIteration: + workbook.close() + raise SystemExit(f"빈 xlsx입니다: {path}") + rows = [dict(zip(headers, row)) for row in iterator] + workbook.close() + return headers, rows + + def load_human_length_band(frame, text_column: str) -> tuple[list[int], int, int]: """human 분포에서 목표 길이 후보와 허용 구간을 뽑는다.""" lengths = sorted( - len(t) for t in frame[text_column].dropna().astype(str) if len(t) >= 120 + len(str(value)) for value in _column_values(frame, text_column) + if value is not None and len(str(value)) >= 120 ) if not lengths: raise SystemExit(f"'{text_column}' 에서 길이를 잴 수 있는 행이 없습니다.") @@ -201,29 +225,29 @@ def main() -> int: if args.concurrency < 1: ap.error("--concurrency는 1 이상이어야 합니다") - import pandas as pd - - frame = pd.read_excel(args.xlsx, sheet_name=0) - lengths, low_chars, high_chars = load_human_length_band(frame, args.text_column) + columns, rows = load_xlsx_rows(args.xlsx) + if args.text_column not in columns: + raise SystemExit(f"길이 분포 컬럼이 없습니다: {args.text_column}") + lengths, low_chars, high_chars = load_human_length_band(rows, args.text_column) logger.info( "human 길이 분포: 중앙값 %d자, 허용 구간 %d~%d자 (n=%d)", lengths[len(lengths) // 2], low_chars, high_chars, len(lengths), ) - requested_meta = args.meta_column or [c for c in META_COLUMNS if c in frame.columns] - missing_meta = [c for c in requested_meta if c not in frame.columns] + requested_meta = args.meta_column or [c for c in META_COLUMNS if c in columns] + missing_meta = [c for c in requested_meta if c not in columns] if missing_meta: raise SystemExit(f"프롬프트 소재 컬럼이 없습니다: {missing_meta}") usable_meta = assert_no_body_columns(requested_meta) if not usable_meta: raise SystemExit( - f"소재로 쓸 컬럼이 없습니다. 기대: {META_COLUMNS} / 실제: {list(frame.columns)}" + f"소재로 쓸 컬럼이 없습니다. 기대: {META_COLUMNS} / 실제: {columns}" ) logger.info("프롬프트 소재 컬럼: %s", usable_meta) group_salt = os.environ.get(args.group_salt_env, "") if args.group_column: - if args.group_column not in frame.columns: + if args.group_column not in columns: raise SystemExit(f"그룹 컬럼이 없습니다: {args.group_column}") if not group_salt: raise SystemExit( @@ -243,7 +267,7 @@ def main() -> int: rng = random.Random(args.seed) # 소재 행을 결정적으로 섞는다. 같은 seed 면 같은 순서가 나온다. - order = sorted(range(len(frame)), key=lambda i: _stable_int(f"{args.seed}:{i}")) + order = sorted(range(len(rows)), key=lambda i: _stable_int(f"{args.seed}:{i}")) client = None if not args.dry_run: @@ -272,15 +296,15 @@ def main() -> int: if prompt_id in done: return None meta = { - c: sanitize_prompt_metadata(str(frame[c].iloc[idx])) + c: sanitize_prompt_metadata(str(rows[idx].get(c))) for c in usable_meta - if isinstance(frame[c].iloc[idx], str) and frame[c].iloc[idx].strip() + if isinstance(rows[idx].get(c), str) and str(rows[idx].get(c)).strip() } if not meta: return None source_group = f"ai-topic:{prompt_id}" if args.group_column: - raw_group = str(frame[args.group_column].iloc[idx] or "").strip() + raw_group = str(rows[idx].get(args.group_column) or "").strip() if raw_group and raw_group.lower() != "nan": source_group = pseudonymous_id(raw_group, group_salt) target = rng.choice(target_pool)