upload codes and model weights

LICENSE

@@ -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.

ListConRanker_ckpt/config.json

@@ -0,0 +1,38 @@
+{
+  "architectures": [
+    "BertModel"
+  ],
+  "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,
+  "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
+}

ListConRanker_ckpt/linear_in_embedding.pt

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

ListConRanker_ckpt/list_transformer.pt

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

ListConRanker_ckpt/model.safetensors

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

ListConRanker_ckpt/special_tokens_map.json

@@ -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
+  }
+}

ListConRanker_ckpt/tokenizer_config.json

@@ -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": 1000000000000000019884624838656,
+  "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]"
+}

README.md

@@ -1,3 +1,164 @@
----
-license: mit
----
+---
+model-index:
+- name: ListConRanker
+  results:
+  - dataset:
+      config: default
+      name: MTEB CMedQAv1-reranking (default)
+      revision: null
+      split: test
+      type: C-MTEB/CMedQAv1-reranking
+    metrics:
+    - type: map
+      value: 90.55366308098787
+    - type: mrr_1
+      value: 87.8
+    - type: mrr_10
+      value: 92.45134920634919
+    - type: mrr_5
+      value: 92.325
+    - type: main_score
+      value: 90.55366308098787
+    task:
+      type: Reranking
+  - dataset:
+      config: default
+      name: MTEB CMedQAv2-reranking (default)
+      revision: null
+      split: test
+      type: C-MTEB/CMedQAv2-reranking
+    metrics:
+    - type: map
+      value: 89.38076135722042
+    - type: mrr_1
+      value: 85.9
+    - type: mrr_10
+      value: 91.28769841269842
+    - type: mrr_5
+      value: 91.08999999999999
+    - type: main_score
+      value: 89.38076135722042
+    task:
+      type: Reranking
+  - dataset:
+      config: default
+      name: MTEB MMarcoReranking (default)
+      revision: null
+      split: dev
+      type: C-MTEB/Mmarco-reranking
+    metrics:
+    - type: map
+      value: 43.881461866703894
+    - type: mrr_1
+      value: 32.0
+    - type: mrr_10
+      value: 44.700793650793656
+    - type: mrr_5
+      value: 43.61666666666667
+    - type: main_score
+      value: 43.881461866703894
+    task:
+      type: Reranking
+  - dataset:
+      config: default
+      name: MTEB T2Reranking (default)
+      revision: null
+      split: dev
+      type: C-MTEB/T2Reranking
+    metrics:
+    - type: map
+      value: 69.16513825032682
+    - type: mrr_1
+      value: 67.41706161137441
+    - type: mrr_10
+      value: 80.0946053776961
+    - type: mrr_5
+      value: 79.71676822387724
+    - type: main_score
+      value: 69.16513825032682
+    task:
+      type: Reranking
+tags:
+- mteb
+---
+
+
+# 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 modules.listconranker import ListConRanker
+
+reranker = ListConRanker('./ListConRanker_ckpt', use_fp16=True, list_transformer_layer=2)
+
+# [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.compute_score(batch)
+print(scores)
+# [[0.5126953125, 0.331298828125, 0.3642578125], [0.63671875, 0.71630859375, 0.42822265625, 0.35302734375]]
+
+# for iterative inferfence, only a batch size of 1 is supported
+# the scores do not indicate similarity but are intended only for ranking
+scores = reranker.iterative_inference(batch[0])
+print(scores)
+# [0.5126953125, 0.331298828125, 0.3642578125]
+```
+
+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/).

eval_listconranker.py

@@ -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="./ListConRanker_ckpt", 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))

eval_listconranker_iterative_inference.py

@@ -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="./ListConRanker_ckpt", 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))

modules/Reranking.py

@@ -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',
+        }

modules/Reranking_loop.py

@@ -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',
+        }

modules/listconranker.py

@@ -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

modules/modeling.py

@@ -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

requirements.txt

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