02 - Vector Embeddings Explained: Theory and Practice

# The embedding of the sentence "The cat sleeps on the couch"
# generated with OpenAI text-embedding-3-small (1536 dimensions)
embedding = [
    0.0231, -0.0456, 0.0891, -0.0123, 0.0567, -0.0234,
    0.0789, -0.0345, 0.0123, -0.0678, 0.0456, -0.0891,
    # ... 1524 more values ...
]
print(f"Type: {type(embedding)}")      # <class 'list'>
print(f"Dimensions: {len(embedding)}") # 1536
    

# Vocabulary: ["cat", "dog", "fish", "car", "bike"]
# Vector size = vocabulary size = 5

cat  = [1, 0, 0, 0, 0]
dog  = [0, 1, 0, 0, 0]
fish = [0, 0, 1, 0, 0]
car  = [0, 0, 0, 1, 0]
bike = [0, 0, 0, 0, 1]

# Problem 1: the distance between "cat" and "dog" equals
# the distance between "cat" and "car"
import numpy as np
dist_cat_dog = np.linalg.norm(
    np.array(cat) - np.array(dog)
)  # sqrt(2) = 1.414
dist_cat_car = np.linalg.norm(
    np.array(cat) - np.array(car)
)  # sqrt(2) = 1.414 -- identical!

# Problem 2: with a vocabulary of 100,000 words,
# each vector has 100,000 dimensions (sparse, inefficient)
    

from sklearn.feature_extraction.text import TfidfVectorizer

documents = [
    "the cat sleeps on the couch",
    "the dog plays in the garden",
    "the automobile drives on the road",
    "the feline rests on the armchair",
]

vectorizer = TfidfVectorizer()
tfidf_matrix = vectorizer.fit_transform(documents)

# Result: sparse matrix (4 documents x N terms)
print(f"Shape: {tfidf_matrix.shape}")  # (4, 14)
print(f"Terms: {vectorizer.get_feature_names_out()}")

# Problem: "cat sleeps" and "feline rests" are distant
# because they use different words, even though meaning is similar
from sklearn.metrics.pairwise import cosine_similarity
sim = cosine_similarity(tfidf_matrix[0:1], tfidf_matrix[3:4])
print(f"Similarity cat-feline: {sim[0][0]:.3f}")  # ~0.07 (low!)
    

from gensim.models import Word2Vec

# Sample corpus (in production: millions of sentences)
sentences = [
    ["the", "cat", "sleeps", "on", "the", "couch"],
    ["the", "dog", "plays", "in", "the", "garden"],
    ["the", "feline", "rests", "on", "the", "armchair"],
    ["the", "dog", "runs", "in", "the", "park"],
]

# Train Word2Vec (Skip-gram)
model = Word2Vec(
    sentences=sentences,
    vector_size=100,    # embedding dimensionality
    window=5,           # context: 5 words before and after
    min_count=1,        # include words with at least 1 occurrence
    sg=1,               # 1 = Skip-gram, 0 = CBOW
    epochs=100
)

# Now "cat" and "feline" are close!
print(model.wv.most_similar("cat", topn=3))
# [('feline', 0.92), ('dog', 0.85), ('sleeps', 0.71)]

# Access to the vector
vector_cat = model.wv["cat"]
print(f"Dimensions: {vector_cat.shape}")  # (100,)
print(f"First 5: {vector_cat[:5]}")
    

import gensim.downloader as api

# Load pre-trained GloVe embeddings
model = api.load("glove-wiki-gigaword-100")

# king - man + woman = ?
result = model.most_similar(
    positive=["king", "woman"],
    negative=["man"],
    topn=3
)
print(result)
# [('queen', 0.7698), ('princess', 0.6450), ('monarch', 0.6345)]

# Other analogies that work:
# Paris - France + Italy = Rome
result2 = model.most_similar(
    positive=["paris", "italy"],
    negative=["france"],
    topn=1
)
print(result2)  # [('rome', 0.8722)]

# good - bad + sad = ?
result3 = model.most_similar(
    positive=["good", "sad"],
    negative=["bad"],
    topn=1
)
print(result3)  # [('happy', 0.6891)]
    

import numpy as np
from scipy.spatial.distance import cosine, euclidean, cityblock

