endee-llamaindex 0.1.2__py3-none-any.whl → 0.1.5__py3-none-any.whl

endee_llamaindex/constants.py

@@ -0,0 +1,70 @@
+"""
+Constants for EndeeVectorStore.
+
+This module contains all constants used by the Endee LlamaIndex integration.
+"""
+
+from llama_index.core.vector_stores.types import FilterOperator
+
+# Endee default constants
+# These may be overridden by importing from endee.constants if available
+MAX_VECTORS_PER_BATCH = 1000
+DEFAULT_EF_SEARCH = 128
+MAX_TOP_K_ALLOWED = 512
+MAX_EF_SEARCH_ALLOWED = 1024
+MAX_DIMENSION_ALLOWED = 10000
+MAX_INDEX_NAME_LENGTH_ALLOWED = 48
+
+# Try to import constants from endee package to stay in sync
+try:
+    from endee.constants import (
+        DEFAULT_EF_SEARCH as _DEFAULT_EF_SEARCH,
+        MAX_DIMENSION_ALLOWED as _MAX_DIMENSION_ALLOWED,
+        MAX_EF_SEARCH_ALLOWED as _MAX_EF_SEARCH_ALLOWED,
+        MAX_INDEX_NAME_LENGTH_ALLOWED as _MAX_INDEX_NAME_LENGTH_ALLOWED,
+        MAX_TOP_K_ALLOWED as _MAX_TOP_K_ALLOWED,
+        MAX_VECTORS_PER_BATCH as _MAX_VECTORS_PER_BATCH,
+    )
+    # Override defaults with values from endee package
+    DEFAULT_EF_SEARCH = _DEFAULT_EF_SEARCH
+    MAX_DIMENSION_ALLOWED = _MAX_DIMENSION_ALLOWED
+    MAX_EF_SEARCH_ALLOWED = _MAX_EF_SEARCH_ALLOWED
+    MAX_INDEX_NAME_LENGTH_ALLOWED = _MAX_INDEX_NAME_LENGTH_ALLOWED
+    MAX_TOP_K_ALLOWED = _MAX_TOP_K_ALLOWED
+    MAX_VECTORS_PER_BATCH = _MAX_VECTORS_PER_BATCH
+except ImportError:
+    pass
+
+# Space types and precision types for index creation
+SPACE_TYPES_VALID = ("cosine", "l2", "ip")
+PRECISION_VALID = ("binary", "float16", "float32", "int16d", "int8d")
+
+# Space type mapping (aliases)
+SPACE_TYPE_MAP = {
+    "cosine": "cosine",
+    "l2": "l2",
+    "ip": "ip",
+    "euclidean": "l2",
+    "inner_product": "ip",
+}
+
+# Vector store keys
+ID_KEY = "id"
+VECTOR_KEY = "values"
+SPARSE_VECTOR_KEY = "sparse_values"
+METADATA_KEY = "metadata"
+
+# Batch size for add(); capped by MAX_VECTORS_PER_BATCH
+DEFAULT_BATCH_SIZE = 100
+
+# Supported filter operations: currently only EQ and IN.
+# Map FilterOperator -> endee/backend filter symbol.
+SUPPORTED_FILTER_OPERATORS = (
+    FilterOperator.EQ,   # eq  -> $eq
+    FilterOperator.IN,   # in  -> $in
+)
+
+REVERSE_OPERATOR_MAP = {
+    FilterOperator.EQ: "$eq",
+    FilterOperator.IN: "$in",
+}

endee_llamaindex/utils.py

