dtunkelang/bag-of-documents

Bag-of-Documents: Product Search Dataset

• —Blog post: Distilling Retrieval Pipelines to a Single Embedding Model
• —Live demo: huggingface.co/spaces/dtunkelang/bag-of-documents-demo
• —Code: github.com/dtunkelang/bag-of-documents

Dataset Description

A large-scale bag-of-documents dataset for e-commerce product search, built on Amazon product data. Each search query is represented as a distribution of relevant products in embedding space, captured by a centroid vector and a specificity score.

• —Centroid: the mean direction in product embedding space, representing what the query means
• —Specificity: the tightness of the distribution (high = narrow query like "hp laptop 16gb ram", low = broad query like "laptop")

Dataset Summary

The bags themselves were computed against the broader 6M source catalog (per the original blog post). The 1.2M ESCI subset is what's loaded in combined_index/index.faiss here — it's the catalog the rerank result table below was evaluated on, and the catalog the live Space demo serves.

Supported Tasks

• —Retrieval model training: fine-tune an embedding model to predict bag centroids from query text, producing a query encoder specialized for product search
• —Reranking: use BoD-trained query encoders as bi-encoder rerankers on top of base FAISS retrieval (see "BoD as reranker" below — the deployable architecture)
• —Specificity prediction: predict whether a query is broad or narrow using kNN on bag centroids
• —Search evaluation: compare retrieval models using bag centroids as ground-truth query representations

Languages

English (US)

Dataset Structure

combined_index/
├── index.faiss                          # FAISS HNSW index over the catalog
├── titles.json                          # Product titles, parallel to FAISS positions
├── rerank_A.vecs.fp16.npy               # cached product embeddings under Reranker A
├── rerank_B.vecs.fp16.npy               # cached product embeddings under Reranker B
├── rerank_G.vecs.fp16.npy               # cached product embeddings under Reranker G
├── tantivy_index/                       # tantivy BM25 index (legacy, en_stem)
├── bm25s_index/                         # bm25s BM25 index (k1=0.3, b=0.6) - SOTA retriever
└── spell_vocab.json                     # NEW: catalog-vocab spell-correction dict (172K tokens)
query_model/                             # Original BoD-as-retriever (cosine-loss MNRL)
query_model_6m_mnrl/                     # Reranker A - 75K queries × 6M corpus, MNRL
query_model_hardneg/                     # Reranker B - qrels-based bags + hard negatives
query_model_esci_supervised/             # Reranker G - ESCI E-vs-I triplets, MNRL
bags.jsonl                               # Bag centroids + specificity for each query
queries.jsonl                            # Source queries
eval/
└── regime_queries.jsonl                 # 45-query per-regime eval harness

Bags (JSONL)

Each bag is a JSON object:

{
  "query": "wireless keyboard",
  "num_results": 42,
  "query_vector": [0.023, -0.051, ...],
  "specificity": 0.95,
  "results": [
    {"title": "Logitech K380 Multi-Device Bluetooth Keyboard"},
    ...
  ]
}

Cached reranker embeddings

combined_index/rerank_A.vecs.fp16.npy and rerank_B.vecs.fp16.npy are 1,215,851 × 384 fp16 numpy arrays, parallel to combined_index/titles.json. Loading them lets the rerank pipeline skip the live candidate-encoding step — only the query is encoded at runtime, candidates are looked up by FAISS-returned index. This is what makes the reranker deployable at sub-100ms latency on commodity hardware.

Architecture progression

The originally-published architecture treats BoD as a retrieval-stage model (single encoder + FAISS, end of story). Iterating on loss function (cosine-to-centroid → MNRL), training scale (full 6M signal), and stage (retrieval vs rerank vs hybrid) produced a series of improvements. The progression below is measured on the full 22,458-query ESCI test set, R@10 with E+S as relevant, nDCG@10 with E=1.0/S=0.1, Kretrieve=100, Keval=10.

Five things to note:

