Spaces:
Running
Running
| # embeddings/embedding_models.py | |
| """ | |
| Multiple embedding model implementations | |
| """ | |
| import numpy as np | |
| from typing import List, Optional | |
| import torch | |
| from sentence_transformers import SentenceTransformer | |
| class EmbeddingManager: | |
| """Manager for multiple embedding models""" | |
| def __init__(self, model_name: str = "all-MiniLM-L6-v2"): | |
| self.model_name = model_name | |
| self.model = None | |
| self._load_model() | |
| def _load_model(self): | |
| """Load the specified embedding model""" | |
| try: | |
| print(f"🔧 Loading embedding model: {self.model_name}") | |
| self.model = SentenceTransformer(self.model_name) | |
| print(f"✅ Model loaded successfully: {self.model_name}") | |
| except Exception as e: | |
| print(f"❌ Failed to load model {self.model_name}: {e}") | |
| # Fallback to default model | |
| self.model_name = "all-MiniLM-L6-v2" | |
| self.model = SentenceTransformer(self.model_name) | |
| print(f"🔄 Using fallback model: {self.model_name}") | |
| def encode(self, texts: List[str], batch_size: int = 32) -> np.ndarray: | |
| """Encode texts into embeddings""" | |
| if not self.model: | |
| raise ValueError("Embedding model not loaded") | |
| if isinstance(texts, str): | |
| texts = [texts] | |
| try: | |
| embeddings = self.model.encode( | |
| texts, | |
| batch_size=batch_size, | |
| show_progress_bar=False, | |
| convert_to_numpy=True | |
| ) | |
| return embeddings | |
| except Exception as e: | |
| print(f"❌ Embedding encoding error: {e}") | |
| raise | |
| def get_embedding_dimensions(self) -> int: | |
| """Get the dimensions of the embeddings""" | |
| # Test encoding to get dimensions | |
| test_embedding = self.encode(["test"]) | |
| return test_embedding.shape[1] | |
| def get_model_info(self) -> dict: | |
| """Get information about the current model""" | |
| return { | |
| "model_name": self.model_name, | |
| "dimensions": self.get_embedding_dimensions(), | |
| "max_sequence_length": getattr(self.model, 'max_seq_length', 512) | |
| } | |
| class MultiEmbeddingManager: | |
| """Manager that can switch between multiple embedding models""" | |
| def __init__(self): | |
| self.models = {} | |
| self.current_model = None | |
| def load_model(self, model_name: str) -> EmbeddingManager: | |
| """Load a specific embedding model""" | |
| if model_name not in self.models: | |
| self.models[model_name] = EmbeddingManager(model_name) | |
| self.current_model = self.models[model_name] | |
| return self.current_model | |
| def get_model(self, model_name: str = None) -> EmbeddingManager: | |
| """Get a model instance""" | |
| if model_name: | |
| return self.load_model(model_name) | |
| elif self.current_model: | |
| return self.current_model | |
| else: | |
| # Load default model | |
| return self.load_model("all-MiniLM-L6-v2") | |
| def list_loaded_models(self) -> List[str]: | |
| """List all currently loaded models""" | |
| return list(self.models.keys()) | |
| # Quick test function | |
| def test_embedding_models(): | |
| """Test all available embedding models""" | |
| from config.vector_config import EMBEDDING_MODELS | |
| multi_manager = MultiEmbeddingManager() | |
| test_texts = [ | |
| "Deep learning for medical image analysis", | |
| "Transformer architectures in genomics", | |
| "AI-driven drug discovery methods" | |
| ] | |
| print("🧪 Testing Embedding Models") | |
| print("=" * 50) | |
| for model_name in EMBEDDING_MODELS.keys(): | |
| try: | |
| print(f"\n🔬 Testing: {model_name}") | |
| manager = multi_manager.load_model(model_name) | |
| info = manager.get_model_info() | |
| print(f" Dimensions: {info['dimensions']}") | |
| print(f" Max sequence length: {info['max_sequence_length']}") | |
| # Test encoding | |
| embeddings = manager.encode(test_texts) | |
| print(f" Embedding shape: {embeddings.shape}") | |
| print(f" ✅ {model_name} working correctly") | |
| except Exception as e: | |
| print(f" ❌ {model_name} failed: {e}") | |
| if __name__ == "__main__": | |
| test_embedding_models() |