# Two sample vectors (normalized)
a = np.array([0.5, 0.3, 0.8, 0.1, 0.6])
b = np.array([0.4, 0.35, 0.75, 0.15, 0.55])

# L2 normalization
a_norm = a / np.linalg.norm(a)
b_norm = b / np.linalg.norm(b)

# 1. Cosine Similarity (1 - cosine distance)
cos_sim = 1 - cosine(a_norm, b_norm)
print(f"Cosine similarity:   {cos_sim:.6f}")   # ~0.999

# 2. Euclidean Distance (L2)
l2_dist = euclidean(a_norm, b_norm)
print(f"L2 distance:         {l2_dist:.6f}")   # ~0.042

# 3. Dot Product (for normalized vectors = cosine similarity)
dot = np.dot(a_norm, b_norm)
print(f"Dot product:         {dot:.6f}")       # ~0.999

# 4. Manhattan Distance (L1)
l1_dist = cityblock(a_norm, b_norm)
print(f"L1 distance:         {l1_dist:.6f}")   # ~0.072

# L2-Cosine relation for normalized vectors:
# d_L2^2 = 2 * (1 - cos_sim)
print(f"\nVerification: L2^2 = {l2_dist**2:.6f}")
print(f"2*(1-cos) = {2*(1-cos_sim):.6f}")     # identical!
    

# pip install sentence-transformers

from sentence_transformers import SentenceTransformer
import numpy as np

# Load the model (downloaded automatically on first use)
model = SentenceTransformer("all-MiniLM-L6-v2")

# Single sentence embedding
sentence = "PostgreSQL is an open-source relational database"
embedding = model.encode(sentence)
print(f"Type: {type(embedding)}")          # numpy.ndarray
print(f"Dimensions: {embedding.shape}")     # (384,)

# Batch embedding (much more efficient)
sentences = [
    "PostgreSQL is an open-source relational database",
    "pgvector adds vector support to PostgreSQL",
    "Machine learning requires large amounts of data",
    "Pizza margherita is a classic Italian dish",
]

embeddings = model.encode(
    sentences,
    batch_size=32,            # process 32 sentences at a time
    show_progress_bar=True,   # show progress for large batches
    normalize_embeddings=True # normalize to L2 norm = 1
)
print(f"Shape: {embeddings.shape}")  # (4, 384)

# Compute pairwise similarity
from sentence_transformers.util import cos_sim
similarities = cos_sim(embeddings, embeddings)
print(f"\nSimilarity matrix:\n{similarities}")
# The first 2 sentences (about PostgreSQL) will have high similarity
# The pizza sentence will be distant from the others
    

# pip install openai

from openai import OpenAI
import os

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

def get_embeddings(
    texts: list[str],
    model: str = "text-embedding-3-small"
) -> list[list[float]]:
    """Generate embeddings for a list of texts."""
    response = client.embeddings.create(
        input=texts,
        model=model,
    )
    return [item.embedding for item in response.data]

# Single embedding
text = "PostgreSQL as vector database for AI"
embedding = get_embeddings([text])[0]
print(f"Dimensions: {len(embedding)}")  # 1536

# Batch embeddings (up to 2048 texts per call)
texts = [
    "How to install pgvector on Docker",
    "Tutorial for similarity search in PostgreSQL",
    "Guide to pasta al forno cooking",
]
embeddings = get_embeddings(texts)
print(f"Embeddings generated: {len(embeddings)}")  # 3

# Reduced dimension with text-embedding-3-small
# You can specify smaller dimensions to save space
response = client.embeddings.create(
    input=["Sample text"],
    model="text-embedding-3-small",
    dimensions=512  # reduced from 1536 to 512
)
emb_small = response.data[0].embedding
print(f"Reduced dimensions: {len(emb_small)}")  # 512
    

# pip install huggingface_hub

from huggingface_hub import InferenceClient
import os

client = InferenceClient(
    token=os.getenv("HF_TOKEN")
)

def get_hf_embeddings(
    texts: list[str],
    model: str = "BAAI/bge-large-en-v1.5"
) -> list[list[float]]:
    """Generate embeddings using HuggingFace Inference API."""
    result = client.feature_extraction(
        text=texts,
        model=model,
    )
    return result

