Vector Search Experiment

This repository contains experiments for dense retrieval fine-tuning with synthetic queries, cross-encoder reranking, and listwise distillation. The work is based on the CADET/Listwise Distillation pipeline and explores whether newer embedding models such as Nomic Embed and ModernBERT Embed benefit from standard InfoNCE fine-tuning or from a combined InfoNCE + listwise distillation objective.

The main experiment flow is:

1. Evaluate a base dense retriever on BEIR/MS MARCO-style retrieval tasks.
2. Generate synthetic queries for each corpus.
3. Retrieve top-k passages for those synthetic queries.
4. Rerank retrieved passages with a cross-encoder teacher such as RankT5-3B.
5. Normalize reranker scores and build train/dev JSONL files.
6. Fine-tune retrievers with InfoNCE or InfoNCE + KL/listwise loss.
7. Re-evaluate with NDCG@10 and Recall@100.

Repository Layout

.
|-- train_dpr.py                         # SentenceTransformers DPR-style training
|-- model-distillation.py                # Teacher/student embedding distillation example
|-- evaluate_beir.py                     # BEIR evaluation with a PLAID index
|-- evaluate_dpr.py                      # MTEB evaluation helper
|-- run_beir_eval.sh                     # Batch wrapper for evaluate_beir.py
|-- run_eval_without_reranker.sh         # Dense retrieval evaluation wrapper
|-- run_eval_with_reranker.sh            # Dense retrieval + reranking evaluation wrapper
|-- setup_dev_env.sh                     # Conda + Pyserini/Anserini dev setup
|-- setup_eval_env.sh                    # CUDA/Linux eval environment setup
|-- setup_mac_eval_env.sh                # macOS eval environment setup
|-- requirements.txt
`-- listwise_distillation/
    |-- encoding/                        # Corpus encoding and dense retrieval evals
    |-- query_generation/                # Synthetic query generation/filtering
    |-- reranker/                        # Reranking and training data creation
    `-- training_scripts/                # InfoNCE and joint-loss training scripts

Environment Setup

For model training, a Python virtual environment is usually enough:

python -m venv venv
source venv/bin/activate
pip install --upgrade pip
pip install -r requirements.txt

Optional dependencies used by some GPU runs:

MAX_JOBS=4 pip install flash-attn --no-build-isolation
wandb login
huggingface-cli login

For retrieval evaluation with Pyserini/Anserini, use the conda setup scripts:

chmod +x install_miniconda.sh
./install_miniconda.sh
source ~/.bashrc

chmod +x setup_dev_env.sh
bash setup_dev_env.sh
conda activate pyserini-dev

Alternative eval environments:

# Linux/CUDA-oriented setup
bash setup_eval_env.sh

# macOS/CPU-oriented setup
bash setup_mac_eval_env.sh

For long training/evaluation jobs, run inside tmux so jobs continue after the terminal disconnects.

Data Preparation

The listwise training scripts expect train/dev files in this format:

final_data/beir.{retriever}.{dataset}.train.generated_queries.listwise.jsonl
final_data/beir.{retriever}.{dataset}.dev.generated_queries.listwise.jsonl

Each JSONL row should contain a query and a passages list. Each passage should include at least:

{
  "query": "synthetic or human query text",
  "passages": [
    {
      "docid": "document-id",
      "text": "document text",
      "score": 0.91
    }
  ]
}

The current training scripts use the top 20 passages per query by default.

Generate Synthetic Queries

Run the bundled query generation wrapper:

chmod +x listwise_distillation/query_generation/query_gen_scripts/run_all_query_generation.sh
./listwise_distillation/query_generation/query_gen_scripts/run_all_query_generation.sh

This writes files such as:

generated_queries/{dataset}_generated_queries_{query_type}.tsv

Supported query types in the current scripts include:

keywords, titles, claims, questions, random, msmarco

Retrieve Top-20 Passages for Generated Queries

Edit dataset, model, model_name, and model_prefix in listwise_distillation/query_generation/model_retrieve_20.sh, then run:

bash listwise_distillation/query_generation/model_retrieve_20.sh

Expected output:

retrieval_runs/run.{model_name}.{dataset}.generated-queries-{query_type}_20.txt

Filter, Rerank, Normalize, and Split

The working pipeline described in notes.md is:

1. Filter generated-query runs to retain examples whose positive passage appears in the top 20.
2. Convert filtered TREC runs to JSONL for reranking.
3. Rerank with the cross-encoder teacher.
4. Filter again using reranker rank/score.
5. Normalize reranker scores.
6. Combine query types and split into train/dev files.

The relevant scripts are:

python listwise_distillation/query_generation/retriever_filtering_step.py
python listwise_distillation/query_generation/generate_jsonl_for_reranking.py
bash listwise_distillation/reranker/reranking_scripts/run_reranking_for_generated_queries.sh
python listwise_distillation/reranker/filter_by_reranker.py
python listwise_distillation/reranker/normalize_scores.py
python listwise_distillation/reranker/create_train_dev_data.py

Most of these scripts currently use hard-coded dataset/model lists. Check the top-level variables in each script before launching a full run.

Training

InfoNCE Training on Synthetic Query Data

Use this for pure contrastive fine-tuning:

python listwise_distillation/training_scripts/new_train_model_infonce_loss.py \
  --dataset scifact

The script reads:

final_data/beir.modernbert-embed-base.scifact.train.generated_queries.listwise.jsonl
final_data/beir.modernbert-embed-base.scifact.dev.generated_queries.listwise.jsonl

and saves the best checkpoint to:

models/modernbert-embed-base_scifact-infonce-loss

Joint InfoNCE + Listwise/KL Training

Use this for the combined objective:

python listwise_distillation/training_scripts/new_train_model_joint_loss.py \
  --dataset scifact

The script reads:

final_data/beir.modernbert-embed-base.scifact.train.generated_queries.listwise.jsonl
final_data/beir.modernbert-embed-base.scifact.dev.generated_queries.listwise.jsonl

and saves the best checkpoint to:

models/modernbert-embed-base_scifact-joint-loss-w0.8

Current defaults for both listwise training scripts live in the Config class:

model_name_or_path = nomic-ai/modernbert-embed-base
retriever = modernbert-embed-base
batch_size = 16
list_length = 20
accumulation_steps = 256
lr = 2e-4
query_maxlength = 64
text_maxlength = 512

To train another retriever or use another loss weight, edit Config before running. For joint training, contrastive_loss_weight = 0.1 means:

loss = 0.1 * InfoNCE + 0.9 * KL/listwise

Common Gotchas

• Several scripts currently use hard-coded model names, dataset lists, and output paths. Check the top of each script before launching expensive runs.
• The listwise training scripts require CUDA and call .cuda() directly.
• Some encoder/evaluation code requires trust_remote_code=True for Nomic and ModernBERT embedding models.
• Pyserini/Anserini evaluation requires Java, Maven, FAISS, and the compiled evaluation helpers.