ByteDance/ListConRanker
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

初始化项目，由ModelHub XC社区提供模型

Browse Source 
Model: ByteDance/ListConRanker
Source: Original Platform
This commit is contained in:
ModelHub XC
2026-05-14 14:04:55 +08:00
commit 77d52fa9fd
19 changed files with 23157 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
.gitattributes vendored Normal file
View File

@@ -0,0 +1,35 @@
*.7z filter=lfs diff=lfs merge=lfs -text
*.arrow filter=lfs diff=lfs merge=lfs -text
*.bin filter=lfs diff=lfs merge=lfs -text
*.bz2 filter=lfs diff=lfs merge=lfs -text
*.ckpt filter=lfs diff=lfs merge=lfs -text
*.ftz filter=lfs diff=lfs merge=lfs -text
*.gz filter=lfs diff=lfs merge=lfs -text
*.h5 filter=lfs diff=lfs merge=lfs -text
*.joblib filter=lfs diff=lfs merge=lfs -text
*.lfs.* filter=lfs diff=lfs merge=lfs -text
*.mlmodel filter=lfs diff=lfs merge=lfs -text
*.model filter=lfs diff=lfs merge=lfs -text
*.msgpack filter=lfs diff=lfs merge=lfs -text
*.npy filter=lfs diff=lfs merge=lfs -text
*.npz filter=lfs diff=lfs merge=lfs -text
*.onnx filter=lfs diff=lfs merge=lfs -text
*.ot filter=lfs diff=lfs merge=lfs -text
*.parquet filter=lfs diff=lfs merge=lfs -text
*.pb filter=lfs diff=lfs merge=lfs -text
*.pickle filter=lfs diff=lfs merge=lfs -text
*.pkl filter=lfs diff=lfs merge=lfs -text
*.pt filter=lfs diff=lfs merge=lfs -text
*.pth filter=lfs diff=lfs merge=lfs -text
*.rar filter=lfs diff=lfs merge=lfs -text
*.safetensors filter=lfs diff=lfs merge=lfs -text
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
*.tar.* filter=lfs diff=lfs merge=lfs -text
*.tar filter=lfs diff=lfs merge=lfs -text
*.tflite filter=lfs diff=lfs merge=lfs -text
*.tgz filter=lfs diff=lfs merge=lfs -text
*.wasm filter=lfs diff=lfs merge=lfs -text
*.xz filter=lfs diff=lfs merge=lfs -text
*.zip filter=lfs diff=lfs merge=lfs -text
*.zst filter=lfs diff=lfs merge=lfs -text
*tfevents* filter=lfs diff=lfs merge=lfs -text

21
LICENSE Normal file
View File

@@ -0,0 +1,21 @@
MIT License
Copyright (c) 2024 data-comment
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

203
README.md Normal file
View File

