ai-parrot 0.17.2__cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

parrot/embeddings/huggingface.py

@@ -0,0 +1,74 @@
+from typing import List, Any
+from enum import Enum
+import numpy as np
+from .base import EmbeddingModel
+from ..conf import HUGGINGFACE_EMBEDDING_CACHE_DIR
+
+
+class ModelType(Enum):
+    """Enumerator for different model types used in embeddings."""
+    MPNET = "sentence-transformers/all-mpnet-base-v2"
+    MINILM = "sentence-transformers/all-MiniLM-L6-v2"
+    BGE_LARGE = "BAAI/bge-large-en-v1.5"
+    GTE_BASE = "thenlper/gte-base"
+
+
+class SentenceTransformerModel(EmbeddingModel):
+    """
+    A wrapper class for HuggingFace sentence-transformers embeddings.
+    """
+    model_name: str = "sentence-transformers/all-mpnet-base-v2"
+
+    def __init__(self, model_name: str, **kwargs):
+        """
+        Initializes the embedding model with the specified model name.
+
+        Args:
+            model_name: The name of the Hugging Face model to load.
+            **kwargs: Additional keyword arguments for SentenceTransformer.
+        """
+        super().__init__(model_name=model_name, **kwargs)
+        self.logger.info(
+            f"Initialized SentenceTransformerModel with model: {self.model_name}"
+        )
+
+    def _create_embedding(self, model_name: str = None, **kwargs) -> Any:
+        """
+        Creates and returns the SentenceTransformer model instance.
+
+        Args:
+            model_name: The name of the Hugging Face model to load.
+
+        Returns:
+            An instance of SentenceTransformer.
+        """
+        from sentence_transformers import SentenceTransformer
+        import torch
+        
+        # Access self.device via property (calls _get_device lazily)
+        device = self.device
+        model_name = model_name or self.model_name
+        
+        self.logger.info(
+            f"Loading embedding model '{model_name}' on device '{device}'"
+        )
+        
+        model = SentenceTransformer(
+            model_name,
+            device=device,
+            cache_folder=HUGGINGFACE_EMBEDDING_CACHE_DIR
+        )
+        
+        # Set dimension after loading model
+        self._dimension = model.get_sentence_embedding_dimension()
+        
+        # Production optimizations
+        model.eval()
+        if str(device) == "cuda":
+            model.half()  # Use FP16 for GPU inference
+            torch.backends.cudnn.benchmark = True
+
+        return model
+
+    async def encode(self, texts: List[str], **kwargs) -> np.ndarray:
+        return self.model.encode(texts, **kwargs)

parrot/embeddings/openai.py

@@ -0,0 +1,84 @@
+from typing import Any, List, Union, Optional
+import asyncio
+import numpy as np
+from navconfig import config
+from openai import AsyncOpenAI
+from .base import EmbeddingModel
+
+class OpenAIEmbeddingModel(EmbeddingModel):
+    """A wrapper class for OpenAI Embedding models.
+    """
+    model_name: str = "text-embedding-3-large"
+
+    def __init__(self, model_name: str = None, dimensions: int = None, **kwargs):
+        self.api_key = kwargs.pop('api_key', config.get('OPENAI_API_KEY'))
+        if model_name:
+            self.model_name = model_name
+        self.dimensions = dimensions
+        self._dimension = dimensions
+        super().__init__(model_name=self.model_name, **kwargs)
+
+    def _create_embedding(self, model_name: str = None, **kwargs) -> Any:
+        """
+        Creates and returns an OpenAI client instance.
+
+        Args:
+            model_name: The name of the OpenAI model to load.
+
+        Returns:
+            An instance of AsyncOpenAI client.
+        """
+        if model_name:
+            self.model_name = model_name
+        
+        self.logger.info(
+            f"Loading embedding model '{self.model_name}'"
+        )
+        # Using AsyncOpenAI as consistent with other parts of the system
+        return AsyncOpenAI(api_key=self.api_key)
+
+    async def encode(self, texts: List[str], **kwargs) -> List[List[float]]:
+        """Generate embeddings for a list of texts using AsyncOpenAI.
+
+        This mirrors the official OpenAI embeddings guide, allowing callers to
+        override the model or dimensions per request while keeping asynchronous
+        semantics.
+        """
+        model = kwargs.pop("model", self.model_name)
+        dimensions = kwargs.pop("dimensions", self.dimensions)
+        call_kwargs = {
+            "model": model,
+            "input": texts,
+            "encoding_format": "float",
+            **kwargs,
+        }
+        if dimensions:
+            call_kwargs["dimensions"] = dimensions
+
+        result = await self.model.embeddings.create(**call_kwargs)
+        embeddings = [data.embedding for data in result.data]
+        if embeddings and not self._dimension:
+            # Keep a record of the dimension reported by the service when not set
+            self._dimension = len(embeddings[0])
+        return embeddings
+
+    async def embed_query(
+        self,
+        text: str,
+        as_nparray: bool = False
+    ) -> Union[List[float], List[np.ndarray]]:
+        """
+        Generates an embedding for a single query string asynchronously.
+        """
+        embeddings = await self.encode([text])
+        embedding = embeddings[0]
+        
+        if as_nparray:
+            return [np.array(embedding)]
+        return embedding
+
+    async def embed_documents(self, texts: List[str], batch_size: Optional[int] = None) -> List[List[float]]:
+        """
+        Generates embeddings for a list of documents asynchronously.
+        """
+        return await self.encode(texts)