# Generate embeddings
texts = [
    "Vector search with PostgreSQL and pgvector",
    "How to create HNSW indexes for fast search",
]
embeddings = get_hf_embeddings(texts)
print(f"Embeddings: {len(embeddings)}")       # 2
print(f"Dimensions: {len(embeddings[0])}")    # 1024 (bge-large)
    

import time
from typing import Generator
from sentence_transformers import SentenceTransformer
import numpy as np

def chunk_list(
    lst: list, chunk_size: int
) -> Generator[list, None, None]:
    """Split a list into fixed-size chunks."""
    for i in range(0, len(lst), chunk_size):
        yield lst[i:i + chunk_size]

def generate_embeddings_batch(
    texts: list[str],
    model_name: str = "all-MiniLM-L6-v2",
    batch_size: int = 256,
    device: str = "cpu"  # "cuda" for GPU
) -> np.ndarray:
    """Generate embeddings in batches with progress tracking."""
    model = SentenceTransformer(model_name, device=device)

    all_embeddings = []
    total_batches = (len(texts) + batch_size - 1) // batch_size
    start = time.time()

    for i, batch in enumerate(chunk_list(texts, batch_size)):
        batch_emb = model.encode(
            batch,
            batch_size=batch_size,
            normalize_embeddings=True,
            show_progress_bar=False
        )
        all_embeddings.append(batch_emb)

        elapsed = time.time() - start
        rate = (i + 1) * batch_size / elapsed
        print(
            f"Batch {i+1}/{total_batches} - "
            f"{rate:.0f} texts/sec"
        )

    return np.vstack(all_embeddings)

# Usage
texts = [f"Document number {i}" for i in range(10_000)]
embeddings = generate_embeddings_batch(
    texts,
    batch_size=256,
    device="cuda"  # use GPU if available
)
print(f"Final shape: {embeddings.shape}")  # (10000, 384)
    

-- Enable pgvector
CREATE EXTENSION IF NOT EXISTS vector;

-- Documents table with embedding
CREATE TABLE documents (
    id BIGSERIAL PRIMARY KEY,
    title VARCHAR(500) NOT NULL,
    content TEXT NOT NULL,
    source VARCHAR(255),
    category VARCHAR(100),
    embedding vector(384),  -- dimension of chosen model
    created_at TIMESTAMPTZ DEFAULT NOW(),
    metadata JSONB DEFAULT '{}'::jsonb
);

-- HNSW index for fast search (cosine distance)
CREATE INDEX idx_documents_embedding
    ON documents
    USING hnsw (embedding vector_cosine_ops)
    WITH (m = 16, ef_construction = 200);

-- Index on category for combined filters
CREATE INDEX idx_documents_category
    ON documents (category);
    

import psycopg2
from psycopg2.extras import execute_values
from sentence_transformers import SentenceTransformer
import numpy as np

# Configuration
DB_CONFIG = {
    "host": "localhost",
    "port": 5432,
    "dbname": "vectordb",
    "user": "admin",
    "password": "secret_password",
}

# 1. Generate embeddings
model = SentenceTransformer("all-MiniLM-L6-v2")

documents = [
    {
        "title": "Introduction to pgvector",
        "content": "pgvector is a PostgreSQL extension for vectors...",
        "source": "blog",
        "category": "database"
    },
    {
        "title": "RAG with LangChain",
        "content": "Retrieval Augmented Generation combines retrieval...",
        "source": "tutorial",
        "category": "ai"
    },
    {
        "title": "Python for Data Science",
        "content": "Python is the most used language for data science...",
        "source": "guide",
        "category": "programming"
    },
]

# Generate embeddings for the content
texts = [d["content"] for d in documents]
embeddings = model.encode(texts, normalize_embeddings=True)

# 2. Save to PostgreSQL
conn = psycopg2.connect(**DB_CONFIG)
cur = conn.cursor()

# Prepare data for batch insert
values = []
for doc, emb in zip(documents, embeddings):
    values.append((
        doc["title"],
        doc["content"],
        doc["source"],
        doc["category"],
        emb.tolist()  # convert numpy array to Python list
    ))

# Efficient batch insert
execute_values(
    cur,
    """INSERT INTO documents
       (title, content, source, category, embedding)
       VALUES %s""",
    values,
    template="(%s, %s, %s, %s, %s::vector)"
)

conn.commit()
print(f"Inserted {len(values)} documents with embeddings")

cur.close()
conn.close()
    

