K-NRM
This is the implementation of the Kernel-based Neural Ranking Model (K-NRM) model from paper End-to-End Neural Ad-hoc Ranking with Kernel Pooling.
If you use this code for your scientific work, please cite it as (bibtex):
C. Xiong, Z. Dai, J. Callan, Z. Liu, and R. Power. End-to-end neural ad-hoc ranking with kernel pooling.
In Proceedings of the 40th International ACM SIGIR Conference on Research & Development in Information Retrieval.
ACM. 2017.
Requirements
- Tensorflow 0.12
- Numpy
- traitlets
Coming soon: K-NRM with Tensorflow 1.0
Guide To Use
Configure: first, configure the model through the config file. Configurable parameters are listed here
Training : pass the config file, training data and validation data as
python ./knrm/model/model_knrm.py config-file\
--train \
--train_file: path to training data\
--validation_file: path to validation data\
--train_size: size of training data (number of training samples)\
--checkpoint_dir: directory to store/load model checkpoints\
--load_model: True or False. Start with a new model or continue training
Testing: pass the config file and testing data as
python ./knrm/model/model_knrm.py config-file\
--test \
--test_file: path to testing data\
--test_size: size of testing data (number of testing samples)\
--checkpoint_dir: directory to load trained model\
--output_score_file: file to output documents score\
Relevance scores will be output to output_score_file, one score per line, in the same order as test_file. We provide a script to convert scores into trec format.
./knrm/tools/gen_trec_from_score.py
Data Preperation
All queries and documents must be mapped into sequences of integer term ids. Term id starts with 1.
-1 indicates OOV or non-existence. Term ids are sepereated by ,
Training Data Format
Each training sample is a tuple of (query, postive document, negative document)
query \t postive_document \t negative_document \t score_difference
Example: 177,705,632 \t 177,705,632,-1,2452,6,98 \t 177,705,632,3,25,14,37,2,146,159, -1 \t 0.119048
If score_difference < 0
, the data generator will swap postive docment and negative document.
If score_difference < lickDataGenerator.min_score_diff
, this training sample will be omitted.
We recommend shuffling the training samples to ease model convergence.
Testing Data Format
Each testing sample is a tuple of (query, document)
q \t document
Example: 177,705,632 \t 177,705,632,-1,2452,6,98
Configurations
Model Configurations
BaseNN.n_bins
: number of kernels (soft bins) (default: 11. One exact match kernel and 10 soft kernels)Knrm.lamb
: defines the guassian kernels' sigma value. sigma = lamb * bin_size (default:0.5 -> sigma=0.1)BaseNN.embedding_size
: embedding dimension (default: 300)BaseNN.max_q_len
: max query length (default: 10)BaseNN.max_d_len
: max document length (default: 50)DataGenerator.max_q_len
: max query length. Should be the same asBaseNN.max_q_len
(default: 10)DataGenerator.max_d_len
: max query length. Should be the same asBaseNN.max_d_len
(default: 50)BaseNN.vocabulary_size
: vocabulary size.DataGenerator.vocabulary_size
: vocabulary size.
Data
Knrm.emb_in
: initial embeddingsDataGenerator.min_score_diff
: minimum score differences between postive documents and negative ones (default: 0)
Training Parameters
BaseNN.bath_size
: batch size (default: 16)BaseNN.max_epochs
: max number of epochs to trainBaseNN.eval_frequency
: evaluate model on validation set very this steps (default: 1000)BaseNN.checkpoint_steps
: save model very this steps (default: 10000)Knrm.learning_rate
: learning rate for Adam Opitmizer (default: 0.001)Knrm.epsilon
: epsilon for Adam Optimizer (default: 0.00001)
Efficiency
During training, it takes about 60ms to process one batch on a single-GPU machine with the following settings:
- batch size: 16
- max_q_len: 10
- max_d_len: 50
- vocabulary_size: 300K
Smaller vocabulary and shorter documents accelerate the training.
Click2Vec
We also provide the click2vec model as described in our paper.
./knrm/click2vec/generate_click_term_pair.py
: generate <query_term, clicked_title_term> pairs./knrm/click2vec/run_word2vec.sh
: call Google's word2vec tool to train click2vec.
Cite the paper
If you use this code for your scientific work, please cite it as:
C. Xiong, Z. Dai, J. Callan, Z. Liu, and R. Power. End-to-end neural ad-hoc ranking with kernel pooling.
In Proceedings of the 40th International ACM SIGIR Conference on Research & Development in Information Retrieval.
ACM. 2017.
@inproceedings{xiong2017neural,
author = {{Xiong}, Chenyan and {Dai}, Zhuyun and {Callan}, Jamie and {Liu}, Zhiyuan and {Power}, Russell},
title = "{End-to-End Neural Ad-hoc Ranking with Kernel Pooling}",
booktitle = {Proceedings of the 40th International ACM SIGIR Conference on Research & Development in Information Retrieval},
organization = {ACM},
year = 2017,
}