Embeddings and Vector Search: Meaning as Numbers, Similarity and Vector Databases

import re, collections, pandas as pd
from sklearn.decomposition import TruncatedSVD
from scipy.sparse import coo_matrix
rv = pd.read_csv("customer_reviews.csv")
docs = [re.findall(r"[a-z]+", t.lower()) for t in rv.review_text]
counts = collections.Counter(w for d in docs for w in d)
vocab = [w for w, c in counts.items() if c >= 5]; idx = {w: i for i, w in enumerate(vocab)}
pairs = collections.Counter()
for d in docs:
    ids = [idx[w] for w in d if w in idx]
    for i, a in enumerate(ids):
        for b in ids[max(0, i - 3): i] + ids[i + 1: i + 4]: pairs[a, b] += 1      # co-occurrence within ±3 words
r, c, v = zip(*[(a, b, n) for (a, b), n in pairs.items()])
M = coo_matrix((v, (r, c)), shape=(len(vocab), len(vocab))).tocsr().astype(float)
total = M.sum(); row = np.asarray(M.sum(1)).ravel(); col = np.asarray(M.sum(0)).ravel()
M = M.tocoo(); pmi = np.log((M.data * total) / (row[M.row] * col[M.col])); keep = pmi > 0
PPMI = coo_matrix((pmi[keep], (M.row[keep], M.col[keep])), shape=M.shape).tocsr()
W = TruncatedSVD(50, random_state=0).fit_transform(PPMI); W /= np.linalg.norm(W, axis=1, keepdims=True)
def similar(word, k=6):
    s = W @ W[idx[word]]; return [vocab[i] for i in np.argsort(-s)[1:k + 1]]
print(len(vocab), "words, 50-dimensional vectors")
for w in ["battery", "delivery", "superb", "slow", "refund"]: print(f"{w:9s} → {similar(w)}")

def cos(a, b): return float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b)))
for a, b in [("battery", "lasts"), ("delivery", "courier"), ("superb", "loved"), ("superb", "slow")]: print(f"cos({a}, {b}) = {cos(W[idx[a]], W[idx[b]]):.2f}")

from sklearn.feature_extraction.text import TfidfVectorizer
tf = TfidfVectorizer(vocabulary=vocab, token_pattern=r"[a-z]+").fit(rv.review_text.str.lower())
def embed(texts):
    T = tf.transform([t.lower() for t in texts]); D = T @ W
    return D / np.maximum(np.linalg.norm(D, axis=1, keepdims=True), 1e-9)
E = embed(rv.review_text.tolist())
def search(query, k=4):
    s = E @ embed([query])[0]; out, seen = [], set()
    for i in np.argsort(-s):                                   # skip duplicate texts
        txt = rv.review_text.iloc[i].rstrip(".!:) ")
        if txt.lower() in seen: continue
        seen.add(txt.lower()); out.append((round(float(s[i]), 2), txt[:80]))
        if len(out) == k: break
    return out
for q in ["charging takes forever", "parcel came late and broken", "support team was very helpful"]:
    print("query:", q); [print("   ", r) for r in search(q)]

import time
rng = np.random.default_rng(0)
N, d = 100_000, 64
centres = rng.normal(size=(2000, d))                                   # real embeddings are clustered by topic, not uniform
X = (centres[rng.integers(0, 2000, N)] + rng.normal(0, .35, (N, d))).astype(np.float32); X /= np.linalg.norm(X, axis=1, keepdims=True)
queries = X[rng.choice(N, 100, replace=False)] + rng.normal(0, .05, (100, d)).astype(np.float32)
t0 = time.perf_counter(); exact = np.argsort(-(queries @ X.T), axis=1)[:, :10]; t_exact = time.perf_counter() - t0
from sklearn.cluster import MiniBatchKMeans
nlist = 256; km = MiniBatchKMeans(nlist, random_state=0, n_init=3, batch_size=4096).fit(X)          # IVF: partition vectors into cells
assign = km.labels_; cells = {c: np.where(assign == c)[0] for c in range(nlist)}
def ivf_search(q, nprobe=8, k=10):
    near = np.argsort(-(km.cluster_centers_ @ q))[:nprobe]; cand = np.concatenate([cells[c] for c in near])
    return cand[np.argsort(-(X[cand] @ q))[:k]]
for nprobe in (1, 4, 16):
    t0 = time.perf_counter(); approx = np.array([ivf_search(q, nprobe) for q in queries]); t = time.perf_counter() - t0
    recall = np.mean([len(set(a) & set(e)) / 10 for a, e in zip(approx, exact)])
    print(f"IVF nprobe={nprobe:2d}: recall@10 {recall:.2f} | {1000 * t / 100:.2f} ms/query (exact: {1000 * t_exact / 100:.2f} ms/query)")

aspect_rx = {"battery": r"batter|charg", "delivery": r"deliver|courier|box arrived|packag", "service": r"support|service|refund|replacement|care"}
labelled = {"phone battery dies quickly": "battery", "my order came late": "delivery", "the customer care was useless": "service"}
from sklearn.feature_extraction.text import TfidfVectorizer as TV
bm = TV(ngram_range=(1, 2), sublinear_tf=True).fit(rv.review_text); B = bm.transform(rv.review_text)
def precision_at_k(scores, aspect, k=10):
    rel = rv.review_text.str.lower().str.contains(aspect_rx[aspect]).to_numpy(); top = np.argsort(-scores)[:k]; return rel[top].mean()
for q, a in labelled.items():
    dense = E @ embed([q])[0]; sparse = (B @ bm.transform([q]).T).toarray().ravel()
    hybrid = 0.5 * dense / dense.max() + 0.5 * sparse / max(sparse.max(), 1e-9)
    print(f"{q:32s} precision@10 dense {precision_at_k(dense, a):.1f} | keyword {precision_at_k(sparse, a):.1f} | hybrid {precision_at_k(hybrid, a):.1f}")

# Voyage AI (recommended by Anthropic for use with Claude) - pip install voyageai
import voyageai
vo = voyageai.Client()                                    # VOYAGE_API_KEY in the environment
docs_emb = vo.embed(texts, model="voyage-3.5", input_type="document").embeddings
query_emb = vo.embed(["phone dies by noon"], model="voyage-3.5", input_type="query").embeddings[0]

# OpenAI - pip install openai
from openai import OpenAI
emb = OpenAI().embeddings.create(model="text-embedding-3-small", input=texts).data

# Local, free - pip install sentence-transformers
from sentence_transformers import SentenceTransformer
model = SentenceTransformer("BAAI/bge-small-en-v1.5")
emb = model.encode(texts, normalize_embeddings=True)

# pgvector example
# CREATE EXTENSION vector;
# CREATE TABLE reviews (id bigserial PRIMARY KEY, text text, product text, embedding vector(1024));
# CREATE INDEX ON reviews USING hnsw (embedding vector_cosine_ops);
# SELECT id, text FROM reviews WHERE product = 'Laptop' ORDER BY embedding <=> $1 LIMIT 5;   -- <=> is cosine distance