Update vector_store.py

vector_store.py

@@ -1,22 +1,49 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-
-"""
-벡터 스토어 모듈: 문서 임베딩 생성 및 벡터 스토어 구축
-배치 처리 적용으로 메모리 사용량 최적화 + 긴 청크 오류 방지
-"""
-
 import os
 import argparse
 import logging
-from tqdm import tqdm
-from langchain_community.vectorstores import FAISS
-from langchain.schema.document import Document
-from langchain_huggingface import HuggingFaceEmbeddings
+import time
+from collections import defaultdict
 
-# 로깅 설정 - 불필요한 경고 메시지 제거
-logging.getLogger().setLevel(logging.ERROR)
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_core.documents import Document
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores import FAISS
 
+# PyMuPDF library
+try:
+    import fitz  # PyMuPDF
+    PYMUPDF_AVAILABLE = True
+    print("✅ PyMuPDF library available")
+except ImportError:
+    PYMUPDF_AVAILABLE = False
+    print("⚠️ PyMuPDF library is not installed. Install with: pip install PyMuPDF")
+
+
+# --------------------------------
+# Log Output
+# --------------------------------
+
+def log(msg):
+    print(f"[{time.strftime('%H:%M:%S')}] {msg}")
+
+# --------------------------------
+# Text Cleaning Function
+# --------------------------------
+
+def clean_text(text):
+    return re.sub(r"[^\uAC00-\uD7A3\u1100-\u11FF\u3130-\u318F\w\s.,!?\"'()$:\-]", "", text)
+
+def apply_corrections(text):
+    corrections = {
+        'º©': 'info', 'Ì': 'of', '½': 'operation', 'Ã': '', '©': '',
+        '’': "'", '“': '"', 'â€': '"'
+    }
+    for k, v in corrections.items():
+        text = text.replace(k, v)
+    return text
+
+# --------------------------------
+# Load the embedding model
 def get_embeddings(model_name="intfloat/multilingual-e5-large-instruct", device="cuda"):
     return HuggingFaceEmbeddings(
         model_name=model_name,
@@ -26,26 +53,26 @@ def get_embeddings(model_name="intfloat/multilingual-e5-large-instruct", device=
 
 def build_vector_store_batch(documents, embeddings, save_path="vector_db", batch_size=16):
     if not documents:
-        raise ValueError("문서가 없습니다. 문서가 올바르게 로드되었는지 확인하세요.")
+        raise ValueError("No documents found. Check if documents are loaded correctly.")
 
     texts = [doc.page_content for doc in documents]
     metadatas = [doc.metadata for doc in documents]
 
-    # 배치로 분할
+    # Split into batches
     batches = [texts[i:i + batch_size] for i in range(0, len(texts), batch_size)]
     metadata_batches = [metadatas[i:i + batch_size] for i in range(0, len(metadatas), batch_size)]
 
     print(f"Processing {len(batches)} batches with size {batch_size}")
     print(f"Initializing vector store with batch 1/{len(batches)}")
 
-    # ✅ from_texts 대신 from_documents 사용 (길이 문제 방지)
+    # Use from_documents instead of from_texts (to prevent length issues)
     first_docs = [
         Document(page_content=text, metadata=meta)
         for text, meta in zip(batches[0], metadata_batches[0])
     ]
     vectorstore = FAISS.from_documents(first_docs, embeddings)
 
-    # 나머지 배치 추가
+    # Add remaining batches
     for i in tqdm(range(1, len(batches)), desc="Processing batches"):
         try:
             docs_batch = [
@@ -76,29 +103,28 @@ def build_vector_store_batch(documents, embeddings, save_path="vector_db", batch
 
 def load_vector_store(embeddings, load_path="vector_db"):
     if not os.path.exists(load_path):
-        raise FileNotFoundError(f"벡터 스토어를 찾을 수 없습니다: {load_path}")
+        raise FileNotFoundError(f"Cannot find vector store: {load_path}")
     return FAISS.load_local(load_path, embeddings, allow_dangerous_deserialization=True)
 
-
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="벡터 스토어 구축")
-    parser.add_argument("--folder", type=str, default="dataset", help="문서가 있는 폴더 경로")
-    parser.add_argument("--save_path", type=str, default="vector_db", help="벡터 스토어 저장 경로")
-    parser.add_argument("--batch_size", type=int, default=16, help="배치 크기")
-    parser.add_argument("--model_name", type=str, default="intfloat/multilingual-e5-large-instruct", help="임베딩 모델 이름")
-    parser.add_argument("--device", type=str, default="cuda", help="사용할 디바이스 ('cuda' 또는 'cpu')")
+    parser = argparse.ArgumentParser(description="Builds a vector store")
+    parser.add_argument("--folder", type=str, default="dataset", help="Path to the folder containing the documents")
+    parser.add_argument("--save_path", type=str, default="vector_db", help="Path to save the vector store")
+    parser.add_argument("--batch_size", type=int, default=16, help="Batch size")
+    parser.add_argument("--model_name", type=str, default="intfloat/multilingual-e5-large-instruct", help="Name of the embedding model")
+    parser.add_argument("--device", type=str, default="cuda", choices=["cuda", "cpu"], help="Device to use ('cuda' or 'cpu')")
 
     args = parser.parse_args()
 
-    # 문서 처리 모듈 import
+    # Import the document processing module
     from document_processor import load_documents, split_documents
 
-    # 문서 로드 및 분할
+    # Load and split documents
     documents = load_documents(args.folder)
     chunks = split_documents(documents, chunk_size=800, chunk_overlap=100)
 
-    # 임베딩 모델 로드
+    # Load the embedding model
     embeddings = get_embeddings(model_name=args.model_name, device=args.device)
 
-    # 벡터 스토어 구축
-    build_vector_store_batch(chunks, embeddings, args.save_path, args.batch_size)
+    # Build the vector store
+    build_vector_store_batch(chunks, embeddings, args.save_path, args.batch_size)