We use cookies to enhance your experience on the site
CodeWorlds

Embeddings and Vector Search

Embeddings are vector representations of text that allow machines to "understand" the meaning of words and sentences. They are like the DNA of every text - a unique sequence of numbers representing its semantics!

What Are Embeddings?

An embedding is a vector of real numbers representing text in semantic space:

1from openai import OpenAI
2
3client = OpenAI()
4
5def get_embedding(text: str, model: str = "text-embedding-3-small") -> list[float]:
6    """Generates an embedding for the given text."""
7    response = client.embeddings.create(
8        model=model,
9        input=text
10    )
11    return response.data[0].embedding
12
13# Example
14text = "Python is a great programming language"
15embedding = get_embedding(text)
16
17print(f"Vector dimension: {len(embedding)}")  # 1536 for text-embedding-3-small
18print(f"First 5 values: {embedding[:5]}")

Embedding Models

Different models for different use cases:

1from sentence_transformers import SentenceTransformer
2
3# Open-source models
4class EmbeddingModels:
5    """Various embedding models."""
6
7    def __init__(self):
8        # Fast and lightweight
9        self.minilm = SentenceTransformer('all-MiniLM-L6-v2')
10
11        # Larger, better for multilingual
12        self.multilingual = SentenceTransformer('paraphrase-multilingual-mpnet-base-v2')
13
14    def embed_with_openai(self, texts: list[str]) -> list[list[float]]:
15        """OpenAI embeddings - highest quality."""
16        from openai import OpenAI
17        client = OpenAI()
18
19        response = client.embeddings.create(
20            model="text-embedding-3-large",  # 3072 dimensions
21            input=texts
22        )
23        return [item.embedding for item in response.data]
24
25    def embed_with_sentence_transformers(self, texts: list[str]) -> list[list[float]]:
26        """Local embeddings - no API costs."""
27        return self.minilm.encode(texts).tolist()
28
29# Model comparison
30models = EmbeddingModels()
31
32texts = [
33    "Machine learning is a field of AI",
34    "Deep learning uses neural networks",
35    "I like eating pizza"
36]
37
38# Open-source embeddings
39embeddings = models.embed_with_sentence_transformers(texts)
40print(f"Local embeddings: {len(embeddings[0])} dimensions")

Similarity Metrics

1import numpy as np
2from typing import Callable
3
4def cosine_similarity(vec1: np.ndarray, vec2: np.ndarray) -> float:
5    """Cosine similarity - most commonly used."""
6    return np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))
7
8def euclidean_distance(vec1: np.ndarray, vec2: np.ndarray) -> float:
9    """Euclidean distance."""
10    return np.linalg.norm(vec1 - vec2)
11
12def dot_product(vec1: np.ndarray, vec2: np.ndarray) -> float:
13    """Dot product."""
14    return np.dot(vec1, vec2)
15
16# Demonstration
17def demonstrate_similarity():
18    """Shows how similarity metrics work."""
19    from sentence_transformers import SentenceTransformer
20
21    model = SentenceTransformer('all-MiniLM-L6-v2')
22
23    sentences = [
24        "The cat sits on the mat",              # 0
25        "The kitten lies on the carpet",         # 1 - similar meaning
26        "Programming in Python",                  # 2 - different meaning
27    ]
28
29    embeddings = model.encode(sentences)
30
31    print("Cosine similarity:")
32    print(f"  Sentence 0 vs 1: {cosine_similarity(embeddings[0], embeddings[1]):.3f}")
33    print(f"  Sentence 0 vs 2: {cosine_similarity(embeddings[0], embeddings[2]):.3f}")
34    print(f"  Sentence 1 vs 2: {cosine_similarity(embeddings[1], embeddings[2]):.3f}")
35
36demonstrate_similarity()

Semantic Search

