vector-inspector 0.2.7__py3-none-any.whl → 0.3.2__py3-none-any.whl

vector_inspector/core/embedding_providers/sentence_transformer_provider.py

@@ -0,0 +1,203 @@
+"""Sentence Transformer embedding provider with lazy loading."""
+
+from typing import List, Union, Optional
+import numpy as np
+
+from .base_provider import (
+    EmbeddingProvider, 
+    EmbeddingMetadata, 
+    Modality, 
+    Normalization
+)
+
+
+class SentenceTransformerProvider(EmbeddingProvider):
+    """Provider for sentence-transformers models.
+    
+    Lazy-loads the sentence-transformers library and model on first use.
+    Supports all models from the sentence-transformers library including:
+    - all-MiniLM-L6-v2
+    - all-mpnet-base-v2
+    - BGE, GTE, E5 families
+    - Multilingual variants
+    """
+    
+    def __init__(self, model_name: str, trust_remote_code: bool = False):
+        """Initialize sentence-transformer provider.
+        
+        Args:
+            model_name: HuggingFace model ID or local path
+            trust_remote_code: Whether to trust remote code (for some models)
+        """
+        super().__init__(model_name)
+        self.trust_remote_code = trust_remote_code
+        self._metadata = None
+    
+    def get_metadata(self) -> EmbeddingMetadata:
+        """Get metadata about the sentence-transformer model.
+        
+        This attempts to extract metadata without loading the full model if possible.
+        """
+        if self._metadata is not None:
+            return self._metadata
+        
+        # Try to get metadata without loading full model
+        try:
+            # Import config utilities
+            from sentence_transformers import SentenceTransformer
+            from transformers import AutoConfig
+            
+            # Try to get config without loading weights
+            try:
+                config = AutoConfig.from_pretrained(self.model_name)
+                dimension = getattr(config, 'hidden_size', None)
+                max_length = getattr(config, 'max_position_embeddings', None)
+            except Exception:
+                # If config fails, we'll need to load the model
+                dimension = None
+                max_length = None
+            
+            # If we couldn't get dimension from config, load model
+            if dimension is None:
+                if not self._is_loaded:
+                    self._load_model()
+                dimension = self._model.get_sentence_embedding_dimension()
+                max_length = self._model.max_seq_length
+            
+            self._metadata = EmbeddingMetadata(
+                name=self.model_name,
+                dimension=int(dimension) if dimension is not None else 768,
+                modality=Modality.TEXT,
+                normalization=Normalization.L2,  # Most sentence-transformers normalize
+                model_type="sentence-transformer",
+                source="huggingface",
+                max_sequence_length=max_length,
+                description=f"Sentence-Transformer model: {self.model_name}"
+            )
+            
+        except ImportError:
+            # sentence-transformers not installed
+            raise ImportError(
+                "sentence-transformers library not installed. "
+                "Install with: pip install sentence-transformers"
+            )
+        except Exception as e:
+            # Fallback metadata if we can't determine dimension
+            self._metadata = EmbeddingMetadata(
+                name=self.model_name,
+                dimension=768,  # Common default
+                modality=Modality.TEXT,
+                normalization=Normalization.L2,
+                model_type="sentence-transformer",
+                source="huggingface",
+                description=f"Sentence-Transformer model: {self.model_name} (dimension not verified)"
+            )
+        
+        return self._metadata
+    
+    def _load_model(self):
+        """Load the sentence-transformer model."""
+        try:
+            from sentence_transformers import SentenceTransformer
+        except ImportError:
+            raise ImportError(
+                "sentence-transformers library not installed. "
+                "Install with: pip install sentence-transformers"
+            )
+        
+        # Load model with optional trust_remote_code
+        self._model = SentenceTransformer(
+            self.model_name,
+            trust_remote_code=self.trust_remote_code
+        )
+    
+    def encode(
+        self,
+        inputs: Union[str, List[str]],
+        normalize: bool = True,
+        show_progress: bool = False
+    ) -> np.ndarray:
+        """Encode text inputs into embeddings.
+        
+        Args:
+            inputs: Single string or list of strings
+            normalize: Whether to L2-normalize embeddings
+            show_progress: Whether to show progress bar
+            
+        Returns:
+            numpy array of embeddings
+        """
+        # Ensure model is loaded
+        if not self._is_loaded:
+            self.warmup()
+        
+        # Convert single string to list
+        if isinstance(inputs, str):
+            inputs = [inputs]
+        
+        # Encode
+        embeddings = self._model.encode(
+            inputs,
+            normalize_embeddings=normalize,
+            show_progress_bar=show_progress,
+            convert_to_numpy=True
+        )
+        
+        return embeddings
+    
+    def encode_batch(
+        self,
+        inputs: List[str],
+        batch_size: int = 32,
+        normalize: bool = True,
+        show_progress: bool = True
+    ) -> np.ndarray:
+        """Encode a large batch of texts efficiently.
+        
+        Args:
+            inputs: List of strings
+            batch_size: Batch size for encoding
+            normalize: Whether to L2-normalize embeddings
+            show_progress: Whether to show progress bar
+            
+        Returns:
+            numpy array of embeddings
+        """
+        if not self._is_loaded:
+            self.warmup()
+        
+        embeddings = self._model.encode(
+            inputs,
+            batch_size=batch_size,
+            normalize_embeddings=normalize,
+            show_progress_bar=show_progress,
+            convert_to_numpy=True
+        )
+        
+        return embeddings
+    
+    def similarity(self, query: Union[str, np.ndarray], corpus: List[str]) -> np.ndarray:
+        """Compute similarity between query and corpus.
+        
+        Args:
+            query: Query string or embedding
+            corpus: List of corpus strings
+            
+        Returns:
+            Similarity scores (cosine similarity if normalized)
+        """
+        if not self._is_loaded:
+            self.warmup()
+        
+        # Get embeddings
+        if isinstance(query, str):
+            query_emb = self.encode(query, normalize=True)
+        else:
+            query_emb = query
+        
+        corpus_emb = self.encode(corpus, normalize=True)
+        
+        # Compute cosine similarity (dot product if normalized)
+        similarities = np.dot(corpus_emb, query_emb.T).squeeze()
+        
+        return similarities

