What do 600 million product reviews actually talk about?
Finding out with Jina embeddings and Elasticsearch
Nobody read all 600 million Amazon Electronics reviews or decided what unsupervised clusters they contained. Nobody labeled them or organized them. That structure just exists, silently, inside the text. The only way to find it is to go looking without knowing what you’re looking for.
That’s the problem this post is about. Not search (where you have a query), not classification (where you have labels), but discovery: figuring out what’s in a large corpus before you can ask the right questions.
The answer turns out to be surprisingly contained. You need task-specific embeddings that push similar documents together, an Elasticsearch index that handles both storage and the nearest-neighbor compute, and a density-probed clustering algorithm that never materializes a full distance matrix. Once those pieces are in place, the structure surfaces on its own. You find a battery drain cluster sitting at 2.3 stars. You find a bluetooth pairing cluster that goes from mostly hardware complaints in 2010 to almost entirely app complaints by 2018. You find these things without writing a single rule.
Here’s the full pipeline before we walk through it step by step:
What you’ll learn:
- Part 1 — Discovery clustering: Why clustering embeddings (not retrieval embeddings) matter for topic discovery without a query, and how to load, embed, and index 10,000 Amazon reviews into Elasticsearch.
- Part 2 — Analyzing what you found: How density-probed centroid classification groups documents by topic using Elasticsearch kNN and batched
msearch, howsignificant_textauto-labels clusters without an ML model, and how to query the results with aggregations and kNN search. - Part 3 — Temporal story chains: How to link yearly clusters across 23 years to trace how review themes shift across product generations, and how to visualize those arcs as a color-coded Sankey diagram.
Required tools
Libraries and services:
- Jina v5 clustering embeddings: Task-specific Low-Rank Adaptation (LoRA) adapters for topic grouping. Jina models are available natively through Elastic Inference Service (EIS).
- Elasticsearch: Scalable kNN,
significant_textlabeling, and vector storage. You'll need a deployment on Elastic Cloud, Elasticsearch Serverless, or self-managed 8.18+/9.0+. Thebbq_diskindex format requires 8.18 or later. - DiskBBQ: A disk-based vector index format that combines Better Binary Quantization (BBQ) with hierarchical k-means partitioning for approximate nearest neighbors (ANN) acceleration. This index partitioning is internal to vector search and separate from the density-probed clustering algorithm used in this post.
bbq_diskstores quantized vectors on disk and keeps only partition metadata in heap, dramatically reducing resource requirements compared tobbq_hnsw, while maintaining high recall.
API keys and data access:
- Jina AI API key: The free tier includes 10 million tokens, which comfortably covers the core clustering pipeline for the sample corpus.
- Amazon Reviews 2023 dataset: Available via Hugging Face. You’ll need to accept the dataset’s terms of use on the dataset page before the
datasetslibrary can download it.
Setting up the environment
Install the dependencies:
pip install elasticsearch jina "datasets==2.17.0" umap-learn plotly pandas numpy requests python-dotenvThe Amazon Reviews 2023 dataset uses a custom loading script that datasets version 2.18 and later refuse to run. Version 2.17.0 is the last release that supports these scripts and works cleanly with Python 3.9.
Create a .env file in your project root with your credentials. Never commit this file to version control:
# .env
ELASTICSEARCH_URL=https://your-cluster-url.elastic.co
ELASTICSEARCH_API_KEY=your_elasticsearch_api_key
JINA_API_KEY=your_jina_api_keyThen load these values at the top of your notebook before any other code:
from dotenv import load_dotenv
import os
load_dotenv() # reads .env from the current directory
ELASTICSEARCH_URL = os.environ["ELASTICSEARCH_URL"]
ELASTICSEARCH_API_KEY = os.environ["ELASTICSEARCH_API_KEY"]
JINA_API_KEY = os.environ["JINA_API_KEY"]
JINA_API_URL = "https://api.jina.ai/v1/embeddings"All subsequent cells reference these variables rather than hardcoded strings.
Part 1: Discovery clustering
Why the embedding task matters more than you’d expect
Here’s a thing that tripped me up: a 5-star review praising battery life and a 1-star review complaining about battery life will end up very close together in a standard retrieval embedding space. Both are about batteries. Both use the same vocabulary. Retrieval embeddings are trained to surface relevant documents for a query, and “battery life” is genuinely relevant to both.
That’s fine for search. For discovery, it’s a problem. If you’re trying to find what topic groups exist in your corpus, you don’t want sentiment mixed into the topical signal. You want reviews about batteries near other reviews about batteries, reviews about sound quality near other sound quality reviews, and complaints about shipping near other shipping complaints.
This is what Jina v5’s task="clustering" adapter is actually doing. LoRA (Low-Rank Adaptation) adds small low-rank updates to targeted internal layers of the base model while keeping most weights frozen. The clustering adapter is trained to maximize separation between topically distinct documents, not to maximize relevance to a query.
TaskTrained forUse caseretrievalQuery-document matchingSearch, RAGclusteringTopic grouping (optimized for tight clusters)Discovery, categorization
Switch tasks, same base model, different geometry. For the rest of this walkthrough, task="clustering" is used throughout.
Loading the dataset
The Electronics category has 23 years of reviews, from 2000 to 2023, and covers a product category that shifted dramatically during that period. Reviews from 2004 are mostly about physical hardware: battery compartments, build quality, manual legibility. Reviews from 2020 are mostly about software: app connectivity, firmware updates, voice assistant integration. That evolution is exactly what the temporal story chains in Part 2 are designed to find.
The dataset is large (around 32GB for Electronics alone), so we stream it rather than download it. The loader collects up to 400 reviews per year and stops once the year buckets are full. A few things are worth noting about this approach: early years like 2001 or 2002 may have far fewer than 400 reviews available, so they’ll contribute whatever they have. Streaming also means the sample is reproducible without a full download, and the per-year quota keeps the corpus temporally balanced, which matters for Part 2.
Streaming connections can time out on slow networks. The loader wraps initialization in a retry with exponential backoff:
from datasets import load_dataset
from collections import defaultdict
import pandas as pd
import time
SAMPLES_PER_YEAR = 400 # ~400 × 25 years ≈ 10k total
MIN_YEAR = 2000 # pre-2000 reviews are very sparse
MAX_YEAR = 2023
def load_dataset_with_retry(max_retries: int = 3, backoff: int = 10):
for attempt in range(max_retries):
try:
return load_dataset(
"McAuley-Lab/Amazon-Reviews-2023",
"raw_review_Electronics",
split="full",
trust_remote_code=True,
streaming=True, # no full download
)
except Exception as e:
if attempt < max_retries - 1:
print(f"Connection failed ({e}), retrying in {backoff}s...")
time.sleep(backoff)
backoff *= 2
else:
raise
dataset = load_dataset_with_retry()
# Stratified streaming: collect SAMPLES_PER_YEAR reviews per year, then stop
REPORT_EVERY = 1_000 # print a status bar every 1k records collected
RED = "\033[31m"
RESET = "\033[0m"
BAR_WIDTH = 30 # character width of the filled bar
def print_status_bar(records_scanned: int, records_collected: int, buckets: dict) -> None:
years_seen = len(buckets)
years_full = sum(1 for v in buckets.values() if len(v) >= SAMPLES_PER_YEAR)
total_target = years_seen * SAMPLES_PER_YEAR if years_seen else 1
filled = int(BAR_WIDTH * records_collected / total_target)
bar = "█" * filled + "░" * (BAR_WIDTH - filled)
pct = 100 * records_collected / total_target
print(
f"{RED}[{bar}] {pct:5.1f}% "
f"{records_collected:,} collected / {records_scanned:,} scanned "
f"{years_full}/{years_seen} years full{RESET}"
)
buckets: dict[int, list] = defaultdict(list)
records_scanned = 0
records_collected = 0
print(
f"Streaming Electronics reviews ({MIN_YEAR}–{MAX_YEAR}), "
f"collecting up to {SAMPLES_PER_YEAR} per year. "
f"Status bar updates every {REPORT_EVERY:,} records collected..."
)
for record in dataset:
records_scanned += 1
if not record.get("text") or not record.get("timestamp"):
continue
year = pd.to_datetime(record["timestamp"], unit="ms").year
if year < MIN_YEAR or year > MAX_YEAR:
continue
if len(buckets[year]) < SAMPLES_PER_YEAR:
buckets[year].append({
"text": record["text"],
"rating": record["rating"],
"timestamp": record["timestamp"],
"verified_purchase": record.get("verified_purchase"),
"year": year,
})
records_collected += 1
# Status bar every REPORT_EVERY records collected
if records_collected % REPORT_EVERY == 0:
print_status_bar(records_scanned, records_collected, buckets)
# Stop once we have enough coverage: all years seen so far are full
# and we've seen at least 70% of the expected year range
years_seen = len(buckets)
years_full = sum(1 for v in buckets.values() if len(v) >= SAMPLES_PER_YEAR)
if years_seen >= (MAX_YEAR - MIN_YEAR) * 0.7 and years_full == years_seen:
print_status_bar(records_scanned, records_collected, buckets)
print(f"{RED}Done. Scanned {records_scanned:,} records, collected {records_collected:,}.{RESET}")
break
df = pd.DataFrame([r for records in buckets.values() for r in records])
df["date"] = pd.to_datetime(df["timestamp"], unit="ms")
df["month"] = df["date"].dt.month
df = df.reset_index(drop=True)
print(f"Loaded {len(df)} reviews across {df['year'].nunique()} years")
print(df.groupby("year").size().to_string())The per-year stratification does one more useful thing: it compensates for review volume skew. Electronics reviews on Amazon grew exponentially through the 2010s. Without a cap, recent years would dominate the corpus and the temporal story chains would have almost no signal from the early years where the hardware-to-software shift actually started.
Embedding with the clustering task
We call the Jina v5 API with task="clustering" for all documents. Embeddings are cached to disk on first run; subsequent runs skip the API entirely:
import math
import numpy as np
import requests
import time
import pickle
from pathlib import Path
# JINA_API_KEY and JINA_API_URL loaded from .env - see "Setting up the environment"
CACHE_PATH = Path("embeddings_cache.pkl")
MAX_RETRIES = 5 # max retry attempts per batch on 429
RETRY_BASE = 2 # base wait in seconds; doubles on each attempt
def embed_texts(texts: list[str], task: str = "clustering", batch_size: int = 256) -> np.ndarray:
# Embed a list of texts using Jina v5 task-specific embeddings.
# Retries each batch up to MAX_RETRIES times on 429 with exponential backoff.
# Returns a (len(texts), 1024) float32 array.
all_embeddings = []
n_batches = math.ceil(len(texts) / batch_size)
print(f"Embedding {len(texts):,} texts in {n_batches:,} batches of {batch_size}...")
for i in range(0, len(texts), batch_size):
batch = texts[i : i + batch_size]
wait = RETRY_BASE
for attempt in range(1, MAX_RETRIES + 1):
response = requests.post(
JINA_API_URL,
headers={
"Authorization": f"Bearer {JINA_API_KEY}",
"Content-Type": "application/json",
},
json={
"model": "jina-embeddings-v5-text-small",
"task": task,
"input": batch,
"dimensions": 1024,
},
)
if response.status_code == 429:
# Honour Retry-After if the API provides it, otherwise back off
retry_after = int(response.headers.get("Retry-After", wait))
print(
f"Rate limited on batch {i // batch_size + 1} "
f"(attempt {attempt}/{MAX_RETRIES}). "
f"Waiting {retry_after}s..."
)
time.sleep(retry_after)
wait *= 2 # exponential backoff for subsequent retries
continue
response.raise_for_status()
break # success - exit retry loop
else:
raise RuntimeError(
f"Batch starting at index {i} failed after {MAX_RETRIES} retries."
)
batch_embeddings = [item["embedding"] for item in response.json()["data"]]
all_embeddings.extend(batch_embeddings)
if i + batch_size < len(texts):
time.sleep(0.1)
if i % 5000 == 0:
print(f"Embedded {i + len(batch)}/{len(texts)} reviews...")
return np.array(all_embeddings, dtype=np.float32)
# v5-text-small supports up to 32,768 tokens, so truncation is rarely needed.
# We still cap very long reviews to keep batch times predictable.
df["text_truncated"] = df["text"].str[:2048]
if CACHE_PATH.exists():
print("Loading embeddings from cache...")
with open(CACHE_PATH, "rb") as f:
embeddings = pickle.load(f)
else:
embeddings = embed_texts(df["text_truncated"].tolist())
with open(CACHE_PATH, "wb") as f:
pickle.dump(embeddings, f)
print(f"Embeddings cached to {CACHE_PATH}")
print(f"Embeddings shape: {embeddings.shape}")
# Embeddings shape: (10000, 1024)One thing worth trying after you run this: re-embed a random 200-document sample with task="retrieval.query" and compare the UMAP projections side by side. The retrieval projection looks like a smooth cloud with gradual density variation. The clustering projection looks like a map with distinct islands and clear water between them. That visual difference is what the adapter is doing in 1024-dimensional space.
Indexing into Elasticsearch with bbq_disk
The full corpus goes into one index (reviews-clustering-all) for the global discovery pass. Yearly indices come later.
A 1024-dimensional float32 vector is 4 KB per document. bbq_disk uses hierarchical k-means to partition vectors into small clusters, binary-quantizes them, and stores the full-precision vectors on disk for rescoring. Only partition metadata lives in heap, so memory requirements stay low even for large corpora. For workloads that can afford more heap, bbq_hnsw builds an HNSW graph for faster lookups at higher resource cost.
from elasticsearch import Elasticsearch, helpers
es = Elasticsearch(
ELASTICSEARCH_URL,
api_key=ELASTICSEARCH_API_KEY,
)
# Use es_bulk for bulk operations that need a longer timeout.
# .options() is the preferred way in elasticsearch-py 8.x+.
es_bulk = es.options(request_timeout=60)
INDEX_NAME = "reviews-clustering-all"
es.indices.create(
index=INDEX_NAME,
body={
"mappings": {
"properties": {
"text": {"type": "text"},
"text_truncated": {"type": "text"},
"rating": {"type": "float"},
"year": {"type": "integer"},
"month": {"type": "integer"},
"date": {"type": "date"},
"embedding": {
"type": "dense_vector",
"dims": 1024,
"index": True,
"similarity": "cosine",
"index_options": {
"type": "bbq_disk", # disk-resident quantized vectors; low heap cost
}
}
}
}
},
ignore=400
)
def generate_actions(df, embeddings, index: str = INDEX_NAME):
for i, (_, row) in enumerate(df.iterrows()):
yield {
"_index": index,
"_id": str(i),
"_source": {
"text": row["text"],
"text_truncated": row["text_truncated"],
"rating": row["rating"],
"year": int(row["year"]),
"month": int(row.get("month", 0)),
"date": row["date"].isoformat(),
"embedding": embeddings[i].tolist(),
}
}
success, _ = helpers.bulk(es_bulk, generate_actions(df, embeddings), chunk_size=500)
print(f"Indexed {success} documents")Part 2: Analyzing what you found
With 10,000 reviews indexed and embedded, the global index is ready to be clustered and labeled. This part covers the algorithm that finds the topic groups, the visualizations that confirm they make sense, the significant_text aggregation that names them, and the Elasticsearch queries that let you explore what's inside.
Clustering: density-probed centroid classification
Standard clustering algorithms like K-means assume you can load the full N×d matrix into memory and run repeated passes over it. For 10,000 documents at 1024 dimensions that’s manageable, but the approach doesn’t extend gracefully to millions of reviews. Rewriting it for larger scale means either sampling aggressively or adding infrastructure the Elasticsearch index already provides.
The density-probed approach flips the architecture: instead of bringing the data to the algorithm, it runs the algorithm through the index. Elasticsearch kNN becomes the compute primitive. Every expensive operation (finding neighbors, measuring similarity, classifying documents against centroids) happens server-side. The only client-side work is selecting which high-density probe points become cluster seeds, which takes about 0.01 seconds.
The four steps:
- Sample 5% of documents as density probes (random sample, minimum 10 for the demo corpus; raise
min_probesfor larger corpora). - Probe density via batched
msearchkNN. Each probe fires a kNN query and records the mean similarity of its neighbors. High mean similarity means a dense region of embedding space.msearchsends multiple search requests in a single HTTP call; this is critical here because density probing generates hundreds of kNN queries, and batching avoids per-request overhead. - Select high-density seeds with diversification. Candidates above median density are sorted by density descending and greedily accepted only when their cosine similarity to every existing seed is below a separation threshold. This is the only client-side compute, and it takes roughly 0.01s for a 10k corpus.
- Classify all documents against centroids via
msearchkNN. Each seed acts as a centroid; a kNN search retrieves nearby documents above a similarity threshold. Each document is assigned to whichever centroid returned it with the highest score. Small clusters are dissolved to noise.
Clusters have natural sizes determined by embedding space density around each centroid, not by a hard k cap. Dense topic regions produce larger clusters; niche topics produce smaller ones.
import random
from collections import Counter
from elasticsearch.helpers import scan
def density_probe_cluster(
es,
index: str,
embedding_field: str = "embedding",
probe_fraction: float = 0.05,
probe_k: int = 15,
min_probes: int = 10, # lowered from 50 for 10k demo; raise for larger corpora
separation_threshold: float = 0.15,
similarity_threshold: float = 0.82,
classify_k: int = 50,
min_cluster_size: int = 10,
msearch_batch_size: int = 50,
) -> tuple[dict, dict]:
# Density-probed centroid classification on an Elasticsearch index.
# Returns (assignments, id_to_embedding):
# assignments: dict mapping document _id -> cluster_id (-1 = noise)
# id_to_embedding: dict mapping _id -> embedding vector
# Fetch all doc IDs and embeddings via scroll (bypasses the 10k result-window limit)
all_ids = []
id_to_embedding = {}
for hit in scan(es, index=index, query={"query": {"match_all": {}}, "_source": [embedding_field]}):
doc_id = hit["_id"]
all_ids.append(doc_id)
id_to_embedding[doc_id] = hit["_source"][embedding_field]
n_docs = len(all_ids)
if n_docs == 0:
print("No documents found in index - skipping.")
return {}, {}, {}
n_probes = min(n_docs, max(min_probes, int(n_docs * probe_fraction)))
probe_ids = random.sample(all_ids, n_probes)
print(f"Probing density on {n_probes} samples from {n_docs} documents...")
# Step 2: batched msearch density probing
def msearch_batch(queries: list[dict]) -> list[dict]:
body = []
for q in queries:
body.append({"index": index})
body.append(q)
return es.msearch(body=body)["responses"]
probe_densities = {}
for batch_start in range(0, len(probe_ids), msearch_batch_size):
batch = probe_ids[batch_start : batch_start + msearch_batch_size]
queries = [
{
"knn": {
"field": embedding_field,
"query_vector": id_to_embedding[pid],
"k": probe_k,
"num_candidates": probe_k * 2,
},
"_source": False,
}
for pid in batch
]
responses = msearch_batch(queries)
for pid, resp in zip(batch, responses):
hits = resp.get("hits", {}).get("hits", [])
probe_densities[pid] = (
sum(h["_score"] for h in hits) / len(hits) if hits else 0.0
)
# Step 3: select high-density seeds with diversification (client-side, ~0.01s)
median_density = sorted(probe_densities.values())[len(probe_densities) // 2]
candidates = sorted(
[(pid, d) for pid, d in probe_densities.items() if d >= median_density],
key=lambda x: -x[1],
)
seeds = []
for pid, _ in candidates:
vec = np.array(id_to_embedding[pid])
too_close = False
for s in seeds:
s_vec = np.array(id_to_embedding[s])
cos_sim = np.dot(vec, s_vec) / (np.linalg.norm(vec) * np.linalg.norm(s_vec))
if cos_sim >= (1 - separation_threshold):
too_close = True
break
if not too_close:
seeds.append(pid)
print(f"Selected {len(seeds)} seeds after diversification")
# Step 4: classify all documents against centroids via kNN
assignments = {doc_id: -1 for doc_id in all_ids}
best_scores = {doc_id: -1.0 for doc_id in all_ids}
for cluster_id, seed_id in enumerate(seeds):
response = es.search(
index=index,
body={
"knn": {
"field": embedding_field,
"query_vector": id_to_embedding[seed_id],
"k": classify_k,
"num_candidates": classify_k * 4,
},
"min_score": similarity_threshold,
"_source": False,
"size": classify_k,
}
)
for hit in response["hits"]["hits"]:
doc_id = hit["_id"]
score = hit["_score"]
if score > best_scores.get(doc_id, -1):
assignments[doc_id] = cluster_id
best_scores[doc_id] = score
# Dissolve small clusters back to noise
cluster_sizes = Counter(v for v in assignments.values() if v != -1)
for doc_id in assignments:
if assignments[doc_id] != -1 and cluster_sizes[assignments[doc_id]] < min_cluster_size:
assignments[doc_id] = -1
n_clusters = len(set(v for v in assignments.values() if v != -1))
n_noise = sum(1 for v in assignments.values() if v == -1)
print(f"Found {n_clusters} clusters, {n_noise} noise documents ({100*n_noise/n_docs:.1f}%)")
return assignments, id_to_embedding, probe_densities
assignments, id_to_embedding, probe_densities = density_probe_cluster(es, INDEX_NAME)The noise fraction will typically land around 25–30%. That’s not a tuning failure. Short reviews, one-off opinions, and reviews that span multiple unrelated topics genuinely don’t belong to any coherent cluster. Forcing them in by lowering similarity_threshold would dilute the clusters that do matter. The noise documents are the algorithm telling you which reviews resist categorization. That's its own signal.
Visualizing probe density
Before writing cluster assignments, it’s worth looking at where the density probes landed. A scatter plot of probe embeddings reduced to 2D with UMAP (Uniform Manifold Approximation and Projection) and colored by density score reveals whether the algorithm found genuine dense regions or distributed uniformly across the space.
Bright peaks correspond to topic islands where a probe’s 15 nearest neighbors all live very close together. The dark scatter between them is the open water: reviews that don’t cluster strongly around any particular theme. A plot with no bright peaks and uniform medium density is a diagnostic signal: try increasing probe_k or probe_fraction.
# If fig.show(renderer="notebook") does not render in your environment,
# try renderer="iframe" (JupyterLab) or renderer="browser" (opens in a new tab).
import umap
import plotly.express as px
import pandas as pd
# Reduce probe embeddings to 2D for visualization
probe_embeddings = np.array(
[id_to_embedding[pid] for pid in probe_densities.keys()],
dtype=np.float32,
)
probe_density_scores = list(probe_densities.values())
reducer = umap.UMAP(n_neighbors=15, min_dist=0.1, metric="cosine", random_state=42)
probe_2d = reducer.fit_transform(probe_embeddings)
probe_df = pd.DataFrame({
"x": probe_2d[:, 0],
"y": probe_2d[:, 1],
"density": probe_density_scores,
})
fig = px.scatter(
probe_df,
x="x",
y="y",
color="density",
color_continuous_scale="Plasma",
title="Probe density scores (UMAP projection) - bright = dense embedding region",
labels={"density": "Density score"},
opacity=0.8,
width=900,
height=600,
)
fig.update_traces(marker=dict(size=6))
fig.write_html("probe_density.html")
fig.show(renderer="notebook")Writing cluster assignments back to Elasticsearch
At this point, assignments is a Python dict mapping each document _id to a cluster integer (or -1 for noise). It lives only in memory. Writing it back to Elasticsearch is what makes the rest of the pipeline possible: the aggregations, the bubble chart, the kNN search within a cluster, and the temporal linking in Part 2 all rely on cluster being a queryable field in the index.
The update is a two-step operation: first, add the cluster field to the index mapping so Elasticsearch knows to treat it as an integer; second, bulk-update every document with its assigned value. Documents labeled -1 are written too, which means you can explicitly filter noise out of any downstream query rather than accidentally including it.
def generate_cluster_updates(assignments: dict):
for doc_id, cluster_id in assignments.items():
yield {
"_op_type": "update",
"_index": INDEX_NAME,
"_id": doc_id,
"doc": {"cluster": cluster_id}
}
es.indices.put_mapping(
index=INDEX_NAME,
body={"properties": {"cluster": {"type": "integer"}}}
)
success, _ = helpers.bulk(es, generate_cluster_updates(assignments), chunk_size=500)
print(f"Updated {success} documents with cluster assignments")Sanity-checking clusters with UMAP
UMAP reduces the 1024-dimensional embeddings to 2D while preserving local neighborhood structure. The goal here isn’t proof: it’s a quick check that the clusters are spatially coherent. A single cluster spanning half the plot, or heavy overlap between clusters, means the density parameters need revisiting.
The thing that tends to surprise people: “battery life is incredible” and “battery life is garbage” sit near each other in the 2D projection even when they’re in different clusters. The neighborhood structure is topical, not polar. Sentiment doesn’t separate in the clustering embedding space; topics do. Two reviews that disagree violently about the same feature are still closer to each other than either is to a review about sound quality.
Noise documents (cluster -1) are plotted in light gray. Well-formed results look like distinct islands surrounded by gray. If you see one giant blob, lower similarity_threshold. If you see nothing but gray with a few dots, raise probe_fraction or probe_k.
# umap and plotly.express already imported above
# Fetch all embeddings and their cluster assignments from Elasticsearch
from elasticsearch.helpers import scan
rows = []
for hit in scan(es, index=INDEX_NAME, query={"query": {"match_all": {}}, "_source": ["embedding", "cluster"]}):
rows.append({
"embedding": hit["_source"]["embedding"],
"cluster": hit["_source"].get("cluster", -1),
"id": hit["_id"],
})
all_embeddings = np.array([r["embedding"] for r in rows], dtype=np.float32)
all_clusters = [r["cluster"] for r in rows]
# Reduce to 2D - subsample for speed if corpus is large
sample_size = min(10_000, len(all_embeddings))
sample_idx = np.random.choice(len(all_embeddings), sample_size, replace=False)
sample_embeddings = all_embeddings[sample_idx]
sample_clusters = [str(all_clusters[i]) for i in sample_idx]
reducer = umap.UMAP(n_neighbors=15, min_dist=0.1, metric="cosine", random_state=42)
embeddings_2d = reducer.fit_transform(sample_embeddings)
import pandas as pd # noqa: re-import guard for standalone cell execution
umap_df = pd.DataFrame({
"x": embeddings_2d[:, 0],
"y": embeddings_2d[:, 1],
"cluster": sample_clusters,
})
# Separate noise from assigned clusters for coloring
umap_df["is_noise"] = umap_df["cluster"] == "-1"
fig = px.scatter(
umap_df,
x="x",
y="y",
color="cluster",
title="Amazon Reviews embeddings (UMAP projection, colored by cluster)",
opacity=0.5,
width=1000,
height=700,
color_discrete_map={"-1": "lightgray"},
)
fig.update_traces(marker=dict(size=3))
fig.write_html("clusters_umap.html")
fig.show(renderer="notebook")Automatic labels with significant_text
Once you have cluster assignments, you need to know what each cluster is actually about. The obvious approach is to look at the most frequent terms. The problem is that frequent terms in a “battery complaints” cluster include words like the, and, it, product, because those words are frequent everywhere. Frequency is the wrong signal.
What you want is statistical surprise: terms that appear more often in this cluster than you’d expect given the full corpus. Elasticsearch’s significant_text aggregation does exactly this using the JLH score, which balances absolute frequency with relative frequency shift. A cluster about Bluetooth connectivity will surface pairing, bluetooth, disconnect because those terms are disproportionately common in that cluster compared to the overall review corpus, not because they're the most common words there.
This also works as a quality gate. Clusters that produce no significant terms have no distinguishing vocabulary. They’re not coherent topics; they’re collection artifacts. Dissolving them to noise before analysis saves time and keeps the results clean.
def label_cluster(es, index: str, cluster_id: int, top_n: int = 8) -> list[str]:
# Use significant_text to auto-label a cluster against the full corpus background.
# Returns significant terms, or an empty list if the cluster is incoherent.
response = es.search(
index=index,
body={
"size": 0,
"query": {"term": {"cluster": cluster_id}},
"aggs": {
"significant_terms": {
"significant_text": {
"field": "text",
"size": top_n,
"filter_duplicate_text": True,
}
}
}
}
)
buckets = response["aggregations"]["significant_terms"]["buckets"]
return [b["key"] for b in buckets]
# Label all discovered clusters and dissolve incoherent ones
cluster_ids = sorted(set(v for v in assignments.values() if v != -1))
cluster_labels = {}
for cluster_id in cluster_ids:
terms = label_cluster(es, INDEX_NAME, cluster_id)
if terms:
cluster_labels[cluster_id] = terms
print(f"Cluster {cluster_id}: {', '.join(terms)}")
else:
print(f"Cluster {cluster_id}: no significant terms - dissolving to noise")
for doc_id in assignments:
if assignments[doc_id] == cluster_id:
assignments[doc_id] = -1When you run this against the Electronics reviews corpus, the labels are immediately legible. One cluster surfaces battery, drain, charge, life. Another surfaces bluetooth, pairing, connect, disconnect. A third surfaces sound, bass, audio, headphone and sits noticeably higher in average rating than the complaint clusters. No training data, no LLM, no hand-written rules.
Querying clusters with Elasticsearch aggregations
With cluster labels indexed, you can ask questions that aren’t possible with a traditional search query. Average rating per cluster immediately shows which topic groups correlate with satisfaction and which with frustration:
K = len(cluster_labels)
response = es.search(
index=INDEX_NAME,
body={
"size": 0,
"aggs": {
"by_cluster": {
"terms": {"field": "cluster", "size": K, "exclude": [-1]},
"aggs": {
"avg_rating": {"avg": {"field": "rating"}},
"by_year": {
"date_histogram": {
"field": "date",
"calendar_interval": "year"
},
"aggs": {
"avg_rating": {"avg": {"field": "rating"}}
}
}
}
}
}
}
)
for bucket in response["aggregations"]["by_cluster"]["buckets"]:
cid = bucket["key"]
label = ", ".join(cluster_labels.get(cid, ["unlabeled"]))
avg = round(bucket["avg_rating"]["value"], 2)
print(f"Cluster {cid} [{label}]: avg rating {avg} ({bucket['doc_count']} docs)")Plotting this as a bubble chart with average rating on the x-axis and document count on the y-axis makes the satisfaction landscape readable at a glance.
Clusters in the top-right (high rating, high volume) are the satisfied majority: unboxing experiences, first impressions, things that just worked. Clusters in the bottom-left (low rating, smaller volume) are concentrated complaint themes worth investigating. Clusters with high volume but middling ratings are mixed-sentiment areas like “connectivity” or “setup experience” where individual reviews vary widely within the same topic.
import plotly.graph_objects as go
# Build a dataframe from the aggregation response
cluster_stats = []
for bucket in response["aggregations"]["by_cluster"]["buckets"]:
cid = bucket["key"]
if cid == -1:
continue
terms = cluster_labels.get(cid, ["unlabeled"])
label = ", ".join(terms[:2])
cluster_stats.append({
"cluster_id": cid,
"label": label,
"avg_rating": round(bucket["avg_rating"]["value"] or 0, 2),
"doc_count": bucket["doc_count"],
})
stats_df = pd.DataFrame(cluster_stats)
fig = go.Figure()
fig.add_trace(go.Scatter(
x=stats_df["avg_rating"],
y=stats_df["doc_count"],
mode="markers+text",
text=stats_df["label"],
textposition="top center",
marker=dict(
size=stats_df["doc_count"] / stats_df["doc_count"].max() * 60 + 10,
color=stats_df["avg_rating"],
colorscale="RdYlGn",
showscale=True,
colorbar=dict(title="Avg rating"),
line=dict(width=1, color="white"),
),
hovertemplate=(
"<b>%{text}</b><br>"
"Avg rating: %{x:.2f}<br>"
"Documents: %{y:,}<extra></extra>"
),
))
fig.update_layout(
title="Cluster landscape: avg rating vs. volume (bubble size = doc count)",
xaxis=dict(title="Average rating", range=[1, 5.2]),
yaxis=dict(title="Document count"),
width=1000,
height=650,
)
fig.write_html("cluster_bubble.html")
fig.show(renderer="notebook")You can also query within a specific cluster using Elasticsearch’s native kNN, with a filter that restricts results to documents already assigned to that cluster:
def find_similar_in_cluster(query_text: str, cluster_id: int, top_k: int = 5):
# Find the most similar reviews within a specific cluster using Elasticsearch kNN.
query_embedding = embed_texts([query_text], task="retrieval.query")[0]
response = es.search(
index=INDEX_NAME,
body={
"knn": {
"field": "embedding",
"query_vector": query_embedding.tolist(),
"k": top_k,
"num_candidates": 100,
"filter": {"term": {"cluster": cluster_id}}
},
"_source": ["text", "rating", "year"]
}
)
return [hit["_source"] for hit in response["hits"]["hits"]]
results = find_similar_in_cluster("battery drains too fast", cluster_id=2)
for r in results:
print(f"[{r['year']} | {r['rating']}★] {r['text'][:200]}")Part 3: Temporal story chains
How review themes change over 23 years
Part 1 found the topics. Part 3 watches them evolve.
The same density-probed clustering runs independently per year on yearly indices, then clusters are linked across adjacent years. Each year produces its own assignments tuned to that year’s content. The 2005 index has its own “setup and installation” cluster with its own vocabulary. The 2018 index has its own “setup and installation” cluster with a completely different vocabulary. The linking step finds that these two independent clusters are thematically continuous across the years between them.
This independence matters. A global clustering over all 23 years would blur the temporal signal. You’d find a “setup” cluster, but you wouldn’t see when its vocabulary pivoted from manual, bracket, screwdriver to app, bluetooth, firmware. Running clusters per year and linking them across time is what surfaces that shift.
# Index reviews per year into separate indices and run clustering per year
assignments_by_year = {}
id_to_embedding_by_year = {} # store per-year embeddings for temporal linking
for year, year_df in df.groupby("year"):
year_index = f"reviews-clustering-{year}"
# Use positional iloc to extract the matching rows from the global embeddings array
year_positions = [df.index.get_loc(i) for i in year_df.index]
year_embeddings = embeddings[year_positions]
year_df_reset = year_df.reset_index(drop=True)
es.indices.create(
index=year_index,
body={
"mappings": {
"properties": {
"text": {"type": "text"},
"text_truncated": {"type": "text"},
"rating": {"type": "float"},
"year": {"type": "integer"},
"date": {"type": "date"},
"cluster": {"type": "integer"},
"embedding": {
"type": "dense_vector",
"dims": 1024,
"index": True,
"similarity": "cosine",
"index_options": {"type": "bbq_disk"}
}
}
}
},
ignore=400
)
helpers.bulk(
es_bulk,
generate_actions(year_df_reset, year_embeddings, index=year_index),
chunk_size=500,
)
# Lower similarity_threshold for sparse yearly slices (fewer docs = looser clusters)
year_assignments, year_id_to_emb, _ = density_probe_cluster(
es, year_index, similarity_threshold=0.75
)
# Write cluster assignments back to the year index so link_yearly_clusters
# can read the cluster field via kNN hits
def _year_cluster_updates(assignments, idx):
for doc_id, cluster_id in assignments.items():
yield {
"_op_type": "update",
"_index": idx,
"_id": doc_id,
"doc": {"cluster": cluster_id},
}
es.indices.put_mapping(index=year_index, body={"properties": {"cluster": {"type": "integer"}}})
helpers.bulk(es, _year_cluster_updates(year_assignments, year_index), chunk_size=500)
# Refresh so the cluster field is visible to subsequent kNN queries
es.indices.refresh(index=year_index)
assignments_by_year[year] = year_assignments
id_to_embedding_by_year[year] = year_id_to_emb
n_real = sum(1 for v in year_assignments.values() if v != -1)
n_clusters = len(set(v for v in year_assignments.values() if v != -1))
print(f"{year}: {n_clusters} clusters, {n_real}/{len(year_assignments)} docs assigned")The linking approach: sample-and-query
For each cluster on year A:
- Sample a few representative documents.
- Run kNN against year B’s index.
- Count how many hits land in each year B cluster.
- If the hit fraction exceeds a threshold (kNN fraction ≥ 0.2 for the demo corpus; raise
link_thresholdfor larger corpora), record a link.
This is fast: only a handful of documents per cluster are queried, not all of them. It uses Elasticsearch’s native kNN with no external dependencies:
def link_yearly_clusters(
es,
year_a: int,
year_b: int,
assignments_a: dict,
id_to_embedding: dict,
n_samples: int = 10, # more samples per cluster for sparse yearly slices
link_threshold: float = 0.2, # lowered from 0.4 for 10k demo corpus
) -> list[dict]:
# Link clusters from year_a to clusters in year_b using sample-and-query kNN.
# Returns a list of link dicts: source cluster, target cluster, kNN fraction.
index_b = f"reviews-clustering-{year_b}"
cluster_ids_a = set(v for v in assignments_a.values() if v != -1)
links = []
for cluster_id_a in cluster_ids_a:
cluster_doc_ids = [
doc_id for doc_id, cid in assignments_a.items() if cid == cluster_id_a
]
sample_ids = random.sample(cluster_doc_ids, min(n_samples, len(cluster_doc_ids)))
target_hits: Counter = Counter()
total_hits = 0
for doc_id in sample_ids:
response = es.search(
index=index_b,
body={
"knn": {
"field": "embedding",
"query_vector": id_to_embedding[doc_id],
"k": 10,
"num_candidates": 50,
},
"_source": ["cluster"],
"size": 10,
}
)
for hit in response["hits"]["hits"]:
target_cluster = hit["_source"].get("cluster", -1)
if target_cluster != -1:
target_hits[target_cluster] += 1
total_hits += 1
if total_hits == 0:
continue
for target_cluster, count in target_hits.most_common():
fraction = count / total_hits
if fraction >= link_threshold:
links.append({
"source_year": year_a,
"source_cluster": cluster_id_a,
"target_year": year_b,
"target_cluster": target_cluster,
"knn_fraction": round(fraction, 2),
})
return links
# Run linking across all consecutive year pairs
# Each year's id_to_embedding is used: doc IDs are local to each year index
years = sorted(df["year"].unique())
all_links = []
for year_a, year_b in zip(years[:-1], years[1:]):
if year_a not in id_to_embedding_by_year or year_b not in id_to_embedding_by_year:
continue
links = link_yearly_clusters(
es, year_a, year_b,
assignments_by_year[year_a],
id_to_embedding_by_year[year_a],
n_samples=10, # more samples per cluster for sparse yearly slices
link_threshold=0.2, # lowered from 0.4 for 10k demo corpus
)
all_links.extend(links)
print(f"{year_a}→{year_b}: {len(links)} links found")Building story chains
Pairwise links tell you that 2010’s “setup” cluster connects to 2011’s. Chains reveal the full arc.
The chain builder is greedy: starting from the earliest unlinked cluster in a chain, it always follows the strongest outgoing link. A chain that spans 10 years means 10 consecutive year-pairs where the embedding neighborhoods overlapped enough to meet the link threshold. That’s a strong signal of thematic continuity.
def build_story_chains(all_links: list[dict]) -> list[list[dict]]:
# Greedily build story chains from pairwise year-to-year links.
# Always follows the strongest outgoing link at each step.
link_index = {}
for link in all_links:
key = (link["source_year"], link["source_cluster"])
if key not in link_index or link["knn_fraction"] > link_index[key]["knn_fraction"]:
link_index[key] = link
targets = {(l["target_year"], l["target_cluster"]) for l in all_links}
sources = {(l["source_year"], l["source_cluster"]) for l in all_links}
chain_starts = sources - targets
chains = []
for start in sorted(chain_starts):
chain = []
current = start
while current in link_index:
link = link_index[current]
chain.append(link)
current = (link["target_year"], link["target_cluster"])
if len(chain) >= 2:
chains.append(chain)
return sorted(chains, key=len, reverse=True)
chains = build_story_chains(all_links)
print("Top story chains by length:")
for chain in chains[:5]:
span = f"{chain[0]['source_year']}→{chain[-1]['target_year']}"
print(f" {span} ({len(chain)+1} years)")Visualizing story chains as a Sankey diagram
Each color represents one story chain end-to-end: nodes and the links between them share the same hue, so you can trace a single theme across years without losing it in the flow. Link opacity is set to 35% so overlapping chains stay legible where they share the same time period.
Color makes three things immediately readable that are hard to see in a table. First, which themes are stable: a single color flowing wide and unbroken across many years means the vocabulary and embedding neighborhood of that topic barely changed across product generations. Second, where divergence happens: a node that fans out into links of two different colors means a previously unified topic split into distinct sub-communities in the following year, often corresponding to a product category fork like corded vs. wireless headphones. Third, where chains merge: a node receiving links of different colors means two previously separate themes converged, typically when a product category consolidates around a dominant design.
import plotly.graph_objects as go
# One distinct color per chain - solid for nodes, semi-transparent for links
CHAIN_COLORS = [
"#e41a1c", "#377eb8", "#4daf4a", "#984ea3",
"#ff7f00", "#a65628", "#f781bf", "#999999",
]
def rgba(hex_color: str, alpha: float) -> str:
"""Convert a hex color to an rgba string with the given alpha."""
h = hex_color.lstrip("#")
r, g, b = int(h[0:2], 16), int(h[2:4], 16), int(h[4:6], 16)
return f"rgba({r},{g},{b},{alpha})"
# Build Sankey nodes and links from the top N chains
TOP_N_CHAINS = 8
top_chains = chains[:TOP_N_CHAINS]
node_labels = []
node_colors = [] # one color per node
node_index = {}
sankey_sources = []
sankey_targets = []
sankey_values = []
sankey_link_labels = []
sankey_link_colors = [] # one color per link
def get_year_cluster_label(year: int, cluster: int) -> str:
"""Fetch significant_text terms for a specific year index + cluster."""
idx = f"reviews-clustering-{year}"
try:
resp = es.search(
index=idx,
body={
"size": 0,
"query": {"term": {"cluster": cluster}},
"aggs": {"sig": {"significant_text": {"field": "text", "size": 2, "filter_duplicate_text": True}}},
},
)
terms = [b["key"] for b in resp["aggregations"]["sig"]["buckets"]]
return ", ".join(terms) if terms else f"cluster {cluster}"
except Exception:
return f"cluster {cluster}"
def node_id(year, cluster, color):
key = (year, cluster)
if key not in node_index:
node_index[key] = len(node_labels)
label = get_year_cluster_label(year, cluster)
node_labels.append(f"{year}: {label}")
node_colors.append(color)
return node_index[key]
for chain_idx, chain in enumerate(top_chains):
color = CHAIN_COLORS[chain_idx % len(CHAIN_COLORS)]
link_color = rgba(color, 0.35)
for link in chain:
src = node_id(link["source_year"], link["source_cluster"], color)
tgt = node_id(link["target_year"], link["target_cluster"], color)
sankey_sources.append(src)
sankey_targets.append(tgt)
sankey_values.append(max(link["knn_fraction"] * 100, 1))
sankey_link_labels.append(f"kNN fraction: {link['knn_fraction']:.0%}")
sankey_link_colors.append(link_color)
fig = go.Figure(go.Sankey(
arrangement="snap",
node=dict(
pad=20,
thickness=15,
label=node_labels,
color=node_colors, # per-node chain color
),
link=dict(
source=sankey_sources,
target=sankey_targets,
value=sankey_values,
label=sankey_link_labels,
color=sankey_link_colors, # per-link chain color, semi-transparent
),
))
fig.update_layout(
title="Review theme story chains across years (each color = one chain, width = kNN fraction)",
font_size=11,
width=1100,
height=700,
)
fig.write_html("story_chains_sankey.html")
fig.show(renderer="notebook")What the temporal shift looks like
The most interesting story chain in this dataset isn’t a single cluster. It’s what happens to the “setup and installation” chain over a decade.
In reviews from 2008 to 2014, this chain’s significant_text terms are dominated by physical setup language: screwdriver, bracket, manual, instructions. From 2016 onward, the same chain's vocabulary shifts to software setup language: app, bluetooth, firmware, update, connect.
The story stayed constant (“setting up a new device is frustrating”) but the vocabulary drifted as the product category evolved from hardware-first to software-first. No classifier would have found this transition without labeled examples that explicitly captured it. The temporal linking found it purely from embedding similarity across yearly indices.
One genuine caveat: cluster labels are vocabulary lists, and interpreting the story arcs is manual work. The significant_text terms give you strong signals, but you'll occasionally find chains that resist a clean narrative. That's not a bug in the method: it means the embedding space found a pattern that doesn't map neatly onto a single human concept.
Cleaning up
Once you’re done testing, delete all indices created during the walkthrough. This removes the global clustering index and all per-year indices in one block:
# Delete the global clustering index
if es.indices.exists(index=INDEX_NAME):
es.indices.delete(index=INDEX_NAME)
print(f"Deleted {INDEX_NAME}")
else:
print(f"{INDEX_NAME} not found — already deleted or never created")
# Delete all yearly indices
years_to_clean = sorted(df["year"].unique()) if "df" in dir() else []
for year in years_to_clean:
year_index = f"reviews-clustering-{year}"
if es.indices.exists(index=year_index):
es.indices.delete(index=year_index)
print(f"Deleted {year_index}")
# Delete the local embeddings cache
from pathlib import Path
cache = Path("embeddings_cache.pkl")
if cache.exists():
cache.unlink()
print(f"Deleted {cache}")
print("Cleanup complete.")If you want to verify nothing was left behind, list all indices matching the prefix:
index_pattern = "reviews-clustering-*"
existing = list(es.indices.get(index=index_pattern).keys())
if existing:
print(f"Indices still present: {existing}")
else:
print("No reviews-clustering-* indices found. Cluster is clean.")Key takeaways
You just built a pipeline that finds structure in a corpus nobody labeled and traced how that structure shifts over 23 years. A few things are worth naming explicitly before you take this further.
Task-specific embeddings are upstream of everything. The choice of task="clustering" over task="retrieval.query" isn't a minor detail. Retrieval embeddings are trained to spread topics broadly so a query can reach documents it isn't identical to: that's the wrong geometry for grouping. Clustering embeddings are trained to pull documents about the same topic into tight, well-separated islands. Swap the task parameter and the rest of the pipeline gets worse in ways that are hard to diagnose: clusters become noisier, significant_text labels become vaguer, and the temporal links weaken. The embedding choice is the one you can't fix downstream.
Elasticsearch is doing more work than you might think. The density-probed clustering algorithm never loads the full embedding matrix into a local process and runs iterative passes over it. Instead, it uses the index as a compute primitive: every probe query, every centroid classification, every temporal kNN link is a request against the cluster. This means the same approach scales to tens of millions of documents without changing a line of code. You add nodes to the cluster, not infrastructure to your pipeline.
The noise label is a feature. Documents that don’t fit any dense cluster are labeled -1 and left unassigned. It's tempting to lower similarity_threshold to assign everything, but the noise documents are telling you something: they're the one-off opinions, very short reviews, and off-topic entries that genuinely don't belong to a coherent theme. Forcing them into a cluster dilutes the clusters that do matter. Keep the noise rate in mind when reading the Sankey: a year with 40% noise and 3 tight clusters is more interesting than a year with 0% noise and 15 vague ones.
significant_text does two jobs: it labels and it quality-gates. The labeling step and the quality gate are the same query: significant_text returns terms if the cluster is coherent, and nothing if it isn't. No terms means no distinguishing vocabulary. Dissolve and move on.
The temporal story chains reward patience. The setup and installation chain that shifts from screwdriver to firmware over a decade is the kind of finding you don't get from a single query or a one-pass analysis. It required running independent clusters per year, linking them, and following the chain. That's also what makes it defensible: the transition wasn't asserted, it was discovered by the model finding its own path through 23 years of embedding similarity.
Ready to run your own experiment?
Here are three places to start:
Change the product category. The Electronics split is one of the largest and noisiest. All_Beauty, Sports_and_Outdoors, and Home_and_Kitchen are smaller, have cleaner temporal coverage, and often surface more legible story chains. Swap "raw_review_Electronics" for any category name in the dataset page and re-run from the embedding step: the embeddings cache means you only pay the API cost once.
Tighten or loosen the cluster boundaries. similarity_threshold=0.82 is conservative for the global index. Lower it to 0.70 and you'll get more clusters with broader membership. Raise it to 0.90 and you'll get fewer, tighter clusters with a higher noise rate. The UMAP plot is the fastest way to see the effect: well-tuned parameters produce compact islands separated by gray; over-aggressive parameters produce a gray sea with a few dots.
Apply this to your own data. The pipeline has no dependency on Amazon reviews. Any corpus of text documents indexed in Elasticsearch works: support tickets, GitHub issues, news articles, research abstracts. The generate_actions function expects a dataframe with text, text_truncated, rating, year, month, and date columns: strip the rating and date fields if your corpus doesn't have them and remove the aggregations and temporal linking steps accordingly. The clustering and labeling work on any dense vector index.
