CodeWeaver/src/vector_db/local_embeddings.py

"""
로컬 임베딩 관리 모듈.

BAAI/bge-m3 모델을 사용해 로컬에서 임베딩을 생성한다.
"""

import logging
from typing import List

from sentence_transformers import SentenceTransformer

logger = logging.getLogger(__name__)


class LocalEmbeddingManager:
    """BAAI/bge-m3 로컬 임베딩 생성기."""

    def __init__(self, model_name: str = "BAAI/bge-m3") -> None:
        logger.info("로컬 임베딩 모델 로딩 중: %s", model_name)
        self.model = SentenceTransformer(model_name)
        dim = self.model.get_sentence_embedding_dimension()
        logger.info("로컬 임베딩 모델 로딩 완료 (차원: %d)", dim)

    def get_embedding(self, text: str) -> List[float]:
        """단일 텍스트를 임베딩."""
        embedding = self.model.encode(text, convert_to_numpy=True)
        return embedding.tolist()

    def get_embeddings_batch(self, texts: List[str]) -> List[List[float]]:
        """배치 텍스트 임베딩."""
        embeddings = self.model.encode(texts, convert_to_numpy=True)
        return embeddings.tolist()