@@ -0,0 +1,203 @@
---
tags:
- mteb
- sentence-transformers
- transformers
pipeline_tag: text-ranking
model-index:
- name: ListConRanker
results:
- task:
type: Reranking
dataset:
type: C-MTEB/CMedQAv1-reranking
name: MTEB CMedQAv1
config: default
split: test
revision: None
metrics:
- type: map
value: 90.55366308098787
- type: mrr_1
value: 87.8
- type: mrr_10
value: 92.45134920634919
- type: mrr_5
value: 92.325
- task:
type: Reranking
dataset:
type: C-MTEB/CMedQAv2-reranking
name: MTEB CMedQAv2
config: default
split: test
revision: None
metrics:
- type: map
value: 89.38076135722042
- type: mrr_1
value: 85.9
- type: mrr_10
value: 91.28769841269842
- type: mrr_5
value: 91.08999999999999
- task:
type: Reranking
dataset:
type: C-MTEB/Mmarco-reranking
name: MTEB MMarcoReranking
config: default
split: dev
revision: None
metrics:
- type: map
value: 43.881461866703894
- type: mrr_1
value: 32.0
- type: mrr_10
value: 44.534126984126985
- type: mrr_5
value: 43.45
- task:
type: Reranking
dataset:
type: C-MTEB/T2Reranking
name: MTEB T2Reranking
config: default
split: dev
revision: None
metrics:
- type: map
value: 69.16513825032682
- type: mrr_1
value: 67.36628300609343
- type: mrr_10
value: 80.06914890758831
- type: mrr_5
value: 79.69137892123675
---
# ListConRanker
## Model
- We propose a **List**wise-encoded **Con**trastive text re**Ranker** (**ListConRanker**), includes a ListTransformer module for listwise encoding. The ListTransformer can facilitate global contrastive information learning between passage features, including the clustering of similar passages, the clustering between dissimilar passages, and the distinction between similar and dissimilar passages. Besides, we propose ListAttention to help ListTransformer maintain the features of the query while learning global comparative information.
- The training loss function is Circle Loss[1]. Compared with cross-entropy loss and ranking loss, it can solve the problems of low data efficiency and unsmooth gradient change.
## Data
The training data consists of approximately 2.6 million queries, each corresponding to multiple passages. The data comes from the training sets of several datasets, including cMedQA1.0, cMedQA2.0, MMarcoReranking, T2Reranking, huatuo, MARC, XL-sum, CSL and so on.
## Training
We trained the model in two stages. In the first stage, we freeze the parameters of embedding model and only train the ListTransformer for 4 epochs with a batch size of 1024. In the second stage, we do not freeze any parameter and train for another 2 epochs with a batch size of 256.
## Inference
Due to the limited memory of GPUs, we input about 20 passages at a time for each query during training. However, during actual use, there may be situations where far more than 20 passages are input at the same time (e.g, MMarcoReranking).
To reduce the discrepancy between training and inference, we propose iterative inference. The iterative inference feeds the passages into the ListConRanker multiple times, and each time it only decides the ranking of the passage at the end of the list.
## Performance
| Model | cMedQA1.0 | cMedQA2.0 | MMarcoReranking | T2Reranking | Avg. |
| :--- | :---: | :---: | :---: | :---: | :---: |
| LdIR-Qwen2-reranker-1.5B | 86.50 | 87.11 | 39.35 | 68.84 | 70.45 |
| zpoint-large-embedding-zh | 91.11 | 90.07 | 38.87 | 69.29 | 72.34 |
| xiaobu-embedding-v2 | 90.96 | 90.41 | 39.91 | 69.03 | 72.58 |
| Conan-embedding-v1 | 91.39 | 89.72 | 41.58 | 68.36 | 72.76 |
| ListConRanker | 90.55 | 89.38 | 43.88 | 69.17 | **73.25** |
| - w/o Iterative Inference | 90.20 | 89.98 | 37.52 | 69.17 | 71.72 |
## How to use
```python
from transfoermers import AutoModelForSequenceClassification, AutoTokenizer
reranker = AutoModelForSequenceClassification.from_pretrained('ByteDance/ListConRanker', trust_remote_code=True)
# [query, passages_1, passage_2, ..., passage_n]
batch = [
[
'皮蛋是寒性的食物吗', # query
'营养医师介绍皮蛋是属于凉性的食物,中医认为皮蛋可治眼疼、牙疼、高血压、耳鸣眩晕等疾病。体虚者要少吃。', # passage_1
'皮蛋这种食品是在中国地域才常见的传统食品,它的生长汗青也是非常的悠长。', # passage_2
'喜欢皮蛋的人会觉得皮蛋是最美味的食物,不喜欢皮蛋的人则觉得皮蛋是黑暗料理,尤其很多外国朋友都不理解我们吃皮蛋的习惯' # passage_3
],
[
'月有阴晴圆缺的意义', # query
'形容的是月所有的状态,晴朗明媚,阴沉混沌,有月圆时,但多数时总是有缺陷。', # passage_1
'人有悲欢离合,月有阴晴圆缺这句话意思是人有悲欢离合的变迁,月有阴晴圆缺的转换。', # passage_2
'既然是诗歌,又哪里会有真正含义呢? 大概可以说:人生有太多坎坷,苦难,从容坦荡面对就好。', # passage_3
'一零七六年苏轼贬官密州,时年四十一岁的他政治上很不得志,时值中秋佳节,非常想念自己的弟弟子由内心颇感忧郁,情绪低沉,有感而发写了这首词。' # passage_4
]
]
# for conventional inference, please manage the batch size by yourself
scores = reranker.multi_passage(batch)
print(scores)
# [0.5126814246177673, 0.33125635981559753, 0.3642643094062805, 0.6367220282554626, 0.7166246175765991, 0.4281482696533203, 0.3530198335647583]
# for iterative inferfence, only a batch size of 1 is supported
# the scores do not carry similarity meaning and are only used for ranking
scores = reranker.multi_passage_in_iterative_inference(batch[0])
print(scores)
# [0.5126813650131226, 0.3312564790248871, 0.3642643094062805]
tokenizer = AutoTokenizer.from_pretrained('ByteDance/ListConRanker')
inputs = tokenizer(
[
[
"皮蛋是寒性的食物吗",
"营养医师介绍皮蛋是属于凉性的食物,中医认为皮蛋可治眼疼、牙疼、高血压、耳鸣眩晕等疾病。体虚者要少吃。",
],
[
"皮蛋是寒性的食物吗",
"皮蛋这种食品是在中国地域才常见的传统食品,它的生长汗青也是非常的悠长。",
],
[
"月有阴晴圆缺的意义",
"形容的是月所有的状态,晴朗明媚,阴沉混沌,有月圆时,但多数时总是有缺陷。",
],
],
return_tensors="pt",
padding=True,
truncation=False
)
# tensor([[0.5070], [0.3334], [0.6294]], device='cuda:0', dtype=torch.float16, grad_fn=<ViewBackward0>)
```
or using the `sentence_transformers` library (We do not recommend using `sentence_transformers`. Because its truncation strategy may not match the model design, which may lead to performance degradation.):
```python
from sentence_transformers import CrossEncoder
model = CrossEncoder('ByteDance/ListConRanker', trust_remote_code=True)
inputs = [
[
"皮蛋是寒性的食物吗",
"营养医师介绍皮蛋是属于凉性的食物,中医认为皮蛋可治眼疼、牙疼、高血压、耳鸣眩晕等疾病。体虚者要少吃。",
],
[
"皮蛋是寒性的食物吗",
"皮蛋这种食品是在中国地域才常见的传统食品,它的生长汗青也是非常的悠长。",
],
[
"月有阴晴圆缺的意义",
"形容的是月所有的状态,晴朗明媚,阴沉混沌,有月圆时,但多数时总是有缺陷。",
],
]
scores = model.predict(inputs)
print(scores)
```
To reproduce the results with iterative inference, please run:
```bash
python3 eval_listconranker_iterative_inference.py
```
To reproduce the results without iterative inference, please run:
```bash
python3 eval_listconranker.py
```
## Reference
1. https://arxiv.org/abs/2002.10857
2. https://github.com/FlagOpen/FlagEmbedding
3. https://arxiv.org/abs/2408.15710
## License
This work is licensed under a [MIT License](https://opensource.org/license/MIT) and the weight of models is licensed under a [Creative Commons Attribution-NonCommercial 4.0 International License](https://creativecommons.org/licenses/by-nc/4.0/).

45
config.json Normal file
View File

@@ -0,0 +1,45 @@
{
"architectures": [
"ListConRanker"
],
"auto_map": {
"AutoConfig": "configuration_listconranker.ListConRankerConfig",
"AutoModelForSequenceClassification": "modeling_listconranker.ListConRankerModel"
},
"attention_probs_dropout_prob": 0.1,
"classifier_dropout": null,
"directionality": "bidi",
"gradient_checkpointing": false,
"hidden_act": "gelu",
"hidden_dropout_prob": 0.1,
"hidden_size": 1024,
"list_con_hidden_size": 1792,
"id2label": {
"0": "LABEL_0"
},
"initializer_range": 0.02,
"intermediate_size": 4096,
"label2id": {
"LABEL_0": 0
},
"layer_norm_eps": 1e-12,
"max_position_embeddings": 512,
"model_type": "bert",
"num_attention_heads": 16,
"num_hidden_layers": 24,
"output_hidden_states": true,
"pad_token_id": 0,
"pooler_fc_size": 768,
"pooler_num_attention_heads": 12,
"pooler_num_fc_layers": 3,
"pooler_size_per_head": 128,
"pooler_type": "first_token_transform",
"position_embedding_type": "absolute",
"torch_dtype": "bfloat16",
"transformers_version": "4.45.2",
"type_vocab_size": 2,
"use_cache": true,
"vocab_size": 21128,
"cls_token_id": 101,
"sep_token_id": 102
}

44
configuration_listconranker.py Normal file
View File

@@ -0,0 +1,44 @@
# Copyright 2024 Bytedance Ltd. and/or its affiliates
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of this software
# and associated documentation files (the "Software"), to deal in the Software without
# restriction, including without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all copies or
# substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
# OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
# ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.
from __future__ import annotations
from transformers import BertConfig
class ListConRankerConfig(BertConfig):
"""Configuration class for ListConRanker model."""
model_type = "ListConRanker"
def __init__(
self,
list_transformer_layers: int = 2,
list_con_hidden_size: int = 1792,
num_labels: int = 1,
cls_token_id: int = 101,
sep_token_id: int = 102,
**kwargs,
):
super().__init__(**kwargs)
self.list_transformer_layers = list_transformer_layers
self.list_con_hidden_size = list_con_hidden_size
self.num_labels = num_labels
self.cls_token_id = cls_token_id
self.sep_token_id = sep_token_id
self.bert_config = BertConfig(**kwargs)
self.bert_config.output_hidden_states = True

45
eval_listconranker.py Normal file
View File

@@ -0,0 +1,45 @@
# Copyright 2024 Bytedance Ltd. and/or its affiliates
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of this software
# and associated documentation files (the “Software”), to deal in the Software without
# restriction, including without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all copies or
# substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
# OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
# ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.
import argparse
from modules.Reranking import *
from mteb import MTEB
from modules.listconranker import ListConRanker
def get_args():
parser = argparse.ArgumentParser()
parser.add_argument('--model_name_or_path', default="./", type=str)
return parser.parse_args()
if __name__ == '__main__':
args = get_args()
model = ListConRanker(args.model_name_or_path, use_fp16=True, list_transformer_layer=2)
dir_name = args.model_name_or_path.split('/')[-2]
if 'checkpoint-' in args.model_name_or_path:
save_name = "_".join(args.model_name_or_path.split('/')[-2:])
dir_name = args.model_name_or_path.split('/')[-3]
else:
save_name = "_".join(args.model_name_or_path.split('/')[-1:])
dir_name = args.model_name_or_path.split('/')[-2]
evaluation = MTEB(task_types=["Reranking"], task_langs=['zh'])
evaluation.run(model, output_folder="reranker_results/{}/{}".format(dir_name, save_name))

45
eval_listconranker_iterative_inference.py Normal file
View File

@@ -0,0 +1,45 @@
# Copyright 2024 Bytedance Ltd. and/or its affiliates
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of this software
# and associated documentation files (the “Software”), to deal in the Software without
# restriction, including without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all copies or
# substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
# OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
# ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.
import argparse
from modules.Reranking_loop import *
from mteb import MTEB
from modules.listconranker import ListConRanker
def get_args():
parser = argparse.ArgumentParser()
parser.add_argument('--model_name_or_path', default="./", type=str)
return parser.parse_args()
if __name__ == '__main__':
args = get_args()
model = ListConRanker(args.model_name_or_path, use_fp16=True, list_transformer_layer=2)
dir_name = args.model_name_or_path.split('/')[-2]
if 'checkpoint-' in args.model_name_or_path:
save_name = "_".join(args.model_name_or_path.split('/')[-2:])
dir_name = args.model_name_or_path.split('/')[-3]
else:
save_name = "_".join(args.model_name_or_path.split('/')[-1:])
dir_name = args.model_name_or_path.split('/')[-2]
evaluation = MTEB(task_types=["Reranking"], task_langs=['zh'])
evaluation.run(model, output_folder="reranker_results/{}/{}".format(dir_name, save_name))

3
linear_in_embedding.pt Normal file
View File

@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:0ae1bb10d5c23c3bdbe50dfee2f37bb243d1606cc1a41e02a8ffb7bf61b71033
size 7348826

3
list_transformer.pt Normal file
View File

@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:cabbabfb04bc1feb6fa859a074f97397e88b0dc6bf14b0ef9ad3a0ddfac1cef5
size 293894397

3
model.safetensors Normal file
View File

@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:15c1c9ebc02bd5255758d0ba5498b3c93fd4cc8dd25845fd6a2cac8b2d12cefc
size 1302134568

530
modeling_listconranker.py Normal file
View File

@@ -0,0 +1,530 @@
# Copyright 2024 Bytedance Ltd. and/or its affiliates
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of this software
# and associated documentation files (the "Software"), to deal in the Software without
# restriction, including without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all copies or
# substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
# OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
# ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.
from __future__ import annotations
import torch
from torch import nn
from torch.nn import functional as F
from transformers import (
PreTrainedModel,
BertModel,
AutoTokenizer,
)
import os
from transformers.modeling_outputs import SequenceClassifierOutput
from typing import Union, List, Optional
from collections import defaultdict
import numpy as np
import math
from huggingface_hub import hf_hub_download
from .configuration_listconranker import ListConRankerConfig
class QueryEmbedding(nn.Module):
def __init__(self, config) -> None:
super().__init__()
self.query_embedding = nn.Embedding(2, config.list_con_hidden_size)
self.layerNorm = nn.LayerNorm(config.list_con_hidden_size)
def forward(self, x, tags):
query_embeddings = self.query_embedding(tags)
x += query_embeddings
x = self.layerNorm(x)
return x
class ListTransformer(nn.Module):
def __init__(self, num_layer, config) -> None:
super().__init__()
self.config = config
self.list_transformer_layer = nn.TransformerEncoderLayer(
config.list_con_hidden_size,
self.config.num_attention_heads,
batch_first=True,
activation=F.gelu,
norm_first=False,
)
self.list_transformer = nn.TransformerEncoder(
self.list_transformer_layer, num_layer
)
self.relu = nn.ReLU()
self.query_embedding = QueryEmbedding(config)
self.linear_score3 = nn.Linear(
config.list_con_hidden_size * 2, config.list_con_hidden_size
)
self.linear_score2 = nn.Linear(
config.list_con_hidden_size * 2, config.list_con_hidden_size
)
self.linear_score1 = nn.Linear(config.list_con_hidden_size * 2, 1)
def forward(
self, pair_features: torch.Tensor, pair_nums: List[int]
) -> torch.Tensor:
batch_pair_features = pair_features.split(pair_nums)
pair_feature_query_passage_concat_list = []
for i in range(len(batch_pair_features)):
pair_feature_query = (
batch_pair_features[i][0].unsqueeze(0).repeat(pair_nums[i] - 1, 1)
)
pair_feature_passage = batch_pair_features[i][1:]
pair_feature_query_passage_concat_list.append(
torch.cat([pair_feature_query, pair_feature_passage], dim=1)
)
pair_feature_query_passage_concat = torch.cat(
pair_feature_query_passage_concat_list, dim=0
)
batch_pair_features = nn.utils.rnn.pad_sequence(
batch_pair_features, batch_first=True
)
query_embedding_tags = torch.zeros(
batch_pair_features.size(0),
batch_pair_features.size(1),
dtype=torch.long,
device=self.device,
)
query_embedding_tags[:, 0] = 1
batch_pair_features = self.query_embedding(
batch_pair_features, query_embedding_tags
)
mask = self.generate_attention_mask(pair_nums)
query_mask = self.generate_attention_mask_custom(pair_nums)
pair_list_features = self.list_transformer(
batch_pair_features, src_key_padding_mask=mask, mask=query_mask
)
output_pair_list_features = []
output_query_list_features = []
pair_features_after_transformer_list = []
for idx, pair_num in enumerate(pair_nums):
output_pair_list_features.append(pair_list_features[idx, 1:pair_num, :])
output_query_list_features.append(pair_list_features[idx, 0, :])
pair_features_after_transformer_list.append(
pair_list_features[idx, :pair_num, :]
)
pair_features_after_transformer_cat_query_list = []
for idx, pair_num in enumerate(pair_nums):
query_ft = (
output_query_list_features[idx].unsqueeze(0).repeat(pair_num - 1, 1)
)
pair_features_after_transformer_cat_query = torch.cat(
[query_ft, output_pair_list_features[idx]], dim=1
)
pair_features_after_transformer_cat_query_list.append(
pair_features_after_transformer_cat_query
)
pair_features_after_transformer_cat_query = torch.cat(
pair_features_after_transformer_cat_query_list, dim=0
)
pair_feature_query_passage_concat = self.relu(
self.linear_score2(pair_feature_query_passage_concat)
)
pair_features_after_transformer_cat_query = self.relu(
self.linear_score3(pair_features_after_transformer_cat_query)
)
final_ft = torch.cat(
[
pair_feature_query_passage_concat,
pair_features_after_transformer_cat_query,
],
dim=1,
)
logits = self.linear_score1(final_ft).squeeze()
return logits, torch.cat(pair_features_after_transformer_list, dim=0)
def generate_attention_mask(self, pair_num):
max_len = max(pair_num)
batch_size = len(pair_num)
mask = torch.zeros(batch_size, max_len, dtype=torch.bool, device=self.device)
for i, length in enumerate(pair_num):
mask[i, length:] = True
return mask
def generate_attention_mask_custom(self, pair_num):
max_len = max(pair_num)
mask = torch.zeros(max_len, max_len, dtype=torch.bool, device=self.device)
mask[0, 1:] = True
return mask
class ListConRankerModel(PreTrainedModel):
"""
ListConRanker model for sequence classification that's compatible with AutoModelForSequenceClassification.
"""
config_class = ListConRankerConfig
base_model_prefix = "listconranker"
def __init__(self, config: ListConRankerConfig):
super().__init__(config)
self.config = config
self.num_labels = config.num_labels
self.hf_model = BertModel(config.bert_config)
self.sigmoid = nn.Sigmoid()
self.linear_in_embedding = nn.Linear(
config.hidden_size, config.list_con_hidden_size
)
self.list_transformer = ListTransformer(
config.list_transformer_layers,
config,
)
def forward(
self,
input_ids: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
token_type_ids: Optional[torch.Tensor] = None,
position_ids: Optional[torch.Tensor] = None,
head_mask: Optional[torch.Tensor] = None,
inputs_embeds: Optional[torch.Tensor] = None,
labels: Optional[torch.Tensor] = None,
output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
**kwargs,
) -> Union[tuple[torch.Tensor], SequenceClassifierOutput]:
if self.training:
raise NotImplementedError("Training not supported; use eval mode.")
device = input_ids.device
self.list_transformer.device = device
# Reorganize by unique queries and their passages
(
reorganized_input_ids,
reorganized_attention_mask,
reorganized_token_type_ids,
pair_nums,
group_indices,
) = self._reorganize_inputs(input_ids, attention_mask, token_type_ids)
out = self.hf_model(
input_ids=reorganized_input_ids,
attention_mask=reorganized_attention_mask,
token_type_ids=reorganized_token_type_ids,
return_dict=True,
)
feats = out.last_hidden_state
pooled = self.average_pooling(feats, reorganized_attention_mask)
embedded = self.linear_in_embedding(pooled)
logits, _ = self.list_transformer(embedded, pair_nums)
probs = self.sigmoid(logits)
# Restore original order
sorted_probs = self._restore_original_order(probs, group_indices)
sorted_logits = self._restore_original_order(logits, group_indices)
if not return_dict:
return (sorted_probs, sorted_logits)
return SequenceClassifierOutput(
loss=None,
logits=sorted_logits,
hidden_states=out.hidden_states,
attentions=out.attentions,
)
def _reorganize_inputs(
self,
input_ids: torch.Tensor,
attention_mask: torch.Tensor,
token_type_ids: Optional[torch.Tensor],
) -> tuple[
torch.Tensor, torch.Tensor, Optional[torch.Tensor], List[int], List[List[int]]
]:
"""
Group inputs by unique queries: for each query, produce [query] + its passages,
then flatten, pad, and return pair sizes and original indices mapping.
"""
batch_size = input_ids.size(0)
# Structure: query_key -> {
# 'query': (seq, mask, tt),
# 'passages': [(seq, mask, tt), ...],
# 'indices': [original_index, ...]
# }
grouped = {}
for idx in range(batch_size):
seq = input_ids[idx]
mask = attention_mask[idx]
token_type_ids[idx] if token_type_ids is not None else torch.zeros_like(seq)
sep_idxs = (seq == self.config.sep_token_id).nonzero(as_tuple=True)[0]
if sep_idxs.numel() == 0:
raise ValueError(f"No SEP in sequence {idx}")
first_sep = sep_idxs[0].item()
second_sep = sep_idxs[1].item()
# Extract query and passage
q_seq = seq[: first_sep + 1]
q_mask = mask[: first_sep + 1]
q_tt = torch.zeros_like(q_seq)
p_seq = seq[first_sep : second_sep + 1]
p_mask = mask[first_sep : second_sep + 1]
p_seq = p_seq.clone()
p_seq[0] = self.config.cls_token_id
p_tt = torch.zeros_like(p_seq)
# Build key excluding CLS/SEP
key = tuple(
q_seq[
(q_seq != self.config.cls_token_id)
& (q_seq != self.config.sep_token_id)
].tolist()
)
# truncation
q_seq = q_seq[: self.config.max_position_embeddings]
q_seq[-1] = self.config.sep_token_id
p_seq = p_seq[: self.config.max_position_embeddings]
p_seq[-1] = self.config.sep_token_id
q_mask = q_mask[: self.config.max_position_embeddings]
p_mask = p_mask[: self.config.max_position_embeddings]
q_tt = q_tt[: self.config.max_position_embeddings]
p_tt = p_tt[: self.config.max_position_embeddings]
if key not in grouped:
grouped[key] = {
"query": (q_seq, q_mask, q_tt),
"passages": [],
"indices": [],
}
grouped[key]["passages"].append((p_seq, p_mask, p_tt))
grouped[key]["indices"].append(idx)
# Flatten according to group insertion order
seqs, masks, tts, pair_nums, group_indices = [], [], [], [], []
for key, data in grouped.items():
q_seq, q_mask, q_tt = data["query"]
passages = data["passages"]
indices = data["indices"]
# record sizes and original positions
pair_nums.append(len(passages) + 1) # +1 for the query
group_indices.append(indices)
# append query then its passages
seqs.append(q_seq)
masks.append(q_mask)
tts.append(q_tt)
for p_seq, p_mask, p_tt in passages:
seqs.append(p_seq)
masks.append(p_mask)
tts.append(p_tt)
# Pad to uniform length
max_len = max(s.size(0) for s in seqs)
padded_seqs, padded_masks, padded_tts = [], [], []
for s, m, t in zip(seqs, masks, tts):
ps = torch.zeros(max_len, dtype=s.dtype, device=s.device)
pm = torch.zeros(max_len, dtype=m.dtype, device=m.device)
pt = torch.zeros(max_len, dtype=t.dtype, device=t.device)
ps[: s.size(0)] = s
pm[: m.size(0)] = m
pt[: t.size(0)] = t
padded_seqs.append(ps)
padded_masks.append(pm)
padded_tts.append(pt)
rid = torch.stack(padded_seqs)
ram = torch.stack(padded_masks)
rtt = torch.stack(padded_tts) if token_type_ids is not None else None
return rid, ram, rtt, pair_nums, group_indices
def _restore_original_order(
self,
logits: torch.Tensor,
group_indices: List[List[int]],
) -> torch.Tensor:
"""
Map flattened logits back so each original index gets its passage score.
"""
out = torch.zeros(logits.size(0), dtype=logits.dtype, device=logits.device)
i = 0
for indices in group_indices:
for idx in indices:
out[idx] = logits[i]
i += 1
return out.reshape(-1, 1)
def average_pooling(self, hidden_state, attention_mask):
extended_attention_mask = (
attention_mask.unsqueeze(-1)
.expand(hidden_state.size())
.to(dtype=hidden_state.dtype)
)
masked_hidden_state = hidden_state * extended_attention_mask
sum_embeddings = torch.sum(masked_hidden_state, dim=1)
sum_mask = extended_attention_mask.sum(dim=1)
return sum_embeddings / sum_mask
@classmethod
def from_pretrained(
cls, model_name_or_path, config: Optional[ListConRankerConfig] = None, **kwargs
):
model = super().from_pretrained(model_name_or_path, config=config, **kwargs)
model.hf_model = BertModel.from_pretrained(
model_name_or_path, config=model.config.bert_config, **kwargs
)
linear_path = hf_hub_download(
repo_id = model_name_or_path,
filename = "linear_in_embedding.pt",
revision = "main",
cache_dir = kwargs['cache_dir'] if 'cache_dir' in kwargs else None
)
list_transformer_path = hf_hub_download(
repo_id = "ByteDance/ListConRanker",
filename = "list_transformer.pt",
revision = "main",
cache_dir = kwargs['cache_dir'] if 'cache_dir' in kwargs else None
)
try:
model.linear_in_embedding.load_state_dict(torch.load(linear_path))
model.list_transformer.load_state_dict(torch.load(list_transformer_path))
except FileNotFoundError as e:
raise e
return model
def multi_passage(
self,
sentences: List[List[str]],
batch_size: int = 32,
tokenizer: AutoTokenizer = AutoTokenizer.from_pretrained(
"ByteDance/ListConRanker"
),
):
"""
Process multiple passages for each query.
:param sentences: List of lists, where each inner list contains sentences for a query.
:return: Tensor of logits for each passage.
"""
pairs = []
for batch in sentences:
if len(batch) < 2:
raise ValueError("Each query must have at least one passage.")
query = batch[0]
passages = batch[1:]
for passage in passages:
pairs.append((query, passage))
total_batches = (len(pairs) + batch_size - 1) // batch_size
total_logits = torch.zeros(len(pairs), dtype=torch.float, device=self.device)
for batch in range(total_batches):
batch_pairs = pairs[batch * batch_size : (batch + 1) * batch_size]
inputs = tokenizer(
batch_pairs,
padding=True,
truncation=False,
return_tensors="pt",
)
for k, v in inputs.items():
inputs[k] = v.to(self.device)
logits = self(**inputs)[0]
total_logits[batch * batch_size : (batch + 1) * batch_size] = (
logits.squeeze(1)
)
return total_logits.tolist()
def multi_passage_in_iterative_inference(
self,
sentences: List[str],
stop_num: int = 20,
decrement_rate: float = 0.2,
min_filter_num: int = 10,
tokenizer: AutoTokenizer = AutoTokenizer.from_pretrained(
"ByteDance/ListConRanker"
),
):
"""
Process multiple passages for one query in iterative inference.
:param sentences: List contains sentences for a query.
:return: Tensor of logits for each passage.
"""
if stop_num < 1:
raise ValueError("stop_num must be greater than 0")
if decrement_rate <= 0 or decrement_rate >= 1:
raise ValueError("decrement_rate must be in (0, 1)")
if min_filter_num < 1:
raise ValueError("min_filter_num must be greater than 0")
query = sentences[0]
passage = sentences[1:]
filter_times = 0
passage2score = defaultdict(list)
while len(passage) > stop_num:
batch = [[query] + passage]
pred_scores = self.multi_passage(
batch, batch_size=len(batch[0]) - 1, tokenizer=tokenizer
)
pred_scores_argsort = np.argsort(
pred_scores
).tolist() # Sort in increasing order
passage_len = len(passage)
to_filter_num = math.ceil(passage_len * decrement_rate)
if to_filter_num < min_filter_num:
to_filter_num = min_filter_num
have_filter_num = 0
while have_filter_num < to_filter_num:
idx = pred_scores_argsort[have_filter_num]
passage2score[passage[idx]].append(pred_scores[idx] + filter_times)
have_filter_num += 1
while (
pred_scores[pred_scores_argsort[have_filter_num - 1]]
== pred_scores[pred_scores_argsort[have_filter_num]]
):
idx = pred_scores_argsort[have_filter_num]
passage2score[passage[idx]].append(pred_scores[idx] + filter_times)
have_filter_num += 1
next_passage = []
next_passage_idx = have_filter_num
while next_passage_idx < len(passage):
idx = pred_scores_argsort[next_passage_idx]
next_passage.append(passage[idx])
next_passage_idx += 1
passage = next_passage
filter_times += 1
batch = [[query] + passage]
pred_scores = self.multi_passage(
batch, batch_size=len(batch[0]) - 1, tokenizer=tokenizer
)
cnt = 0
while cnt < len(passage):
passage2score[passage[cnt]].append(pred_scores[cnt] + filter_times)
cnt += 1
passage = sentences[1:]
final_score = []
for i in range(len(passage)):
p = passage[i]
final_score.append(passage2score[p][0])
return final_score

287
modules/Reranking.py Normal file
View File

@@ -0,0 +1,287 @@
# Copyright 2024 Bytedance Ltd. and/or its affiliates
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of this software
# and associated documentation files (the “Software”), to deal in the Software without
# restriction, including without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all copies or
# substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
# OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
# ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.
import logging
import numpy as np
from mteb import RerankingEvaluator, AbsTaskReranking
from tqdm import tqdm
logger = logging.getLogger(__name__)
class ChineseRerankingEvaluator(RerankingEvaluator):
"""
This class evaluates a SentenceTransformer model for the task of re-ranking.
Given a query and a list of documents, it computes the score [query, doc_i] for all possible
documents and sorts them in decreasing order. Then, MRR@10 and MAP is compute to measure the quality of the ranking.
:param samples: Must be a list and each element is of the form:
- {'query': '', 'positive': [], 'negative': []}. Query is the search query, positive is a list of positive
(relevant) documents, negative is a list of negative (irrelevant) documents.
- {'query': [], 'positive': [], 'negative': []}. Where query is a list of strings, which embeddings we average
to get the query embedding.
"""
def __call__(self, model):
scores = self.compute_metrics(model)
return scores
def compute_metrics(self, model):
return (
self.compute_metrics_batched(model)
if self.use_batched_encoding
else self.compute_metrics_individual(model)
)
def compute_metrics_batched(self, model):
"""
Computes the metrices in a batched way, by batching all queries and
all documents together
"""
if hasattr(model, 'compute_score'):
return self.compute_metrics_batched_from_crossencoder(model)
else:
return self.compute_metrics_batched_from_biencoder(model)
def compute_metrics_batched_from_crossencoder(self, model):
batch_size = 4
all_ap_scores = []
all_mrr_1_scores = []
all_mrr_5_scores = []
all_mrr_10_scores = []
all_scores = []
tmp_pairs = []
for sample in tqdm(self.samples, desc="Evaluating"):
b_pairs = [sample['query']]
for p in sample['positive']:
b_pairs.append(p)
for n in sample['negative']:
b_pairs.append(n)
tmp_pairs.append(b_pairs)
if len(tmp_pairs) == batch_size:
sample_scores = model.compute_score(tmp_pairs)
sample_scores = sum(sample_scores, [])
all_scores += sample_scores
tmp_pairs = []
if len(tmp_pairs) > 0:
sample_scores = model.compute_score(tmp_pairs)
sample_scores = sum(sample_scores, [])
all_scores += sample_scores
all_scores = np.array(all_scores)
start_inx = 0
for sample in tqdm(self.samples, desc="Evaluating"):
is_relevant = [True] * len(sample['positive']) + [False] * len(sample['negative'])
pred_scores = all_scores[start_inx:start_inx + len(is_relevant)]
start_inx += len(is_relevant)
pred_scores_argsort = np.argsort(-pred_scores) # Sort in decreasing order
ap = self.ap_score(is_relevant, pred_scores)
mrr_1 = self.mrr_at_k_score(is_relevant, pred_scores_argsort, 1)
mrr_5 = self.mrr_at_k_score(is_relevant, pred_scores_argsort, 5)
mrr_10 = self.mrr_at_k_score(is_relevant, pred_scores_argsort, 10)
all_mrr_1_scores.append(mrr_1)
all_mrr_5_scores.append(mrr_5)
all_mrr_10_scores.append(mrr_10)
all_ap_scores.append(ap)
mean_ap = np.mean(all_ap_scores)
mean_mrr_1 = np.mean(all_mrr_1_scores)
mean_mrr_5 = np.mean(all_mrr_5_scores)
mean_mrr_10 = np.mean(all_mrr_10_scores)
return {"map": mean_ap, "mrr_1": mean_mrr_1, 'mrr_5': mean_mrr_5, 'mrr_10': mean_mrr_10}
def compute_metrics_batched_from_biencoder(self, model):
all_mrr_scores = []
all_ap_scores = []
logger.info("Encoding queries...")
if isinstance(self.samples[0]["query"], str):
if hasattr(model, 'encode_queries'):
all_query_embs = model.encode_queries(
[sample["query"] for sample in self.samples],
convert_to_tensor=True,
batch_size=self.batch_size,
)
else:
all_query_embs = model.encode(
[sample["query"] for sample in self.samples],
convert_to_tensor=True,
batch_size=self.batch_size,
)
elif isinstance(self.samples[0]["query"], list):
# In case the query is a list of strings, we get the most similar embedding to any of the queries
all_query_flattened = [q for sample in self.samples for q in sample["query"]]
if hasattr(model, 'encode_queries'):
all_query_embs = model.encode_queries(all_query_flattened, convert_to_tensor=True,
batch_size=self.batch_size)
else:
all_query_embs = model.encode(all_query_flattened, convert_to_tensor=True, batch_size=self.batch_size)
else:
raise ValueError(f"Query must be a string or a list of strings but is {type(self.samples[0]['query'])}")
logger.info("Encoding candidates...")
all_docs = []
for sample in self.samples:
all_docs.extend(sample["positive"])
all_docs.extend(sample["negative"])
all_docs_embs = model.encode(all_docs, convert_to_tensor=True, batch_size=self.batch_size)
# Compute scores
logger.info("Evaluating...")
query_idx, docs_idx = 0, 0
for instance in self.samples:
num_subqueries = len(instance["query"]) if isinstance(instance["query"], list) else 1
query_emb = all_query_embs[query_idx: query_idx + num_subqueries]
query_idx += num_subqueries
num_pos = len(instance["positive"])
num_neg = len(instance["negative"])
docs_emb = all_docs_embs[docs_idx: docs_idx + num_pos + num_neg]
docs_idx += num_pos + num_neg
if num_pos == 0 or num_neg == 0:
continue
is_relevant = [True] * num_pos + [False] * num_neg
scores = self._compute_metrics_instance(query_emb, docs_emb, is_relevant)
all_mrr_scores.append(scores["mrr"])
all_ap_scores.append(scores["ap"])
mean_ap = np.mean(all_ap_scores)
mean_mrr = np.mean(all_mrr_scores)
return {"map": mean_ap, "mrr": mean_mrr}
def evaluate(self, model, split="test", **kwargs):
if not self.data_loaded:
self.load_data()
data_split = self.dataset[split]
evaluator = ChineseRerankingEvaluator(data_split, **kwargs)
scores = evaluator(model)
return dict(scores)
AbsTaskReranking.evaluate = evaluate
class T2Reranking(AbsTaskReranking):
@property
def description(self):
return {
'name': 'T2Reranking',
'hf_hub_name': "C-MTEB/T2Reranking",
'description': 'T2Ranking: A large-scale Chinese Benchmark for Passage Ranking',
"reference": "https://arxiv.org/abs/2304.03679",
'type': 'Reranking',
'category': 's2p',
'eval_splits': ['dev'],
'eval_langs': ['zh'],
'main_score': 'map',
}
class T2RerankingZh2En(AbsTaskReranking):
@property
def description(self):
return {
'name': 'T2RerankingZh2En',
'hf_hub_name': "C-MTEB/T2Reranking_zh2en",
'description': 'T2Ranking: A large-scale Chinese Benchmark for Passage Ranking',
"reference": "https://arxiv.org/abs/2304.03679",
'type': 'Reranking',
'category': 's2p',
'eval_splits': ['dev'],
'eval_langs': ['zh2en'],
'main_score': 'map',
}
class T2RerankingEn2Zh(AbsTaskReranking):
@property
def description(self):
return {
'name': 'T2RerankingEn2Zh',
'hf_hub_name': "C-MTEB/T2Reranking_en2zh",
'description': 'T2Ranking: A large-scale Chinese Benchmark for Passage Ranking',
"reference": "https://arxiv.org/abs/2304.03679",
'type': 'Reranking',
'category': 's2p',
'eval_splits': ['dev'],
'eval_langs': ['en2zh'],
'main_score': 'map',
}
class MMarcoReranking(AbsTaskReranking):
@property
def description(self):
return {
'name': 'MMarcoReranking',
'hf_hub_name': "C-MTEB/Mmarco-reranking",
'description': 'mMARCO is a multilingual version of the MS MARCO passage ranking dataset',
"reference": "https://github.com/unicamp-dl/mMARCO",
'type': 'Reranking',
'category': 's2p',
'eval_splits': ['dev'],
'eval_langs': ['zh'],
'main_score': 'map',
}
class CMedQAv1(AbsTaskReranking):
@property
def description(self):
return {
'name': 'CMedQAv1',
"hf_hub_name": "C-MTEB/CMedQAv1-reranking",
'description': 'Chinese community medical question answering',
"reference": "https://github.com/zhangsheng93/cMedQA",
'type': 'Reranking',
'category': 's2p',
'eval_splits': ['test'],
'eval_langs': ['zh'],
'main_score': 'map',
}
class CMedQAv2(AbsTaskReranking):
@property
def description(self):
return {
'name': 'CMedQAv2',
"hf_hub_name": "C-MTEB/CMedQAv2-reranking",
'description': 'Chinese community medical question answering',
"reference": "https://github.com/zhangsheng93/cMedQA2",
'type': 'Reranking',
'category': 's2p',
'eval_splits': ['test'],
'eval_langs': ['zh'],
'main_score': 'map',
}

325
modules/Reranking_loop.py Normal file
View File

@@ -0,0 +1,325 @@
# Copyright 2024 Bytedance Ltd. and/or its affiliates
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of this software
# and associated documentation files (the “Software”), to deal in the Software without
# restriction, including without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all copies or
# substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
# OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
# ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.
import logging
import numpy as np
from mteb import RerankingEvaluator, AbsTaskReranking
from tqdm import tqdm
import math
logger = logging.getLogger(__name__)
class ChineseRerankingEvaluator(RerankingEvaluator):
"""
This class evaluates a SentenceTransformer model for the task of re-ranking.
Given a query and a list of documents, it computes the score [query, doc_i] for all possible
documents and sorts them in decreasing order. Then, MRR@10 and MAP is compute to measure the quality of the ranking.
:param samples: Must be a list and each element is of the form:
- {'query': '', 'positive': [], 'negative': []}. Query is the search query, positive is a list of positive
(relevant) documents, negative is a list of negative (irrelevant) documents.
- {'query': [], 'positive': [], 'negative': []}. Where query is a list of strings, which embeddings we average
to get the query embedding.
"""
def __call__(self, model):
scores = self.compute_metrics(model)
return scores
def compute_metrics(self, model):
return (
self.compute_metrics_batched(model)
if self.use_batched_encoding
else self.compute_metrics_individual(model)
)
def compute_metrics_batched(self, model):
"""
Computes the metrices in a batched way, by batching all queries and
all documents together
"""
if hasattr(model, 'compute_score'):
return self.compute_metrics_batched_from_crossencoder(model)
else:
return self.compute_metrics_batched_from_biencoder(model)
def compute_metrics_batched_from_crossencoder(self, model):
all_ap_scores = []
all_mrr_1_scores = []
all_mrr_5_scores = []
all_mrr_10_scores = []
for sample in tqdm(self.samples, desc="Evaluating"):
query = sample['query']
pos = sample['positive']
neg = sample['negative']
passage = pos + neg
passage2label = {}
for p in pos:
passage2label[p] = True
for p in neg:
passage2label[p] = False
filter_times = 0
passage2score = {}
while len(passage) > 20:
batch = [[query] + passage]
pred_scores = model.compute_score(batch)[0]
# Sort in increasing order
pred_scores_argsort = np.argsort(pred_scores).tolist()
passage_len = len(passage)
to_filter_num = math.ceil(passage_len * 0.2)
if to_filter_num < 10:
to_filter_num = 10
have_filter_num = 0
while have_filter_num < to_filter_num:
idx = pred_scores_argsort[have_filter_num]
if passage[idx] in passage2score:
passage2score[passage[idx]].append(pred_scores[idx] + filter_times)
else:
passage2score[passage[idx]] = [pred_scores[idx] + filter_times]
have_filter_num += 1
while pred_scores[pred_scores_argsort[have_filter_num - 1]] == pred_scores[pred_scores_argsort[have_filter_num]]:
idx = pred_scores_argsort[have_filter_num]
if passage[idx] in passage2score:
passage2score[passage[idx]].append(pred_scores[idx] + filter_times)
else:
passage2score[passage[idx]] = [pred_scores[idx] + filter_times]
have_filter_num += 1
next_passage = []
next_passage_idx = have_filter_num
while next_passage_idx < len(passage):
idx = pred_scores_argsort[next_passage_idx]
next_passage.append(passage[idx])
next_passage_idx += 1
passage = next_passage
filter_times += 1
batch = [[query] + passage]
pred_scores = model.compute_score(batch)[0]
cnt = 0
while cnt < len(passage):
if passage[cnt] in passage2score:
passage2score[passage[cnt]].append(pred_scores[cnt] + filter_times)
else:
passage2score[passage[cnt]] = [pred_scores[cnt] + filter_times]
cnt += 1
passage = list(set(pos + neg))
is_relevant = []
final_score = []
for i in range(len(passage)):
p = passage[i]
is_relevant += [passage2label[p]] * len(passage2score[p])
final_score += passage2score[p]
ap = self.ap_score(is_relevant, final_score)
pred_scores_argsort = np.argsort(-(np.array(final_score)))
mrr_1 = self.mrr_at_k_score(is_relevant, pred_scores_argsort, 1)
mrr_5 = self.mrr_at_k_score(is_relevant, pred_scores_argsort, 5)
mrr_10 = self.mrr_at_k_score(is_relevant, pred_scores_argsort, 10)
all_ap_scores.append(ap)
all_mrr_1_scores.append(mrr_1)
all_mrr_5_scores.append(mrr_5)
all_mrr_10_scores.append(mrr_10)
mean_ap = np.mean(all_ap_scores)
mean_mrr_1 = np.mean(all_mrr_1_scores)
mean_mrr_5 = np.mean(all_mrr_5_scores)
mean_mrr_10 = np.mean(all_mrr_10_scores)
return {"map": mean_ap, "mrr_1": mean_mrr_1, 'mrr_5': mean_mrr_5, 'mrr_10': mean_mrr_10}
def compute_metrics_batched_from_biencoder(self, model):
all_mrr_scores = []
all_ap_scores = []
logger.info("Encoding queries...")
if isinstance(self.samples[0]["query"], str):
if hasattr(model, 'encode_queries'):
all_query_embs = model.encode_queries(
[sample["query"] for sample in self.samples],
convert_to_tensor=True,
batch_size=self.batch_size,
)
else:
all_query_embs = model.encode(
[sample["query"] for sample in self.samples],
convert_to_tensor=True,
batch_size=self.batch_size,
)
elif isinstance(self.samples[0]["query"], list):
# In case the query is a list of strings, we get the most similar embedding to any of the queries
all_query_flattened = [q for sample in self.samples for q in sample["query"]]
if hasattr(model, 'encode_queries'):
all_query_embs = model.encode_queries(all_query_flattened, convert_to_tensor=True,
batch_size=self.batch_size)
else:
all_query_embs = model.encode(all_query_flattened, convert_to_tensor=True, batch_size=self.batch_size)
else:
raise ValueError(f"Query must be a string or a list of strings but is {type(self.samples[0]['query'])}")
logger.info("Encoding candidates...")
all_docs = []
for sample in self.samples:
all_docs.extend(sample["positive"])
all_docs.extend(sample["negative"])
all_docs_embs = model.encode(all_docs, convert_to_tensor=True, batch_size=self.batch_size)
# Compute scores
logger.info("Evaluating...")
query_idx, docs_idx = 0, 0
for instance in self.samples:
num_subqueries = len(instance["query"]) if isinstance(instance["query"], list) else 1
query_emb = all_query_embs[query_idx: query_idx + num_subqueries]
query_idx += num_subqueries
num_pos = len(instance["positive"])
num_neg = len(instance["negative"])
docs_emb = all_docs_embs[docs_idx: docs_idx + num_pos + num_neg]
docs_idx += num_pos + num_neg
if num_pos == 0 or num_neg == 0:
continue
is_relevant = [True] * num_pos + [False] * num_neg
scores = self._compute_metrics_instance(query_emb, docs_emb, is_relevant)
all_mrr_scores.append(scores["mrr"])
all_ap_scores.append(scores["ap"])
mean_ap = np.mean(all_ap_scores)
mean_mrr = np.mean(all_mrr_scores)
return {"map": mean_ap, "mrr": mean_mrr}
def evaluate(self, model, split="test", **kwargs):
if not self.data_loaded:
self.load_data()
data_split = self.dataset[split]
evaluator = ChineseRerankingEvaluator(data_split, **kwargs)
scores = evaluator(model)
return dict(scores)
AbsTaskReranking.evaluate = evaluate
class T2Reranking(AbsTaskReranking):
@property
def description(self):
return {
'name': 'T2Reranking',
'hf_hub_name': "C-MTEB/T2Reranking",
'description': 'T2Ranking: A large-scale Chinese Benchmark for Passage Ranking',
"reference": "https://arxiv.org/abs/2304.03679",
'type': 'Reranking',
'category': 's2p',
'eval_splits': ['dev'],
'eval_langs': ['zh'],
'main_score': 'map',
}
class T2RerankingZh2En(AbsTaskReranking):
@property
def description(self):
return {
'name': 'T2RerankingZh2En',
'hf_hub_name': "C-MTEB/T2Reranking_zh2en",
'description': 'T2Ranking: A large-scale Chinese Benchmark for Passage Ranking',
"reference": "https://arxiv.org/abs/2304.03679",
'type': 'Reranking',
'category': 's2p',
'eval_splits': ['dev'],
'eval_langs': ['zh2en'],
'main_score': 'map',
}
class T2RerankingEn2Zh(AbsTaskReranking):
@property
def description(self):
return {
'name': 'T2RerankingEn2Zh',
'hf_hub_name': "C-MTEB/T2Reranking_en2zh",
'description': 'T2Ranking: A large-scale Chinese Benchmark for Passage Ranking',
"reference": "https://arxiv.org/abs/2304.03679",
'type': 'Reranking',
'category': 's2p',
'eval_splits': ['dev'],
'eval_langs': ['en2zh'],
'main_score': 'map',
}
class MMarcoReranking(AbsTaskReranking):
@property
def description(self):
return {
'name': 'MMarcoReranking',
'hf_hub_name': "C-MTEB/Mmarco-reranking",
'description': 'mMARCO is a multilingual version of the MS MARCO passage ranking dataset',
"reference": "https://github.com/unicamp-dl/mMARCO",
'type': 'Reranking',
'category': 's2p',
'eval_splits': ['dev'],
'eval_langs': ['zh'],
'main_score': 'map',
}
class CMedQAv1(AbsTaskReranking):
@property
def description(self):
return {
'name': 'CMedQAv1',
"hf_hub_name": "C-MTEB/CMedQAv1-reranking",
'description': 'Chinese community medical question answering',
"reference": "https://github.com/zhangsheng93/cMedQA",
'type': 'Reranking',
'category': 's2p',
'eval_splits': ['test'],
'eval_langs': ['zh'],
'main_score': 'map',
}
class CMedQAv2(AbsTaskReranking):
@property
def description(self):
return {
'name': 'CMedQAv2',
"hf_hub_name": "C-MTEB/CMedQAv2-reranking",
'description': 'Chinese community medical question answering',
"reference": "https://github.com/zhangsheng93/cMedQA2",
'type': 'Reranking',
'category': 's2p',
'eval_splits': ['test'],
'eval_langs': ['zh'],
'main_score': 'map',
}

161
modules/listconranker.py Normal file
View File

@@ -0,0 +1,161 @@
# Copyright 2024 Bytedance Ltd. and/or its affiliates
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of this software
# and associated documentation files (the “Software”), to deal in the Software without
# restriction, including without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all copies or
# substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
# OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
# ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.
import math
import torch
import numpy as np
from transformers import AutoTokenizer, is_torch_npu_available
from typing import Union, List
from .modeling import CrossEncoder
import os
os.environ['TOKENIZERS_PARALLELISM'] = 'true'
def sigmoid(x):
return 1 / (1 + np.exp(-x))
class ListConRanker:
def __init__(
self,
model_name_or_path: str = None,
use_fp16: bool = False,
cache_dir: str = None,
device: Union[str, int] = None,
list_transformer_layer = None
) -> None:
self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, cache_dir=cache_dir)
self.model = CrossEncoder.from_pretrained_for_eval(model_name_or_path, list_transformer_layer)
if device and isinstance(device, str):
self.device = torch.device(device)
if device == 'cpu':
use_fp16 = False
else:
if torch.cuda.is_available():
if device is not None:
self.device = torch.device(f"cuda:{device}")
else:
self.device = torch.device("cuda")
elif torch.backends.mps.is_available():
self.device = torch.device("mps")
elif is_torch_npu_available():
self.device = torch.device("npu")
else:
self.device = torch.device("cpu")
use_fp16 = False
if use_fp16:
self.model.half()
self.model = self.model.to(self.device)
self.model.eval()
if device is None:
self.num_gpus = torch.cuda.device_count()
if self.num_gpus > 1:
print(f"----------using {self.num_gpus}*GPUs----------")
self.model = torch.nn.DataParallel(self.model)
else:
self.num_gpus = 1
@torch.no_grad()
def compute_score(self, sentence_pairs: List[List[str]], max_length: int = 512) -> List[List[float]]:
pair_nums = [len(pairs) - 1 for pairs in sentence_pairs]
sentences_batch = sum(sentence_pairs, [])
inputs = self.tokenizer(
sentences_batch,
padding=True,
truncation=True,
return_tensors='pt',
max_length=max_length,
).to(self.device)
inputs['pair_num'] = torch.LongTensor(pair_nums)
scores = self.model(inputs).float()
all_scores = scores.cpu().numpy().tolist()
if isinstance(all_scores, float):
return [all_scores]
result = []
curr_idx = 0
for i in range(len(pair_nums)):
result.append(all_scores[curr_idx: curr_idx + pair_nums[i]])
curr_idx += pair_nums[i]
# return all_scores
return result
@torch.no_grad()
def iterative_inference(self, sentence_pairs: List[str], max_length: int = 512) -> List[float]:
query = sentence_pairs[0]
passage = sentence_pairs[1:]
filter_times = 0
passage2score = {}
while len(passage) > 20:
batch = [[query] + passage]
pred_scores = self.compute_score(batch, max_length)[0]
# Sort in increasing order
pred_scores_argsort = np.argsort(pred_scores).tolist()
passage_len = len(passage)
to_filter_num = math.ceil(passage_len * 0.2)
if to_filter_num < 10:
to_filter_num = 10
have_filter_num = 0
while have_filter_num < to_filter_num:
idx = pred_scores_argsort[have_filter_num]
if passage[idx] in passage2score:
passage2score[passage[idx]].append(pred_scores[idx] + filter_times)
else:
passage2score[passage[idx]] = [pred_scores[idx] + filter_times]
have_filter_num += 1
while pred_scores[pred_scores_argsort[have_filter_num - 1]] == pred_scores[pred_scores_argsort[have_filter_num]]:
idx = pred_scores_argsort[have_filter_num]
if passage[idx] in passage2score:
passage2score[passage[idx]].append(pred_scores[idx] + filter_times)
else:
passage2score[passage[idx]] = [pred_scores[idx] + filter_times]
have_filter_num += 1
next_passage = []
next_passage_idx = have_filter_num
while next_passage_idx < len(passage):
idx = pred_scores_argsort[next_passage_idx]
next_passage.append(passage[idx])
next_passage_idx += 1
passage = next_passage
filter_times += 1
batch = [[query] + passage]
pred_scores = self.compute_score(batch, max_length)[0]
cnt = 0
while cnt < len(passage):
if passage[cnt] in passage2score:
passage2score[passage[cnt]].append(pred_scores[cnt] + filter_times)
else:
passage2score[passage[cnt]] = [pred_scores[cnt] + filter_times]
cnt += 1
passage = sentence_pairs[1:]
final_score = []
for i in range(len(passage)):
p = passage[i]
final_score += passage2score[p]
return final_score

