GenRT: List-aware Reranking-Truncation Joint Model

Paper: List-aware Reranking-Truncation Joint Model for Search and Retrieval-augmented Generation
Authors: Shicheng Xu, Liang Pang, Jun Xu, Huawei Shen, Xueqi Cheng
Venue: WWW '24: Proceedings of the ACM Web Conference 2024

🗂️ Project Structure

GenRT/
├── README.md                   # This file
├── requirements.txt            # Python dependencies
├── .gitignore                  # Git ignore rules
│
├── app.py                      # Flask web application (Phase 4)
├── train_web10k.py             # Training script (Phase 2)
├── test_inference.py           # Inference testing (Phase 3)
├── test_gpu_fit.py             # GPU memory verification
│
├── app/                        # Web application assets
│   ├── templates/
│   │   ├── index.html          # Search interface
│   │   └── result.html         # Results display
│   └── static/
│       └── style.css           # Styling
│
├── model/                      # GenRT model implementation
│   ├── modeling_t5_add_pos_eof.py      # Main model (training)
│   ├── modeling_t5_add_pos_eof_predict.py  # Model (inference)
│   └── embedding_layer.py      # Feature embedding layers
│
├── services/                   # Inference service layer
│   ├── __init__.py
│   └── inference_service.py    # GenRTInference class
│
├── setup/                      # Data preparation scripts
│   ├── create_data_web10k.py   # Process LETOR → pickle
│   ├── merge_tables.py         # Merge feature tables
│   ├── download_t5_clean.py    # Download T5-small
│   ├── fix_t5_model.py         # Fix model format issues
│   └── verify_data.py          # Verify processed data
│
├── data/                       # Raw dataset (LETOR format)
│   └── MSLR-WEB10K/
│       └── Fold1/
│           ├── train.txt       # Training data
│           ├── train.predict   # LambdaMART scores
│           ├── test.txt        # Test data
│           ├── test.predict    # LambdaMART scores
│           ├── vali.txt        # Validation data
│           └── my_model.txt    # Trained LambdaMART model
│
├── dataset/                    # Processed data (pickle files)
│   └── MSLR-WEB10K/
│       └── FOLD1/
│           └── processed_data/
│               ├── train_top40.pkl
│               ├── train_top40_table.pkl
│               ├── test_top40.pkl
│               ├── test_top40_table.pkl
│               └── feature_top40_table.pkl
│
├── checkpoints/                # Trained model weights
│   ├── checkpoint_best.pth     # Best validation checkpoint
│   └── checkpoint_latest.pth   # Latest checkpoint
│
├── T5-small/                   # T5-small base model
│   ├── config.json
│   ├── pytorch_model.bin
│   └── spiece.model
│
└── mytransformers_emb_sim/     # Modified HuggingFace Transformers
    └── src/transformers/
        └── models/t5/          # Custom T5 with GenRT modifications

---

🚀 Quick Start

Prerequisites

Requirement	Version	Notes
Python	3.8+	Tested on 3.10, 3.11
PyTorch	1.10+	CUDA support recommended
GPU	8GB+ VRAM	RTX 4060 or better
Java	JRE 8+	For RankLib

Installation

# 1. Navigate to project directory
cd GenRT

# 2. Create virtual environment (recommended)
python -m venv venv
source venv/bin/activate  # Linux/Mac
# or: venv\Scripts\activate  # Windows

# 3. Install dependencies
pip install -r requirements.txt

# 4. Verify GPU
python test_gpu_fit.py

---

📊 Complete Pipeline

Phase 1: Data Preparation

Step 1.1: Download MSLR-WEB10K Dataset

1. Visit Microsoft Learning to Rank Datasets
2. Download MSLR-WEB10K.zip (~1GB)
3. Extract to data/MSLR-WEB10K/

Expected structure:

data/MSLR-WEB10K/
├── Fold1/
│   ├── train.txt    # ~235MB, 723K lines
│   ├── test.txt     # ~78MB, 241K lines
│   └── vali.txt     # ~78MB, 241K lines
├── Fold2/
│   └── ...
└── Fold5/