1from dataclasses import dataclass
2import numpy as np
3
4@dataclass
5class SearchResult:
6    """Semantic search result."""
7    text: str
8    score: float
9    index: int
10
11class SemanticSearch:
12    """Simple semantic search implementation."""
13
14    def __init__(self, model_name: str = 'all-MiniLM-L6-v2'):
15        from sentence_transformers import SentenceTransformer
16        self.model = SentenceTransformer(model_name)
17        self.documents: list[str] = []
18        self.embeddings: np.ndarray | None = None
19
20    def index_documents(self, documents: list[str]) -> None:
21        """Indexes documents."""
22        self.documents = documents
23        self.embeddings = self.model.encode(documents)
24        print(f"Indexed {len(documents)} documents")
25
26    def search(self, query: str, top_k: int = 5) -> list[SearchResult]:
27        """Searches for the most similar documents."""
28        if self.embeddings is None:
29            raise ValueError("No indexed documents!")
30
31        query_embedding = self.model.encode([query])[0]
32
33        # Compute similarities
34        similarities = np.dot(self.embeddings, query_embedding)
35        similarities /= np.linalg.norm(self.embeddings, axis=1)
36        similarities /= np.linalg.norm(query_embedding)
37
38        # Top-K results
39        top_indices = np.argsort(similarities)[::-1][:top_k]
40
41        results = []
42        for idx in top_indices:
43            results.append(SearchResult(
44                text=self.documents[idx],
45                score=float(similarities[idx]),
46                index=int(idx)
47            ))
48
49        return results
50
51# Usage example
52search = SemanticSearch()
53
54documents = [
55    "Python is a general-purpose programming language",
56    "Machine Learning uses algorithms to learn from data",
57    "RAG combines retrieval with text generation",
58    "FastAPI is a modern framework for building APIs",
59    "Docker containerizes applications",
60]
61
62search.index_documents(documents)
63
64results = search.search("How to build web applications?")
65for r in results:
66    print(f"Score: {r.score:.3f} | {r.text}")

Batch Processing

1def batch_embed(texts: list[str], batch_size: int = 100) -> list[list[float]]:
2    """Embeds large amounts of text in batches."""
3    from openai import OpenAI
4    client = OpenAI()
5
6    all_embeddings = []
7
8    for i in range(0, len(texts), batch_size):
9        batch = texts[i:i + batch_size]
10
11        response = client.embeddings.create(
12            model="text-embedding-3-small",
13            input=batch
14        )
15
16        batch_embeddings = [item.embedding for item in response.data]
17        all_embeddings.extend(batch_embeddings)
18
19        print(f"Processed {min(i + batch_size, len(texts))}/{len(texts)}")
20
21    return all_embeddings
22
23# Caching embeddings
24import hashlib
25import json
26from pathlib import Path
27
28class EmbeddingCache:
29    """Cache for embeddings."""
30
31    def __init__(self, cache_dir: str = ".embedding_cache"):
32        self.cache_dir = Path(cache_dir)
33        self.cache_dir.mkdir(exist_ok=True)
34
35    def _get_cache_key(self, text: str, model: str) -> str:
36        """Generates a cache key."""
37        content = f"{model}:{text}"
38        return hashlib.md5(content.encode()).hexdigest()
39
40    def get(self, text: str, model: str) -> list[float] | None:
41        """Retrieves an embedding from the cache."""
42        key = self._get_cache_key(text, model)
43        cache_file = self.cache_dir / f"{key}.json"
44
45        if cache_file.exists():
46            return json.loads(cache_file.read_text())
47        return None
48
49    def set(self, text: str, model: str, embedding: list[float]) -> None:
50        """Saves an embedding to the cache."""
51        key = self._get_cache_key(text, model)
52        cache_file = self.cache_dir / f"{key}.json"
53        cache_file.write_text(json.dumps(embedding))

Embeddings are the foundation of every RAG system. In the next lesson you will learn about vector databases - specialized databases for storing and searching vectors!

Go to CodeWorlds