parrot/embeddings/processor.py

@@ -0,0 +1,88 @@
+from typing import List, Dict, Any, Tuple
+import numpy as np
+from .huggingface import SentenceTransformerModel
+
+
+class LateChunkingProcessor:
+    """Processor for handling late chunking of documents using embeddings.
+    This class processes documents by generating embeddings for the entire document first,
+    then splitting it into semantic chunks while preserving boundaries.
+    It uses the SentenceTransformerModel to generate embeddings and chunk metadata.
+    """
+    def __init__(self, embedding_generator: SentenceTransformerModel, chunk_size: int = 8192):
+        self.embedding_generator = embedding_generator
+        self.chunk_size = chunk_size
+
+    async def process_document_late_chunking(
+        self,
+        document_text: str,
+        document_id: str
+    ) -> Tuple[List[np.ndarray], List[Dict[str, Any]]]:
+        """Process document with late chunking strategy"""
+
+        # Step 1: Generate full-document token embeddings
+        full_embeddings = await self.embedding_generator.encode([document_text])
+
+        # Step 2: Split into semantic chunks while preserving boundaries
+        chunks = self._semantic_chunk_split(document_text)
+
+        # Step 3: Map chunks to their corresponding token ranges
+        chunk_embeddings = []
+        chunk_metadata = []
+
+        for chunk_idx, chunk_text in enumerate(chunks):
+            # Calculate token range for this chunk
+            start_pos = document_text.find(chunk_text)
+            end_pos = start_pos + len(chunk_text)
+
+            # Extract embeddings for this chunk's token range
+            chunk_embedding = await self._extract_chunk_embedding(
+                full_embeddings[0], chunk_text, document_text
+            )
+
+            chunk_embeddings.append(chunk_embedding)
+            chunk_metadata.append({
+                'document_id': document_id,
+                'chunk_index': chunk_idx,
+                'start_position': start_pos,
+                'end_position': end_pos,
+                'chunk_text': chunk_text[:200] + '...' if len(chunk_text) > 200 else chunk_text
+            })
+
+        return chunk_embeddings, chunk_metadata
+
+    def _semantic_chunk_split(self, text: str) -> List[str]:
+        """Split text preserving sentence and paragraph boundaries"""
+        # Split by double newlines (paragraphs) first
+        paragraphs = text.split('\n\n')
+        chunks = []
+        current_chunk = []
+        current_size = 0
+
+        for paragraph in paragraphs:
+            paragraph_size = len(paragraph)
+
+            if current_size + paragraph_size > self.chunk_size and current_chunk:
+                chunks.append('\n\n'.join(current_chunk))
+                current_chunk = [paragraph]
+                current_size = paragraph_size
+            else:
+                current_chunk.append(paragraph)
+                current_size += paragraph_size
+
+        if current_chunk:
+            chunks.append('\n\n'.join(current_chunk))
+
+        return chunks
+
+    async def _extract_chunk_embedding(
+        self,
+        full_embedding: np.ndarray,
+        chunk_text: str,
+        full_text: str
+    ) -> np.ndarray:
+        """Extract contextualized embedding for chunk"""
+        # For now, re-embed the chunk with full context
+        # In practice, you might use token-level alignments
+        chunk_embedding = await self.embedding_generator.encode([chunk_text])
+        return chunk_embedding[0]