metadata
license: mit
datasets:
- philipphager/baidu-ultr-pretrain
- philipphager/baidu-ultr_uva-mlm-ctr
metrics:
- dcg@1
- dcg@3
- dcg@5
- dcg@10
- ndcg@10
- mrr@10
Naive Listwise MonoBERT trained on Baidu-ULTR
A flax-based MonoBERT cross encoder trained on the Baidu-ULTR dataset with a listwise softmax cross-entropy loss on clicks. The loss is called "naive" as we use user clicks as a signal of relevance without any additional position bias correction. For more info, read our paper here.
Usage
import jax.numpy as jnp
from src.model import ListwiseCrossEncoder
model = ListwiseCrossEncoder.from_pretrained(
"philipphager/baidu-ultr_uva-bert_naive-listwise",
)
# Mock batch following Baidu-ULTR with 4 documents, each with 8 tokens
batch = {
# Query_id for each document
"query_id": jnp.array([1, 1, 1, 1]),
# Document position in SERP
"positions": jnp.array([1, 2, 3, 4]),
# Token ids for: [CLS] Query [SEP] Document
"tokens": jnp.array([
[2, 21448, 21874, 21436, 1, 20206, 4012, 2860],
[2, 21448, 21874, 21436, 1, 16794, 4522, 2082],
[2, 21448, 21874, 21436, 1, 20206, 10082, 9773],
[2, 21448, 21874, 21436, 1, 2618, 8520, 2860],
]),
# Specify if a token id belongs to the query (0) or document (1)
"token_types": jnp.array([
[0, 0, 0, 0, 1, 1, 1, 1],
[0, 0, 0, 0, 1, 1, 1, 1],
[0, 0, 0, 0, 1, 1, 1, 1],
[0, 0, 0, 0, 1, 1, 1, 1],
]),
# Marks if a token should be attended to (True) or ignored, e.g., padding tokens (False):
"attention_mask": jnp.array([
[True, True, True, True, True, True, True, True],
[True, True, True, True, True, True, True, True],
[True, True, True, True, True, True, True, True],
[True, True, True, True, True, True, True, True],
]),
}
outputs = model(batch)
print(outputs)
Test Results on Baidu-ULTR Expert Annotations
| Model | log-likelihood | DCG@1 | DCG@3 | DCG@5 | DCG@10 | nDCG@10 | MRR@10 |
|---|---|---|---|---|---|---|---|
| Naive Pointwise | 0.2272 | 1.6836 | 3.5616 | 4.8822 | 7.4244 | 0.3640 | 0.6096 |
| Naive Listwise | - | 1.9738 | 4.1609 | 5.6861 | 8.5432 | 0.4091 | 0.6436 |