AIyaodianzhushou/data-pipeline/processors/embedder.py

"""
BGE-M3 Embedding 生成器
将文本 chunk 向量化写入 Milvus
"""
import logging
from typing import Optional
from dataclasses import dataclass

from app.core.config import get_settings

settings = get_settings()
logger = logging.getLogger(__name__)


class Embedder:
    def __init__(self, model_name: Optional[str] = None, device: Optional[str] = None):
        self.model_name = model_name or settings.embedding_model
        self.device = device or settings.embedding_device
        self.model = None
        self._loaded = False

    def load_model(self):
        """Phase 2 实现：加载 BGE-M3 模型"""
        if self._loaded:
            return
        logger.info(f"Loading embedding model: {self.model_name}")
        self._loaded = True

    def encode(self, texts: list[str], batch_size: int = 32) -> list[list[float]]:
        """
        批量文本向量化。
        BGE-M3 支持文本前加 instruction 前缀：
          查询： "Represent this sentence for searching relevant passages: "
          文档： ""

        Phase 2 实现。
        """
        return [[0.0] * settings.embedding_dim] * len(texts)

    def encode_queries(self, queries: list[str]) -> list[list[float]]:
        """查询向量化（带 instruction 前缀）"""
        prefixed = [
            f"Represent this sentence for searching relevant passages: {q}"
            for q in queries
        ]
        return self.encode(prefixed)

    def encode_documents(self, documents: list[str]) -> list[list[float]]:
        """文档向量化"""
        return self.encode(documents)

    def embed_and_store(
        self,
        chunks: list[str],
        metadata_list: list[dict],
        collection_name: Optional[str] = None,
    ):
        """
        向量化并写入 Milvus。
        Phase 2 实现。
        """
        vectors = self.encode_documents(chunks)
        logger.info(f"Generated {len(vectors)} embeddings for {len(chunks)} chunks")
        return len(vectors)