Step 1.2: Generate LambdaMART Predictions

Download RankLib-2.18.jar and place in data/MSLR-WEB10K/Fold1/.

cd data/MSLR-WEB10K/Fold1

# Train LambdaMART model (takes ~10-30 minutes)
java -Xmx8g -jar RankLib.jar \
    -train train.txt \
    -ranker 6 \
    -metric2t NDCG@10 \
    -tree 1000 \
    -leaf 20 \
    -save my_model.txt

# Generate prediction scores for each split
java -jar RankLib.jar -load my_model.txt -rank train.txt -score train.predict
java -jar RankLib.jar -load my_model.txt -rank test.txt -score test.predict
java -jar RankLib.jar -load my_model.txt -rank vali.txt -score vali.predict

Step 1.3: Download T5-small Model

cd setup
python download_t5_clean.py

If you encounter safetensors issues:

python fix_t5_model.py

Step 1.4: Process Data to Pickle Format

cd setup

# Convert LETOR format to pickle (creates train/test/vali pkl files)
python create_data_web10k.py

# Merge all feature tables into one
python merge_tables.py

# Verify everything is ready
python verify_data.py

---

Phase 2: Training

💡 Pre-trained Checkpoints Available!
You can skip training by downloading our pre-trained checkpoints:
📥 Download Checkpoints from Google Drive
Extract and place the files in the checkpoints/ directory.

# From project root
python train_web10k.py

Training Options

# Custom hyperparameters
python train_web10k.py --batch_size 32 --epochs 15 --lr 5e-5

# Resume from checkpoint
python train_web10k.py --resume latest --epochs 20

# Reduced memory usage
python train_web10k.py --batch_size 8 --grad_accum 2

Command Line Arguments

Argument	Default	Description
--batch_size	16	Training batch size
--eval_batch_size	32	Evaluation batch size
--epochs	10	Number of training epochs
--lr	1e-4	Learning rate
--weight_decay	0.0005	Weight decay
--lr_decay	0.95	LR decay per epoch
--resume	None	'latest', 'best', or checkpoint path
--no_amp	False	Disable mixed precision
--grad_accum	1	Gradient accumulation steps
--early_stop	5	Early stopping patience

Expected Output

🚀 GenRT Trainer Initialization
========================================
✅ GPU: NVIDIA GeForce RTX 4060
   Memory: 8.0 GB

📂 Loading data...
   ✅ Train: 10,000 queries
   ✅ Test: 3,000 queries

🤖 Initializing model...
   ✅ Model parameters: 12.5M

📈 Training Progress:
Epoch 1/10: 100%|██████████| 625/625 [05:32<00:00]
   Train Loss: 2.3456
   Test NDCG@10: 0.4521
   ✅ Best model saved!

---

Phase 3: Inference Testing

python test_inference.py

Expected output:

🧪 Testing GenRT Inference Service
========================================
🔧 Initializing inference service...
   ✅ Loaded 3000 queries

📋 Available sample QIDs: [1, 2, 3, ...]

🔍 Testing prediction for QID: 1
✅ Prediction successful!

📊 Results:
   NDCG@10 Baseline: 0.4231
   NDCG@10 GenRT:    0.4687
   Improvement:      +0.0456

---

Phase 4: Web Application

python app.py

Open browser: http://localhost:5000

Features

• Enter Query ID to see reranking results
• Side-by-side comparison: Baseline vs GenRT
• Visual indicators for promoted/demoted documents
• NDCG metrics display
• Truncation point visualization

---

🔧 API Reference

Python API

from services.inference_service import GenRTInference

# Initialize service (loads model once)
inference = GenRTInference(
    checkpoint_path="./checkpoints/checkpoint_best.pth",
    device="cuda"  # or "cpu"
)

# Get available query IDs
all_qids = inference.get_available_qids()
sample_qids = inference.get_sample_qids(n=10)

# Predict for a single query
result = inference.predict(qid=12345)