174
modules/modeling.py Normal file
View File

@@ -0,0 +1,174 @@
# Copyright 2024 Bytedance Ltd. and/or its affiliates
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of this software
# and associated documentation files (the “Software”), to deal in the Software without
# restriction, including without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all copies or
# substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
# OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
# ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.
import logging
import torch
from torch import nn
from transformers import AutoModel, PreTrainedModel
from torch.nn import functional as F
logger = logging.getLogger(__name__)
class ListTransformer(nn.Module):
def __init__(self, num_layer, config, device) -> None:
super().__init__()
self.config = config
self.device = device
self.list_transformer_layer = nn.TransformerEncoderLayer(1792, self.config.num_attention_heads, batch_first=True, activation=F.gelu, norm_first=False)
self.list_transformer = nn.TransformerEncoder(self.list_transformer_layer, num_layer)
self.relu = nn.ReLU()
self.query_embedding = QueryEmbedding(config, device)
self.linear_score3 = nn.Linear(1792 * 2, 1792)
self.linear_score2 = nn.Linear(1792 * 2, 1792)
self.linear_score1 = nn.Linear(1792 * 2, 1)
def forward(self, pair_features, pair_nums):
pair_nums = [x + 1 for x in pair_nums]
batch_pair_features = pair_features.split(pair_nums)
pair_feature_query_passage_concat_list = []
for i in range(len(batch_pair_features)):
pair_feature_query = batch_pair_features[i][0].unsqueeze(0).repeat(pair_nums[i] - 1, 1)
pair_feature_passage = batch_pair_features[i][1:]
pair_feature_query_passage_concat_list.append(torch.cat([pair_feature_query, pair_feature_passage], dim=1))
pair_feature_query_passage_concat = torch.cat(pair_feature_query_passage_concat_list, dim=0)
batch_pair_features = nn.utils.rnn.pad_sequence(batch_pair_features, batch_first=True)
query_embedding_tags = torch.zeros(batch_pair_features.size(0), batch_pair_features.size(1), dtype=torch.long, device=self.device)
query_embedding_tags[:, 0] = 1
batch_pair_features = self.query_embedding(batch_pair_features, query_embedding_tags)
mask = self.generate_attention_mask(pair_nums)
query_mask = self.generate_attention_mask_custom(pair_nums)
pair_list_features = self.list_transformer(batch_pair_features, src_key_padding_mask=mask, mask=query_mask)
output_pair_list_features = []
output_query_list_features = []
pair_features_after_transformer_list = []
for idx, pair_num in enumerate(pair_nums):
output_pair_list_features.append(pair_list_features[idx, 1:pair_num, :])
output_query_list_features.append(pair_list_features[idx, 0, :])
pair_features_after_transformer_list.append(pair_list_features[idx, :pair_num, :])
pair_features_after_transformer_cat_query_list = []
for idx, pair_num in enumerate(pair_nums):
query_ft = output_query_list_features[idx].unsqueeze(0).repeat(pair_num - 1, 1)
pair_features_after_transformer_cat_query = torch.cat([query_ft, output_pair_list_features[idx]], dim=1)
pair_features_after_transformer_cat_query_list.append(pair_features_after_transformer_cat_query)
pair_features_after_transformer_cat_query = torch.cat(pair_features_after_transformer_cat_query_list, dim=0)
pair_feature_query_passage_concat = self.relu(self.linear_score2(pair_feature_query_passage_concat))
pair_features_after_transformer_cat_query = self.relu(self.linear_score3(pair_features_after_transformer_cat_query))
final_ft = torch.cat([pair_feature_query_passage_concat, pair_features_after_transformer_cat_query], dim=1)
logits = self.linear_score1(final_ft).squeeze()
return logits, torch.cat(pair_features_after_transformer_list, dim=0)
def generate_attention_mask(self, pair_num):
max_len = max(pair_num)
batch_size = len(pair_num)
mask = torch.zeros(batch_size, max_len, dtype=torch.bool, device=self.device)
for i, length in enumerate(pair_num):
mask[i, length:] = True
return mask
def generate_attention_mask_custom(self, pair_num):
max_len = max(pair_num)
mask = torch.zeros(max_len, max_len, dtype=torch.bool, device=self.device)
mask[0, 1:] = True
return mask
class QueryEmbedding(nn.Module):
def __init__(self, config, device) -> None:
super().__init__()
self.query_embedding = nn.Embedding(2, 1792)
self.layerNorm = nn.LayerNorm(1792)
def forward(self, x, tags):
query_embeddings = self.query_embedding(tags)
x += query_embeddings
x = self.layerNorm(x)
return x
class CrossEncoder(nn.Module):
def __init__(self, hf_model: PreTrainedModel, list_transformer_layer_4eval: int=None):
super().__init__()
self.hf_model = hf_model
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.sigmoid = nn.Sigmoid()
self.config = self.hf_model.config
self.config.output_hidden_states = True
self.linear_in_embedding = nn.Linear(1024, 1792)
self.list_transformer_layer = list_transformer_layer_4eval
self.list_transformer = ListTransformer(self.list_transformer_layer, self.config, self.device)
def forward(self, batch):
if 'pair_num' in batch:
pair_nums = batch.pop('pair_num').tolist()
if self.training:
pass
else:
split_batch = 400
input_ids = batch['input_ids']
attention_mask = batch['attention_mask']
if sum(pair_nums) > split_batch:
last_hidden_state_list = []
input_ids_list = input_ids.split(split_batch)
attention_mask_list = attention_mask.split(split_batch)
for i in range(len(input_ids_list)):
last_hidden_state = self.hf_model(input_ids=input_ids_list[i], attention_mask=attention_mask_list[i], return_dict=True).hidden_states[-1]
last_hidden_state_list.append(last_hidden_state)
last_hidden_state = torch.cat(last_hidden_state_list, dim=0)
else:
ranker_out = self.hf_model(**batch, return_dict=True)
last_hidden_state = ranker_out.last_hidden_state
pair_features = self.average_pooling(last_hidden_state, attention_mask)
pair_features = self.linear_in_embedding(pair_features)
logits, pair_features_after_list_transformer = self.list_transformer(pair_features, pair_nums)
logits = self.sigmoid(logits)
return logits
@classmethod
def from_pretrained_for_eval(cls, model_name_or_path, list_transformer_layer):
hf_model = AutoModel.from_pretrained(model_name_or_path)
reranker = cls(hf_model, list_transformer_layer)
reranker.linear_in_embedding.load_state_dict(torch.load(model_name_or_path + '/linear_in_embedding.pt'))
reranker.list_transformer.load_state_dict(torch.load(model_name_or_path + '/list_transformer.pt'))
return reranker
def average_pooling(self, hidden_state, attention_mask):
extended_attention_mask = attention_mask.unsqueeze(-1).expand(hidden_state.size()).to(dtype=hidden_state.dtype)
masked_hidden_state = hidden_state * extended_attention_mask
sum_embeddings = torch.sum(masked_hidden_state, dim=1)
sum_mask = extended_attention_mask.sum(dim=1)
return sum_embeddings / sum_mask

