retrieval/vector_retriever.py

"""
벡터 검색 구현 모듈
"""

import os
import numpy as np
from typing import List, Dict, Any, Optional, Union, Tuple
import logging
from sentence_transformers import SentenceTransformer
from .base_retriever import BaseRetriever

logger = logging.getLogger(__name__)

class VectorRetriever(BaseRetriever):
    """
    임베딩 기반 벡터 검색 구현
    """
    
    def __init__(
        self, 
        embedding_model: Optional[Union[str, SentenceTransformer]] = "paraphrase-multilingual-MiniLM-L12-v2",
        documents: Optional[List[Dict[str, Any]]] = None,
        embedding_field: str = "text",
        embedding_device: str = "cpu"
    ):
        """
        VectorRetriever 초기화
        
        Args:
            embedding_model: 임베딩 모델 이름 또는 SentenceTransformer 인스턴스
            documents: 초기 문서 목록 (선택 사항)
            embedding_field: 임베딩할 문서 필드 이름
            embedding_device: 임베딩 모델 실행 장치 ('cpu' 또는 'cuda')
        """
        self.embedding_field = embedding_field
        self.model_name = None
        
        # 임베딩 모델 로드
        if isinstance(embedding_model, str):
            logger.info(f"임베딩 모델 '{embedding_model}' 로드 중...")
            self.model_name = embedding_model
            self.embedding_model = SentenceTransformer(embedding_model, device=embedding_device)
        else:
            self.embedding_model = embedding_model
            # 모델이 이미 로드된 인스턴스일 경우 이름 추출
            if hasattr(embedding_model, '_modules') and 'modules' in embedding_model._modules:
                self.model_name = "loaded_sentence_transformer"
        
        # 문서 저장소 초기화
        self.documents = []
        self.document_embeddings = None
        
        # 초기 문서가 제공된 경우 추가
        if documents:
            self.add_documents(documents)
    
    def add_documents(self, documents: List[Dict[str, Any]]) -> None:
        """
        검색기에 문서를 추가하고 임베딩 생성
        
        Args:
            documents: 추가할 문서 목록
        """
        if not documents:
            logger.warning("추가할 문서가 없습니다.")
            return
        
        # 문서 추가
        document_texts = []
        for doc in documents:
            if self.embedding_field not in doc:
                logger.warning(f"문서에 필드 '{self.embedding_field}'가 없습니다. 건너뜁니다.")
                continue
            
            self.documents.append(doc)
            document_texts.append(doc[self.embedding_field])
        
        if not document_texts:
            logger.warning(f"임베딩할 텍스트가 없습니다. 모든 문서에 '{self.embedding_field}' 필드가 있는지 확인하세요.")
            return
        
        # 문서 임베딩 생성
        logger.info(f"{len(document_texts)}개 문서의 임베딩 생성 중...")
        new_embeddings = self.embedding_model.encode(document_texts, show_progress_bar=True)
        
        # 기존 임베딩과 병합
        if self.document_embeddings is None:
            self.document_embeddings = new_embeddings
        else:
            self.document_embeddings = np.vstack([self.document_embeddings, new_embeddings])
        
        logger.info(f"총 {len(self.documents)}개 문서, {self.document_embeddings.shape[0]}개 임베딩 저장됨")
    
    def search(self, query: str, top_k: int = 5, **kwargs) -> List[Dict[str, Any]]:
        """
        쿼리에 대한 벡터 검색 수행
        
        Args:
            query: 검색 쿼리
            top_k: 반환할 상위 결과 수
            **kwargs: 추가 검색 매개변수
        
        Returns:
            관련성 점수와 함께 검색된 문서 목록
        """
        if not self.documents or self.document_embeddings is None:
            logger.warning("검색할 문서가 없습니다.")
            return []
        
        # 쿼리 임베딩 생성
        query_embedding = self.embedding_model.encode(query)
        
        # 코사인 유사도 계산
        scores = np.dot(self.document_embeddings, query_embedding) / (
            np.linalg.norm(self.document_embeddings, axis=1) * np.linalg.norm(query_embedding)
        )
        
        # 상위 결과 선택
        top_indices = np.argsort(scores)[-top_k:][::-1]
        
        # 결과 형식화
        results = []
        for idx in top_indices:
            doc = self.documents[idx].copy()
            doc["score"] = float(scores[idx])
            results.append(doc)
        
        return results
    
    def save(self, directory: str) -> None:
        """
        검색기 상태를 디스크에 저장
        
        Args:
            directory: 저장할 디렉토리 경로
        """
        import pickle
        import json
        
        os.makedirs(directory, exist_ok=True)
        
        # 문서 저장
        with open(os.path.join(directory, "documents.json"), "w", encoding="utf-8") as f:
            json.dump(self.documents, f, ensure_ascii=False, indent=2)
        
        # 임베딩 저장
        if self.document_embeddings is not None:
            np.save(os.path.join(directory, "embeddings.npy"), self.document_embeddings)
        
        # 모델 정보 저장
        model_info = {
            "model_name": self.model_name or "paraphrase-multilingual-MiniLM-L12-v2",  # 기본값 설정
            "embedding_dim": self.embedding_model.get_sentence_embedding_dimension() if hasattr(self.embedding_model, 'get_sentence_embedding_dimension') else 384 
        }
        
        with open(os.path.join(directory, "model_info.json"), "w") as f:
            json.dump(model_info, f)
        
        logger.info(f"검색기 상태를 '{directory}'에 저장했습니다.")
    
    @classmethod
    def load(cls, directory: str, embedding_model: Optional[Union[str, SentenceTransformer]] = None) -> "VectorRetriever":
        """
        디스크에서 검색기 상태를 로드
        
        Args:
            directory: 로드할 디렉토리 경로
            embedding_model: 사용할 임베딩 모델 (제공되지 않으면 저장된 정보 사용)
        
        Returns:
            로드된 VectorRetriever 인스턴스
        """
        import json
        
        # 모델 정보 로드
        with open(os.path.join(directory, "model_info.json"), "r") as f:
            model_info = json.load(f)
        
        # 임베딩 모델 인스턴스화
        if embedding_model is None:
            # 모델 이름을 사용하여 모델 인스턴스화
            if "model_name" in model_info and isinstance(model_info["model_name"], str):
                embedding_model = model_info["model_name"]
            else:
                # 안전장치: 모델 이름이 없거나 정수인 경우(이전 버전 호환성) 기본 모델 사용
                logger.warning("유효한 모델 이름을 찾을 수 없습니다. 기본 모델을 사용합니다.")
                embedding_model = "paraphrase-multilingual-MiniLM-L12-v2"
        
        # 검색기 인스턴스 생성 (문서 없이)
        retriever = cls(embedding_model=embedding_model)
        
        # 문서 로드
        with open(os.path.join(directory, "documents.json"), "r", encoding="utf-8") as f:
            retriever.documents = json.load(f)
        
        # 임베딩 로드
        embeddings_path = os.path.join(directory, "embeddings.npy")
        if os.path.exists(embeddings_path):
            retriever.document_embeddings = np.load(embeddings_path)
        
        logger.info(f"검색기 상태를 '{directory}'에서 로드했습니다.")
        return retriever