def similarity_search(
    query: str,
    top_k: int = 5,
    category: str = None,
    threshold: float = 0.3
) -> list[dict]:
    """Find documents similar to the query."""
    # Generate query embedding
    query_embedding = model.encode(
        query, normalize_embeddings=True
    ).tolist()

    conn = psycopg2.connect(**DB_CONFIG)
    cur = conn.cursor()

    # Query with optional filter
    if category:
        cur.execute("""
            SELECT id, title, content, category,
                   1 - (embedding <=> %s::vector) AS similarity
            FROM documents
            WHERE category = %s
              AND 1 - (embedding <=> %s::vector) > %s
            ORDER BY embedding <=> %s::vector
            LIMIT %s
        """, (
            query_embedding, category,
            query_embedding, threshold,
            query_embedding, top_k
        ))
    else:
        cur.execute("""
            SELECT id, title, content, category,
                   1 - (embedding <=> %s::vector) AS similarity
            FROM documents
            WHERE 1 - (embedding <=> %s::vector) > %s
            ORDER BY embedding <=> %s::vector
            LIMIT %s
        """, (
            query_embedding,
            query_embedding, threshold,
            query_embedding, top_k
        ))

    results = []
    for row in cur.fetchall():
        results.append({
            "id": row[0],
            "title": row[1],
            "content": row[2][:200],  # truncated
            "category": row[3],
            "similarity": round(row[4], 4),
        })

    cur.close()
    conn.close()
    return results

# Example usage
results = similarity_search(
    "how to use vectors in a database",
    top_k=3,
    category="database"
)
for r in results:
    print(f"[{r['similarity']}] {r['title']}")
    

-- HNSW (recommended for production)
-- m: connections per node (16-64, default 16)
-- ef_construction: build quality (64-512, default 64)
CREATE INDEX idx_hnsw_cosine
    ON documents
    USING hnsw (embedding vector_cosine_ops)
    WITH (m = 16, ef_construction = 200);

-- For L2 distance
CREATE INDEX idx_hnsw_l2
    ON documents
    USING hnsw (embedding vector_l2_ops)
    WITH (m = 16, ef_construction = 200);

-- IVFFlat (faster to build)
-- lists: number of clusters (sqrt(N) as a base rule)
CREATE INDEX idx_ivfflat_cosine
    ON documents
    USING ivfflat (embedding vector_cosine_ops)
    WITH (lists = 100);  -- for ~10K documents

-- Query parameters to control recall vs speed
SET hnsw.ef_search = 100;       -- default 40, increase for more recall
SET ivfflat.probes = 10;        -- default 1, increase for more recall
    

# pip install transformers pillow

from transformers import CLIPProcessor, CLIPModel
from PIL import Image
import torch
import numpy as np

# Load CLIP
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

# Image embedding
image = Image.open("cat_photo.jpg")
inputs = processor(images=image, return_tensors="pt")
with torch.no_grad():
    image_embedding = model.get_image_features(**inputs)
image_emb = image_embedding[0].numpy()
print(f"Image embedding: {image_emb.shape}")  # (512,)

# Text embedding (in the SAME space!)
text_inputs = processor(
    text=["a sleeping cat", "a playing dog"],
    return_tensors="pt",
    padding=True
)
with torch.no_grad():
    text_embeddings = model.get_text_features(**text_inputs)
text_embs = text_embeddings.numpy()

# Compute cross-modal similarity
from numpy.linalg import norm
for i, text in enumerate(["a sleeping cat", "a playing dog"]):
    sim = np.dot(image_emb, text_embs[i]) / (
        norm(image_emb) * norm(text_embs[i])
    )
    print(f"Similarity '{text}': {sim:.4f}")
# "a sleeping cat" will have higher similarity with cat_photo.jpg
    

from sentence_transformers import SentenceTransformer
from sentence_transformers.evaluation import (
    InformationRetrievalEvaluator
)

# Prepare evaluation dataset
queries = {
    "q1": "how to install pgvector",
    "q2": "what is similarity search",
    "q3": "image embeddings with CLIP",
}

corpus = {
    "d1": "Guide to installing pgvector on Ubuntu",
    "d2": "pgvector for PostgreSQL: Docker setup",
    "d3": "Vector similarity search",
    "d4": "CLIP: multimodal model for images and text",
    "d5": "Spaghetti carbonara recipe",
}