vector_inspector/core/embedding_utils.py

@@ -3,33 +3,35 @@
 from typing import Optional, Union, Tuple
 from sentence_transformers import SentenceTransformer
 
+from .model_registry import get_model_registry
 
-# Mapping of vector dimensions to appropriate models
-# Format: dimension -> list of (model_name, model_type, description)
-# Listed in order of preference for ambiguous cases
-DIMENSION_TO_MODEL = {
-    384: [
-        ("all-MiniLM-L6-v2", "sentence-transformer", "Fast text embeddings"),
-    ],
-    512: [
-        ("openai/clip-vit-base-patch32", "clip", "Multi-modal (text + images)"),
-        ("paraphrase-albert-small-v2", "sentence-transformer", "Text-only paraphrase"),
-    ],
-    768: [
-        ("all-mpnet-base-v2", "sentence-transformer", "High quality text embeddings"),
-    ],
-    1024: [
-        ("all-roberta-large-v1", "sentence-transformer", "Large text embeddings"),
-    ],
-    1536: [
-        ("gtr-t5-large", "sentence-transformer", "Very large text embeddings"),
-    ],
-}
 
 # Default model to use when dimension is unknown or not mapped
 DEFAULT_MODEL = ("all-MiniLM-L6-v2", "sentence-transformer")
 
 
+def _get_dimension_to_model_dict():
+    """Build dimension->models dictionary from registry.
+    
+    Returns:
+        Dict mapping dimension to list of (name, type, description) tuples
+    """
+    registry = get_model_registry()
+    dimension_map = {}
+    
+    for dimension in registry.get_all_dimensions():
+        models = registry.get_models_by_dimension(dimension)
+        dimension_map[dimension] = [
+            (m.name, m.type, m.description) for m in models
+        ]
+    
+    return dimension_map
+
+
+# For backward compatibility - dynamically loads from registry
+DIMENSION_TO_MODEL = _get_dimension_to_model_dict()
+
+
 def get_model_for_dimension(dimension: int, prefer_multimodal: bool = True) -> Tuple[str, str]:
     """
     Get the appropriate embedding model name and type for a given vector dimension.
@@ -42,30 +44,36 @@ def get_model_for_dimension(dimension: int, prefer_multimodal: bool = True) -> T
     Returns:
         Tuple of (model_name, model_type) where model_type is "sentence-transformer" or "clip"
     """
