|
--- |
|
language: |
|
- en |
|
tags: |
|
- search |
|
- ranking |
|
- vespa |
|
license: mit |
|
|
|
base_model: "cross-encoder/ms-marco-MiniLM-L-6-v2" |
|
--- |
|
|
|
# MS Marco Ranking with ColBERT on Vespa.ai |
|
|
|
Model is based on [ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT](https://arxiv.org/abs/2004.12832). |
|
This BERT model is based on [cross-encoder/ms-marco-MiniLM-L-6-v2](https://huggingface.co/cross-encoder/ms-marco-MiniLM-L-6-v2) and trained using the |
|
original [ColBERT training routine](https://github.com/stanford-futuredata/ColBERT/). |
|
|
|
This model has 22.3M trainable parameters and is approximately 2x faster than |
|
[vespa-engine/colbert-medium](https://huggingface.co/vespa-engine/colbert-medium) and with better or on pair MRR@10 on dev. |
|
|
|
The model weights have been tuned by training using a randomized sample of MS Marco training triplets |
|
[MSMARCO-Passage-Ranking](https://github.com/microsoft/MSMARCO-Passage-Ranking). |
|
|
|
To use this model with vespa.ai for MS Marco Passage Ranking, see |
|
[MS Marco Ranking using Vespa.ai sample app](https://github.com/vespa-engine/sample-apps/tree/master/msmarco-ranking). |
|
|
|
# MS Marco Passage Ranking |
|
|
|
| MS Marco Passage Ranking Query Set | MRR@10 ColBERT on Vespa.ai | |
|
|------------------------------------|----------------| |
|
| Dev | 0.364 | |
|
|
|
Recall@k On Dev (6980 queries) |
|
|
|
|K | Recall@K | |
|
|------------------------------------|----------------| |
|
| 50 | 0.816 | |
|
| 200 | 0.905 | |
|
| 1000 | 0.939 | |
|
|
|
|
|
The MRR@10 on dev is achieved by re-ranking 1K retrieved by a dense retriever based on |
|
[sentence-transformers/msmarco-MiniLM-L-6-v3](https://huggingface.co/sentence-transformers/msmarco-MiniLM-L-6-v3). Re-ranking the original top 1000 dev |
|
is 0.354 MRR@10 (Recall@1K 0.82). |
|
|
|
The official baseline BM25 ranking model MRR@10 0.16 on eval and 0.167 on dev question set. |
|
See [MS Marco Passage Ranking Leaderboard](https://microsoft.github.io/msmarco/). |
|
|
|
## Export ColBERT query encoder to ONNX |
|
We represent the ColBERT query encoder in the Vespa runtime, to map the textual query representation to the tensor representation. For this |
|
we use Vespa's support for running ONNX models. One can use the following snippet to export the model for serving. |
|
|
|
```python |
|
from transformers import BertModel |
|
from transformers import BertPreTrainedModel |
|
from transformers import BertConfig |
|
import torch |
|
import torch.nn as nn |
|
|
|
class VespaColBERT(BertPreTrainedModel): |
|
|
|
def __init__(self,config): |
|
super().__init__(config) |
|
self.bert = BertModel(config) |
|
self.linear = nn.Linear(config.hidden_size, 32, bias=False) |
|
self.init_weights() |
|
|
|
def forward(self, input_ids, attention_mask): |
|
Q = self.bert(input_ids,attention_mask=attention_mask)[0] |
|
Q = self.linear(Q) |
|
return torch.nn.functional.normalize(Q, p=2, dim=2) |
|
|
|
colbert_query_encoder = VespaColBERT.from_pretrained("vespa-engine/col-minilm") |
|
|
|
#Export model to ONNX for serving in Vespa |
|
|
|
input_names = ["input_ids", "attention_mask"] |
|
output_names = ["contextual"] |
|
#input, max 32 query term |
|
input_ids = torch.ones(1,32, dtype=torch.int64) |
|
attention_mask = torch.ones(1,32,dtype=torch.int64) |
|
args = (input_ids, attention_mask) |
|
torch.onnx.export(colbert_query_encoder, |
|
args=args, |
|
f="query_encoder_colbert.onnx", |
|
input_names = input_names, |
|
output_names = output_names, |
|
dynamic_axes = { |
|
"input_ids": {0: "batch"}, |
|
"attention_mask": {0: "batch"}, |
|
"contextual": {0: "batch"}, |
|
}, |
|
opset_version=11) |
|
``` |
|
|
|
# Representing the model on Vespa.ai |
|
See [Ranking with ONNX models](https://docs.vespa.ai/documentation/onnx.html) and the [MS Marco Ranking sample app](https://github.com/vespa-engine/sample-apps/tree/master/msmarco-ranking) |
|
|
|
|
