docsif 0.1.0__py3-none-any.whl

sif/embedding/embedder.py

@@ -0,0 +1,328 @@
+"""Embedding model implementation."""
+
+from __future__ import annotations
+
+import hashlib
+from pathlib import Path
+
+import numpy as np
+
+from sif.core.models import Embedder
+from sif.models.download import ModelDownloader
+from sif.utils.logging import get_logger, is_quiet, suppress_output
+
+
+logger = get_logger(__name__)
+
+
+class SentenceTransformerEmbedder(Embedder):
+    """Embedder using sentence-transformers."""
+
+    def __init__(
+        self,
+        model_name: str = "all-MiniLM-L6-v2",
+        device: str | None = None,
+        cache_dir: str | None = None,
+    ) -> None:
+        """Initialize sentence transformer embedder.
+
+        Args:
+            model_name: HuggingFace model name
+            device: Device to use (cpu, cuda, etc.)
+            cache_dir: Cache directory for models
+        """
+        try:
+            from sentence_transformers import SentenceTransformer  # noqa: PLC0415
+        except ImportError:
+            logger.error(
+                "sentence-transformers not installed. "
+                "Install with: pip install sentence-transformers"
+            )
+            raise
+
+        self.model_name = model_name
+
+        if device is None:
+            import torch  # noqa: PLC0415
+
+            if torch.cuda.is_available():
+                device = "cuda"
+            elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+                device = "mps"
+            else:
+                device = "cpu"
+
+        logger.info(f"Loading embedding model: {model_name} on {device}")
+
+        if is_quiet():
+            with suppress_output():
+                self.model = SentenceTransformer(
+                    model_name,
+                    device=device,
+                    cache_folder=cache_dir,
+                )
+        else:
+            self.model = SentenceTransformer(
+                model_name,
+                device=device,
+                cache_folder=cache_dir,
+            )
+
+        self._dimension = self.model.get_sentence_embedding_dimension()
+        logger.info(f"Embedding dimension: {self._dimension}")
+
+    def embed(self, text: str) -> list[float]:
+        """Embed a single text.
+
+        Args:
+            text: Text to embed
+
+        Returns:
+            Embedding vector
+        """
+        embedding = self.model.encode(text, normalize_embeddings=True)
+        return embedding.tolist()
+
+    def embed_batch(self, texts: list[str]) -> list[list[float]]:
+        """Embed multiple texts.
+
+        Args:
+            texts: List of texts to embed
+
+        Returns:
+            List of embedding vectors
+        """
+        embeddings = self.model.encode(
+            texts,
+            normalize_embeddings=True,
+            show_progress_bar=len(texts) > 100,
+        )
+        return embeddings.tolist()
+
+    @property
+    def dimension(self) -> int:
+        """Get embedding dimension."""
+        return self._dimension
+
+
+class LlamaCppEmbedder(Embedder):
+    """Embedder using llama-cpp-python."""
+
+    def __init__(
+        self,
+        model_path: str | Path,
+        n_ctx: int = 8192,
+        n_threads: int | None = None,
+        verbose: bool = False,
+    ) -> None:
+        """Initialize llama.cpp embedder.
+
+        Args:
+            model_path: Path to GGUF model file
+            n_ctx: Context size
+            n_threads: Number of threads (auto if None)
+            verbose: Enable verbose output
+        """
+        try:
+            from llama_cpp import Llama  # noqa: PLC0415
+        except ImportError:
+            logger.error(
+                "llama-cpp-python not installed. Install with: pip install llama-cpp-python"
+            )
+            raise
+
+        self.model_path = Path(model_path)
+
+        if n_threads is None:
+            import os  # noqa: PLC0415
+
+            n_threads = os.cpu_count() or 4
+
+        logger.info(f"Loading llama.cpp model: {self.model_path}")
+
+        self.model = Llama(
+            model_path=str(self.model_path),
+            n_ctx=n_ctx,
+            n_threads=n_threads,
+            embedding=True,
+            verbose=verbose,
+        )
+
+        self._dimension = self.model.n_embd()
+        logger.info(f"Embedding dimension: {self._dimension}")
+
+    def embed(self, text: str) -> list[float]:
+        """Embed a single text."""
+        embedding = self.model.embed(text)
+        # Normalize
+        norm = np.linalg.norm(embedding)
+        if norm > 0:
+            embedding = embedding / norm
+        return embedding.tolist()
+
+    def embed_batch(self, texts: list[str]) -> list[list[float]]:
+        """Embed multiple texts."""
+        embeddings = []
+        for text in texts:
+            embedding = self.embed(text)
+            embeddings.append(embedding)
+        return embeddings
+
+    @property
+    def dimension(self) -> int:
+        """Get embedding dimension."""
+        return self._dimension
+
+
+class ModelScopeEmbedder(Embedder):
+    """Embedder that downloads from ModelScope and uses sentence-transformers."""
+
+    def __init__(
+        self,
+        model_id: str = "iic/gte_Qwen2-7B-instruct",
+        device: str | None = None,
+        cache_dir: str | None = None,
+        force_download: bool = False,
+    ) -> None:
+        """Initialize ModelScope embedder.
+
+        Args:
+            model_id: ModelScope model ID
+            device: Device to use
+            cache_dir: Cache directory
+            force_download: Force re-download
+        """
+        # Download model from ModelScope
+        downloader = ModelDownloader(cache_dir)
+        model_path = downloader.download(model_id, force=force_download)
+
+        # Find the actual model directory (snapshot_download creates subdirectories)
+        model_dirs = list(model_path.glob("*"))
+        if model_dirs:
+            # Use the first subdirectory if exists
+            actual_model_path = model_dirs[0] if model_dirs[0].is_dir() else model_path
+        else:
+            actual_model_path = model_path
+
+        # Initialize with sentence-transformers
+        try:
+            from sentence_transformers import SentenceTransformer  # noqa: PLC0415
+        except ImportError:
+            logger.error("sentence-transformers not installed")
+            raise
+
+        if device is None:
+            import torch  # noqa: PLC0415
+
+            if torch.cuda.is_available():
+                device = "cuda"
+            elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+                device = "mps"
+            else:
+                device = "cpu"
+
+        logger.info(f"Loading ModelScope model: {model_id} on {device}")
+
+        if is_quiet():
+            with suppress_output():
+                self.model = SentenceTransformer(
+                    str(actual_model_path),
+                    device=device,
+                )
+        else:
+            self.model = SentenceTransformer(
+                str(actual_model_path),
+                device=device,
+            )
+
+        self._dimension = self.model.get_sentence_embedding_dimension()
+        self.model_id = model_id
+        logger.info(f"Embedding dimension: {self._dimension}")
+
+    def embed(self, text: str) -> list[float]:
+        """Embed a single text."""
+        embedding = self.model.encode(text, normalize_embeddings=True)
+        return embedding.tolist()
+
+    def embed_batch(self, texts: list[str]) -> list[list[float]]:
+        """Embed multiple texts."""
+        embeddings = self.model.encode(
+            texts,
+            normalize_embeddings=True,
+            show_progress_bar=len(texts) > 100,
+        )
+        return embeddings.tolist()
+
+    @property
+    def dimension(self) -> int:
+        """Get embedding dimension."""
+        return self._dimension
+
+
+class SimpleEmbedder(Embedder):
+    """Simple embedder using TF-IDF (fallback when no ML models available)."""
+
+    def __init__(self, dimension: int = 384) -> None:
+        """Initialize simple TF-IDF embedder.
+
+        Args:
+            dimension: Embedding dimension
+        """
+        self._dimension = dimension
+        self.vocabulary: dict = {}
+        self._doc_count = 0
+
+    def embed(self, text: str) -> list[float]:
+        """Embed text using simple hashing."""
+        # Simple hash-based embedding
+
+        # Create a deterministic embedding based on text hash
+        hash_bytes = hashlib.sha256(text.encode()).digest()
+
+        # Expand to desired dimension
+        embedding = []
+        for i in range(self._dimension):
+            # Use different parts of hash
+            idx = i % len(hash_bytes)
+            val = (hash_bytes[idx] / 255.0) * 2 - 1  # Scale to [-1, 1]
+            embedding.append(val)
+
+        # Normalize
+        norm = np.linalg.norm(embedding)
+        if norm > 0:
+            embedding = [v / norm for v in embedding]
+
+        return embedding
+
+    def embed_batch(self, texts: list[str]) -> list[list[float]]:
+        """Embed multiple texts."""
+        return [self.embed(text) for text in texts]
+
+    @property
+    def dimension(self) -> int:
+        """Get embedding dimension."""
+        return self._dimension
+
+
+def create_embedder(
+    embedder_type: str = "sentence_transformer",
+    **kwargs,
+) -> Embedder:
+    """Factory function to create embedders.
+
+    Args:
+        embedder_type: Type of embedder to create
+        **kwargs: Additional arguments for the embedder
+
+    Returns:
+        Embedder instance
+    """
+    if embedder_type == "sentence_transformer":
+        return SentenceTransformerEmbedder(**kwargs)
+    if embedder_type == "llama_cpp":
+        return LlamaCppEmbedder(**kwargs)
+    if embedder_type == "modelscope":
+        return ModelScopeEmbedder(**kwargs)
+    if embedder_type == "simple":
+        return SimpleEmbedder(**kwargs)
+    raise ValueError(f"Unknown embedder type: {embedder_type}")