-    if dimension in DIMENSION_TO_MODEL:
-        models = DIMENSION_TO_MODEL[dimension]
-        if len(models) == 1:
-            return (models[0][0], models[0][1])
-        
-        # Multiple models available - apply preference
-        if prefer_multimodal:
-            # Prefer CLIP/multimodal
-            for model_name, model_type, desc in models:
-                if model_type == "clip":
-                    return (model_name, model_type)
-        
-        # Default to first option
-        return (models[0][0], models[0][1])
+    registry = get_model_registry()
+    models = registry.get_models_by_dimension(dimension)
+    
+    if not models:
+        # Find the closest dimension if exact match not found
+        closest_dim = registry.find_closest_dimension(dimension)
+        if closest_dim:
+            models = registry.get_models_by_dimension(closest_dim)
+    
+    if not models:
+        return DEFAULT_MODEL
+    
+    if len(models) == 1:
+        return (models[0].name, models[0].type)
     
-    # Find the closest dimension if exact match not found
-    closest_dim = min(DIMENSION_TO_MODEL.keys(), key=lambda x: abs(x - dimension))
-    models = DIMENSION_TO_MODEL[closest_dim]
-    return (models[0][0], models[0][1])
+    # Multiple models available - apply preference
+    if prefer_multimodal:
+        # Prefer CLIP/multimodal
+        for model in models:
+            if model.modality == "multimodal" or model.type == "clip":
+                return (model.name, model.type)
+    
+    # Default to first option
+    return (models[0].name, models[0].type)
 
 
 def get_available_models_for_dimension(dimension: int) -> list:
     """
     Get all available model options for a given dimension.
+    Includes both predefined (from registry) and custom user-added models.
     
     Args:
         dimension: The vector dimension size
@@ -73,9 +81,28 @@ def get_available_models_for_dimension(dimension: int) -> list:
     Returns:
         List of tuples: [(model_name, model_type, description), ...]
     """
-    if dimension in DIMENSION_TO_MODEL:
-        return DIMENSION_TO_MODEL[dimension]
-    return []
+    # Start with models from registry
+    registry = get_model_registry()
+    registry_models = registry.get_models_by_dimension(dimension)
+    models = [(m.name, m.type, m.description) for m in registry_models]
+    
+    # Add custom models from settings
+    try:
+        from ..services.settings_service import SettingsService
+        settings = SettingsService()
+        custom_models = settings.get_custom_embedding_models(dimension)
+        
+        for model in custom_models:
+            # Format: (model_name, model_type, description)
+            models.append((
+                model["name"],
+                model["type"],
+                f"{model['description']} (custom)"
+            ))
+    except Exception as e:
+        print(f"Warning: Could not load custom models: {e}")
+    
+    return models
 
 
 def load_embedding_model(model_name: str, model_type: str) -> Union[SentenceTransformer, any]:

vector_inspector/core/model_registry.py