# Result structure
{
    'qid': 12345,
    'baseline_ranking': [
        {'doc_id': 0, 'lambdamart_score': 1.234, 'relevance_label': 2},
        ...
    ],
    'genrt_ranking': [
        {'doc_id': 5, 'lambdamart_score': 0.987, 'relevance_label': 3, 'original_rank': 6},
        ...
    ],
    'truncation_point': 15,
    'metrics': {
        'ndcg@5_baseline': 0.412,
        'ndcg@5_genrt': 0.456,
        'ndcg@10_baseline': 0.423,
        'ndcg@10_genrt': 0.468,
        'improvement': 0.045
    },
    'num_docs_baseline': 40,
    'num_docs_genrt': 15
}

# Batch prediction
results = inference.predict_batch(qids=[1, 2, 3, 4, 5])

REST API Endpoints

Endpoint	Method	Description
/	GET	Home page with search interface
/search	POST	Submit query and get results
/api/predict/<qid>	GET	Get prediction JSON for QID
/api/qids	GET	Get list of available QIDs

# Get prediction for QID
curl http://localhost:5000/api/predict/12345

# Get available QIDs (optional: ?n=100)
curl http://localhost:5000/api/qids?n=50

---

📈 Model Architecture

Overview

Input: [Query Features] + [Doc1 Features] + [Doc2 Features] + ... + [DocN Features]
                                    ↓
                    ┌───────────────────────────────┐
                    │   Global Dependency Encoder   │
                    │   (Multi-Head Self-Attention) │
                    └───────────────────────────────┘
                                    ↓
                    ┌───────────────────────────────┐
                    │  Sequential Dependency Decoder │
                    │   (Step-by-step Generation)   │
                    └───────────────────────────────┘
                                    ↓
Output: Reranked List + Truncation Point

Global Dependency Encoder

• Input Embedding: 136 LETOR features + 1 LambdaMART score = 137-dim vector
• Transfer Layer: MLP with Swish activation → hidden size
• Context Modeling: Multi-Head Self-Attention captures list-level dependencies
• Residual Connection: Preserves original features while adding context

Sequential Dependency Decoder

• Step-by-step Generation: Outputs one document per step
• Truncation Decision: Binary classification at each step
• Cross-Attention: Attends to encoder output for document selection
• Position Encoding: Relative position encoding for order awareness

Loss Functions

1. Reranking Loss ($\mathcal{L}_R$):

   • Step-adaptive attention loss (weights top positions more)
   • Step-by-step lambda loss (pairwise ranking)

2. Truncation Loss ($\mathcal{L}_T$):

   • Binary classification with TDCG-based soft labels
   • Reward Augmented Maximum Likelihood (RAML)

---

📚 Citation

If you use this code, please cite the original paper:

@inproceedings{xu2024genrt,
  title={List-aware Reranking-Truncation Joint Model for Search and Retrieval-augmented Generation},
  author={Xu, Shicheng and Pang, Liang and Xu, Jun and Shen, Huawei and Cheng, Xueqi},
  booktitle={Proceedings of the ACM Web Conference 2024},
  pages={1028--1037},
  year={2024},
  organization={ACM}
}

---

🐛 Troubleshooting

Common Issues

Issue	Solution
CUDA out of memory	Reduce --batch_size to 8 or 4; Use --grad_accum 2
T5 model loading error	Run python setup/fix_t5_model.py; Ensure pytorch_model.bin exists
Missing .predict files	Follow Step 1.2 to generate with RankLib
Feature table not found	Run python setup/merge_tables.py
Web app fails to start	Check checkpoint and processed data exist
Java not found	Install JRE 8+; Add to PATH
Import errors	Ensure virtual environment is activated

Verification Commands

# Check GPU
python -c "import torch; print(torch.cuda.is_available())"

# Check data files
python setup/verify_data.py

# Check model loads
python test_gpu_fit.py

---

📝 License

This project is for educational and research purposes only.

---

🙏 Acknowledgments

• Original GenRT paper authors
• Microsoft Research for MSLR datasets
• HuggingFace for Transformers library
• RankLib for LambdaMART implementation