ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter3/dense-embedding/embedding_service.py

@@ -0,0 +1,218 @@
+"""Embedding service using BGE-M3 model."""
+
+import time
+import numpy as np
+from typing import List, Dict, Optional
+from FlagEmbedding import BGEM3FlagModel
+from logger import VectorSearchLogger, log_execution_time
+import logging
+
+
+class EmbeddingService:
+    """Service for generating embeddings using BGE-M3 model."""
+    
+    def __init__(self, model_name: str = "BAAI/bge-m3", use_fp16: bool = True, 
+                 max_seq_length: int = 512, logger: Optional[VectorSearchLogger] = None):
+        """
+        Initialize the embedding service with BGE-M3 model.
+        
+        Args:
+            model_name: Name of the BGE-M3 model
+            use_fp16: Whether to use FP16 for inference
+            max_seq_length: Maximum sequence length
+            logger: Logger instance for educational output
+        """
+        self.model_name = model_name
+        self.use_fp16 = use_fp16
+        self.max_seq_length = max_seq_length
+        self.logger = logger
+        self.std_logger = logging.getLogger("vector_search")
+        
+        # Initialize the model
+        self._initialize_model()
+    
+    def _initialize_model(self):
+        """Initialize the BGE-M3 model."""
+        start_time = time.time()
+        
+        if self.logger:
+            self.logger.logger.info(f"🚀 Initializing BGE-M3 model: {self.model_name}")
+            self.logger.logger.debug(f"  - Using FP16: {self.use_fp16}")
+            self.logger.logger.debug(f"  - Max sequence length: {self.max_seq_length}")
+        
+        try:
+            self.model = BGEM3FlagModel(
+                self.model_name, 
+                use_fp16=self.use_fp16
+            )
+            
+            # Get embedding dimension by encoding a test sentence
+            test_embedding = self.model.encode(["test"])
+            if isinstance(test_embedding, dict):
+                self.embedding_dim = test_embedding['dense_vecs'].shape[1]
+            else:
+                self.embedding_dim = test_embedding.shape[1]
+            
+            load_time = time.time() - start_time
+            
+            if self.logger:
+                self.logger.logger.info(f"✅ Model loaded successfully in {load_time:.2f} seconds")
+                self.logger.logger.debug(f"  - Embedding dimension: {self.embedding_dim}")
+                self.logger.logger.debug(f"  - Model supports: dense, sparse, and multi-vector retrieval")
+                
+        except Exception as e:
+            if self.logger:
+                self.logger.logger.error(f"Failed to load model: {e}")
+            raise
+    
+    @log_execution_time()
+    def encode_text(self, text: str, return_sparse: bool = False, 
+                   return_colbert: bool = False) -> Dict[str, np.ndarray]:
+        """
+        Encode a single text into embeddings.
+        
+        Args:
+            text: Input text to encode
+            return_sparse: Whether to return sparse embeddings
+            return_colbert: Whether to return ColBERT embeddings
+        
+        Returns:
+            Dictionary containing different types of embeddings
+        """
+        start_time = time.time()
+        
+        if self.logger:
+            self.logger.logger.debug(f"📝 Encoding text (length: {len(text)} chars)")
+            self.logger.logger.debug(f"  Text preview: {text[:100]}..." if len(text) > 100 else f"  Text: {text}")
+        
+        # Encode the text
+        embeddings = self.model.encode(
+            [text],
+            return_dense=True,
+            return_sparse=return_sparse,
+            return_colbert_vecs=return_colbert
+        )
+        
+        # Extract dense embeddings
+        dense_vec = embeddings['dense_vecs'][0]
+        
+        result = {
+            'dense': dense_vec,
+            'dimension': len(dense_vec)
+        }
+        
+        # Add sparse embeddings if requested
+        if return_sparse and 'lexical_weights' in embeddings:
+            result['sparse'] = embeddings['lexical_weights'][0]
+            if self.logger:
+                num_tokens = len(result['sparse'])
+                self.logger.logger.debug(f"  Sparse embedding: {num_tokens} non-zero tokens")
+        
+        # Add ColBERT embeddings if requested  
+        if return_colbert and 'colbert_vecs' in embeddings:
+            result['colbert'] = embeddings['colbert_vecs'][0]
+            if self.logger:
+                colbert_shape = result['colbert'].shape
+                self.logger.logger.debug(f"  ColBERT embedding shape: {colbert_shape}")
+        
+        encoding_time = time.time() - start_time
+        
+        if self.logger:
+            self.logger.logger.debug(f"✅ Encoding completed in {encoding_time:.4f} seconds")
+            self.logger.log_embedding_vector(dense_vec, sample_size=10)
+        
+        return result
+    
+    @log_execution_time()
+    def encode_batch(self, texts: List[str], return_sparse: bool = False, 
+                    return_colbert: bool = False) -> Dict[str, np.ndarray]:
+        """
+        Encode multiple texts into embeddings.
+        
+        Args:
+            texts: List of input texts to encode
+            return_sparse: Whether to return sparse embeddings
+            return_colbert: Whether to return ColBERT embeddings
+        
+        Returns:
+            Dictionary containing different types of embeddings for all texts
+        """
+        start_time = time.time()
+        
+        if self.logger:
+            self.logger.logger.info(f"📚 Batch encoding {len(texts)} texts")
+            total_chars = sum(len(t) for t in texts)
+            self.logger.logger.debug(f"  Total characters: {total_chars}")
+            avg_len = total_chars / len(texts) if texts else 0.0
+            self.logger.logger.debug(f"  Average text length: {avg_len:.1f} chars")
+        
+        # Encode all texts
+        embeddings = self.model.encode(
+            texts,
+            return_dense=True,
+            return_sparse=return_sparse,
+            return_colbert_vecs=return_colbert
+        )
+        
+        result = {
+            'dense': embeddings['dense_vecs'],
+            'dimension': embeddings['dense_vecs'].shape[1],
+            'num_texts': len(texts)
+        }
+        
+        # Add sparse embeddings if requested
+        if return_sparse and 'lexical_weights' in embeddings:
+            result['sparse'] = embeddings['lexical_weights']
+        
+        # Add ColBERT embeddings if requested
+        if return_colbert and 'colbert_vecs' in embeddings:
+            result['colbert'] = embeddings['colbert_vecs']
+        
+        encoding_time = time.time() - start_time
+        
+        if self.logger:
+            self.logger.logger.info(f"✅ Batch encoding completed in {encoding_time:.4f} seconds")
+            avg_time = encoding_time / len(texts) if texts else 0.0
+            self.logger.logger.debug(f"  Average time per text: {avg_time:.4f} seconds")
+        
+        return result
+    
+    def get_embedding_dimension(self) -> int:
+        """Get the dimension of the embeddings."""
+        return self.embedding_dim
+    
+    def compute_similarity(self, vec1: np.ndarray, vec2: np.ndarray, 
+                          metric: str = "cosine") -> float:
+        """
+        Compute similarity between two vectors.
+        
+        Args:
+            vec1: First vector
+            vec2: Second vector
+            metric: Similarity metric ('cosine', 'euclidean', 'dot')
+        
+        Returns:
+            Similarity score
+        """
+        if metric == "cosine":
+            # Cosine similarity
+            dot_product = np.dot(vec1, vec2)
+            norm1 = np.linalg.norm(vec1)
+            norm2 = np.linalg.norm(vec2)
+            similarity = dot_product / (norm1 * norm2)
+            
+        elif metric == "euclidean":
+            # Euclidean distance (negative for similarity)
+            similarity = -np.linalg.norm(vec1 - vec2)
+            
+        elif metric == "dot":
+            # Dot product
+            similarity = np.dot(vec1, vec2)
+            
+        else:
+            raise ValueError(f"Unknown metric: {metric}")
+        
+        if self.logger:
+            self.logger.logger.debug(f"  Similarity ({metric}): {similarity:.6f}")
+        
+        return float(similarity)