sif/embedding/factory.py

@@ -0,0 +1,92 @@
+"""Embedding model factory."""
+
+from sif.embedding.embedder import (
+    LlamaCppEmbedder,
+    ModelScopeEmbedder,
+    SentenceTransformerEmbedder,
+)
+from sif.embedding.model import EmbeddingModel, ModelType
+from sif.utils.logging import get_logger
+
+
+logger = get_logger(__name__)
+
+
+class EmbeddingModelFactory:
+    """Factory for creating embedding model instances."""
+
+    def create_model(
+        self,
+        model_type: ModelType,
+        model_path: str | None,
+        model_name: str,
+        **kwargs: dict[str, any],
+    ) -> EmbeddingModel:
+        """Create an embedding model instance."""
+        if model_type == ModelType.SENTENCE_TRANSFORMERS:
+            return self._create_sentence_transformers_model(model_name, **kwargs)
+        if model_type == ModelType.GGUF:
+            return self._create_gguf_model(model_path, **kwargs)
+        if model_type == ModelType.OPENAI:
+            return self._create_openai_model(model_name, **kwargs)
+        if model_type == ModelType.HUGGINGFACE:
+            return self._create_huggingface_model(model_name, **kwargs)
+        if model_type == ModelType.MODELSCOPE:
+            return self._create_modelscope_model(model_name, **kwargs)
+        raise ValueError(f"Unsupported model type: {model_type}")
+
+    def _create_gguf_model(
+        self,
+        model_path: str | None,
+        **kwargs: dict[str, any],
+    ) -> EmbeddingModel:
+        """Create a GGUF model using llama-cpp-python."""
+        if not model_path:
+            raise ValueError("model_path is required for GGUF models")
+
+        return LlamaCppEmbedder(
+            model_path=model_path,
+            n_ctx=kwargs.get("n_ctx", 2048),
+            n_threads=kwargs.get("n_threads"),
+            verbose=kwargs.get("verbose", False),
+        )
+
+    def _create_sentence_transformers_model(
+        self,
+        model_name: str,
+        **kwargs: dict[str, any],
+    ) -> EmbeddingModel:
+        """Create a Sentence Transformers model."""
+        return SentenceTransformerEmbedder(
+            model_name=model_name,
+            device=kwargs.get("device"),
+            cache_dir=kwargs.get("cache_dir"),
+        )
+
+    def _create_openai_model(
+        self,
+        model_name: str,
+        **kwargs: dict[str, any],
+    ) -> EmbeddingModel:
+        """Create an OpenAI API model."""
+        raise NotImplementedError("OpenAI models not yet implemented")
+
+    def _create_huggingface_model(
+        self,
+        model_name: str,
+        **kwargs: dict[str, any],
+    ) -> EmbeddingModel:
+        """Create a HuggingFace Transformers model."""
+        raise NotImplementedError("HuggingFace models not yet implemented")
+
+    def _create_modelscope_model(
+        self,
+        model_name: str,
+        **kwargs: dict[str, any],
+    ) -> EmbeddingModel:
+        """Create a ModelScope embedder."""
+        return ModelScopeEmbedder(
+            model_id=model_name,
+            device=kwargs.get("device"),
+            cache_dir=kwargs.get("cache_dir"),
+        )