# Mapping query -> relevant documents
relevant = {
    "q1": {"d1": 1, "d2": 1},  # d1 and d2 relevant for q1
    "q2": {"d3": 1},
    "q3": {"d4": 1},
}

# Evaluate the model
model = SentenceTransformer("all-MiniLM-L6-v2")
evaluator = InformationRetrievalEvaluator(
    queries=queries,
    corpus=corpus,
    relevant_docs=relevant,
    name="custom-eval"
)
results = evaluator(model)
print(f"NDCG@10: {results['custom-eval_ndcg@10']:.4f}")
print(f"MAP@10:  {results['custom-eval_map@10']:.4f}")
    

# pip install umap-learn matplotlib

import umap
import matplotlib.pyplot as plt
import numpy as np
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("all-MiniLM-L6-v2")

# Texts with different categories
texts = [
    # Databases
    "PostgreSQL is a relational database",
    "MongoDB is a NoSQL database",
    "Redis is an in-memory database",
    # AI/ML
    "Deep learning uses deep neural networks",
    "GPT is a language model",
    "Linear regression is a simple algorithm",
    # Food
    "Pizza is baked in a wood-fired oven",
    "Tiramisu is an Italian dessert",
    "Pasta carbonara uses eggs and guanciale",
]
categories = ["DB"]*3 + ["AI"]*3 + ["Food"]*3
colors = ["blue"]*3 + ["red"]*3 + ["green"]*3

# Generate embeddings (384 dimensions)
embeddings = model.encode(texts, normalize_embeddings=True)

# Reduce to 2D with UMAP
reducer = umap.UMAP(n_components=2, random_state=42)
emb_2d = reducer.fit_transform(embeddings)

# Visualize
plt.figure(figsize=(10, 8))
for i, (x, y) in enumerate(emb_2d):
    plt.scatter(x, y, c=colors[i], s=100, zorder=5)
    plt.annotate(texts[i][:30], (x, y), fontsize=8, ha='left')
plt.title("Embeddings in 2D (UMAP)")
plt.savefig("embeddings_umap.png", dpi=150)
plt.show()
    

from openai import OpenAI
import numpy as np
import os

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

text = "PostgreSQL as a vector database for AI applications"

# Generate the same embedding at different dimensions
for dim in [256, 512, 1024, 1536]:
    response = client.embeddings.create(
        input=[text],
        model="text-embedding-3-small",
        dimensions=dim
    )
    emb = response.data[0].embedding
    print(f"Dimensions: {dim}, norm: {np.linalg.norm(emb):.4f}")

# In PostgreSQL: use columns with appropriate dimensions
# CREATE TABLE docs_compact (
#     id BIGSERIAL PRIMARY KEY,
#     content TEXT,
#     embedding vector(256)  -- more compact, 6x less storage
# );
    

import hashlib
import psycopg2

def should_reindex(conn, source_path: str, content: str) -> bool:
    """Returns True if the document needs to be re-indexed."""
    content_hash = hashlib.md5(content.encode()).hexdigest()

    with conn.cursor() as cur:
        cur.execute("""
            SELECT content_hash FROM documents
            WHERE source_path = %s
            LIMIT 1
        """, (source_path,))
        row = cur.fetchone()

    if row is None:
        return True  # new document
    return row[0] != content_hash  # True if changed

# Save the hash in the table schema:
# CREATE TABLE documents (
#     ...
#     content_hash TEXT,  -- MD5 of the content
#     embedding_model TEXT NOT NULL,  -- which model was used
#     ...
# )
    

import psycopg2
import time
from openai import OpenAI
from sentence_transformers import SentenceTransformer

# ===================================
# Evaluation framework for embedding models
# ===================================

