Vector Embeddings in Production

import os
import logging
from typing import List
import torch
from sentence_transformers import SentenceTransformer
import openai
 
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("EmbeddingGenerator")
 
class DualEmbeddingEngine:
    def __init__(self, openai_api_key: Optional[str] = None):
        # 1. Initialize local model with device acceleration detection
        if torch.cuda.is_available():
            self.device = "cuda"
        elif torch.backends.mps.is_available():
            self.device = "mps"
        else:
            self.device = "cpu"
        
        logger.info(f"Initializing local sentence-transformers on device: {self.device}")
        self.local_model = SentenceTransformer("BAAI/bge-small-en-v1.5", device=self.device)
 
        # 2. Initialize OpenAI client
        api_key = openai_api_key or os.getenv("OPENAI_API_KEY")
        if api_key:
            self.openai_client = openai.OpenAI(api_key=api_key)
        else:
            self.openai_client = None
            logger.warning("OpenAI API key not found. Hosted API model is disabled.")
 
    def generate_local(self, texts: List[str]) -> List[List[float]]:
        """Generates embeddings locally using BGE on CUDA/MPS/CPU."""
        logger.info(f"Generating {len(texts)} local embeddings...")
        embeddings = self.local_model.encode(texts, normalize_embeddings=True)
        return embeddings.tolist()
 
    def generate_api(self, texts: List[str]) -> List[List[float]]:
        """Generates embeddings using OpenAI's text-embedding-3-small."""
        if not self.openai_client:
            raise ValueError("OpenAI client is not initialized. Provide an API key.")
        
        logger.info(f"Generating {len(texts)} hosted API embeddings...")
        response = self.openai_client.embeddings.create(
            model="text-embedding-3-small",
            input=texts
        )
        # OpenAI returns embeddings in index order; extract and return
        return [data.embedding for data in response.data]
 
# Example Usage:
# engine = DualEmbeddingEngine()
# local_vectors = engine.generate_local(["AI Engineering is systems engineering."])

import time
import random
from typing import List, Generator
 
def batch_texts(texts: List[str], batch_size: int = 32) -> Generator[List[str], None, None]:
    """Yields consecutive batches of a specified size."""
    for i in range(0, len(texts), batch_size):
        yield texts[i:i + batch_size]
 
def generate_embeddings_with_backoff(
    client: openai.OpenAI, 
    texts: List[str], 
    max_retries: int = 5
) -> List[List[float]]:
    """Generates embeddings for a list of texts using exponential backoff with jitter."""
    all_embeddings = []
    
    for batch in batch_texts(texts, batch_size=64):
        retries = 0
        while True:
            try:
                response = client.embeddings.create(
                    model="text-embedding-3-small",
                    input=batch
                )
                all_embeddings.extend([d.embedding for d in response.data])
                break
            except openai.RateLimitError as e:
                if retries >= max_retries:
                    raise RuntimeError("Max retries exceeded on OpenAI rate limiting.") from e
                
                # Exponential backoff with jitter: 2^retries + random float
                sleep_time = (2 ** retries) + random.uniform(0.1, 1.0)
                logger.warning(f"Rate limited. Retrying in {sleep_time:.2f} seconds...")
                time.sleep(sleep_time)
                retries += 1
                
    return all_embeddings

import numpy as np
from sklearn.decomposition import PCA
 
def reduce_dimensions_pca(embeddings: List[List[float]], target_dim: int = 128) -> List[List[float]]:
    """
    Compresses high-dimensional embeddings to a target dimension using PCA.
    Note: PCA should be fit on a representative corpus first.
    """
    matrix = np.array(embeddings)
    n_samples, n_features = matrix.shape
    
    # PCA target components cannot exceed sample size
    actual_target = min(n_samples, target_dim)
    
    pca = PCA(n_components=actual_target)
    compressed_matrix = pca.fit_transform(matrix)
    
    # Re-normalize to unit length for Dot Product search compatibility
    norms = np.linalg.norm(compressed_matrix, axis=1, keepdims=True)
    normalized_matrix = compressed_matrix / np.maximum(norms, 1e-12)
    
    return normalized_matrix.tolist()

[Sentence 1] ➔ (Measure Similarity) ➔ [Sentence 2] ➔ (Similarity Drop) ➔ [New Chunk]

import re
from typing import List, Dict
import numpy as np
 
def split_into_sentences(text: str) -> List[str]:
    """Splits raw text into sentences using simple regex patterns."""
    # Split on punctuation followed by whitespace and a capital letter
    sentences = re.split(r'(?<=[\.\?\!])\s+(?=[A-Z])', text)
    return [s.strip() for s in sentences if s.strip()]
 
def cosine_similarity(v1: np.ndarray, v2: np.ndarray) -> float:
    """Computes the cosine similarity between two 1D arrays."""
    dot = np.dot(v1, v2)
    norm_v1 = np.linalg.norm(v1)
    norm_v2 = np.linalg.norm(v2)
    if norm_v1 == 0 or norm_v2 == 0:
        return 0.0
    return float(dot / (norm_v1 * norm_v2))
 
def create_semantic_chunks(
    text: str, 
    embedding_engine: DualEmbeddingEngine, 
    threshold_percentile: float = 60.0
) -> List[str]:
    # 1. Split text into sentences
    sentences = split_into_sentences(text)
    if len(sentences) < 2:
        return sentences
 
    # 2. Generate embeddings for all sentences
    embeddings = embedding_engine.generate_local(sentences)
    vectors = [np.array(emb) for emb in embeddings]
 
    # 3. Calculate distance shifts between consecutive sentences
    similarities = []
    for idx in range(len(vectors) - 1):
        sim = cosine_similarity(vectors[idx], vectors[idx + 1])
        similarities.append(sim)
 
    # Convert similarity to distance shifts (1 - similarity)
    distances = [1.0 - sim for sim in similarities]
    
    # 4. Calculate dynamic threshold based on percentiles
    threshold = np.percentile(distances, threshold_percentile)
 
    # 5. Partition sentences into chunks based on threshold crossings
    chunks = []
    current_chunk_sentences = [sentences[0]]
 
    for idx, dist in enumerate(distances):
        if dist > threshold:
            # Distance exceeds threshold; finalize chunk and start new one
            chunks.append(" ".join(current_chunk_sentences))
            current_chunk_sentences = [sentences[idx + 1]]
        else:
            current_chunk_sentences.append(sentences[idx + 1])
 
    # Add final chunk
    if current_chunk_sentences:
        chunks.append(" ".join(current_chunk_sentences))
 
    return chunks