@@ -0,0 +1,160 @@
+import logging
+from typing import Any, Callable, List, Optional
+
+_logger = logging.getLogger(__name__)
+
+# Supported sparse embedding models
+SUPPORTED_SPARSE_MODELS = {
+    "splade_pp": "prithivida/Splade_PP_en_v1",
+    "splade_cocondenser": "naver/splade-cocondenser-ensembledistil",
+    "bert_base": "bert-base-uncased",
+    "distilbert": "distilbert-base-uncased",
+    "minilm": "sentence-transformers/all-MiniLM-L6-v2",
+    "mpnet": "sentence-transformers/all-mpnet-base-v2",
+    "roberta": "roberta-base",
+    "xlm_roberta": "xlm-roberta-base",
+}
+
+
+def _import_endee() -> Any:
+    """Import endee module."""
+    try:
+        import endee
+    except ImportError as e:
+        raise ImportError(
+            "Could not import endee python package. "
+            "Please install it with `pip install endee`."
+        ) from e
+    return endee
+
+
+def _initialize_sparse_encoder_fastembed(
+    model_name: str,
+    batch_size: int = 256,
+    cache_dir: Optional[str] = None,
+    threads: Optional[int] = None,
+) -> Callable:
+    """
+    Initialize a sparse encoder using FastEmbed (recommended for SPLADE models).
+    """
+    try:
+        from fastembed.sparse.sparse_text_embedding import SparseTextEmbedding
+    except ImportError as e:
+        raise ImportError(
+            "FastEmbed is required for hybrid search but not installed.\n"
+            "Install options:\n"
+            "  - CPU:  pip install endee-llamaindex[hybrid]\n"
+            "  - GPU:  pip install endee-llamaindex[hybrid-gpu]\n"
+            "  - Or:   pip install fastembed\n"
+            "For dense-only search, create vector store without sparse_dim/model_name."
+        ) from e
+
+    resolved_model_name = SUPPORTED_SPARSE_MODELS.get(model_name, model_name)
+    
+    try:
+        model = SparseTextEmbedding(
+            resolved_model_name,
+            cache_dir=cache_dir,
+            threads=threads,
+            providers=["CUDAExecutionProvider"],
+        )
+        _logger.info(f"Initialized sparse encoder '{resolved_model_name}' on GPU")
+    except Exception:
+        model = SparseTextEmbedding(
+            resolved_model_name, 
+            cache_dir=cache_dir, 
+            threads=threads
+        )
+        _logger.info(f"Initialized sparse encoder '{resolved_model_name}' on CPU")
+
+    def compute_vectors(texts: List[str]) -> tuple:
+        """Compute sparse vectors (indices, values) for a list of texts."""
+        embeddings = model.embed(texts, batch_size=batch_size)
+        indices = []
+        values = []
+        for embedding in embeddings:
+            indices.append(embedding.indices.tolist())
+            values.append(embedding.values.tolist())
+        return indices, values
+
+    return compute_vectors
+
+
+def _initialize_sparse_encoder_transformers(model_name: str) -> Callable:
+    """
+    Initialize a sparse encoder using Transformers library.
+    """
+    try:
+        import torch
+        from transformers import AutoModelForMaskedLM, AutoTokenizer
+    except ImportError as e:
+        raise ImportError(
+            "Could not import transformers library. "
+            'Please install transformers with `pip install "transformers[torch]"`'
+        ) from e
+
+    resolved_model_name = SUPPORTED_SPARSE_MODELS.get(model_name, model_name)
+    
+    tokenizer = AutoTokenizer.from_pretrained(resolved_model_name)
+    model = AutoModelForMaskedLM.from_pretrained(resolved_model_name)
+    
+    if torch.cuda.is_available():
+        model = model.to("cuda")
+        _logger.info(f"Initialized sparse encoder '{resolved_model_name}' on GPU")
+    else:
+        _logger.info(f"Initialized sparse encoder '{resolved_model_name}' on CPU")
+
+    def compute_vectors(texts: List[str]) -> tuple:
+        """Compute sparse vectors from logits."""
+        tokens = tokenizer(
+            texts, 
+            truncation=True, 
+            padding=True, 
+            max_length=512, 
+            return_tensors="pt"
+        )
+        
+        if torch.cuda.is_available():
+            tokens = tokens.to("cuda")
+
+        with torch.no_grad():
+            output = model(**tokens)
+            logits, attention_mask = output.logits, tokens.attention_mask
+            relu_log = torch.log(1 + torch.relu(logits))
+            weighted_log = relu_log * attention_mask.unsqueeze(-1)
+            tvecs, _ = torch.max(weighted_log, dim=1)
+
+        indices = []
+        values = []
+        for batch in tvecs:
+            nz_indices = batch.nonzero(as_tuple=True)[0].tolist()
+            indices.append(nz_indices)
+            values.append(batch[nz_indices].tolist())
+
+        return indices, values
+
+    return compute_vectors
+
+
+def get_sparse_encoder(
+    model_name: Optional[str] = None,
+    use_fastembed: bool = True,
+    batch_size: int = 256,
+    cache_dir: Optional[str] = None,
+    threads: Optional[int] = None,
+) -> Optional[Callable]:
+    """
+    Get a sparse encoder function for the specified model.
+    """
+    if model_name is None:
+        return None
+    
+    if use_fastembed:
+        return _initialize_sparse_encoder_fastembed(
+            model_name=model_name,
+            batch_size=batch_size,
+            cache_dir=cache_dir,
+            threads=threads,
+        )
+    else:
+        return _initialize_sparse_encoder_transformers(model_name=model_name)