class EmbeddingEvaluator:
    """
    Evaluate embedding models on your specific domain data.

    Methodology:
    1. Create a test set of (query, relevant_document) pairs
    2. Embed all documents with each candidate model
    3. For each query, measure Recall@K and MRR (Mean Reciprocal Rank)
    4. Compare latency and cost
    """

    def __init__(self, conn, test_pairs: list[tuple]):
        """
        Args:
            conn: PostgreSQL connection
            test_pairs: List of (query, expected_source_path) tuples
        """
        self.conn = conn
        self.test_pairs = test_pairs  # ground truth

    def evaluate_model(self, model_name: str, embedder, k: int = 5) -> dict:
        """Run full evaluation for one embedding model."""
        recall_scores = []
        mrr_scores = []
        latencies = []

        for query, expected_source in self.test_pairs:
            # Embed query
            t_start = time.perf_counter()
            query_vec = embedder.embed(query)
            embed_time = (time.perf_counter() - t_start) * 1000

            # Search in PostgreSQL
            t_start = time.perf_counter()
            with self.conn.cursor() as cur:
                cur.execute("""
                    SELECT source_path
                    FROM documents
                    WHERE embedding_model = %s
                    ORDER BY embedding <=> %s::vector
                    LIMIT %s
                """, (model_name, query_vec, k))
                results = [r[0] for r in cur.fetchall()]
            search_time = (time.perf_counter() - t_start) * 1000

            # Calculate Recall@K
            recall = 1.0 if expected_source in results else 0.0
            recall_scores.append(recall)

            # Calculate MRR (position of first correct result)
            mrr = 0.0
            if expected_source in results:
                position = results.index(expected_source) + 1
                mrr = 1.0 / position
            mrr_scores.append(mrr)

            latencies.append(embed_time + search_time)

        return {
            "model": model_name,
            f"recall_at_{k}": round(sum(recall_scores) / len(recall_scores), 4),
            "mrr": round(sum(mrr_scores) / len(mrr_scores), 4),
            "avg_latency_ms": round(sum(latencies) / len(latencies), 2),
            "n_queries": len(self.test_pairs)
        }

# Example: compare OpenAI small vs local MiniLM
test_pairs = [
    ("How do I create an HNSW index?", "docs/pgvector_guide.md"),
    ("What is cosine similarity?", "docs/embeddings_intro.md"),
    ("PgBouncer connection pooling configuration", "docs/production_guide.md"),
    # ... add 20-50 pairs for statistical significance
]

evaluator = EmbeddingEvaluator(conn, test_pairs)

# Results example (your domain may differ):
# model=text-embedding-3-small: recall@5=0.91, mrr=0.78, latency=45ms
# model=all-MiniLM-L6-v2:      recall@5=0.84, mrr=0.71, latency=12ms
# Conclusion: small quality gap, 3.7x latency difference for local model

# The simplest domain adaptation: prefix prompting
# Some models (like text-embedding-3-small) respond well to task-specific prefixes

def embed_with_task_prefix(text: str, task: str = "search_document") -> list[float]:
    """
    Add task-specific prefix to improve embedding quality for specific tasks.
    Tested with E5 and similar models that support instruction-following.

    task values:
    - "search_document": for documents to be retrieved
    - "search_query": for user queries
    - "classification": for classification tasks
    - "clustering": for clustering tasks
    """
    prefixed = f"Represent this {task}: {text}"
    return embedder.embed(prefixed)

# Example: E5-large-instruct with task prefix
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("intfloat/e5-large-instruct")

# For ingestion (documents):
doc_embedding = model.encode(
    "Represent this document for retrieval: " + document_text,
    normalize_embeddings=True
)

# For queries:
query_embedding = model.encode(
    "Represent this query for retrieving relevant documents: " + user_query,
    normalize_embeddings=True
)

# The asymmetric approach (different prefixes for docs vs queries)
# often improves Recall@10 by 5-15% compared to symmetric embeddings

from sentence_transformers import SentenceTransformer, InputExample, losses
from torch.utils.data import DataLoader

def fine_tune_embedding_model(
    base_model: str,
    training_pairs: list[tuple],  # (query, relevant_doc, irrelevant_doc) triplets
    output_path: str,
    epochs: int = 3,
    batch_size: int = 32
):
    """
    Fine-tune an embedding model on domain-specific data using triplet loss.

    training_pairs: List of (anchor, positive, negative) tuples where:
    - anchor: a user query from your domain
    - positive: a relevant document for that query
    - negative: an irrelevant document (hard negative = challenging)

    Hard negatives are documents that are superficially similar but not relevant.
    Mining hard negatives is the most important step for good fine-tuning.
    """
    model = SentenceTransformer(base_model)

    # Convert to InputExamples
    examples = [
        InputExample(texts=[anchor, positive, negative])
        for anchor, positive, negative in training_pairs
    ]

    # DataLoader
    dataloader = DataLoader(examples, shuffle=True, batch_size=batch_size)

    # Triplet loss: maximize distance(anchor, negative) - distance(anchor, positive)
    loss = losses.TripletLoss(model=model)

    # Training
    model.fit(
        train_objectives=[(dataloader, loss)],
        epochs=epochs,
        warmup_steps=int(0.1 * len(dataloader) * epochs),
        output_path=output_path,
        show_progress_bar=True
    )

    return model

