Initial commit

train/classifier_dataset.py

@@ -0,0 +1,183 @@
+"""
+files/json/{LABEL}/{파일명}/page_xxx.json 구조를 읽어
+페이지 단위로 샘플을 생성하여 캐시 저장
+문서 1개 = 샘플 N개 (페이지 수만큼)
+"""
+import json
+import os
+import pickle
+from pathlib import Path
+
+from loguru import logger
+
+import config
+
+# JSON_DIR      = os.getenv("CLASSIFIER_JSON_DIR",   "./files/json")
+# CACHE_PATH    = os.getenv("CLASSIFIER_CACHE_PATH", "./files/classifier_dataset.pkl")
+# MIN_WORDS     = int(os.getenv("CLASSIFIER_MIN_WORDS", "20"))  # 페이지당 최소 단어 수
+
+JSON_DIR   = config.CLASSIFIER_JSON_DIR
+CACHE_PATH = config.CLASSIFIER_CACHE_PATH
+MIN_WORDS  = config.CLASSIFIER_MIN_WORDS
+
+_status: dict = {"status": "idle", "message": "", "total": 0}
+
+
+def get_status() -> dict:
+    return _status
+
+
+# def build() -> None:
+#     global _status
+#     try:
+#         _status = {"status": "building", "message": "JSON 파일 스캔 중...", "total": 0}
+#         json_path = Path(JSON_DIR)
+#
+#         if not json_path.exists():
+#             raise FileNotFoundError(f"JSON 디렉토리 없음: {JSON_DIR}")
+#
+#         data        = []
+#         skip_count  = 0
+#
+#         for label_dir in sorted(json_path.iterdir()):
+#             if not label_dir.is_dir():
+#                 continue
+#             label      = label_dir.name
+#             label_count = 0
+#
+#             for doc_dir in sorted(label_dir.iterdir()):
+#                 if not doc_dir.is_dir():
+#                     continue
+#
+#                 for page_file in sorted(doc_dir.glob("page_*.json")):
+#                     try:
+#                         with open(page_file, encoding="utf-8") as f:
+#                             d = json.load(f)
+#                         words = d.get("words", [])
+#                     except Exception as e:
+#                         logger.warning(f"파일 읽기 실패 {page_file}: {e}")
+#                         continue
+#
+#                     # 단어 수 미달 페이지 제외
+#                     if len(words) < MIN_WORDS:
+#                         print(f"skip file : {len(words)} : {page_file}")
+#                         skip_count += 1
+#                         continue
+#
+#                     data.append({
+#                         "text":  " ".join(words),
+#                         "label": label,
+#                         "src":   str(page_file),   # 디버깅용
+#                     })
+#                     label_count += 1
+#
+#             _status["message"] = f"{label}: {label_count}페이지 샘플 추가 (누적 {len(data)}개)"
+#             logger.info(_status["message"])
+#
+#         if not data:
+#             raise ValueError("수집된 데이터가 없습니다.")
+#
+#         os.makedirs(os.path.dirname(CACHE_PATH), exist_ok=True)
+#         with open(CACHE_PATH, "wb") as f:
+#             pickle.dump(data, f)
+#
+#         _status = {
+#             "status":  "done",
+#             "message": f"데이터셋 생성 완료 (단어 {MIN_WORDS}개 미만 {skip_count}페이지 제외)",
+#             "total":   len(data),
+#         }
+#         logger.info(f"데이터셋 저장: {CACHE_PATH} ({len(data)}개, 제외 {skip_count}개)")
+#
+#     except Exception as e:
+#         _status = {"status": "error", "message": str(e), "total": 0}
+#         logger.exception(f"데이터셋 생성 실패: {e}")
+
+
+def _process_page_file(page_file: Path, label: str, data: list, skip_count: int) -> tuple[int, int]:
+    """단일 page_*.json 파일 처리. (label_count 증가분, skip_count) 반환"""
+    try:
+        with open(page_file, encoding="utf-8") as f:
+            d = json.load(f)
+        words = d.get("words", [])
+    except Exception as e:
+        logger.warning(f"파일 읽기 실패 {page_file}: {e}")
+        return 0, skip_count
+
+    if len(words) < MIN_WORDS:
+        print(f"skip file : {len(words)} : {page_file}")
+        return 0, skip_count + 1
+
+    data.append({"text": " ".join(words), "label": label, "src": str(page_file)})
+    return 1, skip_count
+
+
+def _process_label_dir(label_dir: Path, data: list, skip_count: int) -> tuple[int, int]:
+    """label 폴더 처리. label_count, skip_count 반환"""
+    label = label_dir.name
+    label_count = 0
+
+    # label 폴더 직속 page_*.json
+    for page_file in sorted(label_dir.glob("*.json")):
+        added, skip_count = _process_page_file(page_file, label, data, skip_count)
+        label_count += added
+
+    # 하위 doc_dir 폴더
+    for doc_dir in sorted(label_dir.iterdir()):
+        if not doc_dir.is_dir():
+            continue
+        for page_file in sorted(doc_dir.glob("*.json")):
+            added, skip_count = _process_page_file(page_file, label, data, skip_count)
+            label_count += added
+
+    return label_count, skip_count
+
+
+def build() -> None:
+    global _status
+    try:
+        _status = {"status": "building", "message": "JSON 파일 스캔 중...", "total": 0}
+        json_path = Path(JSON_DIR)
+
+        if not json_path.exists():
+            raise FileNotFoundError(f"JSON 디렉토리 없음: {JSON_DIR}")
+
+        data       = []
+        skip_count = 0
+
+        for label_dir in sorted(json_path.iterdir()):
+            if not label_dir.is_dir():
+                continue
+
+            label_count, skip_count = _process_label_dir(label_dir, data, skip_count)
+
+            _status["message"] = f"{label_dir.name}: {label_count}페이지 샘플 추가 (누적 {len(data)}개)"
+            logger.info(_status["message"])
+
+        if not data:
+            raise ValueError("수집된 데이터가 없습니다.")
+
+        os.makedirs(os.path.dirname(CACHE_PATH), exist_ok=True)
+        with open(CACHE_PATH, "wb") as f:
+            pickle.dump(data, f)
+
+        _status = {
+            "status":  "done",
+            "message": f"데이터셋 생성 완료 (단어 {MIN_WORDS}개 미만 {skip_count}페이지 제외)",
+            "total":   len(data),
+        }
+        logger.info(f"데이터셋 저장: {CACHE_PATH} ({len(data)}개, 제외 {skip_count}개)")
+
+    except Exception as e:
+        _status = {"status": "error", "message": str(e), "total": 0}
+        logger.exception(f"데이터셋 생성 실패: {e}")
+
+
+def load() -> list[dict]:
+    if not os.path.exists(CACHE_PATH):
+        return []
+    try:
+        with open(CACHE_PATH, "rb") as f:
+            return pickle.load(f)
+    except Exception as e:
+        logger.error(f"데이터셋 로드 실패: {e}")
+        return []