""" Embedding Model — Sentence-Transformers for vector embeddings (RAG). Uses paraphrase-multilingual-MiniLM-L12-v2 (supports Vietnamese). """ import gc import time import torch import logging import numpy as np from typing import List, Union import config logger = logging.getLogger(__name__) class EmbeddingModel: """Manages loading, inference, and unloading of the embedding model.""" def __init__(self, model_id: str, device: str = "cpu", cache_dir: str = None): self.model_id = model_id self.device = device self.cache_dir = cache_dir self.model = None def load(self): """Load embedding model into memory.""" if self.model is not None: logger.info("Embedding model already loaded, skipping.") return start = time.time() logger.info(f"Loading embedding model: {self.model_id} ...") from sentence_transformers import SentenceTransformer self.model = SentenceTransformer( self.model_id, device=self.device, cache_folder=self.cache_dir, ) elapsed = time.time() - start logger.info(f"✅ Embedding model loaded in {elapsed:.1f}s (device={self.device})") def unload(self): """Free memory.""" if self.model is not None: del self.model self.model = None if self.device == "cuda": torch.cuda.empty_cache() gc.collect() logger.info("🗑️ Embedding model unloaded from memory.") def embed(self, text: Union[str, List[str]], prefix: str = "passage: ") -> Union[List[float], List[List[float]]]: """ Generate embedding for a single text or a list of texts (batch). Prepends the required E5 prefix (passage: or query:). """ if self.model is None: raise RuntimeError("Embedding model not loaded. Call load() first.") if isinstance(text, list): # Batch processing: prepend prefix to each element if not present input_texts = [ (t if t.startswith(prefix) else prefix + t) for t in text ] else: # Single text processing input_texts = text if text.startswith(prefix) else prefix + text with torch.no_grad(): embeddings = self.model.encode( input_texts, convert_to_tensor=True, show_progress_bar=False, device=self.device ) if isinstance(text, list): return embeddings.tolist() return embeddings.cpu().numpy().tolist() def embed_query(self, query: str) -> list: """Helper to embed a query with 'query: ' prefix.""" return self.embed(query, prefix="query: ") def embed_chunks(self, chunks: list) -> list: """Helper to embed multiple chunks with 'passage: ' prefix.""" if self.model is None: raise RuntimeError("Embedding model not loaded. Call load() first.") prefixed_chunks = [ (c if c.startswith("passage: ") else "passage: " + c) for c in chunks ] embeddings = self.model.encode(prefixed_chunks, convert_to_numpy=True, batch_size=8) return [emb.tolist() for emb in embeddings] def embed_batch(self, texts: list) -> list: """ Generate embeddings for multiple texts. Args: texts: List of input texts Returns: List of embedding vectors """ if self.model is None: raise RuntimeError("Embedding model not loaded. Call load() first.") embeddings = self.model.encode(texts, convert_to_numpy=True, batch_size=8) return [emb.tolist() for emb in embeddings] class SemanticChunker: """ Split text into chunks based on semantic similarity. Works on text WITH or WITHOUT punctuation. Algorithm: 1. Split text into small word-groups (~15 words each) 2. Embed each word-group 3. Compare embeddings of adjacent groups 4. When similarity drops below threshold → that's a topic boundary → cut 5. Merge small chunks into neighbors """ def __init__(self, embed_model: EmbeddingModel, threshold: float = 0.65, min_chunk_words: int = 150, max_chunk_words: int = 350, group_size: int = 25): self.embed_model = embed_model self.threshold = threshold self.min_chunk_words = min_chunk_words self.max_chunk_words = max_chunk_words self.group_size = group_size # Number of words per micro-group (soi kỹ hơn) def chunk(self, text: str) -> list: """ Split text into semantic chunks with overlap. """ words = text.split() if len(words) <= self.max_chunk_words: return [text] # 1. Group words into micro-groups groups = [] for i in range(0, len(words), self.group_size): groups.append(" ".join(words[i:i + self.group_size])) if len(groups) < 2: return [" ".join(words)] # 2. Embed all groups (Using batch processing for speed) group_embeddings = self.embed_model.embed_chunks(groups) # 3. Calculate similarities between adjacent groups similarities = [] for i in range(len(group_embeddings) - 1): sim = self._cosine_similarity(group_embeddings[i], group_embeddings[i+1]) similarities.append(sim) # 4. Build chunks with overlap chunks = [] current_groups = [] overlap_size = 2 # Number of micro-groups to overlap (~20 words) for i in range(len(groups)): current_groups.append(groups[i]) # Decide to cut should_cut = False current_word_count = sum(len(g.split()) for g in current_groups) if i < len(groups) - 1: sim = similarities[i] if sim < self.threshold and current_word_count >= self.min_chunk_words: should_cut = True if current_word_count >= self.max_chunk_words: should_cut = True if should_cut: chunks.append(" ".join(current_groups)) # Create overlap: keep the last N groups for the next chunk current_groups = current_groups[-overlap_size:] if len(current_groups) > overlap_size else [] # Last chunk if current_groups and len(current_groups) > overlap_size: chunks.append(" ".join(current_groups)) # Filter out tiny chunks return [c for c in chunks if len(c.split()) >= self.min_chunk_words/2] def _cosine_similarity(self, v1, v2): v1 = np.array(v1) v2 = np.array(v2) return np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2)) # Global instance for easy access across the models package embedder = EmbeddingModel( config.EMBEDDING_MODEL_ID, device=config.DEVICE, cache_dir=config.HF_CACHE_DIR )