# Minimum data requirements for meaningful fine-tuning:
# - At least 1000 triplets for initial improvement
# - 5000+ triplets for significant gains
# - Use existing production query logs + relevant document pairs
# - Generate hard negatives with BM25 (retrieve by keyword, use as false positives)

import asyncio
import aiohttp
from typing import AsyncGenerator
import time

class AsyncEmbeddingPipeline:
    """
    Production-grade async embedding pipeline with:
    - Rate limiting (tokens per minute)
    - Automatic retry with exponential backoff
    - Progress tracking
    - Batch size optimization
    """

    def __init__(self,
                 api_key: str,
                 model: str = "text-embedding-3-small",
                 max_tpm: int = 1_000_000,  # tokens per minute limit
                 max_retries: int = 3):
        self.api_key = api_key
        self.model = model
        self.max_tpm = max_tpm
        self.max_retries = max_retries
        self._token_count = 0
        self._window_start = time.time()

    async def _check_rate_limit(self, tokens: int):
        """Wait if we are approaching the rate limit."""
        elapsed = time.time() - self._window_start
        if elapsed >= 60:
            # Reset window
            self._token_count = 0
            self._window_start = time.time()
        elif self._token_count + tokens > self.max_tpm:
            # Wait until the minute window resets
            wait_time = 60 - elapsed
            print(f"Rate limit reached. Waiting {wait_time:.1f}s...")
            await asyncio.sleep(wait_time)
            self._token_count = 0
            self._window_start = time.time()

        self._token_count += tokens

    async def embed_batch_async(self,
                                session: aiohttp.ClientSession,
                                texts: list[str],
                                attempt: int = 0) -> list[list[float]]:
        """Embed a batch of texts with retry logic."""
        # Estimate token count (rough: 1 token ~= 4 chars)
        estimated_tokens = sum(len(t) // 4 for t in texts)
        await self._check_rate_limit(estimated_tokens)

        try:
            async with session.post(
                "https://api.openai.com/v1/embeddings",
                headers={"Authorization": f"Bearer {self.api_key}"},
                json={"input": texts, "model": self.model},
                timeout=aiohttp.ClientTimeout(total=30)
            ) as response:
                if response.status == 429:
                    # Rate limited: exponential backoff
                    if attempt < self.max_retries:
                        wait = (2 ** attempt) * 5
                        await asyncio.sleep(wait)
                        return await self.embed_batch_async(session, texts, attempt + 1)
                response.raise_for_status()
                data = await response.json()
                return [item["embedding"] for item in sorted(data["data"], key=lambda x: x["index"])]

        except Exception as e:
            if attempt < self.max_retries:
                await asyncio.sleep(2 ** attempt)
                return await self.embed_batch_async(session, texts, attempt + 1)
            raise

    async def process_documents(self,
                                texts: list[str],
                                batch_size: int = 100) -> list[list[float]]:
        """Process all documents with concurrent batches."""
        all_embeddings = [None] * len(texts)

        async with aiohttp.ClientSession() as session:
            # Process in batches, with max 3 concurrent requests
            semaphore = asyncio.Semaphore(3)

            async def process_batch(batch_texts, start_idx):
                async with semaphore:
                    embeddings = await self.embed_batch_async(session, batch_texts)
                    for i, emb in enumerate(embeddings):
                        all_embeddings[start_idx + i] = emb

            tasks = []
            for i in range(0, len(texts), batch_size):
                batch = texts[i:i + batch_size]
                tasks.append(process_batch(batch, i))

            await asyncio.gather(*tasks)

        return all_embeddings

# Usage:
pipeline = AsyncEmbeddingPipeline(api_key=openai_api_key)
embeddings = asyncio.run(pipeline.process_documents(all_texts, batch_size=100))
print(f"Generated {len(embeddings)} embeddings")