v3:추출 파이프라인_260122

converters/pipeline/step5_rag.py

@@ -0,0 +1,141 @@
+# -*- 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\survey_test\process")
+OUTPUT_ROOT = Path(r"D:\for python\survey_test\output")
+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()