• —*MNRL-trained BoD beats base as a retriever*** (B vs A: +2.50pp R@10). The original cosine-distilled BoD-as-retriever loses on this stricter benchmark; the MNRL-trained variant doesn't.
• —BM25 alone is competitive with the dense rerank stack (H ≈ E, within rounding). On entity-heavy product queries, lexical matching does most of the work.
• —MNRL retrieval is dead weight in the SOTA pipeline. Adding MNRL retrieval to the candidate pool dilutes BM25's lexically-anchored hits with semantically-near-but-irrelevant ones. The deployable architecture has no dense retrieval lane in the inference path.
• —BM25 hyperparameter tuning matters. Default Lucene/tantivy params (k1=1.2, b=0.75) assume long natural-language documents. Amazon product titles are short and keyword-stuffed. A sweep finds (k1=0.3, b=0.6) optimal — early term-frequency saturation, moderate length normalization. The bm25s rows show the +0.83pp / +0.16pp / +0.29pp lift the parameter swap brings on top of every downstream rerank stage.
• —Cross-encoder fusion is the biggest single lift at the quality tier. The fast SOTA (CC3-50 + spell) → CC4-100 (medium, + LiYuan CE @ w=0.25) is +0.49pp R@10 / +2.32pp E@1. CC4-100 → CC5-100 (quality, + BGE-reranker-v2-m3 fused 3-way mean) is +1.00pp R@10 / +2.95pp E@1 — both deltas statistically significant via 1000-resample paired bootstrap. BGE-reranker (XLM-RoBERTa-large, ~568M params, BEIR-tested) is ~6× slower than LiYuan but adds orthogonal signal worth the cost.

The bi-encoder rerank fuses three encoders (query_model_6m_mnrl, query_model_hardneg, query_model_esci_supervised) by averaging their cosine similarities. With cached product embeddings (rerank_A.vecs.fp16.npy, rerank_B.vecs.fp16.npy, rerank_G.vecs.fp16.npy), only the query is encoded live; candidate vectors are looked up by index. The retrieval lane uses bm25s (k1=0.3, b=0.6) with optional pre-BM25 catalog-vocab spell correction (spell_vocab.json). Cross-encoders (LiYuan/Amazon-Cup-Cross-Encoder-Regression and BAAI/bge-reranker-v2-m3) are loaded directly from HF Hub at runtime — no dataset artifact.

Dataset Creation

Source Data

• —Products: Amazon Reviews 2023 (McAuley Lab, UCSD; data collected 1996-2023). 20% random sample of the full catalog across all 33 categories (~6M of ~30M unique products).
• —Queries: All 75K US-locale queries from the Amazon Shopping Queries Dataset (ESCI, KDD Cup 2022) — real Amazon search queries spanning all product categories.

Bag Construction Pipeline

Query text
  -> Hybrid retrieval: keyword (tantivy AND with relaxation) + FAISS embedding similarity
  -> Cross-encoder scoring: ESCI RoBERTa CE scores ALL candidates, threshold 0.3
  -> Top 50 passing candidates -> encode -> bag centroid + specificity

The cross-encoder is LiYuan/Amazon-Cup-Cross-Encoder-Regression, a RoBERTa model trained on ESCI data for the KDD Cup 2022 competition.

Fine-Tuning the Reranker Models

• —query_model_6m_mnrl: trained with MultipleNegativesRanking loss on bags from the full 6M-product corpus
• —query_model_hardneg: trained with MNRL on qrels-derived bags + hard-mined negative products
• —Both share the same all-MiniLM-L6-v2 base model

The cached vec files were produced by encoding all 1.2M ESCI products under each reranker (precompute_rerank_vecs.py in the code repo) and saving as fp16 numpy.

Citation

@misc{tunkelang2026bagdocs,
  title={Bag-of-Documents: Product Search Dataset},
  author={Daniel Tunkelang and Aritra Mandal},
  year={2026},
  url={https://huggingface.co/datasets/dtunkelang/bag-of-documents}
}

Related Work

• —Tunkelang, D. Distilling Retrieval Pipelines to a Single Embedding Model. 2026.
• —Tunkelang, D. Modeling Queries as Bags of Documents. 2024.
• —Reddy, C.K. et al. Shopping Queries Dataset: A Large-Scale ESCI Benchmark for Improving Product Search. KDD Cup 2022.
• —McAuley Lab. Amazon Reviews 2023.

License

MIT