Working with Embeddings
Embeddings are the foundation of modern AI search, recommendation, and RAG systems. They convert text into dense numeric vectors — lists of 768 to 3072 floating-point numbers — where semantically similar texts produce geometrically nearby vectors. This lesson covers how to generate embeddings, measure similarity, and build an efficient retrieval pipeline.
What Are Embeddings?
An embedding model takes a piece of text and outputs a fixed-size vector of floats. The key property: semantically similar texts produce similar vectors. "How do I fix a Python ImportError?" and "Python module import failing" produce vectors that are very close in the embedding space, even though they share few exact words.
from openai import OpenAI
client = OpenAI()
def embed_text(text: str, model: str = "text-embedding-3-small") -> list[float]:
"""Generate an embedding for a single text."""
response = client.embeddings.create(
model=model,
input=text,
)
return response.data[0].embedding
def embed_batch(texts: list[str], model: str = "text-embedding-3-small") -> list[list[float]]:
"""Generate embeddings for multiple texts in one API call."""
response = client.embeddings.create(
model=model,
input=texts, # batch — more efficient than multiple single calls
)
# Response is ordered to match input
return [item.embedding for item in sorted(response.data, key=lambda x: x.index)]
# Example
query = "How do language models generate text?"
query_embedding = embed_text(query)
print(f"Embedding dimensions: {len(query_embedding)}") # 1536 for text-embedding-3-small
print(f"First 5 values: {query_embedding[:5]}")Vector Similarity: The Math
The most common similarity metric for embeddings is cosine similarity — it measures the angle between two vectors, independent of their magnitude. A score of 1.0 means identical direction (semantically the same), 0.0 means orthogonal (unrelated), -1.0 means opposite.
import numpy as np
from typing import Callable
def cosine_similarity(a: list[float], b: list[float]) -> float:
"""
Cosine similarity between two vectors.
Range: -1.0 (opposite) to 1.0 (identical).
"""
a_arr = np.array(a)
b_arr = np.array(b)
dot = np.dot(a_arr, b_arr)
magnitude = np.linalg.norm(a_arr) * np.linalg.norm(b_arr)
if magnitude == 0:
return 0.0
return float(dot / magnitude)
def euclidean_distance(a: list[float], b: list[float]) -> float:
"""Euclidean distance — smaller = more similar. Used by some vector DBs."""
return float(np.linalg.norm(np.array(a) - np.array(b)))
def dot_product(a: list[float], b: list[float]) -> float:
"""
Dot product — fast when embeddings are normalized.
If vectors are unit-normalized, dot product == cosine similarity.
"""
return float(np.dot(np.array(a), np.array(b)))
def normalize(vector: list[float]) -> list[float]:
"""L2-normalize a vector — makes it a unit vector."""
arr = np.array(vector)
norm = np.linalg.norm(arr)
return (arr / norm).tolist() if norm > 0 else arr.tolist()
# Demo: semantic similarity
texts = [
"Python is a programming language",
"Python programming tutorial",
"Snake is a reptile",
"Machine learning with neural networks",
]
embeddings = embed_batch(texts)
query_emb = embed_text("Learn Python coding")
scores = [
(text, cosine_similarity(query_emb, emb))
for text, emb in zip(texts, embeddings)
]
scores.sort(key=lambda x: x[1], reverse=True)
for text, score in scores:
print(f"{score:.3f} {text}")Building an In-Memory Semantic Search
For small document collections (under 100K chunks), you don't need a vector database — NumPy can handle similarity search efficiently in memory:
import numpy as np
from dataclasses import dataclass, field
@dataclass
class InMemoryVectorStore:
"""Simple in-memory vector store for prototyping and small datasets."""
_documents: list[str] = field(default_factory=list)
_metadata: list[dict] = field(default_factory=list)
_embeddings: list[list[float]] = field(default_factory=list)
_embed_fn: Callable = embed_text
def add_documents(
self,
documents: list[str],
metadata: list[dict] | None = None,
) -> None:
"""Add documents and compute their embeddings."""
if metadata is None:
metadata = [{} for _ in documents]
# Batch embed for efficiency
embeddings = embed_batch(documents)
self._documents.extend(documents)
self._metadata.extend(metadata)
self._embeddings.extend(embeddings)
def search(
self,
query: str,
top_k: int = 5,
min_score: float = 0.0,
) -> list[dict]:
"""Find the most similar documents to a query."""
if not self._embeddings:
return []
query_embedding = np.array(self._embed_fn(query))
# Matrix multiplication for all similarities at once — O(n*d) not O(n*d*n)
emb_matrix = np.array(self._embeddings) # shape: (n_docs, dim)
# Normalize both query and document embeddings
query_norm = query_embedding / np.linalg.norm(query_embedding)
doc_norms = emb_matrix / np.linalg.norm(emb_matrix, axis=1, keepdims=True)
# All cosine similarities in one operation
similarities = np.dot(doc_norms, query_norm) # shape: (n_docs,)
# Get top_k indices sorted by similarity
top_indices = np.argsort(similarities)[::-1][:top_k]
results = []
for idx in top_indices:
score = float(similarities[idx])
if score >= min_score:
results.append({
"document": self._documents[idx],
"metadata": self._metadata[idx],
"score": score,
})
return results
# Usage
store = InMemoryVectorStore()
# Index your corpus
store.add_documents(
documents=[
"Python asyncio enables concurrent I/O operations",
"Pydantic validates data using Python type hints",
"Vector databases store and search embeddings efficiently",
"RAG combines retrieval with language model generation",
],
metadata=[
{"source": "lesson_1.3"},
{"source": "lesson_1.2"},
{"source": "lesson_3.3"},
{"source": "lesson_3.1"},
],
)
# Search
results = store.search("How do I handle many API calls at once?", top_k=2)
for r in results:
print(f"Score: {r['score']:.3f} | {r['document']}")Embedding Models Comparison
| Model | Dimensions | Max Tokens | Speed | Best For |
|---|---|---|---|---|
| text-embedding-3-small | 1536 | 8191 | Fast | Production RAG, general use |
| text-embedding-3-large | 3072 | 8191 | Slower | High-accuracy retrieval |
| sentence-transformers/all-MiniLM-L6-v2 | 384 | 256 | Very fast | Local/offline, high volume |
| BAAI/bge-large-en-v1.5 | 1024 | 512 | Moderate | Open-source, competitive quality |
Batch Embedding with Rate Limit Handling
Embedding a whole corpus means thousands of API calls, so doing it one at a time is painfully slow and doing it all at once trips rate limits. The production pattern batches multiple texts per request (the embeddings API accepts a list), runs several batches concurrently with a semaphore to cap parallelism, and retries on failure. This combines the concurrency from Lesson 1.3 with the resilience from Lesson 2.1.
import asyncio
import openai
from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
async_openai = openai.AsyncOpenAI()
@retry(
retry=retry_if_exception_type(openai.RateLimitError),
wait=wait_exponential(multiplier=2, min=2, max=30),
stop=stop_after_attempt(5),
)
async def embed_batch_async(
texts: list[str],
model: str = "text-embedding-3-small",
) -> list[list[float]]:
"""Async batch embedding with retry."""
response = await async_openai.embeddings.create(model=model, input=texts)
return [item.embedding for item in sorted(response.data, key=lambda x: x.index)]
async def embed_corpus(
documents: list[str],
batch_size: int = 100,
max_concurrent: int = 5,
) -> list[list[float]]:
"""Embed a large corpus concurrently in batches."""
semaphore = asyncio.Semaphore(max_concurrent)
async def embed_one_batch(batch: list[str]) -> list[list[float]]:
async with semaphore:
return await embed_batch_async(batch)
batches = [documents[i:i+batch_size] for i in range(0, len(documents), batch_size)]
batch_results = await asyncio.gather(*[embed_one_batch(b) for b in batches])
# Flatten batch results
return [emb for batch in batch_results for emb in batch]OpenAI's text-embedding-3-* models support dimension reduction via the dimensions parameter. You can request 512 or 256 dimensions instead of 1536, trading some accuracy for 3–6x faster similarity search and smaller storage. This is often the right trade-off for large-scale production RAG — benchmark your specific use case before committing.
Key Takeaways
- Embeddings are fixed-size float vectors where similar texts produce geometrically nearby vectors — the basis of semantic search
- Always batch embed (pass multiple texts in one API call) — it's 5–10x more efficient than single-text calls
- Cosine similarity is the standard metric: 1.0 = identical direction, 0.0 = unrelated. Always normalize vectors first for dot-product search
- For small datasets (<100K chunks), NumPy matrix multiplication gives you fast similarity search without a vector database
- Use async embedding with a Semaphore for large corpora — respect the API's rate limits while maximizing throughput
- Consider dimension reduction for large-scale production: OpenAI's models can output 256/512 dimensions instead of 1536 for faster search with modest quality trade-off
Comments
Loading comments…