sif/embedding/manager.py

@@ -0,0 +1,200 @@
+"""Embedding manager for model lifecycle and embedding generation."""
+
+import time
+
+from sif.config.settings import Settings
+from sif.core.models import Embedder
+from sif.embedding.cache import EmbeddingCache
+from sif.embedding.factory import EmbeddingModelFactory
+from sif.models.embedding import EmbeddingConfig, EmbeddingResponse, ModelType
+from sif.utils.logging import get_logger
+
+
+logger = get_logger(__name__)
+
+
+class EmbeddingManager:
+    """Manager for embedding model lifecycle and operations.
+
+    This class provides a high-level interface for:
+    - Loading and unloading embedding models
+    - Generating embeddings for text
+    - Caching embeddings for reuse
+    - Managing model configuration
+    """
+
+    def __init__(
+        self,
+        config: EmbeddingConfig | None = None,
+        factory: EmbeddingModelFactory | None = None,
+        cache: EmbeddingCache | None = None,
+    ) -> None:
+        """Initialize embedding manager.
+
+        Args:
+            config: Embedding configuration
+            factory: Model factory (uses default if not provided)
+            cache: Embedding cache (creates default if not provided)
+        """
+        self._config = config or EmbeddingConfig()
+        self._factory = factory or EmbeddingModelFactory()
+        self._cache = cache
+        self._model: Embedder | None = None
+
+    @classmethod
+    def from_settings(cls, settings: Settings) -> "EmbeddingManager":
+        """Create an embedding manager from settings.
+
+        Args:
+            settings: Application settings
+
+        Returns:
+            Configured embedding manager
+        """
+        config = EmbeddingConfig(
+            model_type=ModelType(settings.model_type),
+            model_name=settings.model_name,
+            model_path=str(settings.model_path) if settings.model_path else None,
+            embedding_dim=settings.embedding_dim,
+            max_tokens=settings.max_tokens,
+            batch_size=settings.batch_size,
+            n_gpu_layers=settings.n_gpu_layers,
+            api_key=settings.api_key,
+            api_base=settings.api_base,
+            cache_embeddings=settings.cache_embeddings,
+            cache_dir=str(settings.get_cache_dir()) if settings.cache_embeddings else None,
+        )
+
+        cache = None
+        if config.cache_embeddings and config.cache_dir:
+            cache = EmbeddingCache(config.cache_dir)
+
+        return cls(config=config, cache=cache)
+
+    def load_model(self) -> None:
+        """Load the embedding model."""
+        if self._model is not None:
+            return
+
+        logger.info(f"Loading embedding model: {self._config.model_name}")
+
+        self._model = self._factory.create_model(
+            model_type=self._config.model_type,
+            model_path=self._config.model_path,
+            model_name=self._config.model_name,
+            embedding_dim=self._config.embedding_dim,
+            max_tokens=self._config.max_tokens,
+            n_gpu_layers=self._config.n_gpu_layers,
+            n_ctx=self._config.n_ctx,
+            api_key=self._config.api_key,
+            api_base=self._config.api_base,
+            cache_dir=self._config.cache_dir,
+        )
+
+        logger.info("Embedding model loaded successfully")
+
+    def unload_model(self) -> None:
+        """Unload the embedding model to free memory."""
+        if self._model is not None:
+            self._model = None
+            logger.info("Embedding model unloaded")
+
+    def embed(
+        self,
+        texts: list[str],
+        normalize: bool = True,  # noqa: ARG002
+        use_cache: bool = True,
+    ) -> EmbeddingResponse:
+        """Generate embeddings for texts.
+
+        Args:
+            texts: Texts to embed
+            normalize: Whether to normalize embeddings
+            use_cache: Whether to use caching
+
+        Returns:
+            Embedding response with embeddings and metadata
+        """
+        self.load_model()
+
+        start_time = time.time()
+
+        # Check cache for cached embeddings
+        embeddings: list[list[float] | None] = [None] * len(texts)
+        texts_to_embed: list[tuple[int, str]] = []
+
+        if use_cache and self._cache:
+            for i, text in enumerate(texts):
+                cached = self._cache.get(text)
+                if cached is not None:
+                    embeddings[i] = cached
+                else:
+                    texts_to_embed.append((i, text))
+        else:
+            texts_to_embed = list(enumerate(texts))
+
+        # Generate embeddings for uncached texts
+        if texts_to_embed and self._model:
+            indices, to_embed = zip(*texts_to_embed)
+            new_embeddings = self._model.embed_batch(list(to_embed))
+
+            # Store in cache
+            if use_cache and self._cache:
+                for _idx, text, emb in zip(indices, to_embed, new_embeddings):
+                    self._cache.set(text, emb)
+
+            # Fill in results
+            for idx, emb in zip(indices, new_embeddings):
+                embeddings[idx] = emb
+
+        processing_time = (time.time() - start_time) * 1000
+
+        # Count tokens (rough estimate)
+        total_tokens = sum(len(t) // 4 for t in texts)
+
+        actual_dim = self._model.dimension if self._model else self._config.embedding_dim
+        return EmbeddingResponse(
+            embeddings=[e for e in embeddings if e is not None],
+            model_id=self._config.model_name,
+            dimensions=actual_dim,
+            total_tokens=total_tokens,
+            processing_time_ms=processing_time,
+        )
+
+    def embed_single(
+        self,
+        text: str,
+        normalize: bool = True,
+        use_cache: bool = True,
+    ) -> list[float]:
+        """Generate embedding for a single text.
+
+        Args:
+            text: Text to embed
+            normalize: Whether to normalize the embedding
+            use_cache: Whether to use caching
+
+        Returns:
+            Embedding vector
+        """
+        response = self.embed([text], normalize=normalize, use_cache=use_cache)
+        return response.embeddings[0]
+
+    def get_model_info(self) -> dict:
+        """Get information about the loaded model.
+
+        Returns:
+            Model information dictionary
+        """
+        if not self._model:
+            return {
+                "loaded": False,
+                "model_name": self._config.model_name,
+            }
+
+        return {
+            "loaded": True,
+            "model_id": self._config.model_name,
+            "embedding_dim": self._model.dimension,
+            "max_tokens": self._config.max_tokens,
+        }