Welcome to the rankers
repository! This package provides tooling for training neural retrievers using various ranking models. The package is designed to mimic the underlying interface of transformers
to get you started quickly.
- Flexible Model Training: Use the
RankerTrainer
class to train common ranking architectures. - Data Handling:
ir_datasets
integration and ir-specific data processing in a familiar format. - Transformers Integration: Built on top of the Hugging Face
transformers
library for easy model loading and training. - Experiment Tracking: Integrated with
wandb
for tracking experiments and hyperparameters.
pip install rankers
pip install -U git+https://github.com/Parry-Parry/rankers.git
To get started, check out the examples in the examples
directory. Here's a brief overview of what you can find:
train.bert.cat.py
: An example script for training a BERT-based ranking model. It includes setting up data arguments, model arguments, and training arguments, and demonstrates how to train and save the model.
Clone the repository and install the required dependencies:
git clone https://github.com/Parry-Parry/rankers.git
cd rankers
pip install -r requirements.txt
Models can be instantiated like transformers
models, for example, a bi-encoder (dot)
model = Dot.from_pretrained("bert-base-uncased")
Models are converted into rankers in PyTerrier:
ranker = model.to_pyterrier()
Run the example training script:
python examples/train.bert.cat.py --model_name_or_path bert-base-uncased --training_data path/to/data --output_dir path/to/save/model
We welcome contributions! Please read our contributing guidelines for more details.
This project is licensed under the Apache 2.0 License. See the LICENSE file for details.
Happy training!