1. Preprocess

Download and prepare datasets, then optionally clean the text.

# Download datasets (LongEmbed, LitSearch, MLDR)
python prepare.py

# (Optional) Clean text data
python -m utils.preprocess \
  --input datasets/coliee_task1/corpus.json \
  --output datasets/coliee_task1/corpus_clean.json \
  --text_field "text" \
  --remove_structural \
  --remove_metadata \
  --remove_dates

2. Dense Retrieval

Build a retrieval index (BM25, BGE, or E5) over the corpus.

# BM25 (lexical)
python -m dense_retrieval.retrieval.build_index \
  --index_type bm25 \
  --key 2025 \
  --dataset_path coliee_task1

# BGE (dense, with chunking)
python -m dense_retrieval.retrieval.build_index \
  --index_type bge \
  --key 2025 \
  --dataset_path coliee_task1 \
  --batch_size 32 \
  --chunk \
  --chunk_method maxsim

# E5 (dense, with chunking)
python -m dense_retrieval.retrieval.build_index \
  --index_type e5 \
  --key 2025 \
  --dataset_path coliee_task1 \
  --batch_size 32 \
  --chunk \
  --chunk_method average

Indices are saved to retrieval_indices/.

3. Training T5 model

Train T5 model with

# Minimun (only --data_path)
python reranking.t5_train.py --data_path ./data/t5_train_data.csv

# Full cònig
python reranking.t5_train.py \
    --model_name google/flan-t5-large \
    --data_path ./data/t5_train_data.csv \
    --output_dir ./models/my_model \
    --num_train_epochs 5 \
    --learning_rate 1e-4 \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 32 \
    --seed 42

4. Reranking

# Extract feature from each case
python reranking.file_code.py --all

# Build pairs of query case and target cases
python reranking.pairs_code.py --all --files_pickle /path/to/files.pkl