Update vector_store_test.py

vector_store_test.py

@@ -1,26 +1,23 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-
-"""
-벡터 스토어 모듈: 문서 임베딩 생성 및 벡터 스토어 구축
-배치 처리 적용 + 청크 길이 확인 추가
-"""
-
 import os
+import re
+import glob
+import time
 import argparse
 import logging
-from tqdm import tqdm
+from collections import defaultdict
+
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_core.documents import Document
 from langchain_community.vectorstores import FAISS
-from langchain.schema.document import Document
-from langchain_huggingface import HuggingFaceEmbeddings
-from e5_embeddings import E5Embeddings
+from langchain_community.embeddings import HuggingFaceEmbeddings
 
-# 로깅 설정
+# Logging Configuration
 logging.getLogger().setLevel(logging.ERROR)
 
+# Embedding model loading
 def get_embeddings(model_name="intfloat/multilingual-e5-large-instruct", device="cuda"):
-    print(f"[INFO] 임베딩 모델 디바이스: {device}")
-    return E5Embeddings(
+    print(f"[INFO] Embedding model device: {device}")
+    return HuggingFaceEmbeddings(
         model_name=model_name,
         model_kwargs={'device': device},
         encode_kwargs={'normalize_embeddings': True}
@@ -28,31 +25,32 @@ def get_embeddings(model_name="intfloat/multilingual-e5-large-instruct", device=
 
 def build_vector_store_batch(documents, embeddings, save_path="vector_db", batch_size=4):
     if not documents:
-        raise ValueError("문서가 없습니다. 문서가 올바르게 로드되었는지 확인하세요.")
+        raise ValueError("No documents found. Check if documents were loaded correctly.")
 
     texts = [doc.page_content for doc in documents]
     metadatas = [doc.metadata for doc in documents]
 
-    # 청크 길이 출력
+    # Print chunk lengths
     lengths = [len(t) for t in texts]
-    print(f"💡 청크 수: {len(texts)}")
-    print(f"💡 가장 긴 청크 길이: {max(lengths)} chars")
-    print(f"💡 평균 청크 길이: {sum(lengths) // len(lengths)} chars")
+    print(f"💡 Number of chunks: {len(texts)}")
+    print(f"💡 Longest chunk length: {max(lengths)} chars")
+    print(f"💡 Average chunk length: {sum(lengths) // len(lengths)} chars")
 
-    # 배치로 나누기
+    # 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_documents 사용
+    # Use from_documents
     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 = [
@@ -83,39 +81,37 @@ 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="final_dataset", help="문서가 있는 폴더 경로")
-    parser.add_argument("--save_path", type=str, default="vector_db", help="벡터 스토어 저장 경로")
-    parser.add_argument("--batch_size", type=int, default=4, 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.add_argument("--device", type=str, default="cuda", help="사용할 디바이스 ('cuda' 또는 'cpu' 또는 'cuda:1')")
-
+    parser = argparse.ArgumentParser(description="Builds a vector store")
+    parser.add_argument("--folder", type=str, default="final_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=4, 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", help="Device to use ('cuda' or 'cpu' or 'cuda:0')") #Ermöglicht cuda:0
+    
     args = parser.parse_args()
 
-    # 문서 처리 모듈 import
+    # Import the document processing module
     from document_processor_image_test import load_documents, split_documents
 
     documents = load_documents(args.folder)
     chunks = split_documents(documents, chunk_size=800, chunk_overlap=100)
 
-    print(f"[DEBUG] 문서 로딩 및 청크 분할 완료, 임베딩 단계 진입 전")
-    print(f"[INFO] 선택된 디바이스: {args.device}")
+    print(f"[DEBUG] Document loading and chunk splitting complete, entering embedding stage")
+    print(f"[INFO] Selected device: {args.device}")
 
     try:
         embeddings = get_embeddings(
             model_name=args.model_name,
             device=args.device
         )
-        print(f"[DEBUG] 임베딩 모델 생성 완료")
+        print(f"[DEBUG] Embedding model created")
     except Exception as e:
-        print(f"[ERROR] 임베딩 모델 생성 중 에러 발생: {e}")
+        print(f"[ERROR] Error creating embedding model: {e}")
         import traceback; traceback.print_exc()
         exit(1)
 
-    build_vector_store_batch(chunks, embeddings, args.save_path, args.batch_size)
-
+    build_vector_store_batch(chunks, embeddings, args.save_path, args.batch_size)