142 lines
3.8 KiB
Python
142 lines
3.8 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
build_rag.py
|
|
|
|
기능:
|
|
- chunk_and_summary.py 에서 생성된 output/rag/*_chunks.json 파일들을 읽어서
|
|
text + summary 를 임베딩(text-embedding-3-small)한다.
|
|
- FAISS IndexFlatIP 인덱스를 구축하여
|
|
output/rag/faiss.index, meta.json, vectors.npy 를 생성한다.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import faiss
|
|
from openai import OpenAI
|
|
from api_config import API_KEYS
|
|
|
|
# ===== 경로 설정 =====
|
|
DATA_ROOT = Path(r"D:\for python\geulbeot-light\geulbeot-light\00.test\hwpx\out")
|
|
OUTPUT_ROOT = Path(r"D:\for python\geulbeot-light\geulbeot-light\00.test\hwpx\out\out") # 출력 위치
|
|
RAG_DIR = OUTPUT_ROOT / "rag"
|
|
LOG_DIR = OUTPUT_ROOT / "logs"
|
|
|
|
for d in [RAG_DIR, LOG_DIR]:
|
|
d.mkdir(parents=True, exist_ok=True)
|
|
|
|
# ===== OpenAI 설정 (구조 유지) =====
|
|
OPENAI_API_KEY = API_KEYS.get('GPT_API_KEY', '')
|
|
GPT_MODEL = "gpt-5-2025-08-07"
|
|
EMBED_MODEL = "text-embedding-3-small"
|
|
|
|
client = OpenAI(api_key=OPENAI_API_KEY)
|
|
|
|
|
|
def log(msg: str):
|
|
print(msg, flush=True)
|
|
with (LOG_DIR / "build_rag_log.txt").open("a", encoding="utf-8") as f:
|
|
f.write(msg + "\n")
|
|
|
|
|
|
def embed_texts(texts):
|
|
if not texts:
|
|
return np.zeros((0, 1536), dtype="float32")
|
|
embs = []
|
|
B = 96
|
|
for i in range(0, len(texts), B):
|
|
batch = texts[i:i+B]
|
|
resp = client.embeddings.create(model=EMBED_MODEL, input=batch)
|
|
for d in resp.data:
|
|
embs.append(np.array(d.embedding, dtype="float32"))
|
|
return np.vstack(embs)
|
|
|
|
|
|
def _build_embed_input(u: dict) -> str:
|
|
"""
|
|
text + summary 를 합쳐 임베딩 입력을 만든다.
|
|
- text, summary 중 없는 것은 생략
|
|
- 공백 정리
|
|
- 최대 길이 제한
|
|
"""
|
|
sum_ = (u.get("summary") or "").strip()
|
|
txt = (u.get("text") or "").strip()
|
|
|
|
if txt and sum_:
|
|
merged = txt + "\n\n요약: " + sum_[:1000]
|
|
else:
|
|
merged = txt or sum_
|
|
|
|
merged = " ".join(merged.split())
|
|
if not merged:
|
|
return ""
|
|
if len(merged) > 4000:
|
|
merged = merged[:4000]
|
|
return merged
|
|
|
|
|
|
def build_faiss_index():
|
|
docs = []
|
|
metas = []
|
|
|
|
rag_files = list(RAG_DIR.glob("*_chunks.json"))
|
|
if not rag_files:
|
|
log("RAG 파일(*_chunks.json)이 없습니다. 먼저 chunk_and_summary.py를 실행해야 합니다.")
|
|
sys.exit(1)
|
|
|
|
for f in rag_files:
|
|
try:
|
|
units = json.loads(f.read_text(encoding="utf-8", errors="ignore"))
|
|
except Exception as e:
|
|
log(f"[WARN] RAG 파일 읽기 실패: {f.name} | {e}")
|
|
continue
|
|
|
|
for u in units:
|
|
embed_input = _build_embed_input(u)
|
|
if not embed_input:
|
|
continue
|
|
if len(embed_input) < 40:
|
|
continue
|
|
docs.append(embed_input)
|
|
metas.append({
|
|
"source": u.get("source", ""),
|
|
"chunk": int(u.get("chunk", 0)),
|
|
"folder_context": u.get("folder_context", "")
|
|
})
|
|
|
|
if not docs:
|
|
log("임베딩할 텍스트가 없습니다.")
|
|
sys.exit(1)
|
|
|
|
log(f"임베딩 대상 텍스트 수: {len(docs)}")
|
|
|
|
E = embed_texts(docs)
|
|
if E.shape[0] != len(docs):
|
|
log(f"[WARN] 임베딩 수 불일치: E={E.shape[0]}, docs={len(docs)}")
|
|
|
|
faiss.normalize_L2(E)
|
|
index = faiss.IndexFlatIP(E.shape[1])
|
|
index.add(E)
|
|
|
|
np.save(str(RAG_DIR / "vectors.npy"), E)
|
|
(RAG_DIR / "meta.json").write_text(
|
|
json.dumps(metas, ensure_ascii=False, indent=2),
|
|
encoding="utf-8"
|
|
)
|
|
faiss.write_index(index, str(RAG_DIR / "faiss.index"))
|
|
|
|
log(f"FAISS 인덱스 구축 완료: 벡터 수={len(metas)}")
|
|
|
|
|
|
def main():
|
|
log("=== FAISS RAG 인덱스 구축 시작 ===")
|
|
build_faiss_index()
|
|
log("=== FAISS RAG 인덱스 구축 종료 ===")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|