Arithmetic Operations: Embedding Algebra

Addition & Subtraction

import numpy as np
from openai import OpenAI

client = OpenAI()

def get_emb(text):
    return np.array(client.embeddings.create(
        model="text-embedding-3-small",
        input=text
    ).data[0].embedding)

# Arithmetic
emb_king = get_emb("king")
emb_man = get_emb("man")
emb_woman = get_emb("woman")

# Analogy
result = emb_king - emb_man + emb_woman

# Find the closest word
candidates = ["queen", "princess", "prince", "king"]
for word in candidates:
    emb = get_emb(word)
    sim = np.dot(result, emb) / (np.linalg.norm(result) * np.linalg.norm(emb))
    print(f"{word}: {sim:.4f}")

Output:

queen: 0.92    ← Closest!
princess: 0.85
prince: 0.78
king: 0.75

---

Classic analogies

# Paris : France = Berlin : ?
emb_paris = get_emb("Paris")
emb_france = get_emb("France")
emb_berlin = get_emb("Berlin")

result = emb_berlin + emb_france - emb_paris
# Result close to "Germany"

---

Query Expansion

def expand_query(query, expansion_terms, alpha=0.7):
    """Expand a query with related terms"""
    query_emb = get_emb(query)
    
    # Average of the expansion terms
    expansion_embs = [get_emb(term) for term in expansion_terms]
    expansion_avg = np.mean(expansion_embs, axis=0)
    
    # Blend
    expanded = alpha * query_emb + (1 - alpha) * expansion_avg
    
    return expanded

# Example
query = "Python"
expansions = ["programming", "language", "code"]
expanded_emb = expand_query(query, expansions, alpha=0.7)

# Search with expanded_emb (better recall)

---

Summary

• ✅ Addition/Subtraction: Semantic analogies
• ✅ Query expansion: Improves recall
• ✅ Classic: king - man + woman = queen

---

Module 6 - Capsule 02