Spaces:
Configuration error
Configuration error
| """ | |
| 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 | |
| ) | |