Upload reg_embedding_system_v02.py

reg_embedding_system_v02.py

@@ -0,0 +1,453 @@
+import gc
+import json
+import sqlite3
+from pathlib import Path
+from typing import Optional, Tuple, Any, Dict, List, Set, Union
+import numpy as np
+
+import faiss
+
+# [수정됨] 패키지 구조 변경 반영 및 EnsembleRetriever 제거
+from langchain_community.retrievers import BM25Retriever
+from langchain_core.documents import Document
+from langchain_community.vectorstores import FAISS
+from sentence_transformers import SentenceTransformer
+
+# 런타임에 Embeddings 클래스를 찾기 위한 로직
+try:
+    from langchain_core.embeddings import Embeddings
+except ImportError:
+    try:
+        from langchain.embeddings.base import Embeddings
+    except ImportError:
+        Embeddings = object
+
+# --- SQLite 헬퍼 함수 ---
+SQLITE_DB_NAME = "metadata_mapping.db"
+
+# === IDSelector 클래스 정의 ===
+class MetadataIDSelector(faiss.IDSelectorBatch):
+    def __init__(self, allowed_ids: Set[int]):
+        super().__init__(list(allowed_ids))
+
+def get_db_connection(persist_directory: str) -> sqlite3.Connection:
+    """FAISS 저장 경로를 기반으로 SQLite 연결을 설정하고 반환합니다."""
+    db_path = Path(persist_directory) / SQLITE_DB_NAME
+    conn = sqlite3.connect(db_path)
+    return conn
+
+def _create_and_populate_sqlite_db(chunks: List[Document], persist_directory: str):
+    """
+    문서 청크를 기반으로 SQLite DB를 생성하고 채웁니다.
+    [업데이트 반영] 메타데이터 구조: regulation, chapter, section, standard
+    """
+    # 1. 입력 데이터 확인 (가장 중요한 체크 포인트)
+    if not chunks:
+        print("🚨 [오류] _create_and_populate_sqlite_db 함수에 전달된 chunks 리스트가 비어 있습니다!")
+        print("   -> load_chunks_from_jsonl 함수가 정상적으로 파일을 읽었는지 확인해주세요.")
+        return
+
+    # 2. 저장 경로 확인 및 생성
+    save_dir = Path(persist_directory)
+    save_dir.mkdir(parents=True, exist_ok=True)
+    
+    conn = get_db_connection(persist_directory)
+    
+    try:
+        cursor = conn.cursor()
+        
+        # 3. 테이블 생성 (기존 테이블 삭제 후 재생성 옵션 고려)
+        # 스키마가 변경되었으므로 기존 테이블이 있다면 충돌날 수 있습니다. 
+        # 안전하게 지우고 다시 만드는 방법을 추천합니다. (개발 단계)
+        cursor.execute("DROP TABLE IF EXISTS documents") 
+        
+        cursor.execute("""
+        CREATE TABLE documents (
+            faiss_id INTEGER PRIMARY KEY,
+            source TEXT,
+            regulation TEXT,
+            chapter TEXT,
+            section TEXT,
+            standard TEXT,
+            json_metadata TEXT 
+        )
+        """)
+        # 테이블 생성 직후 커밋 (파일에 스키마 기록)
+        conn.commit() 
+        print(f"📂 DB 테이블 생성 완료 (경로: {save_dir}/{SQLITE_DB_NAME})")
+
+        # 4. 데이터 채우기
+        inserted_count = 0
+        for i, doc in enumerate(chunks):
+            faiss_id = i 
+            metadata_json = json.dumps(doc.metadata, ensure_ascii=False)
+            
+            source_val = doc.metadata.get('source', '')
+            regulation_val = doc.metadata.get('regulation', '')
+            chapter_val = doc.metadata.get('chapter', '')
+            section_val = doc.metadata.get('section', '')
+            standard_val = doc.metadata.get('standard', '')
+
+            if isinstance(regulation_val, list): regulation_val = ', '.join(map(str, regulation_val))
+            if isinstance(chapter_val, list): chapter_val = ', '.join(map(str, chapter_val))
+            if isinstance(section_val, list): section_val = ', '.join(map(str, section_val))
+            if isinstance(standard_val, list): standard_val = ', '.join(map(str, standard_val))
+            
+            doc.metadata['faiss_id'] = faiss_id 
+            
+            cursor.execute(
+                """
+                INSERT OR REPLACE INTO documents 
+                (faiss_id, source, regulation, chapter, section, standard, json_metadata) 
+                VALUES (?, ?, ?, ?, ?, ?, ?)
+                """,
+                (faiss_id, source_val, regulation_val, chapter_val, section_val, standard_val, metadata_json)
+            )
+            inserted_count += 1
+
+        # 5. 최종 커밋
+        conn.commit()
+        print(f"✅ SQLite 데이터 저장 완료: 총 {inserted_count}행이 삽입되었습니다.")
+
+    except Exception as e:
+        print(f"🚨 [DB 저장 중 에러 발생] {e}")
+        # 에러가 나도 traceback을 볼 수 있게 함
+        import traceback
+        traceback.print_exc()
+        
+    finally:
+        # 6. 연결 확실히 종료
+        conn.close()
+
+# --- LocalSentenceTransformerEmbeddings ---
+class LocalSentenceTransformerEmbeddings(Embeddings):
+    def __init__(self, st_model, normalize_embeddings: bool = True, encode_batch_size: int = 32):
+        self.model = st_model
+        self.normalize = normalize_embeddings
+        self.encode_batch_size = encode_batch_size
+
+    def embed_documents(self, texts):
+        vecs = self.model.encode(
+            texts,
+            batch_size=self.encode_batch_size,
+            show_progress_bar=False,
+            normalize_embeddings=self.normalize,
+            convert_to_numpy=True,
+        )
+        return vecs.tolist()
+
+    def embed_query(self, text: str):
+        vec = self.model.encode(
+            [text],
+            batch_size=self.encode_batch_size,
+            show_progress_bar=False,
+            normalize_embeddings=self.normalize,
+            convert_to_numpy=True,
+        )[0]
+        return vec.tolist()
+
+def load_chunks_from_jsonl(file_paths: Union[str, List[str]]):
+    """
+    JSONL 파일 로드 함수
+    """
+    if isinstance(file_paths, str):
+        file_paths = [file_paths]
+
+    restored_documents = []
+    print(f" 총 {len(file_paths)}개의 파일 병합 로드를 시작합니다...")
+
+    for file_path in file_paths:
+        try:
+            file_doc_count = 0 
+            with open(file_path, 'r', encoding='utf-8') as f:
+                for line_number, line in enumerate(f):
+                    line = line.strip()
+                    if not line: continue
+                    data = json.loads(line)
+                    doc = Document(
+                        page_content=data.get('page_content', ""),
+                        metadata=data.get('metadata', {})
+                    )
+                    restored_documents.append(doc)
+                    file_doc_count += 1
+            print(f"  - [성공] {file_path}: {file_doc_count}개 Chunk")
+
+        except Exception as e:
+            print(f" [실패] 오류 ({file_path}): {e}")
+            continue
+
+    print(f"✅ 전체 로드 완료: 총 {len(restored_documents)}개의 Chunk가 복원되었습니다.")
+    return restored_documents
+
+# --- save_embedding_system (수정됨: Ensemble 제거 및 개별 반환) ---
+def save_embedding_system(
+    chunks, 
+    persist_directory: str = r"D:/Project AI/RAG", 
+    batch_size: int = 32,
+    device: str = 'cuda'
+):
+    Path(persist_directory).mkdir(parents=True, exist_ok=True)
+    
+    # 1) SQLite DB 저장
+    _create_and_populate_sqlite_db(chunks, persist_directory)
+
+    # 2) 모델 로드
+    model = SentenceTransformer(
+        'nomic-ai/nomic-embed-text-v2-moe', 
+        trust_remote_code=True,
+        device=device
+    )
+
+    embeddings = LocalSentenceTransformerEmbeddings(
+        st_model=model,
+        normalize_embeddings=True,
+        encode_batch_size=batch_size
+    )
+
+    # 3) FAISS 생성
+    vectorstore = None
+    for i in range(0, len(chunks), batch_size):
+        batch = chunks[i:i + batch_size]
+        if vectorstore is None:
+            vectorstore = FAISS.from_documents(documents=batch, embedding=embeddings)
+        else:
+            vectorstore.add_documents(documents=batch)
+        gc.collect()
+
+    # 4) BM25 생성 (Ensemble 없이 독립 생성)
+    bm25_retriever = BM25Retriever.from_documents(chunks)
+    bm25_retriever.k = 5
+    
+    # 5) 저장
+    vectorstore.save_local(persist_directory)
+    
+    # 6) 연결 반환 (개별 요소 반환)
+    sqlite_conn = get_db_connection(persist_directory)
+    gc.collect()
+
+    return bm25_retriever, vectorstore, sqlite_conn
+
+# --- load_embedding_from_faiss (수정됨: Ensemble 제거 및 개별 반환) ---
+def load_embedding_from_faiss(
+    persist_directory: str = r"D:/Project AI/RAG",
+    top_k: int = 10,
+    bm25_k: int = 10,
+    embeddings: Optional[Any] = None,
+    device: str = 'cpu'
+) -> Tuple[Any, FAISS, sqlite3.Connection]:
+    
+    if embeddings is None:
+        st_model = SentenceTransformer(
+            'nomic-ai/nomic-embed-text-v2-moe',
+            trust_remote_code=True,
+            device=device
+        )
+        embeddings = LocalSentenceTransformerEmbeddings(
+            st_model=st_model,
+            normalize_embeddings=True,
+            encode_batch_size=32
+        )
+
+    persist_dir = Path(persist_directory)
+    if not persist_dir.exists():
+        raise FileNotFoundError(f"FAISS 경로가 없습니다: {persist_dir}")
+
+    # FAISS 로드
+    vectorstore = FAISS.load_local(
+        folder_path=str(persist_dir),
+        embeddings=embeddings,
+        allow_dangerous_deserialization=True
+    )
+
+    # BM25 복원 (저장된 문서로부터 재생성)
+    bm25_retriever = None
+    docs = []
+    try:
+        if hasattr(vectorstore, "docstore") and hasattr(vectorstore.docstore, "_dict"):
+            docs = list(vectorstore.docstore._dict.values())
+            if docs:
+                bm25_retriever = BM25Retriever.from_documents(docs)
+                bm25_retriever.k = bm25_k
+            else:
+                print("[경고] 저장된 문서를 찾을 수 없어 BM25를 생성하지 못했습니다.")
+    except Exception as e:
+        print(f"[경고] 저장된 문서를 읽는 중 문제가 발생했습니다: {e}")
+
+    sqlite_conn = get_db_connection(persist_directory)
+
+    return bm25_retriever, vectorstore, sqlite_conn
+
+# --- search_vectorstore (단순 벡터 검색 헬퍼) ---
+def search_vectorstore(bm25_retriever, vectorstore, query, k=5):
+    """
+    vectorstore와 bm25_retriever를 받아 앙상블(Hybrid) 검색을 수행하는 함수.
+    EnsembleRetriever(weights=[0.6, 0.4])와 유사한 결과를 반환합니다.
+    """
+    weights=[0.6, 0.4]
+    # 1. 벡터 검색 수행 (Vector Search)
+    #    FAISS를 리트리버로 변환하여 검색
+    vec_retriever = vectorstore.as_retriever(search_kwargs={"k": k})
+    vec_docs = vec_retriever.invoke(query)
+    
+    # 2. 키워드 검색 수행 (BM25 Search)
+    #    검색 개수를 k개로 맞춰서 실행
+    bm25_docs = bm25_retriever.invoke(query, config={"search_kwargs": {"k": k}})
+    
+    # 3. 랭킹 퓨전 (Weighted Reciprocal Rank Fusion)
+    #    두 리스트의 순위를 기반으로 가중치를 적용해 점수를 매깁니다.
+    
+    doc_scores = {}  # 문서 내용(또는 ID) -> 점수
+    doc_map = {}     # 문서 내용 -> 문서 객체 저장 (나중에 반환하기 위해)
+
+    # 내부 함수: 순위에 따른 점수 계산 (Weight / (Rank + 1))
+    def apply_rank_score(docs, weight):
+        for rank, doc in enumerate(docs):
+            # 고유 키 생성 (page_content가 고유하다고 가정하거나, doc_id가 있다면 사용)
+            doc_key = doc.page_content
+            doc_map[doc_key] = doc
+            
+            if doc_key not in doc_scores:
+                doc_scores[doc_key] = 0.0
+            
+            # 순위가 높을수록(rank가 작을수록) 점수가 높음
+            score = weight / (rank + 1)
+            doc_scores[doc_key] += score
+
+    # 벡터 검색 결과 점수 반영 (가중치 0.6)
+    apply_rank_score(vec_docs, weights[0])
+    
+    # BM25 검색 결과 점수 반영 (가중치 0.4)
+    apply_rank_score(bm25_docs, weights[1])
+
+    # 4. 점수순 정렬 (높은 점수가 상위)
+    #    점수(item[1])를 기준으로 내림차순 정렬
+    sorted_docs = sorted(doc_scores.items(), key=lambda item: item[1], reverse=True)
+    
+    # 5. Top-K 추출 및 문서 객체 반환
+    final_results = [doc_map[key] for key, score in sorted_docs[:k]]
+    
+    return final_results
+
+# --- search_with_metadata_filter (수정됨: 수동 병합 로직 구현) ---
+def search_with_metadata_filter(
+    bm25_retriever: Any,      # [변경] Ensemble 대신 BM25를 직접 받음
+    vectorstore: FAISS,
+    query: str,
+    k: int = 5,
+    metadata_filter: Optional[Dict[str, Any]] = None,
+    sqlite_conn: Optional[sqlite3.Connection] = None
+) -> List[Document]:
+    """
+    SQLite 사전 필터링 -> FAISS 벡터 검색 + BM25 검색 -> 결과 병합
+    """
+    
+    # === 1. SQLite에서 필터링된 FAISS ID 추출 ===
+    filtered_ids = None
+    if metadata_filter and sqlite_conn:
+        cursor = sqlite_conn.cursor()
+        where_clauses = []
+        params = []
+
+        for key, value in metadata_filter.items():
+            if isinstance(value, list):
+                if not value: continue
+                placeholders = ', '.join(['?'] * len(value))
+                where_clauses.append(f"{key} IN ({placeholders})")
+                params.extend(value)
+            else:
+                where_clauses.append(f"{key} = ?")
+                params.append(value)
+
+        if where_clauses:
+            where_sql = " OR ".join(where_clauses)
+            sql_query = f"SELECT faiss_id FROM documents WHERE {where_sql}"
+
+            try:
+                cursor.execute(sql_query, params)
+                filtered_ids = {row[0] for row in cursor.fetchall()}
+                print(f"[사전 필터링] {len(filtered_ids)}개 ID 획득 → FAISS 검색 제한")
+            except Exception as e:
+                print(f"[경고] SQLite 필터링 실패: {e}")
+                filtered_ids = None
+        else:
+            print("[안내] 필터 조건 없음 → 전체 검색")
+    else:
+        print("[안내] 필터 또는 DB 없음 → 전체 검색")
+
+    # === 2. FAISS 벡터 검색 ===
+    vector_docs = []
+    if filtered_ids and len(filtered_ids) > 0:
+        selector = MetadataIDSelector(filtered_ids)
+        index: faiss.Index = vectorstore.index
+        
+        query_embedding = np.array(vectorstore.embeddings.embed_query(query)).astype('float32')
+        query_embedding = query_embedding.reshape(1, -1)
+
+        search_params = faiss.SearchParametersIVF(sel=selector, nprobe=20)
+        _k = max(k * 10, 100)
+        D, I = index.search(query_embedding, _k, params=search_params)
+
+        valid_indices = [i for i in I[0] if i != -1]
+        for idx in valid_indices[:k]:
+            doc_id = vectorstore.index_to_docstore_id[idx]
+            doc = vectorstore.docstore.search(doc_id)
+            if isinstance(doc, Document):
+                vector_docs.append(doc)
+        print(f"[벡터 검색] {len(valid_indices)}개 후보 → {len(vector_docs)}개 유효")
+    else:
+        # 전체 검색
+        vector_retriever = vectorstore.as_retriever(search_kwargs={"k": k})
+        vector_docs = vector_retriever.invoke(query)
+        print(f"[벡터 검색] 전체 검색 → {len(vector_docs)}개 후보")
+
+    # === 3. BM25 검색 ===
+    bm25_docs = []
+    if bm25_retriever:
+        search_k = k * 5
+        candidates = bm25_retriever.invoke(query, config={"search_kwargs": {"k": search_k}})
+        if filtered_ids:
+            bm25_docs = [d for d in candidates if d.metadata.get('faiss_id') in filtered_ids]
+        else:
+            bm25_docs = candidates
+        
+        # Top K 자르기
+        bm25_docs = bm25_docs[:k]
+        print(f"[BM25 검색] {len(candidates)}개 후보 → {len(bm25_docs)}개 필터링 후")
+
+    # === 4. 병합 (Vector 우선 + 중복 제거) ===
+    combined = {id(d): d for d in (vector_docs + bm25_docs)}.values()
+    final_results = list(combined)[:k]
+
+    print(f"[최종 결과] {len(final_results)}개 문서 반환")
+    return final_results
+
+# --- get_unique_metadata_values (빠진 함수 추가) ---
+def get_unique_metadata_values(
+    sqlite_conn: sqlite3.Connection,
+    key_name: str,
+    partial_match: Optional[str] = None
+) -> List[str]:
+    """
+    고유 값 검색 함수.
+    key_name 인자로 'part', 'subpart', 'section', 'source' 등을 사용할 수 있습니다.
+    """
+    if not sqlite_conn:
+        return []
+
+    cursor = sqlite_conn.cursor()
+    # 안전을 위해 key_name은 컬럼명으로 직접 사용 (SQL Injection 주의: 내부 사용 전제)
+    # 실제 프로덕션에서는 key_name을 화이트리스트로 검증하는 것이 좋습니다.
+    sql_query = f"SELECT DISTINCT `{key_name}` FROM documents"
+    params = []
+    
+    if partial_match:
+        sql_query += f" WHERE `{key_name}` LIKE ?"
+        params.append(f"%{partial_match}%")
+
+    try:
+        cursor.execute(sql_query, params)
+        unique_values = [row[0] for row in cursor.fetchall() if row[0] is not None]
+        return unique_values
+    except Exception as e:
+        print(f"[에러] 고유 값 검색 실패 ({key_name}): {e}")
+        return []