DOCSEARCHbyAlgolia
Submit Docs

allRank Framework

repository·master·Indexed 21 days ago

https://github.com/allegro/allrank

A PyTorch-based framework for training neural Learning-to-Rank (LTR) models. It supports pointwise, pairwise, and listwise loss functions, including ListNet, LambdaRank, and NeuralNDCG, as well as FC and Transformer scoring functions. The framework provides tools for training via config.json, implementing custom loss functions, simulating click-through data, and preprocessing MSLR-WEB30K datasets.

Tokens
2.5K
Snippets
9
Records
14
Agent score
77%

What's inside allRank

  1. Rescale NDCG results for different IDCG == 0 treatments

    master

    The allRank implementation assumes NDCG = 1 for queries with no relevant items (where IDCG == 0), following XGBoost and LightGBM defaults.

    If you need to rescale these results to reflect an NDCG = 0 treatment for blank queries, use the following formula:

    avg_ndcg - (num_blank / num_all)

    Below is the reference table for the number of blank queries across MSLR-WEB30K folds:

    FoldTrain (blank/all)Test (blank/all)Validation (blank/all)
    1602 / 18919189 / 6306191 / 6306
    2587 / 18918206 / 6307189 / 6306
    3581 / 18918195 / 6306206 / 6307
    4586 / 18919201 / 6306195 / 6306
    5590 / 18919191 / 6306201 / 6306

  2. Preprocess MSLR-WEB30K data using normalize_features.py

    master

    To reproduce the results from the referenced papers, you must preprocess each MSLR-WEB30K fold separately using the normalize_features.py script. The script performs log-transformation and standardisation on features.

    Run the script with the following parameters:

    • --ds_path: The path to the specific MSLR-WEB30K fold you wish to preprocess.
    • --features_without_logarithm: A list of features to be standardised without a prior log-transform.
    • --features_negative: A list of features that require a constant shift to ensure strict positivity before the log-transform is applied.

    Note: The script contains default values for --features_without_logarithm and --features_negative that match the settings used in the original papers.

    python normalize_features.py --ds_path /path/to/fold --features_without_logarithm [list] --features_negative [list]

  3. Apply a click-model to trained models

    master

    After training an allRank model, you can simulate click-through data using allrank/rank_and_click.py. This process ranks all slates from the dataset specified in your config and applies a click model (configured in the config) to generate a new libsvm dataset.

    This generated dataset can then be used as input for training subsequent models.

    Arguments:

    • --input-model-path: Path to the trained model weights file.
    • --roles: A comma-separated list of dataset roles to process (e.g., train,valid).
    • --config_file_name: Path to your configuration file.
    • --run_id: The name of your experiment.
    • --job_dir: The directory to save results.
    python allrank/rank_and_click.py --input-model-path <path_to_the_model_weights_file> --roles <comma_separated_list_of_ds_roles_to_process> --config_file_name allrank/config.json --run_id <the_name_of_your_experiment> --job_dir <the_place_to_save_results>

  4. Implement a custom loss function

    master

    To add a custom loss function to allRank:

    1. Implement a function that accepts two tensors (model prediction and ground truth) as input.
    2. Place this function in the losses package and ensure it is exposed at the package level.
    3. Reference the loss in your config.json using the loss key with the function's name and any required args.

    Example configuration for a custom loss:

    "loss": {
    "name": "yourLoss",
    "args": {
        "arg1": val1,
        "arg2": val2
    }
}

  5. Select the correct architecture version (GPU vs CPU)

    master

    Because PyTorch binaries differ between GPU and CPU, you must select the appropriate architecture version when building the Docker image.

    Use the arch_version build argument during the Docker build process. When running the run_example.sh script, pass gpu or cpu as the first command-line argument. If no argument is provided, cpu is the default.

    # Building the image
docker build --build-arg arch_version=${ARCH_VERSION} .

# Running the example with GPU
./run_example.sh gpu

  6. Run the allRank quickstart example

    master

    To quickly see how allRank works, you can use the provided run_example.sh script. This script requires Docker. It generates dummy ranking data in libsvm format and trains a Transformer model using a default config.json.

    After execution, dummy data is stored in the dummy_data directory and experiment results are in the test_run directory.

    ./run_example.sh [gpu|cpu]

  7. Train a custom model using config.json

    master

    To train your own model, you must define your hyperparameters (model definition, data location, loss, metrics, etc.) in a config.json file.

    Data Requirements:

    • Your training data must be in libsvm format.
    • Following MSLR-WEB30K convention, the training file should be named train.txt.
    • You can specify the names of validation or test datasets in the config.

    Execution: Run the allrank/main.py script with the following arguments:

    • --config_file_name: Path to your configuration file.
    • --run_id: A unique name for your experiment.
    • --job_dir: The directory where results will be saved.

    Results are saved to <job_dir>/results/<run_id>.

    python allrank/main.py --config_file_name allrank/config.json --run_id <the_name_of_your_experiment> --job_dir <the_place_to_save_results>

  8. Configure model and loss function settings

    master

    Configuration files for reproducing specific research papers are located in the configs/ directory.

    Available Configurations

    For Context-Aware Learning to Rank with Self-Attention:

    • MLP and context-aware rankers using ordinal loss (best-performing context-aware).
    • MLP and context-aware rankers using NDCG_Loss2++ (best-performing MLP).

    For NeuralNDCG: Direct Optimisation of a Ranking Metric via Differentiable Relaxation of Sorting:

    • Context-aware rankers using NeuralNDCG@max (best-performing NeuralNDCG variant).
    • Context-aware rankers using ApproxNDCG (direct NDCG optimization).
    • Context-aware rankers using LambdaRank@max (best-performing loss function).

    Setup Requirements

    1. Dataset Path: You must manually update the dataset path in the configuration files to point to your local data.
    2. Customization: These files can be adapted for other model/loss combinations by modifying the model and loss function parameters.

  11. run

    master

    The main execution function that orchestrates the training pipeline. It performs the following steps:

    1. Reproducibility: Sets seeds for torch and numpy.
    2. Path Management: Initializes PathsContainer and creates output directories.
    3. Configuration: Loads the JSON config via Config.from_json and saves a copy to the output directory.
    4. Data Loading: Loads LibSVM datasets and creates PyTorch DataLoaders.
    5. Model Setup: Instantiates the model using make_model and wraps it in CustomDataParallel if multiple GPUs are available.
    6. Optimization: Dynamically instantiates the optimizer, loss function, and LR scheduler based on the config.
    7. Training: Executes the training loop using fit.
    8. Post-processing: Dumps experiment results, uploads to GS if necessary, and asserts expected metrics.
    from allrank.main import run

# This is typically called via the CLI, but can be invoked directly
run()

  12. parse_args

    master

    Parses command-line arguments for the allRank training pipeline. It returns a Namespace object containing job_dir, run_id, and config_file_name.

    from allrank.main import parse_args

args = parse_args()
# args contains: job_dir, run_id, config_file_name