183 lines
6.2 KiB
Python
183 lines
6.2 KiB
Python
"""
|
|
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 [] |