Views
No views yet
| MS Marco Passage Ranking Query Set | MRR@10 ColBERT on Vespa.ai |
|---|---|
| Dev | 0.364 |
| K | Recall@K |
|---|---|
| 50 | 0.816 |
| 200 | 0.905 |
| 1000 | 0.939 |
1from transformers import BertModel
2from transformers import BertPreTrainedModel
3from transformers import BertConfig
4import torch
5import torch.nn as nn
6
7class VespaColBERT(BertPreTrainedModel):
8
9 def __init__(self,config):
10 super().__init__(config)
11 self.bert = BertModel(config)
12 self.linear = nn.Linear(config.hidden_size, 32, bias=False)
13 self.init_weights()
14
15 def forward(self, input_ids, attention_mask):
16 Q = self.bert(input_ids,attention_mask=attention_mask)[0]
17 Q = self.linear(Q)
18 return torch.nn.functional.normalize(Q, p=2, dim=2)
19
20colbert_query_encoder = VespaColBERT.from_pretrained("vespa-engine/col-minilm")
21
22#Export model to ONNX for serving in Vespa
23
24input_names = ["input_ids", "attention_mask"]
25output_names = ["contextual"]
26#input, max 32 query term
27input_ids = torch.ones(1,32, dtype=torch.int64)
28attention_mask = torch.ones(1,32,dtype=torch.int64)
29args = (input_ids, attention_mask)
30torch.onnx.export(colbert_query_encoder,
31 args=args,
32 f="query_encoder_colbert.onnx",
33 input_names = input_names,
34 output_names = output_names,
35 dynamic_axes = {
36 "input_ids": {0: "batch"},
37 "attention_mask": {0: "batch"},
38 "contextual": {0: "batch"},
39 },
40 opset_version=11)