""" Semantic text segmentation and embedding generation module. This module implements the core logic for splitting preprocessed (shielded) text into sentences and generating dense vector embeddings using multilingual Bi-Encoder models with long-context support. CRITICAL: All models must be pre-downloaded to local storage. Mathematical Foundations ------------------------ 1. Sentence Segmentation: - Rule-based tokenization via PySBD (Python Sentence Boundary Disambiguation) - Time complexity: O(n) where n = text length - Reference: Sadvilkar & Neumann, "PySBD: Pragmatic Sentence Boundary Disambiguation", EMNLP 2020 [1] 2. Embedding Generation (Bi-Encoder Architecture): - Input: tokenized sentence s = [t₁, t₂, ..., tₖ] - Output: embedding v ∈ ℝᵈ where d = embedding dimension (384 for MiniLM) - Forward pass: v = MeanPool(BERT(s)) ∈ ℝᵈ - Normalization: v̂ = v / ||v||₂ (L2 normalization for cosine similarity) - Reference: Reimers & Gurevych, "Sentence-BERT", EMNLP 2019 [2] 3. Cosine Similarity for Semantic Distance: - sim(u, v) = (u · v) / (||u||₂ · ||v||₂) ∈ [-1, 1] - For L2-normalized vectors: sim(u, v) = u · v - Distance metric: d(u, v) = 1 - sim(u, v) ∈ [0, 2] - Reference: Manning & Schütze, "Foundations of Statistical NLP" [3] 4. ONNX Runtime Optimization: - Graph optimization: operator fusion, constant folding, layout optimization - Execution providers: CUDA (GPU) → CPU fallback chain - Speedup: 2-4x inference acceleration vs. PyTorch eager mode - Reference: Microsoft ONNX Runtime Documentation [4] References ---------- [1] Sadvilkar, N., & Neumann, M. (2020). PySBD: Pragmatic sentence boundary disambiguation for 20+ languages. EMNLP 2020 System Demonstrations. https://github.com/nipunsadvilkar/pySBD [2] Reimers, N., & Gurevych, I. (2019). Sentence-BERT: Sentence embeddings using Siamese BERT-networks. EMNLP-IJCNLP 2019. https://github.com/UKPLab/sentence-transformers https://arxiv.org/abs/1908.10084 [3] Manning, C. D., & Schütze, H. (1999). Foundations of Statistical Natural Language Processing. MIT Press. Chapter 15: Vector Space Models. [4] Microsoft. (2024). ONNX Runtime: High-performance inference engine. https://github.com/microsoft/onnxruntime https://onnxruntime.ai/docs/execution-providers/ Performance Characteristics --------------------------- - Model loading: ~50-200ms (ONNX from local), ~100-400ms (PyTorch from local) - Sentence segmentation: O(n) with negligible constant factor (~0.1ms/KB) - Embedding inference: * ONNX + CUDA: ~2-5ms per sentence (batch=32) * ONNX + CPU: ~10-30ms per sentence (batch=32) * PyTorch + CUDA: ~5-15ms per sentence (batch=32) - Memory footprint: ~90MB (MiniLM-L6-v2), ~45MB with FP16 on CUDA Thread Safety ------------- - Singleton instance uses double-checked locking for thread-safe lazy init - spaCy/langdetect model caches are protected by class-level locks - ONNX Runtime sessions are thread-safe by design - All instance methods are reentrant; no mutable shared state after init Author: IntelliDeep Labs Team License: BSL 1.1 """ from __future__ import annotations import asyncio import logging import os import threading import time from dataclasses import dataclass from enum import Enum from pathlib import Path from typing import Dict, List, Optional, Tuple, Union import numpy as np import torch from langdetect import detect, LangDetectException from pysbd import Segmenter as PySBDSegmenter from scipy.spatial.distance import cdist # Conditional imports for optional backends try: from optimum.onnxruntime import ORTModelForFeatureExtraction from transformers import AutoTokenizer _ONNX_AVAILABLE = True except ImportError: _ONNX_AVAILABLE = False ORTModelForFeatureExtraction = None # type: ignore AutoTokenizer = None # type: ignore try: from sentence_transformers import SentenceTransformer _SENTENCE_TRANSFORMERS_AVAILABLE = True except ImportError: _SENTENCE_TRANSFORMERS_AVAILABLE = False SentenceTransformer = None # type: ignore # Configure module logger logger = logging.getLogger(__name__) class EmbeddingBackend(str, Enum): """ Enumeration of supported embedding model backends. - ONNX: Optimized inference via ONNX Runtime (recommended for production) - PYTORCH: Native PyTorch inference via SentenceTransformers (flexible for dev) """ ONNX = "onnx" PYTORCH = "pytorch" @dataclass(frozen=True) class SegmentationConfig: """ Immutable configuration for sentence segmentation behavior. Attributes ---------- clean : bool Whether to strip whitespace and filter empty segments (default: False). char_span : bool Whether to return character spans for each sentence (default: False). """ clean: bool = False char_span: bool = False class SemanticSegmenter: """ High-performance semantic text segmenter and embedding generator. This class implements a two-stage pipeline: 1. Language-aware sentence boundary detection via PySBD 2. Dense vector embedding generation via Bi-Encoder (MiniLM or compatible) CRITICAL: Local-Only Model Loading --------------------------------- This module does NOT support automatic HuggingFace downloads. All models must be pre-downloaded to the local filesystem using: python -m nlproxy download_models.py Expected directory structure: models/ ├── all-MiniLM-L6-v2/ │ ├── model.onnx # ONNX format (preferred) │ ├── config.json │ ├── tokenizer.json │ └── ... └── {other-model}/ └── ... Key Features ------------ - Multilingual support: Auto-detects or accepts explicit language codes - Backend flexibility: ONNX Runtime (production) or PyTorch (development) - FP16 optimization: Optional half-precision inference on CUDA devices - Async support: Non-blocking inference via asyncio.to_thread() - Caching: Per-language PySBD segmenters cached at instance level Mathematical Foundations ------------------------ 1. Embedding Normalization (L2): Given raw embedding v ∈ ℝᵈ: ||v||₂ = √(Σᵢ vᵢ²) v̂ = v / (||v||₂ + ε) where ε = 1e-9 for numerical stability Properties: - ||v̂||₂ = 1 (unit norm) - sim(u, v) = u · v = cos(θ) where θ = angle between vectors 2. Mean Pooling for Sentence Embeddings: Given token embeddings T = [t₁, t₂, ..., tₖ] ∈ ℝᵏˣᵈ: v = (1/k) Σᵢ tᵢ ∈ ℝᵈ Alternative: Attention-weighted pooling (not implemented; future extension) 3. Batch Processing Efficiency: For batch size B and sequence length L: Time: O(B · L · d²) for transformer forward pass Memory: O(B · L · d) for activations + gradients (training only) Practical guidance: - Use batch_size=32-128 for optimal GPU utilization - Truncate sequences to max_seq_length=512 for MiniLM compatibility References ---------- [2] Reimers & Gurevych (2019). Sentence-BERT. [4] Microsoft ONNX Runtime Documentation. Performance Notes ----------------- - Model loading is deferred until first inference call (lazy initialization) - FP16 mode reduces memory by ~50% with <1% accuracy loss on STS benchmarks - Warm-up inference call eliminates first-call JIT compilation overhead Usage Example ------------- >>> # Basic usage with local model path >>> segmenter = SemanticSegmenter( ... model_name="all-MiniLM-L6-v2", ... models_dir=Path("models"), ... batch_size=32 ... ) >>> sentences, embeddings = segmenter.segment_and_encode("Hello world. How are you?") >>> print(f"Generated {len(sentences)} embeddings of shape {embeddings.shape}") >>> # Async usage (non-blocking) >>> sentences, embeddings = await segmenter.segment_and_encode_async(text) >>> # Singleton pattern (shared instance) >>> segmenter = SemanticSegmenter.get_instance(batch_size=64) """ # Singleton instance management (thread-safe lazy initialization) _instance: Optional[SemanticSegmenter] = None _singleton_lock: threading.Lock = threading.Lock() # Class-level caches (shared across instances for efficiency) _segmenter_cache: Dict[str, PySBDSegmenter] = {} _segmenter_lock: threading.Lock = threading.Lock() # Model path constants (configurable via environment in production) DEFAULT_MODEL_NAME: str = "all-MiniLM-L6-v2" DEFAULT_MODELS_DIR: Path = Path(os.getenv("NLPROXY_MODELS_DIR") or str(Path(__file__).resolve().parent.parent / "models")) DEFAULT_MAX_SEQ_LENGTH: int = 512 # MiniLM context window DEFAULT_EMBEDDING_DIM: int = 384 # MiniLM output dimension # Numerical constants _L2_NORMALIZATION_EPSILON: float = 1e-9 _WARMUP_SENTENCE: str = "warmup" # Supported languages for PySBD (ISO 639-1 codes) _SUPPORTED_LANGUAGES: Tuple[str, ...] = ( 'en', 'es', 'de', 'fr', 'it', 'pt', 'nl', 'ru', 'zh', 'ja', 'ar', 'hi', 'ko', 'tr', 'pl', 'sv', 'da', 'no', 'fi' ) # Required files for model validation _ONNX_REQUIRED_FILES: Tuple[str, ...] = ("model.onnx", "config.json", "tokenizer.json") _PYTORCH_REQUIRED_FILES: Tuple[str, ...] = ("config.json", "pytorch_model.bin") def __init__( self, model_name: str = DEFAULT_MODEL_NAME, models_dir: Optional[Union[str, Path]] = None, use_fp16: bool = True, device: Optional[str] = None, max_seq_length: int = DEFAULT_MAX_SEQ_LENGTH, batch_size: int = 32, language: Optional[str] = None, backend: Optional[EmbeddingBackend] = None, onnx_int8: bool = False, ) -> None: """ Initialize the SemanticSegmenter with local model loading. Parameters ---------- model_name : str, optional Name of the model directory under models_dir (default: all-MiniLM-L6-v2). Example: models_dir="models", model_name="all-MiniLM-L6-v2" → loads from "models/all-MiniLM-L6-v2/" models_dir : Optional[Union[str, Path]], optional Base directory containing pre-downloaded models (default: "models"). Must contain subdirectories for each model with required files. use_fp16 : bool, optional Enable half-precision inference on CUDA devices (default: True). device : Optional[str], optional Explicit device specification ("cuda", "cpu", "cuda:0"). If None, auto-detects CUDA availability. max_seq_length : int, optional Maximum token sequence length for truncation (default: 512). batch_size : int, optional Number of sentences to process per inference call (default: 32). language : Optional[str], optional ISO 639-1 language code to force segmentation language. If None, auto-detects via langdetect with fallback to 'en'. backend : Optional[EmbeddingBackend], optional Explicit backend selection. If None, auto-detects ONNX availability. onnx_int8 : bool, optional Enable support for CPU INT8 quantized ONNX models when available. Raises ------ FileNotFoundError If the specified model directory or required files are missing. ImportError If neither ONNX Runtime nor SentenceTransformers is available. ValueError If specified backend is unavailable or device is invalid. Thread Safety ------------- - Model loading is protected by instance-level state; safe for single-threaded use - For multi-threaded scenarios, prefer `get_instance()` singleton pattern - PySBD segmenter cache is protected by class-level lock Example ------- >>> from pathlib import Path >>> segmenter = SemanticSegmenter( ... model_name="all-MiniLM-L6-v2", ... models_dir=Path("/opt/nlproxy/models"), ... device="cuda", ... batch_size=64 ... ) """ # Validate dependencies if not _ONNX_AVAILABLE and not _SENTENCE_TRANSFORMERS_AVAILABLE: raise ImportError( "Either 'optimum[onnxruntime]' or 'sentence-transformers' " "must be installed. Install with: " "pip install nlproxy[onnx] or pip install nlproxy[dev]" ) # Resolve models directory and default to model-specific folder under nlproxy/models self.model_name = model_name if models_dir: candidate = Path(models_dir) # If caller provided a directory that already points to the model folder, use it if candidate.exists() and candidate.name == model_name: self.model_path = candidate self.models_dir = candidate else: # Treat provided value as base models root self.models_dir = candidate self.model_path = self.models_dir / model_name else: # Default to nlproxy/models/{model_name} self.model_path = self.DEFAULT_MODELS_DIR / model_name self.models_dir = self.model_path # Validate model directory exists if not self.model_path.exists(): raise FileNotFoundError( f"Model directory not found: {self.model_path}\n" f"Please download models using: python -m nlproxy download_models" ) # Store configuration self.use_fp16 = use_fp16 self.batch_size = batch_size self.max_seq_length = max_seq_length self.language = language self.onnx_int8 = onnx_int8 self._backend = backend self._embedding_dim = self.DEFAULT_EMBEDDING_DIM # Resolve device if device is None: self.device = "cuda" if torch.cuda.is_available() else "cpu" else: # Validate device string if device.startswith("cuda") and not torch.cuda.is_available(): logger.warning(f"CUDA requested but not available; falling back to CPU") self.device = "cpu" else: self.device = device # Resolve backend preference (ONNX preferred if available and files exist) if self._backend is None: if _ONNX_AVAILABLE and self._has_onnx_files(): self._backend = EmbeddingBackend.ONNX elif _SENTENCE_TRANSFORMERS_AVAILABLE: self._backend = EmbeddingBackend.PYTORCH else: raise ImportError( f"No compatible backend available for model {model_name}. " f"ONNX files: {self._has_onnx_files()}, " f"PyTorch available: {_SENTENCE_TRANSFORMERS_AVAILABLE}" ) logger.debug(f"Auto-selected backend: {self._backend.value}") elif self._backend == EmbeddingBackend.ONNX and not _ONNX_AVAILABLE: logger.warning("ONNX backend requested but unavailable; falling back to PyTorch") self._backend = EmbeddingBackend.PYTORCH # Lazy initialization flags self._model_loaded: bool = False self._embedding_model: Optional[Union[ORTModelForFeatureExtraction, SentenceTransformer]] = None self._tokenizer: Optional[AutoTokenizer] = None self._is_onnx: bool = (self._backend == EmbeddingBackend.ONNX) self._loading_lock: threading.Lock = threading.Lock() # Per-instance PySBD cache (language -> segmenter) self._segmenters: Dict[str, PySBDSegmenter] = {} logger.info( f"SemanticSegmenter initialized: model={model_name}, " f"path={self.model_path}, backend={self._backend.value}, " f"device={self.device}, fp16={self.use_fp16 and self.device == 'cuda'}, " f"batch_size={self.batch_size}" ) def _has_onnx_files(self) -> bool: """ Check if required ONNX model files exist in the model directory. Returns ------- bool True if all required ONNX files are present. """ return all((self.model_path / f).exists() for f in self._ONNX_REQUIRED_FILES) def _has_pytorch_files(self) -> bool: """ Check if required PyTorch model files exist in the model directory. Returns ------- bool True if all required PyTorch files are present. """ return all((self.model_path / f).exists() for f in self._PYTORCH_REQUIRED_FILES) @classmethod def get_instance( cls, model_name: str = DEFAULT_MODEL_NAME, models_dir: Optional[Union[str, Path]] = None, use_fp16: bool = True, device: Optional[str] = None, max_seq_length: int = DEFAULT_MAX_SEQ_LENGTH, batch_size: int = 32, language: Optional[str] = None, backend: Optional[EmbeddingBackend] = None ) -> SemanticSegmenter: """ Get or create the singleton instance of SemanticSegmenter. Thread-safe implementation using double-checked locking pattern. Recommended for applications requiring shared embedding model state. Parameters ---------- model_name : str, optional Model directory name (only used on first creation). models_dir : Optional[Union[str, Path]], optional Base directory for models (only used on first creation). use_fp16 : bool, optional Enable FP16 inference (only used on first creation). device : Optional[str], optional Device specification (only used on first creation). max_seq_length : int, optional Maximum sequence length (only used on first creation). batch_size : int, optional Inference batch size (only used on first creation). language : Optional[str], optional Default language for segmentation (only used on first creation). backend : Optional[EmbeddingBackend], optional Backend preference (only used on first creation). Returns ------- SemanticSegmenter The singleton instance. Note ---- Subsequent calls return the existing instance regardless of parameters. To change configuration, use `reset_instance()` and call again. """ if cls._instance is None: with cls._singleton_lock: if cls._instance is None: cls._instance = cls( model_name=model_name, models_dir=models_dir, use_fp16=use_fp16, device=device, max_seq_length=max_seq_length, batch_size=batch_size, language=language, backend=backend ) return cls._instance @classmethod def reset_instance(cls) -> None: """Reset the singleton instance (useful for testing).""" with cls._singleton_lock: cls._instance = None def _load_model(self) -> None: """ Thread-safe model loading with ONNX fallback to PyTorch. Includes proper warm-up handling and error recovery. """ # Fast path: already loaded if self._model_loaded: return # Double-checked locking for thread safety with self._loading_lock: if self._model_loaded: return start = time.time() logger.info(f"Starting model loading sequence for {self.model_name}...") if self._backend == EmbeddingBackend.ONNX: if not self._has_onnx_files(): logger.warning(f"ONNX files missing in {self.model_path}. Switching to PyTorch.") self._backend = EmbeddingBackend.PYTORCH else: try: onnx_model_source = self.model_path if self.onnx_int8 and self.device.startswith("cpu"): int8_model_path = self.model_path / "model_int8.onnx" if int8_model_path.exists(): logger.info( f"Loading CPU INT8 quantized ONNX model from {int8_model_path}..." ) onnx_model_source = int8_model_path else: logger.warning( "INT8 quantization requested but quantized model file " f"not found at {int8_model_path}. Falling back to standard ONNX model." ) else: logger.info(f"Loading ONNX model from {self.model_path}...") import onnxruntime available_providers = onnxruntime.get_available_providers() providers = [] if self.device.startswith("cuda") and torch.cuda.is_available() and "CUDAExecutionProvider" in available_providers: providers.append("CUDAExecutionProvider") providers.append("CPUExecutionProvider") self._embedding_model = ORTModelForFeatureExtraction.from_pretrained( str(onnx_model_source), provider=providers[0] if len(providers) == 1 else providers, export=False ) self._tokenizer = AutoTokenizer.from_pretrained(str(self.model_path)) self._is_onnx = True logger.info("✅ ONNX model loaded successfully.") self._model_loaded = True try: logger.debug("Running ONNX warm-up inference...") self.encode_batch([self._WARMUP_SENTENCE]) logger.debug("✅ ONNX warm-up completed.") except Exception as warmup_error: logger.warning(f"ONNX warm-up failed ({warmup_error}). Falling back to PyTorch.") self._backend = EmbeddingBackend.PYTORCH self._embedding_model = None self._tokenizer = None self._is_onnx = False except (AttributeError, ImportError, RuntimeError) as e: if "int4" in str(e).lower() or "torch" in str(e).lower(): logger.error(f"ONNX load failed (PyTorch compatibility: {e}). Falling back to PyTorch.") else: logger.error(f"ONNX load failed ({type(e).__name__}): {e}. Falling back to PyTorch.") self._backend = EmbeddingBackend.PYTORCH except Exception as e: logger.error(f"Unexpected ONNX error ({type(e).__name__}): {e}. Falling back to PyTorch.") self._backend = EmbeddingBackend.PYTORCH if self._backend == EmbeddingBackend.PYTORCH: if not self._has_pytorch_files(): raise FileNotFoundError( f"PyTorch model files not found in {self.model_path}. " f"Required: {self._PYTORCH_REQUIRED_FILES}. " f"Run: python -m nlproxy download_models" ) logger.info(f"Loading PyTorch model from {self.model_path} on {self.device}...") try: self._embedding_model = SentenceTransformer( str(self.model_path), device=self.device, trust_remote_code=True ) if self.use_fp16 and self.device == "cuda": self._embedding_model = self._embedding_model.half() logger.debug("Enabled FP16 precision") self._embedding_model.max_seq_length = self.max_seq_length self._is_onnx = False logger.info("✅ PyTorch model loaded successfully.") try: self.encode_batch([self._WARMUP_SENTENCE]) logger.debug("✅ PyTorch warm-up completed.") except Exception as e: logger.warning(f"PyTorch warm-up warning: {e}") except Exception as e: raise RuntimeError(f"PyTorch model loading failed: {e}") from e if self._embedding_model is not None: elapsed = time.time() - start logger.info(f"✅ Model fully loaded in {elapsed:.2f}s (Backend: {'ONNX' if self._is_onnx else 'PyTorch'})") self._model_loaded = True else: raise RuntimeError("Model loading failed: no backend successfully initialized") def _get_segmenter(self, text: str, language: Optional[str] = None) -> PySBDSegmenter: """ Retrieve or create a PySBD segmenter for the detected language. Parameters ---------- text : str Input text for language detection (if language not specified). language : Optional[str], optional Explicit ISO 639-1 language code. If None, auto-detect from text. Returns ------- PySBDSegmenter Configured segmenter instance for the target language. Thread Safety ------------- Access to _segmenters cache is protected by _segmenter_lock. """ # Resolve language if language is not None: lang = language elif self.language: lang = self.language else: try: lang = detect(text) if text else "en" if lang not in self._SUPPORTED_LANGUAGES: lang = "en" # Fallback for unsupported languages except (LangDetectException, Exception): lang = "en" # Retrieve or create segmenter (thread-safe) with self._segmenter_lock: if lang not in self._segmenters: try: self._segmenters[lang] = PySBDSegmenter(language=lang, clean=False) except (ValueError, KeyError, Exception): logger.warning(f"PySBD does not support language '{lang}', falling back to 'en'") if "en" not in self._segmenters: self._segmenters["en"] = PySBDSegmenter(language="en", clean=False) self._segmenters[lang] = self._segmenters["en"] logger.debug(f"Created PySBD segmenter for language '{lang}'") return self._segmenters[lang] def split_sentences(self, text: str, language: Optional[str] = None) -> List[str]: """ Split input text into sentences using language-aware segmentation. Parameters ---------- text : str Input text to segment. language : Optional[str], optional ISO 639-1 language code. If None, auto-detect from text. Returns ------- List[str] List of cleaned, non-empty sentences. Raises ------ TypeError If input is not a string. Complexity ---------- Time: O(n) where n = text length Space: O(k) where k = number of sentences """ if not isinstance(text, str): raise TypeError(f"Expected str input, got {type(text).__name__}") segmenter = self._get_segmenter(text, language=language) raw_sentences = segmenter.segment(text) # Filter and clean sentences clean = [s.strip() for s in raw_sentences if s and s.strip()] return clean def encode_batch(self, sentences: List[str], normalize: bool = True) -> np.ndarray: """ Generate dense embeddings for a batch of sentences. Parameters ---------- sentences : List[str] List of sentences to encode. normalize : bool, optional Apply L2 normalization to embeddings (default: True). Returns ------- np.ndarray Array of shape (len(sentences), embedding_dim) with float32 embeddings. Raises ------ ValueError If sentences list is empty. RuntimeError If model is not loaded or inference fails. ONNX Backend ------------ - Tokenization: AutoTokenizer with padding/truncation - Inference: ORTModelForFeatureExtraction with mean pooling - Normalization: L2 normalization with epsilon for stability PyTorch Backend --------------- - Delegates to SentenceTransformer.encode() with batch processing - Automatic device placement (CPU/CUDA) - Optional FP16 precision on CUDA devices Performance ----------- Time: O(B · L · d²) where B=batch_size, L=avg_seq_length, d=embedding_dim Space: O(B · L · d) for intermediate activations """ if not sentences: raise ValueError("Cannot encode empty sentence list") # Ensure model is loaded (lazy initialization) if not self._model_loaded: self._load_model() if self._is_onnx: if self._tokenizer is None: raise RuntimeError("Tokenizer not initialized for ONNX backend") # Tokenize with padding and truncation inputs = self._tokenizer( sentences, padding=True, truncation=True, max_length=self.max_seq_length, return_tensors="np" ) # Forward pass through ONNX model outputs = self._embedding_model(**inputs) # Mean pooling over token embeddings embeddings = outputs.last_hidden_state.mean(axis=1) # Optional L2 normalization if normalize: norms = np.linalg.norm(embeddings, axis=1, keepdims=True) embeddings = embeddings / (norms + self._L2_NORMALIZATION_EPSILON) return embeddings.astype(np.float32) else: # PyTorch backend via SentenceTransformers return self._embedding_model.encode( sentences, batch_size=self.batch_size, show_progress_bar=False, convert_to_numpy=True, normalize_embeddings=normalize, device=self.device ).astype(np.float32) def segment_and_encode(self, text: str) -> Tuple[List[str], np.ndarray]: """ Execute the complete segmentation + embedding pipeline. Parameters ---------- text : str Input text to process. Returns ------- Tuple[List[str], np.ndarray] - List of segmented sentences - Array of corresponding embeddings (shape: [n_sentences, embedding_dim]) Pipeline Stages --------------- 1. Sentence segmentation via PySBD (language-aware) 2. Batch embedding generation via Bi-Encoder 3. Optional L2 normalization for cosine similarity compatibility Debug Output ------------ When logger level is DEBUG, outputs: - First sentence sent to model - Cosine distance matrix for first 5 embeddings (if applicable) - Mean pairwise distance for quality verification Example ------- >>> sentences, embeddings = segmenter.segment_and_encode("Hello. World.") >>> assert len(sentences) == embeddings.shape[0] >>> assert embeddings.shape[1] == 384 # MiniLM dimension """ # Stage 1: Segmentation sentences = self.split_sentences(text) if not sentences: logger.warning("No valid sentences found in input text") return [], np.array([], dtype=np.float32).reshape(0, self._embedding_dim) # Stage 2: Embedding generation embeddings = self.encode_batch(sentences) # Debug logging (only at DEBUG level) if logger.isEnabledFor(logging.DEBUG): logger.debug("--- Embedding Verification ---") logger.debug(f"First sentence: {sentences[0][:100]}...") if embeddings.shape[0] >= 5: # Compute pairwise cosine distances for first 5 embeddings dist_matrix = cdist(embeddings[:5], embeddings[:5], metric='cosine') logger.debug(f"Cosine distance matrix (5x5):\n{dist_matrix}") logger.debug(f"Mean pairwise distance: {dist_matrix.mean():.4f}") return sentences, embeddings async def encode_batch_async(self, sentences: List[str], normalize: bool = True) -> np.ndarray: """ Asynchronous version of encode_batch (non-blocking event loop). Parameters ---------- sentences : List[str] List of sentences to encode. normalize : bool, optional Apply L2 normalization (default: True). Returns ------- np.ndarray Embedding array. Note ---- Uses asyncio.to_thread() to offload CPU-bound inference to worker thread. Does not provide true parallelism but prevents event loop blocking. """ return await asyncio.to_thread(self.encode_batch, sentences, normalize) async def segment_and_encode_async(self, text: str) -> Tuple[List[str], np.ndarray]: """ Asynchronous version of segment_and_encode. Parameters ---------- text : str Input text to process. Returns ------- Tuple[List[str], np.ndarray] Segmented sentences and their embeddings. Note ---- - Sentence segmentation is synchronous (fast, CPU-bound) - Embedding generation is offloaded via encode_batch_async - Suitable for high-concurrency async applications """ # Segmentation is fast and CPU-bound; no async overhead needed sentences = self.split_sentences(text) if not sentences: return [], np.array([], dtype=np.float32).reshape(0, self._embedding_dim) # Offload embedding generation to worker thread embeddings = await self.encode_batch_async(sentences) return sentences, embeddings @property def embedding_dim(self) -> int: """Return the embedding dimension of the loaded model.""" return self._embedding_dim @property def is_onnx(self) -> bool: """Return True if using ONNX backend, False if using PyTorch.""" return self._is_onnx def get_model_info(self) -> Dict[str, any]: """ Return diagnostic information about the loaded model. Returns ------- Dict[str, any] Dictionary with model metadata for debugging/monitoring. """ return { "model_name": self.model_name, "model_path": str(self.model_path), "backend": self._backend.value if self._backend else None, "is_onnx": self._is_onnx, "device": self.device, "embedding_dim": self._embedding_dim, "max_seq_length": self.max_seq_length, "batch_size": self.batch_size, "use_fp16": self.use_fp16 and self.device == "cuda", "model_loaded": self._model_loaded, "supported_languages": list(self._SUPPORTED_LANGUAGES) }