4
requirements.txt Normal file
View File

@@ -0,0 +1,4 @@
mteb==1.1.1
torch==2.1.2
tqdm==4.67.0
transformers==4.46.2

37
special_tokens_map.json Normal file
View File

@@ -0,0 +1,37 @@
{
"cls_token": {
"content": "[CLS]",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false
},
"mask_token": {
"content": "[MASK]",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false
},
"pad_token": {
"content": "[PAD]",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false
},
"sep_token": {
"content": "[SEP]",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false
},
"unk_token": {
"content": "[UNK]",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false
}
}

64
tokenizer_config.json Normal file
View File

@@ -0,0 +1,64 @@
{
"added_tokens_decoder": {
"0": {
"content": "[PAD]",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"100": {
"content": "[UNK]",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"101": {
"content": "[CLS]",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"102": {
"content": "[SEP]",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"103": {
"content": "[MASK]",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
}
},
"clean_up_tokenization_spaces": true,
"cls_token": "[CLS]",
"do_basic_tokenize": true,
"do_lower_case": true,
"mask_token": "[MASK]",
"max_length": 512,
"model_max_length": 512,
"never_split": null,
"pad_to_multiple_of": null,
"pad_token": "[PAD]",
"pad_token_type_id": 0,
"padding_side": "right",
"sep_token": "[SEP]",
"stride": 0,
"strip_accents": null,
"tokenize_chinese_chars": true,
"tokenizer_class": "BertTokenizer",
"truncation_side": "right",
"truncation_strategy": "longest_first",
"unk_token": "[UNK]"
}

21128
vocab.txt Normal file
View File

File diff suppressed because it is too large Load Diff