@@ -0,0 +1,205 @@
+"""Model registry for loading and managing known embedding models."""
+
+import json
+from pathlib import Path
+from typing import List, Dict, Optional, Tuple
+from dataclasses import dataclass
+
+
+@dataclass
+class ModelInfo:
+    """Information about an embedding model."""
+    name: str
+    type: str
+    dimension: int
+    modality: str
+    normalization: str
+    source: str
+    description: str
+    
+    def to_dict(self) -> Dict:
+        """Convert to dictionary."""
+        return {
+            "name": self.name,
+            "type": self.type,
+            "dimension": self.dimension,
+            "modality": self.modality,
+            "normalization": self.normalization,
+            "source": self.source,
+            "description": self.description
+        }
+    
+    @classmethod
+    def from_dict(cls, data: Dict) -> 'ModelInfo':
+        """Create from dictionary."""
+        return cls(
+            name=data["name"],
+            type=data["type"],
+            dimension=data["dimension"],
+            modality=data["modality"],
+            normalization=data["normalization"],
+            source=data["source"],
+            description=data["description"]
+        )
+
+
+class EmbeddingModelRegistry:
+    """Registry of known embedding models loaded from JSON."""
+    
+    _instance = None
+    _models: List[ModelInfo] = []
+    _dimension_index: Dict[int, List[ModelInfo]] = {}
+    _name_index: Dict[str, ModelInfo] = {}
+    
+    def __new__(cls):
+        """Singleton pattern."""
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+            cls._instance._load_registry()
+        return cls._instance
+    
+    def _load_registry(self):
+        """Load models from JSON file."""
+        registry_path = Path(__file__).parent.parent / "config" / "known_embedding_models.json"
+        
+        if not registry_path.exists():
+            print(f"Warning: Model registry not found at {registry_path}")
+            return
+
+        try:
+            with open(registry_path, 'r', encoding='utf-8') as f:
+                data = json.load(f)
+
+            # Parse models
+            for model_data in data.get("models", []):
+                model_info = ModelInfo.from_dict(model_data)
+                self._models.append(model_info)
+                
+                # Index by dimension
+                if model_info.dimension not in self._dimension_index:
+                    self._dimension_index[model_info.dimension] = []
+                self._dimension_index[model_info.dimension].append(model_info)
+                
+                # Index by name
+                self._name_index[model_info.name.lower()] = model_info
+            
+            print(f"Loaded {len(self._models)} models from registry")
+            #...
+        except Exception as e:
+            print(f"Error loading model registry: {e}")
+    
+    def get_models_by_dimension(self, dimension: int) -> List[ModelInfo]:
+        """Get all models for a specific dimension.
+        
+        Args:
+            dimension: Vector dimension
+            
+        Returns:
+            List of ModelInfo objects
+        """
+        return self._dimension_index.get(dimension, [])
+    
+    def get_model_by_name(self, name: str) -> Optional[ModelInfo]:
+        """Get model info by name (case-insensitive).
+        
+        Args:
+            name: Model name
+            
+        Returns:
+            ModelInfo or None if not found
+        """
+        return self._name_index.get(name.lower())
+    
+    def get_all_models(self) -> List[ModelInfo]:
+        """Get all registered models.
+        
+        Returns:
+            List of all ModelInfo objects
+        """
+        return self._models.copy()
+    
+    def get_all_dimensions(self) -> List[int]:
+        """Get all available dimensions.
+        
+        Returns:
+            Sorted list of dimensions
+        """
+        return sorted(self._dimension_index.keys())
+    
+    def find_closest_dimension(self, target_dimension: int) -> Optional[int]:
+        """Find the closest available dimension.
+        
+        Args:
+            target_dimension: Target dimension to match
+            
+        Returns:
+            Closest dimension or None if no models exist
+        """
+        if not self._dimension_index:
+            return None
+        
+        return min(self._dimension_index.keys(), key=lambda x: abs(x - target_dimension))
+    
+    def get_models_by_type(self, model_type: str) -> List[ModelInfo]:
+        """Get all models of a specific type.
+        
+        Args:
+            model_type: Model type (e.g., "sentence-transformer", "clip")
+            
+        Returns:
+            List of ModelInfo objects
+        """
+        return [m for m in self._models if m.type == model_type]
+    
+    def get_models_by_source(self, source: str) -> List[ModelInfo]:
+        """Get all models from a specific source.
+        
+        Args:
+            source: Model source (e.g., "huggingface", "openai-api")
+            
+        Returns:
+            List of ModelInfo objects
+        """
+        return [m for m in self._models if m.source == source]
+    
+    def search_models(self, query: str) -> List[ModelInfo]:
+        """Search models by name or description.
+        
+        Args:
+            query: Search query (case-insensitive)
+            
+        Returns:
+            List of matching ModelInfo objects
+        """
+        query_lower = query.lower()
+        results = []
+        
+        for model in self._models:
+            if (query_lower in model.name.lower() or 
+                query_lower in model.description.lower()):
+                results.append(model)
+        
+        return results
+    
+    def reload(self):
+        """Reload the registry from disk."""
+        self._models.clear()
+        self._dimension_index.clear()
+        self._name_index.clear()
+        self._load_registry()
+
+
+# Global registry instance
+_registry = None
+
+
+def get_model_registry() -> EmbeddingModelRegistry:
+    """Get the global model registry instance.
+    
+    Returns:
+        EmbeddingModelRegistry singleton
+    """
+    global _registry
+    if _registry is None:
+        _registry = EmbeddingModelRegistry()
+    return _registry

vector_inspector/services/backup_restore_service.py

@@ -189,12 +189,28 @@ class BackupRestoreService:
                 return True
             else:
                 print("Failed to restore collection")
+                # Clean up partially created collection
+                try:
+                    if restore_collection_name in connection.list_collections():
+                        print(f"Cleaning up failed restore: deleting collection '{restore_collection_name}'")
+                        connection.delete_collection(restore_collection_name)
+                except Exception as cleanup_error:
+                    print(f"Warning: Failed to clean up collection: {cleanup_error}")
                 return False
                 
         except Exception as e:
             print(f"Restore failed: {e}")
             import traceback
             traceback.print_exc()
+            
+            # Clean up partially created collection
+            try:
+                if restore_collection_name in connection.list_collections():
+                    print(f"Cleaning up failed restore: deleting collection '{restore_collection_name}'")
+                    connection.delete_collection(restore_collection_name)
+            except Exception as cleanup_error:
+                print(f"Warning: Failed to clean up collection: {cleanup_error}")
+            
             return False
     
     @staticmethod