#!/usr/bin/env python3
"""
RAG 인덱스 빌드 스크립트 (성능 최적화 버전)
JSON 형식의 용어집 및 코드 데이터를 읽어, 검색 성능에 최적화된 FAISS 인덱스를 생성합니다.
- 핵심 원리: 불필요한 상용구를 제거하고 핵심 키워드 중심으로 임베딩 텍스트를 구성합니다.
- 모델: 한국어-영어 혼합 데이터에 강력한 `jhgan/ko-sroberta-multitask` 모델을 사용합니다.
- 처리 방식: 파일명이 아닌 데이터의 구조(Key)를 기반으로 동적으로 처리합니다.
Usage:
python knowledge/build_index.py
"""
import json
import logging
import pickle
import shutil
from pathlib import Path
from typing import List, Dict, Any, Tuple
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
# --- 로깅 설정 ---
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
# --- 상수 정의 ---
# 한국어와 영어를 모두 잘 처리하는 고성능 모델
MODEL_NAME = "jhgan/ko-sroberta-multitask"
DATA_PATH = Path("ecom-doc/dict")
INDEX_BASE_PATH = Path("knowledge/index")
class RAGBuilder:
"""RAG 인덱스 생성을 위한 모든 로직을 캡슐화한 클래스"""
def __init__(self, model_name: str):
self.model = SentenceTransformer(model_name)
logger.info(f"✅ Sentence Transformer 모델 로드 완료: {model_name}")
def _create_glossary_texts(self, term: Dict[str, Any]) -> Tuple[str, str]:
"""'인코텀즈/견적' 용어 데이터로부터 표시용 텍스트와 임베딩용 텍스트를 생성합니다."""
display_text = (
f"### {term.get('abbreviation', '')} - {term.get('term_kr', '')}\n\n"
f"**카테고리**: {term.get('category', 'N/A')}\n"
f"**영문명**: {term.get('term_en', 'N/A')}\n"
f"**설명**: {term.get('description', 'N/A')}"
)
# 검색 성능을 극대화하기 위해 핵심 키워드만 조합
embedding_text = (
f"{term.get('term_kr', '')} {term.get('abbreviation', '')} {term.get('term_en', '')} "
f"{term.get('description', '')}"
)
return display_text, embedding_text.strip().replace("\n", " ")
def _create_package_code_texts(self, term: Dict[str, Any]) -> Tuple[str, str]:
"""'패키지 단위 코드' 데이터로부터 텍스트를 생성합니다."""
display_text = (
f"### 패키지 단위 코드: {term.get('code', '')}\n\n"
f"**영문 단위**: {term.get('unit_en', 'N/A')}\n"
f"**한글 단위**: {term.get('unit_kr', 'N/A')}\n"
f"**카테고리**: {term.get('category', 'N/A')}"
)
# 사용자가 검색할 만한 모든 키워드를 공백으로 연결
keywords = " ".join(term.get('keywords', []))
embedding_text = (
f"{term.get('unit_kr', '')} {term.get('unit_en', '')} {term.get('code', '')} {keywords}"
)
return display_text, embedding_text.strip().replace("\n", " ")
def load_and_process_single_file(self, file_path: Path) -> List[Dict[str, Any]]:
"""
단일 JSON 파일을 로드하고, 구조에 따라 적절히 처리하여
RAG에 사용될 문서 목록을 생성합니다.
"""
try:
with open(file_path, "r", encoding="utf-8") as f:
data = json.load(f)
if not isinstance(data, list) or not data:
logger.warning(f"'{file_path.name}'의 내용이 비어있거나 리스트 형태가 아닙니다.")
return []
# 데이터의 첫 항목의 키를 보고 종류를 판별
first_item_keys = data[0].keys()
processor_func = None
if 'abbreviation' in first_item_keys and 'term_kr' in first_item_keys:
processor_func = self._create_glossary_texts
doc_type = "용어집"
elif 'code' in first_item_keys and 'unit_kr' in first_item_keys:
processor_func = self._create_package_code_texts
doc_type = "패키지 코드"
else:
logger.warning(f"'{file_path.name}'에 대한 처리 로직을 찾을 수 없습니다.")
return []
file_documents = []
for i, item in enumerate(data):
display_text, embedding_text = processor_func(item)
file_documents.append({
"text": display_text,
"embedding_text": embedding_text,
"source": str(file_path.name),
"chunk_id": i,
"metadata": item # 원본 데이터 전체를 메타데이터로 저장
})
logger.info(f"로드 완료: {file_path.name} ({len(file_documents)}개 항목, 유형: {doc_type})")
return file_documents
except json.JSONDecodeError as e:
logger.error(f"JSON 파싱 실패 {file_path}: {e}")
return []
except Exception as e:
logger.error(f"파일 처리 실패 {file_path}: {e}")
return []
def build_index(self, documents: List[Dict[str, Any]]) -> Tuple[faiss.Index, List[Dict]]:
"""문서 목록으로부터 FAISS 인덱스를 생성합니다."""
if not documents:
raise ValueError("인덱스를 생성할 문서가 없습니다.")
texts_to_embed = [doc["embedding_text"] for doc in documents]
logger.info(f"임베딩 생성 중... (문서 {len(texts_to_embed)}개)")
embeddings = self.model.encode(
texts_to_embed, convert_to_numpy=True, show_progress_bar=True
)
embeddings_array = np.array(embeddings, dtype=np.float32)
logger.info(f"임베딩 배열 형태: {embeddings_array.shape}")
dimension = embeddings_array.shape[1]
# FAISS 인덱스 생성 (내적 기반)
index = faiss.IndexFlatIP(dimension)
# 코사인 유사도 검색을 위해 벡터를 정규화
faiss.normalize_L2(embeddings_array)
index.add(embeddings_array)
logger.info(f"✅ FAISS 인덱스 생성 완료: {index.ntotal}개 벡터")
return index, documents
@staticmethod
def save_artifacts(index: faiss.Index, documents: List[Dict], output_dir: Path):
"""생성된 인덱스와 문서를 파일로 저장합니다."""
try:
# 출력 디렉토리가 없으면 생성
output_dir.mkdir(parents=True, exist_ok=True)
index_path = output_dir / "rag_index.faiss"
documents_path = output_dir / "rag_documents.pkl"
faiss.write_index(index, str(index_path))
logger.info(f"💾 FAISS 인덱스 저장 완료: {index_path}")
with open(documents_path, "wb") as f:
pickle.dump(documents, f)
logger.info(f"💾 문서 데이터 저장 완료: {documents_path}")
return True
except Exception as e:
logger.error(f"파일 저장 실패: {e}")
return False
def main():
"""메인 실행 함수 - 각 JSON 파일마다 별도의 인덱스를 생성합니다."""
logger.info("🚀 RAG 인덱스 빌드 시작 (파일별 분리 버전)")
# JSON 파일들 찾기
json_files = list(DATA_PATH.glob("*.json"))
if not json_files:
logger.error(f"❌ {DATA_PATH}에서 JSON 파일을 찾을 수 없습니다.")
return False
logger.info(f"발견된 JSON 파일 수: {len(json_files)}")
# 기존 인덱스 디렉토리 정리
if INDEX_BASE_PATH.exists():
shutil.rmtree(INDEX_BASE_PATH)
logger.info(f"기존 인덱스 디렉토리 삭제: {INDEX_BASE_PATH}")
try:
builder = RAGBuilder(model_name=MODEL_NAME)
success_count = 0
for json_file in json_files:
logger.info(f"\n📄 처리 중: {json_file.name}")
# 각 파일에 대한 문서 로드
documents = builder.load_and_process_single_file(json_file)
if not documents:
logger.warning(f"⚠️ {json_file.name}에서 처리할 문서가 없습니다. 건너뜁니다.")
continue
# 인덱스 생성
index, processed_documents = builder.build_index(documents)
# 파일명에서 확장자 제거하여 디렉토리명 생성
file_stem = json_file.stem # 확장자 없는 파일명
output_dir = INDEX_BASE_PATH / file_stem
# 인덱스와 문서 저장
if builder.save_artifacts(index, processed_documents, output_dir):
success_count += 1
logger.info(f"✅ {json_file.name} 인덱스 생성 완료: {output_dir}")
else:
logger.error(f"❌ {json_file.name} 인덱스 저장 실패")
if success_count > 0:
logger.info(f"🎉 RAG 인덱스 빌드 완료! 성공: {success_count}/{len(json_files)} 파일")
return True
else:
logger.error("❌ 모든 파일의 인덱스 생성에 실패했습니다.")
return False
except Exception as e:
logger.error(f"❌ 인덱스 빌드 중 치명적 오류 발생: {e}", exc_info=True)
return False
if __name__ == "__main__":
if main():
print("\n✨ RAG 인덱스가 성공적으로 재생성되었습니다.")
else:
print("\n💥 RAG 인덱스 생성에 실패했습니다. 로그를 확인해주세요.")
exit(1)
