Upload 14 files

src/cross_rerank/collator.py

@@ -0,0 +1,30 @@
+import torch
+import transformers
+from dataclasses import dataclass
+from typing import List, Dict, Any
+from transformers import BatchEncoding, DataCollatorWithPadding
+
+
+@dataclass
+class CrossEncoderCollator(DataCollatorWithPadding):
+
+    def __call__(self, features: List[Dict[str, Any]]) -> BatchEncoding:
+        unpack_features = []
+        for ex in features:
+            keys = list(ex.keys())
+            # assert all(len(ex[k]) == 8 for k in keys)
+            for idx in range(len(ex[keys[0]])):
+                unpack_features.append({k: ex[k][idx] for k in keys})
+
+        old_level = transformers.logging.get_verbosity()
+        transformers.logging.set_verbosity_error()
+        collated_batch_dict = self.tokenizer.pad(
+            unpack_features,
+            padding=self.padding,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            return_tensors=self.return_tensors)
+        transformers.logging.set_verbosity(old_level)
+        
+        collated_batch_dict['labels'] = torch.zeros(len(features), dtype=torch.long)
+
+        return collated_batch_dict

src/cross_rerank/config.py

@@ -0,0 +1,230 @@
+import os
+import torch
+from dataclasses import dataclass, field
+from typing import Optional
+from transformers import TrainingArguments
+
+
+from .logger_config import logger
+
+
+@dataclass
+class Arguments(TrainingArguments):
+    model_name_or_path: str = field(
+        default='bert-base-uncased',
+        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
+    )
+    
+    corpus_file: str = field(
+        default=None, metadata={"help": "Path to corpus file"}
+    )
+
+    data_dir: str = field(
+        default=None, metadata={"help": "Path to train directory"}
+    )
+    task_type: str = field(
+        default='ir', metadata={"help": "task type: ir / qa"}
+    )
+    train_file: Optional[str] = field(
+        default=None, metadata={"help": "The input training data file (a jsonlines file)."}
+    )
+    validation_file: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "An optional input evaluation data file to evaluate the metrics on (a jsonlines file)."
+        },
+    )
+
+    train_n_passages: int = field(
+        default=8,
+        metadata={"help": "number of passages for each example (including both positive and negative passages)"}
+    )
+    share_encoder: bool = field(
+        default=True,
+        metadata={"help": "no weight sharing between qry passage encoders"}
+    )
+    use_first_positive: bool = field(
+        default=False,
+        metadata={"help": "Always use the first positive passage"}
+    )
+    use_first_negative: bool = field(
+        default=False,
+        metadata={"help": "Always use the first positive passage"}
+    )
+    use_scaled_loss: bool = field(
+        default=True,
+        metadata={"help": "Use scaled loss or not"}
+    )
+    loss_scale: float = field(
+        default=-1.,
+        metadata={"help": "loss scale, -1 will use world_size"}
+    )
+    add_pooler: bool = field(default=False)
+    out_dimension: int = field(
+        default=768,
+        metadata={"help": "output dimension for pooler"}
+    )
+    t: float = field(default=0.05, metadata={"help": "temperature of biencoder training"})
+    l2_normalize: bool = field(default=True, metadata={"help": "L2 normalize embeddings or not"})
+    t_warmup: bool = field(default=False, metadata={"help": "warmup temperature"})
+    full_contrastive_loss: bool = field(default=True, metadata={"help": "use full contrastive loss or not"})
+
+    # following arguments are used for encoding documents
+    do_encode: bool = field(default=False, metadata={"help": "run the encoding loop"})
+    encode_in_path: str = field(default=None, metadata={"help": "Path to data to encode"})
+    encode_save_dir: str = field(default=None, metadata={"help": "where to save the encode"})
+    encode_shard_size: int = field(default=int(2 * 10**6))
+    encode_batch_size: int = field(default=256)
+
+    # used for index search
+    do_search: bool = field(default=False, metadata={"help": "run the index search loop"})
+    search_split: str = field(default='dev', metadata={"help": "which split to search"})
+    search_batch_size: int = field(default=128, metadata={"help": "query batch size for index search"})
+    search_topk: int = field(default=200, metadata={"help": "return topk search results"})
+    search_out_dir: str = field(default='', metadata={"help": "output directory for writing search results"})
+
+    # used for reranking
+    do_rerank: bool = field(default=False, metadata={"help": "run the reranking loop"})
+    rerank_max_length: int = field(default=256, metadata={"help": "max length for rerank inputs"})
+    rerank_in_path: str = field(default='', metadata={"help": "Path to predictions for rerank"})
+    rerank_out_path: str = field(default='', metadata={"help": "Path to write rerank results"})
+    rerank_split: str = field(default='dev', metadata={"help": "which split to rerank"})
+    rerank_batch_size: int = field(default=128, metadata={"help": "rerank batch size"})
+    rerank_depth: int = field(default=1000, metadata={"help": "rerank depth, useful for debugging purpose"})
+    rerank_forward_factor: int = field(
+        default=1,
+        metadata={"help": "forward n passages, then select top n/factor passages for backward"}
+    )
+    rerank_use_rdrop: bool = field(default=False, metadata={"help": "use R-Drop regularization for re-ranker"})
+
+    # used for knowledge distillation
+    do_kd_gen_score: bool = field(default=False, metadata={"help": "run the score generation for distillation"})
+    kd_gen_score_split: str = field(default='dev', metadata={
+        "help": "Which split to use for generation of teacher score"
+    })
+    kd_gen_score_batch_size: int = field(default=128, metadata={"help": "batch size for teacher score generation"})
+    kd_gen_score_n_neg: int = field(default=30, metadata={"help": "number of negatives to compute teacher scores"})
+
+    do_kd_biencoder: bool = field(default=False, metadata={"help": "knowledge distillation to biencoder"})
+    kd_mask_hn: bool = field(default=True, metadata={"help": "mask out hard negatives for distillation"})
+    kd_cont_loss_weight: float = field(default=1.0, metadata={"help": "weight for contrastive loss"})
+
+    rlm_generator_model_name: Optional[str] = field(
+        default='google/electra-base-generator',
+        metadata={"help": "generator for replace LM pre-training"}
+    )
+    rlm_freeze_generator: Optional[bool] = field(
+        default=True,
+        metadata={'help': 'freeze generator params or not'}
+    )
+    rlm_generator_mlm_weight: Optional[float] = field(
+        default=0.2,
+        metadata={'help': 'weight for generator MLM loss'}
+    )
+    all_use_mask_token: Optional[bool] = field(
+        default=False,
+        metadata={'help': 'Do not use 80:10:10 mask, use [MASK] for all places'}
+    )
+    rlm_num_eval_samples: Optional[int] = field(
+        default=4096,
+        metadata={"help": "number of evaluation samples pre-training"}
+    )
+    rlm_max_length: Optional[int] = field(
+        default=144,
+        metadata={"help": "max length for MatchLM pre-training"}
+    )
+    rlm_decoder_layers: Optional[int] = field(
+        default=2,
+        metadata={"help": "number of transformer layers for MatchLM decoder part"}
+    )
+    rlm_encoder_mask_prob: Optional[float] = field(
+        default=0.3,
+        metadata={'help': 'mask rate for encoder'}
+    )
+    rlm_decoder_mask_prob: Optional[float] = field(
+        default=0.5,
+        metadata={'help': 'mask rate for decoder'}
+    )
+
+    q_max_len: int = field(
+        default=32,
+        metadata={
+            "help": "The maximum total input sequence length after tokenization for query."
+        },
+    )
+    p_max_len: int = field(
+        default=144,
+        metadata={
+            "help": "The maximum total input sequence length after tokenization for passage."
+        },
+    )
+    
+    chunk_size: int = field(
+        default=8,
+        metadata={
+            "help": "The maximum total chunk"
+        },
+    )
+    
+    max_train_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "For debugging purposes or quicker training, truncate the number of training examples to this "
+                    "value if set."
+        },
+    )
+    dry_run: Optional[bool] = field(
+        default=False,
+        metadata={'help': 'Set dry_run to True for debugging purpose'}
+    )
+
+    def __post_init__(self):
+        assert os.path.exists(self.data_dir)
+        assert torch.cuda.is_available(), 'Only support running on GPUs'
+        assert self.task_type in ['ir', 'qa']
+
+        if self.dry_run:
+            self.logging_steps = 1
+            self.max_train_samples = self.max_train_samples or 128
+            self.num_train_epochs = 1
+            self.per_device_train_batch_size = min(2, self.per_device_train_batch_size)
+            self.train_n_passages = min(4, self.train_n_passages)
+            self.rerank_forward_factor = 1
+            self.gradient_accumulation_steps = 1
+            self.rlm_num_eval_samples = min(256, self.rlm_num_eval_samples)
+            self.max_steps = 30
+            self.save_steps = self.eval_steps = 30
+            logger.warning('Dry run: set logging_steps=1')
+
+        if self.do_encode:
+            assert self.encode_save_dir
+            os.makedirs(self.encode_save_dir, exist_ok=True)
+            assert os.path.exists(self.encode_in_path)
+
+        if self.do_search:
+            assert os.path.exists(self.encode_save_dir)
+            assert self.search_out_dir
+            os.makedirs(self.search_out_dir, exist_ok=True)
+
+        if self.do_rerank:
+            assert os.path.exists(self.rerank_in_path)
+            logger.info('Rerank result will be written to {}'.format(self.rerank_out_path))
+            assert self.train_n_passages > 1, 'Having positive passages only does not make sense for training re-ranker'
+            assert self.train_n_passages % self.rerank_forward_factor == 0
+
+        if self.do_kd_gen_score:
+            assert os.path.exists('{}/{}.jsonl'.format(self.data_dir, self.kd_gen_score_split))
+
+        if self.do_kd_biencoder:
+            if self.use_scaled_loss:
+                assert not self.kd_mask_hn, 'Use scaled loss only works with not masking out hard negatives'
+
+        if torch.cuda.device_count() <= 1:
+            self.logging_steps = min(50, self.logging_steps)
+
+        super(Arguments, self).__post_init__()
+
+        if self.output_dir:
+            os.makedirs(self.output_dir, exist_ok=True)
+
+        self.label_names = ['labels']

src/cross_rerank/data_loader.py

@@ -0,0 +1,106 @@
+import os.path
+import random
+import pandas as pd
+import transformers
+from typing import Tuple, Dict, List, Optional
+from datasets import load_dataset, DatasetDict, Dataset
+from transformers.file_utils import PaddingStrategy
+from transformers import PreTrainedTokenizerFast, Trainer
+
+from .config import Arguments
+from .logger_config import logger
+from .loader_utils import group_doc_ids
+
+
+class CrossEncoderDataLoader:
+    def __init__(self, args: Arguments, tokenizer):
+        self.args = args
+        self.negative_size = args.train_n_passages - 1
+        assert self.negative_size > 0
+        #self.hard_neg_size = args.hard_neg_size if args.hard_neg_size < self.negative_size else self.negative_size
+        #self.rand_neg_size = self.negative_size - self.hard_neg_size if self.hard_neg_size < self.negative_size else 0
+        self.tokenizer = tokenizer
+        #corpus_path = os.path.join(args.data_dir, 'passages.jsonl.gz')
+        #self.corpus: Dataset = load_dataset('json', data_files=corpus_path)['train']
+        self.corpus = pd.read_csv(args.corpus_file)
+        #self.corpus = dcorpus['tokenized_text'].to_list()
+        #self.corpus_bad_ids = [idx for idx in range(len(self.corpus)) if (type(self.corpus['bm25_text'][idx]) is not str)]
+        self.train_dataset, self.eval_dataset = self._get_transformed_datasets()
+
+        # use its state to decide which positives/negatives to sample
+        self.trainer: Optional[Trainer] = None
+
+    def _transform_func(self, examples: Dict[str, List]) -> Dict[str, List]:
+        current_epoch = int(self.trainer.state.epoch or 0)
+
+        input_doc_ids = group_doc_ids(
+            examples=examples,
+            negative_size=self.negative_size,
+            offset=current_epoch + self.args.seed,
+            use_first_positive=self.args.use_first_positive,
+            use_first_negative=self.args.use_first_negative
+        )
+        assert len(input_doc_ids) == len(examples['query']) * self.args.train_n_passages
+
+        input_queries, input_docs = [], []
+        for idx, doc_id in enumerate(input_doc_ids):
+            #prefix = ''
+            #if self.corpus[doc_id].get('title', ''):
+            #    prefix = self.corpus[doc_id]['title'] + ': '
+
+            input_docs.append(self.corpus['tokenized_text'][doc_id])
+            input_queries.append(examples['query'][idx // self.args.train_n_passages])
+
+        old_level = transformers.logging.get_verbosity()
+        transformers.logging.set_verbosity_error()
+        batch_dict = self.tokenizer(input_queries,
+                                    text_pair=input_docs,
+                                    max_length=self.args.rerank_max_length,
+                                    padding=PaddingStrategy.DO_NOT_PAD,
+                                    truncation=True)
+        transformers.logging.set_verbosity(old_level)
+        
+        packed_batch_dict = {}
+        for k in batch_dict:
+            packed_batch_dict[k] = []
+            assert len(examples['query']) * self.args.train_n_passages == len(batch_dict[k])
+            for idx in range(len(examples['query'])):
+                start = idx * self.args.train_n_passages
+                packed_batch_dict[k].append(batch_dict[k][start:(start + self.args.train_n_passages)])
+
+        return packed_batch_dict
+
+    def _get_transformed_datasets(self) -> Tuple:
+        #data_files = {}
+        #if self.args.train_file is not None:
+        #    data_files["train"] = self.args.train_file.split(',')
+        #if self.args.validation_file is not None:
+        #    data_files["validation"] = self.args.validation_file
+        #raw_datasets: DatasetDict = load_dataset('json', data_files=data_files)
+
+        train_dataset, eval_dataset = None, None
+
+        if self.args.do_train:
+            #if "train" not in raw_datasets:
+            try:
+                train_dataset = load_dataset('json', data_files = os.path.join(self.args.data_dir, 'train.jsonl'))['train']
+            except:
+                raise ValueError("--do_train requires a train dataset")
+            #train_dataset = raw_datasets["train"]
+            if self.args.max_train_samples is not None:
+                train_dataset = train_dataset.select(range(self.args.max_train_samples))
+            # Log a few random samples from the training set:
+            for index in random.sample(range(len(train_dataset)), 1):
+                logger.info(f"Sample {index} of the training set: {train_dataset[index]}.")
+            train_dataset.set_transform(self._transform_func)
+
+        if self.args.do_eval:
+            #if "validation" not in raw_datasets:
+            try:
+                eval_dataset = load_dataset('json', data_files = os.path.join(self.args.data_dir, 'eval.jsonl'))['train']
+            except:
+                raise ValueError("--do_eval requires a validation dataset")
+            #eval_dataset = raw_datasets["validation"]
+            eval_dataset.set_transform(self._transform_func)
+
+        return train_dataset, eval_dataset

src/cross_rerank/data_utils.py

@@ -0,0 +1,211 @@
+import os
+import random
+import tqdm
+import json
+
+from typing import Dict, List, Any
+from datasets import load_dataset, Dataset
+from dataclasses import dataclass, field
+
+from .logger_config import logger
+from .config import Arguments
+from .utils import save_json_to_file
+
+
+@dataclass
+class ScoredDoc:
+    qid: str
+    pid: str
+    rank: int
+    score: float = field(default=-1)
+
+
+def load_qrels(path: str) -> Dict[str, Dict[str, int]]:
+    assert path.endswith('.txt')
+
+    # qid -> pid -> score
+    qrels = {}
+    for line in open(path, 'r', encoding='utf-8'):
+        qid, _, pid, score = line.strip().split('\t')
+        if qid not in qrels:
+            qrels[qid] = {}
+        qrels[qid][pid] = int(score)
+
+    logger.info('Load {} queries {} qrels from {}'.format(len(qrels), sum(len(v) for v in qrels.values()), path))
+    return qrels
+
+
+def load_queries(path: str, task_type: str = 'ir') -> Dict[str, str]:
+    assert path.endswith('.tsv')
+
+    if task_type == 'qa':
+        qid_to_query = load_query_answers(path)
+        qid_to_query = {k: v['query'] for k, v in qid_to_query.items()}
+    elif task_type == 'ir':
+        qid_to_query = {}
+        for line in open(path, 'r', encoding='utf-8'):
+            qid, query = line.strip().split('\t')
+            qid_to_query[qid] = query
+    else:
+        raise ValueError('Unknown task type: {}'.format(task_type))
+
+    logger.info('Load {} queries from {}'.format(len(qid_to_query), path))
+    return qid_to_query
+
+
+def normalize_qa_text(text: str) -> str:
+    # TriviaQA has some weird formats
+    # For example: """What breakfast food gets its name from the German word for """"stirrup""""?"""
+    while text.startswith('"') and text.endswith('"'):
+        text = text[1:-1].replace('""', '"')
+    return text
+
+
+def get_question_key(question: str) -> str:
+    # For QA dataset, we'll use normalized question strings as dict key
+    return question
+
+
+def load_query_answers(path: str) -> Dict[str, Dict[str, Any]]:
+    assert path.endswith('.tsv')
+
+    qid_to_query = {}
+    for line in open(path, 'r', encoding='utf-8'):
+        query, answers = line.strip().split('\t')
+        query = normalize_qa_text(query)
+        answers = normalize_qa_text(answers)
+        qid = get_question_key(query)
+        if qid in qid_to_query:
+            logger.warning('Duplicate question: {} vs {}'.format(query, qid_to_query[qid]['query']))
+            continue
+
+        qid_to_query[qid] = {}
+        qid_to_query[qid]['query'] = query
+        qid_to_query[qid]['answers'] = list(eval(answers))
+
+    logger.info('Load {} queries from {}'.format(len(qid_to_query), path))
+    return qid_to_query
+
+
+def load_corpus(path: str) -> Dataset:
+    assert path.endswith('.jsonl') or path.endswith('.jsonl.gz')
+
+    # two fields: id, contents
+    corpus = load_dataset('json', data_files=path)['train']
+    logger.info('Load {} documents from {} with columns {}'.format(len(corpus), path, corpus.column_names))
+    logger.info('A random document: {}'.format(random.choice(corpus)))
+    return corpus
+
+
+def load_msmarco_predictions(path: str) -> Dict[str, List[ScoredDoc]]:
+    assert path.endswith('.txt')
+
+    qid_to_scored_doc = {}
+    for line in tqdm.tqdm(open(path, 'r', encoding='utf-8'), desc='load prediction', mininterval=3):
+        fs = line.strip().split('\t')
+        qid, pid, rank = fs[:3]
+        rank = int(rank)
+        score = round(1 / rank, 4) if len(fs) == 3 else float(fs[3])
+
+        if qid not in qid_to_scored_doc:
+            qid_to_scored_doc[qid] = []
+        scored_doc = ScoredDoc(qid=qid, pid=pid, rank=rank, score=score)
+        qid_to_scored_doc[qid].append(scored_doc)
+
+    qid_to_scored_doc = {qid: sorted(scored_docs, key=lambda sd: sd.rank)
+                         for qid, scored_docs in qid_to_scored_doc.items()}
+
+    logger.info('Load {} query predictions from {}'.format(len(qid_to_scored_doc), path))
+    return qid_to_scored_doc
+
+
+def save_preds_to_msmarco_format(preds: Dict[str, List[ScoredDoc]], out_path: str):
+    with open(out_path, 'w', encoding='utf-8') as writer:
+        for qid in preds:
+            for idx, scored_doc in enumerate(preds[qid]):
+                writer.write('{}\t{}\t{}\t{}\n'.format(qid, scored_doc.pid, idx + 1, round(scored_doc.score, 3)))
+    logger.info('Successfully saved to {}'.format(out_path))
+
+
+def save_to_readable_format(in_path: str, corpus: Dataset):
+    out_path = '{}/readable_{}'.format(os.path.dirname(in_path), os.path.basename(in_path))
+    dataset: Dataset = load_dataset('json', data_files=in_path)['train']
+
+    max_to_keep = 5
+
+    def _create_readable_field(samples: Dict[str, List]) -> List:
+        readable_ex = []
+        for idx in range(min(len(samples['doc_id']), max_to_keep)):
+            doc_id = samples['doc_id'][idx]
+            readable_ex.append({'doc_id': doc_id,
+                                'title': corpus[int(doc_id)].get('title', ''),
+                                'contents': corpus[int(doc_id)]['contents'],
+                                'score': samples['score'][idx]})
+        return readable_ex
+
+    def _mp_func(ex: Dict) -> Dict:
+        ex['positives'] = _create_readable_field(ex['positives'])
+        ex['negatives'] = _create_readable_field(ex['negatives'])
+        return ex
+    dataset = dataset.map(_mp_func, num_proc=8)
+
+    dataset.to_json(out_path, force_ascii=False, lines=False, indent=4)
+    logger.info('Done convert {} to readable format in {}'.format(in_path, out_path))
+
+
+def get_rerank_shard_path(args: Arguments, worker_idx: int) -> str:
+    return '{}_shard_{}'.format(args.rerank_out_path, worker_idx)
+
+
+def merge_rerank_predictions(args: Arguments, gpu_count: int):
+    from metrics import trec_eval, compute_mrr
+
+    qid_to_scored_doc: Dict[str, List[ScoredDoc]] = {}
+    for worker_idx in range(gpu_count):
+        path = get_rerank_shard_path(args, worker_idx)
+        for line in tqdm.tqdm(open(path, 'r', encoding='utf-8'), 'merge results', mininterval=3):
+            fs = line.strip().split('\t')
+            qid, pid, _, score = fs
+            score = float(score)
+
+            if qid not in qid_to_scored_doc:
+                qid_to_scored_doc[qid] = []
+            scored_doc = ScoredDoc(qid=qid, pid=pid, rank=-1, score=score)
+            qid_to_scored_doc[qid].append(scored_doc)
+
+    qid_to_scored_doc = {k: sorted(v, key=lambda sd: sd.score, reverse=True) for k, v in qid_to_scored_doc.items()}
+
+    ori_preds = load_msmarco_predictions(path=args.rerank_in_path)
+    for query_id in list(qid_to_scored_doc.keys()):
+        remain_scored_docs = ori_preds[query_id][args.rerank_depth:]
+        for idx, sd in enumerate(remain_scored_docs):
+            # make sure the order is not broken
+            sd.score = qid_to_scored_doc[query_id][-1].score - idx - 1
+        qid_to_scored_doc[query_id] += remain_scored_docs
+        assert len(set([sd.pid for sd in qid_to_scored_doc[query_id]])) == len(qid_to_scored_doc[query_id])
+
+    save_preds_to_msmarco_format(qid_to_scored_doc, out_path=args.rerank_out_path)
+
+    path_qrels = '{}/{}_qrels.txt'.format(args.data_dir, args.rerank_split)
+    if os.path.exists(path_qrels):
+        qrels = load_qrels(path=path_qrels)
+        all_metrics = trec_eval(qrels=qrels, predictions=qid_to_scored_doc)
+        all_metrics['mrr'] = compute_mrr(qrels=qrels, predictions=qid_to_scored_doc)
+
+        logger.info('{} trec metrics = {}'.format(args.rerank_split, json.dumps(all_metrics, ensure_ascii=False, indent=4)))
+        metrics_out_path = '{}/metrics_rerank_{}.json'.format(os.path.dirname(args.rerank_out_path), args.rerank_split)
+        save_json_to_file(all_metrics, metrics_out_path)
+    else:
+        logger.warning('No qrels found for {}'.format(args.rerank_split))
+
+    # cleanup some intermediate results
+    for worker_idx in range(gpu_count):
+        path = get_rerank_shard_path(args, worker_idx)
+        os.remove(path)
+
+
+if __name__ == '__main__':
+    load_qrels('./data/msmarco/dev_qrels.txt')
+    load_queries('./data/msmarco/dev_queries.tsv')
+    corpus = load_corpus('./data/msmarco/passages.jsonl.gz')
+    preds = load_msmarco_predictions('./data/bm25.msmarco.txt')

src/cross_rerank/loader_utils.py

@@ -0,0 +1,46 @@
+from typing import List, Dict
+
+
+def _slice_with_mod(elements: List, offset: int, cnt: int) -> List:
+    return [elements[(offset + idx) % len(elements)] for idx in range(cnt)]
+
+
+def group_doc_ids(examples: Dict[str, List],
+                  negative_size: int,
+                  offset: int,
+                  use_first_positive: bool = False,
+                  use_first_negative: bool = True) -> List[int]:
+    pos_doc_ids: List[int] = []
+    positives: List[Dict[str, List]] = examples['positives']
+    for idx, ex_pos in enumerate(positives):
+        all_pos_doc_ids = ex_pos['doc_id']
+
+        if use_first_positive:
+            # keep positives that has higher score than all negatives
+            all_pos_doc_ids = [doc_id for p_idx, doc_id in enumerate(all_pos_doc_ids)
+                               if ex_pos['score'][p_idx] == max(ex_pos['score'])]
+
+        cur_pos_doc_id = _slice_with_mod(all_pos_doc_ids, offset=offset, cnt=1)[0]
+        pos_doc_ids.append(int(cur_pos_doc_id))
+
+    neg_doc_ids: List[List[int]] = []
+    negatives: List[Dict[str, List]] = examples['negatives']
+    for ex_neg in negatives:
+        if use_first_negative:
+            cur_neg_doc_ids = ex_neg['doc_id'][:negative_size]
+        else:
+            cur_neg_doc_ids = _slice_with_mod(ex_neg['doc_id'],
+                                            offset=offset * negative_size,
+                                            cnt=negative_size)
+        cur_neg_doc_ids = [int(doc_id) for doc_id in cur_neg_doc_ids]
+        neg_doc_ids.append(cur_neg_doc_ids)
+
+    assert len(pos_doc_ids) == len(neg_doc_ids), '{} != {}'.format(len(pos_doc_ids), len(neg_doc_ids))
+    assert all(len(doc_ids) == negative_size for doc_ids in neg_doc_ids)
+
+    input_doc_ids: List[int] = []
+    for pos_doc_id, neg_ids in zip(pos_doc_ids, neg_doc_ids):
+        input_doc_ids.append(pos_doc_id)
+        input_doc_ids += neg_ids
+
+    return input_doc_ids

src/cross_rerank/logger_config.py

@@ -0,0 +1,32 @@
+import os
+import logging
+
+from transformers.trainer_callback import TrainerCallback
+
+
+def _setup_logger():
+    log_format = logging.Formatter("[%(asctime)s %(levelname)s] %(message)s")
+    logger = logging.getLogger()
+    logger.setLevel(logging.INFO)
+
+    console_handler = logging.StreamHandler()
+    console_handler.setFormatter(log_format)
+
+    data_dir = './data/'
+    os.makedirs(data_dir, exist_ok=True)
+    file_handler = logging.FileHandler('{}/log.txt'.format(data_dir))
+    file_handler.setFormatter(log_format)
+
+    logger.handlers = [console_handler, file_handler]
+
+    return logger
+
+
+logger = _setup_logger()
+
+
+class LoggerCallback(TrainerCallback):
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        _ = logs.pop("total_flos", None)
+        if state.is_world_process_zero:
+            logger.info(logs)

src/cross_rerank/loss.py

@@ -0,0 +1,36 @@
+import torch
+import torch.nn.functional as F
+
+class CrossEncoderNllLoss(object):
+    def __init__(self,
+                 score_type="dot"):
+        self.score_type = score_type
+        
+    def calc(
+        self,
+        logits,
+        labels):
+        """
+        Computes nll loss for the given lists of question and ctx vectors.
+        Return: a tuple of loss value and amount of correct predictions per batch
+        """
+
+        #if len(q_vectors.size()) > 1:
+        #    q_num = q_vectors.size(0)
+        #    scores = scores.view(q_num, -1)
+        #    positive_idx_per_question = [i for i in range(q_num)]
+
+        softmax_scores = F.log_softmax(logits, dim=1)
+        #print("softmax", softmax_scores)
+        #print(softmax_scores.size())
+        #print(labels.size())
+        loss = F.nll_loss(
+            softmax_scores,
+            labels,
+            reduction="mean",
+        )
+        #print(loss)
+        #max_score, max_idxs = torch.max(softmax_scores, 1)
+        #correct_predictions_count = (max_idxs == torch.tensor(positive_idx_per_question).to(max_idxs.device)).sum()
+        return loss#, correct_predictions_count
+

src/cross_rerank/metrics.py

@@ -0,0 +1,105 @@
+import torch
+import pytrec_eval
+
+from typing import List, Dict, Tuple
+
+from .data_utils import ScoredDoc
+from .logger_config import logger
+
+
+def trec_eval(qrels: Dict[str, Dict[str, int]],
+              predictions: Dict[str, List[ScoredDoc]],
+              k_values: Tuple[int] = (10, 50, 100, 200, 1000)) -> Dict[str, float]:
+    ndcg, _map, recall = {}, {}, {}
+
+    for k in k_values:
+        ndcg[f"NDCG@{k}"] = 0.0
+        _map[f"MAP@{k}"] = 0.0
+        recall[f"Recall@{k}"] = 0.0
+
+    map_string = "map_cut." + ",".join([str(k) for k in k_values])
+    ndcg_string = "ndcg_cut." + ",".join([str(k) for k in k_values])
+    recall_string = "recall." + ",".join([str(k) for k in k_values])
+
+    results: Dict[str, Dict[str, float]] = {}
+    for query_id, scored_docs in predictions.items():
+        results.update({query_id: {sd.pid: sd.score for sd in scored_docs}})
+
+    evaluator = pytrec_eval.RelevanceEvaluator(qrels, {map_string, ndcg_string, recall_string})
+    scores = evaluator.evaluate(results)
+
+    for query_id in scores:
+        for k in k_values:
+            ndcg[f"NDCG@{k}"] += scores[query_id]["ndcg_cut_" + str(k)]
+            _map[f"MAP@{k}"] += scores[query_id]["map_cut_" + str(k)]
+            recall[f"Recall@{k}"] += scores[query_id]["recall_" + str(k)]
+
+    def _normalize(m: dict) -> dict:
+        return {k: round(v / len(scores), 5) for k, v in m.items()}
+
+    ndcg = _normalize(ndcg)
+    _map = _normalize(_map)
+    recall = _normalize(recall)
+
+    all_metrics = {}
+    for mt in [ndcg, _map, recall]:
+        all_metrics.update(mt)
+
+    return all_metrics
+
+
+@torch.no_grad()
+def accuracy(output: torch.tensor, target: torch.tensor, topk=(1,)) -> List[float]:
+    """Computes the accuracy over the k top predictions for the specified values of k"""
+    maxk = max(topk)
+    batch_size = target.size(0)
+
+    _, pred = output.topk(maxk, 1, True, True)
+    pred = pred.t()
+    correct = pred.eq(target.view(1, -1).expand_as(pred))
+
+    res = []
+    for k in topk:
+        correct_k = correct[:k].contiguous().view(-1).float().sum(0, keepdim=True)
+        res.append(correct_k.mul_(100.0 / batch_size).item())
+    return res
+
+
+@torch.no_grad()
+def batch_mrr(output: torch.tensor, target: torch.tensor) -> float:
+    assert len(output.shape) == 2
+    assert len(target.shape) == 1
+    sorted_score, sorted_indices = torch.sort(output, dim=-1, descending=True)
+    _, rank = torch.nonzero(sorted_indices.eq(target.unsqueeze(-1)).long(), as_tuple=True)
+    assert rank.shape[0] == output.shape[0]
+
+    rank = rank + 1
+    mrr = torch.sum(100 / rank.float()) / rank.shape[0]
+    return mrr.item()
+
+
+def get_rel_threshold(qrels: Dict[str, Dict[str, int]]) -> int:
+    # For ms-marco passage ranking, score >= 1 is relevant
+    # for trec dl 2019 & 2020, score >= 2 is relevant
+    rel_labels = set()
+    for q_id in qrels:
+        for doc_id, label in qrels[q_id].items():
+            rel_labels.add(label)
+
+    logger.info('relevance labels: {}'.format(rel_labels))
+    return 2 if max(rel_labels) >= 3 else 1
+
+
+def compute_mrr(qrels: Dict[str, Dict[str, int]],
+                predictions: Dict[str, List[ScoredDoc]],
+                k: int = 10) -> float:
+    threshold = get_rel_threshold(qrels)
+    mrr = 0
+    for qid in qrels:
+        scored_docs = predictions.get(qid, [])
+        for idx, scored_doc in enumerate(scored_docs[:k]):
+            if scored_doc.pid in qrels[qid] and qrels[qid][scored_doc.pid] >= threshold:
+                mrr += 1 / (idx + 1)
+                break
+
+    return round(mrr / len(qrels) * 100, 4)

src/cross_rerank/model.py

@@ -0,0 +1,58 @@
+import torch
+import torch.nn as nn
+
+from typing import Optional, Dict
+from transformers import (
+    PreTrainedModel,
+    AutoModelForSequenceClassification
+)
+from transformers.modeling_outputs import SequenceClassifierOutput
+
+from .config import Arguments
+from .loss import CrossEncoderNllLoss
+        
+class Reranker(nn.Module):
+    def __init__(self, hf_model: PreTrainedModel, args: Arguments):
+        super().__init__()
+        self.hf_model = hf_model
+        self.args = args
+        self._keys_to_ignore_on_save = None
+
+        self.cross_entropy = nn.CrossEntropyLoss(reduction='mean')
+        #self.contrastive = CrossEncoderNllLoss()
+        #self.kl_loss_fn = torch.nn.KLDivLoss(reduction="batchmean", log_target=True)
+
+    def forward(self, input_ids, attention_mask, token_type_ids) -> SequenceClassifierOutput:
+        #n_psg_per_query = self.args.train_n_passages // self.args.rerank_forward_factor
+
+        outputs: SequenceClassifierOutput = self.hf_model(input_ids, attention_mask, token_type_ids, return_dict=True)
+        #outputs.logits = outputs.logits.view(-1, n_psg_per_query)
+        #loss = self.cross_entropy(outputs.logits, labels)
+
+        return outputs#, loss
+
+    @classmethod
+    def from_pretrained(cls, all_args: Arguments, *args, **kwargs):
+        hf_model = AutoModelForSequenceClassification.from_pretrained(*args, **kwargs)
+        return cls(hf_model, all_args)
+
+    def save_pretrained(self, output_dir: str):
+        self.hf_model.save_pretrained(output_dir)
+
+
+class RerankerForInference(nn.Module):
+    def __init__(self, model_checkpoint, hf_model: Optional[PreTrainedModel] = None):
+        super().__init__()
+        self.hf_model = hf_model
+        if hf_model is None:
+            self.hf_model = AutoModelForSequenceClassification.from_pretrained(model_checkpoint)
+        self.hf_model.eval()
+
+    @torch.no_grad()
+    def forward(self, batch) -> SequenceClassifierOutput:
+        return self.hf_model(**batch)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str):
+        hf_model = AutoModelForSequenceClassification.from_pretrained(pretrained_model_name_or_path)
+        return cls(hf_model)

src/cross_rerank/trainer.py

@@ -0,0 +1,247 @@
+import os
+import torch
+from torch import nn
+from torch.utils.checkpoint import get_device_states, set_device_states
+#from typing import Optional, Union, Dict, Any
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
+from transformers.trainer import Trainer
+from transformers.modeling_outputs import SequenceClassifierOutput
+from collections.abc import Mapping
+from .logger_config import logger
+from .metrics import accuracy
+from .utils import AverageMeter
+
+def nested_detach(tensors):
+    "Detach `tensors` (even if it's a nested list/tuple/dict of tensors)."
+    if isinstance(tensors, (list, tuple)):
+        return type(tensors)(nested_detach(t) for t in tensors)
+    elif isinstance(tensors, Mapping):
+        return type(tensors)({k: nested_detach(t) for k, t in tensors.items()})
+    return tensors.detach()
+
+class RandContext:
+    def __init__(self, *tensors):
+        self.fwd_cpu_state = torch.get_rng_state()
+        self.fwd_gpu_devices, self.fwd_gpu_states = get_device_states(*tensors)
+
+    def __enter__(self):
+        self._fork = torch.random.fork_rng(
+            devices=self.fwd_gpu_devices,
+            enabled=True
+        )
+        self._fork.__enter__()
+        torch.set_rng_state(self.fwd_cpu_state)
+        set_device_states(self.fwd_gpu_devices, self.fwd_gpu_states)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self._fork.__exit__(exc_type, exc_val, exc_tb)
+        self._fork = None
+        
+class RerankerTrainer(Trainer):
+    def __init__(self, *pargs, **kwargs):
+        super(RerankerTrainer, self).__init__(*pargs, **kwargs)
+
+        self.acc_meter = AverageMeter('acc', round_digits=2)
+        self.last_epoch = 0
+
+    def _save(self, output_dir: Optional[str] = None, state_dict=None):
+        output_dir = output_dir if output_dir is not None else self.args.output_dir
+        os.makedirs(output_dir, exist_ok=True)
+        logger.info("Saving model checkpoint to {}".format(output_dir))
+
+        self.model.save_pretrained(output_dir)
+
+        if self.tokenizer is not None and self.is_world_process_zero():
+            self.tokenizer.save_pretrained(output_dir)
+            
+    def compute_loss(self, model, inputs, return_outputs=False):
+        n_psg_per_query = self.args.train_n_passages // self.args.rerank_forward_factor
+        input_ids = inputs['input_ids']
+        attention_mask = inputs['attention_mask']
+        token_type_ids = inputs['token_type_ids']
+        labels = inputs['labels']
+        outputs = model(input_ids, attention_mask, token_type_ids)
+        outputs.logits = outputs.logits.view(-1, n_psg_per_query)
+        loss = self.model.cross_entropy(outputs.logits, labels)
+
+        if self.model.training:
+            step_acc = accuracy(output=outputs.logits.detach(), target=labels)[0]
+            self.acc_meter.update(step_acc)
+            if self.state.global_step > 0 and self.state.global_step % self.args.logging_steps == 0:
+                logger.info('step: {}, {}'.format(self.state.global_step, self.acc_meter))
+
+            self._reset_meters_if_needed()
+
+        return (loss, outputs) if return_outputs else loss
+            
+    def training_step(self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]]) -> torch.Tensor:
+        model.train()
+        inputs = self._prepare_inputs(inputs)
+
+        with self.compute_loss_context_manager():
+            loss = self.compute_loss_train(model, inputs)
+
+        return loss.detach() / self.args.gradient_accumulation_steps
+
+    def compute_loss_train(self, model, inputs, return_outputs=False):
+        #print(inputs)
+        #print(inputs['input_ids'].size())
+        n_psg_per_query = self.args.train_n_passages // self.args.rerank_forward_factor
+        input_ids = inputs['input_ids']
+        attention_mask = inputs['attention_mask']
+        token_type_ids = inputs['token_type_ids']
+        labels = inputs['labels']
+        
+        all_reps, rnds = [], []
+            
+        id_chunks = input_ids.split(self.args.chunk_size)
+        attn_mask_chunks = attention_mask.split(self.args.chunk_size)
+            
+        type_ids_chunks = token_type_ids.split(self.args.chunk_size)
+            
+        for id_chunk, attn_chunk, type_chunk in zip(id_chunks, attn_mask_chunks, type_ids_chunks):
+            rnds.append(RandContext(id_chunk, attn_chunk, type_chunk))
+            with torch.no_grad():
+                chunk_reps = self.model(id_chunk, attn_chunk, type_chunk).logits
+            all_reps.append(chunk_reps)
+        all_reps = torch.cat(all_reps)
+        all_reps = all_reps.view(-1, n_psg_per_query)
+            
+        all_reps = all_reps.float().detach().requires_grad_()
+        loss = self.model.cross_entropy(all_reps, labels)
+        
+        if self.args.n_gpu > 1:
+            loss = loss.mean()  # mean() to average on multi-gpu parallel training
+
+        self.accelerator.backward(loss)
+        #if self.args.gradient_accumulation_steps > 1:
+        #    loss = loss / self.args.gradient_accumulation_steps
+        #loss.backward()
+        #temp = all_reps.view(-1,1)      
+        grads = all_reps.grad.split(int(self.args.chunk_size/n_psg_per_query))
+            
+        for id_chunk, attn_chunk, type_chunk, grad, rnd in zip(id_chunks, attn_mask_chunks, type_ids_chunks, grads, rnds):
+            #print(id_chunk.size())
+            with rnd:
+                chunk_reps = self.model(id_chunk, attn_chunk, type_chunk).logits
+                #print(chunk_reps.size())
+                #print(grad.size())
+                surrogate = torch.dot(chunk_reps.flatten().float(), grad.flatten())
+                
+            self.accelerator.backward(surrogate)
+                
+        #outputs, loss = model(input_ids, attention_mask, token_type_ids, labels)
+
+        if self.model.training:
+            step_acc = accuracy(all_reps, target=labels)[0]
+            #print(step_acc)
+            self.acc_meter.update(step_acc)
+            if self.state.global_step > 0 and self.state.global_step % self.args.logging_steps == 0:
+                logger.info('step: {}, {}'.format(self.state.global_step, self.acc_meter))
+
+            self._reset_meters_if_needed()
+
+        return (loss, all_reps) if return_outputs else loss
+    
+    '''def compute_loss_pred(self, model, inputs, return_outputs=False):
+        #print(inputs)
+        #print(inputs['input_ids'].size())
+        n_psg_per_query = self.args.train_n_passages // self.args.rerank_forward_factor
+        input_ids = inputs['input_ids']
+        attention_mask = inputs['attention_mask']
+        token_type_ids = inputs['token_type_ids']
+        labels = inputs['labels']
+        
+        all_reps, rnds = [], []
+            
+        id_chunks = input_ids.split(self.args.chunk_size)
+        attn_mask_chunks = attention_mask.split(self.args.chunk_size)
+            
+        type_ids_chunks = token_type_ids.split(self.args.chunk_size)
+            
+        for id_chunk, attn_chunk, type_chunk in zip(id_chunks, attn_mask_chunks, type_ids_chunks):
+            rnds.append(RandContext(id_chunk, attn_chunk, type_chunk))
+            with torch.no_grad():
+                chunk_reps = self.model(id_chunk, attn_chunk, type_chunk).logits
+            all_reps.append(chunk_reps)
+        all_reps = torch.cat(all_reps)
+        all_reps = all_reps.view(-1, n_psg_per_query)
+        loss = self.model.cross_entropy(all_reps, labels)
+
+        if self.model.training:
+            step_acc = accuracy(all_reps, target=labels)[0]
+            #print(step_acc)
+            self.acc_meter.update(step_acc)
+            if self.state.global_step > 0 and self.state.global_step % self.args.logging_steps == 0:
+                logger.info('step: {}, {}'.format(self.state.global_step, self.acc_meter))
+
+            self._reset_meters_if_needed()
+
+        return (loss, all_reps) if return_outputs else loss
+    
+    def prediction_step(
+        self,
+        model: nn.Module,
+        inputs: Dict[str, Union[torch.Tensor, Any]],
+        prediction_loss_only: bool,
+        ignore_keys: Optional[List[str]] = None,
+    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
+        has_labels = False if len(self.label_names) == 0 else all(inputs.get(k) is not None for k in self.label_names)
+        # For CLIP-like models capable of returning loss values.
+        # If `return_loss` is not specified or being `None` in `inputs`, we check if the default value of `return_loss`
+        # is `True` in `model.forward`.
+        return_loss = inputs.get("return_loss", None)
+        if return_loss is None:
+            return_loss = self.can_return_loss
+        loss_without_labels = True if len(self.label_names) == 0 and return_loss else False
+
+        inputs = self._prepare_inputs(inputs)
+        if ignore_keys is None:
+            if hasattr(self.model, "config"):
+                ignore_keys = getattr(self.model.config, "keys_to_ignore_at_inference", [])
+            else:
+                ignore_keys = []
+
+        # labels may be popped when computing the loss (label smoothing for instance) so we grab them first.
+        if has_labels or loss_without_labels:
+            labels = nested_detach(tuple(inputs.get(name) for name in self.label_names))
+            if len(labels) == 1:
+                labels = labels[0]
+        else:
+            labels = None
+
+        with torch.no_grad():
+            if has_labels or loss_without_labels:
+                with self.compute_loss_context_manager():
+                    loss, outputs = self.compute_loss(model, inputs, return_outputs=True)
+                loss = loss.mean().detach()
+
+                if isinstance(outputs, dict):
+                    logits = tuple(v for k, v in outputs.items() if k not in ignore_keys + ["loss"])
+                else:
+                    logits = outputs[1:]
+            else:
+                loss = None
+                with self.compute_loss_context_manager():
+                    outputs = model(**inputs)
+                if isinstance(outputs, dict):
+                    logits = tuple(v for k, v in outputs.items() if k not in ignore_keys)
+                else:
+                    logits = outputs
+                # TODO: this needs to be fixed and made cleaner later.
+                if self.args.past_index >= 0:
+                    self._past = outputs[self.args.past_index - 1]
+
+        if prediction_loss_only:
+            return (loss, None, None)
+
+        logits = nested_detach(logits)
+        if len(logits) == 1:
+            logits = logits[0]
+
+        return (loss, logits, labels)'''
+
+    def _reset_meters_if_needed(self):
+        if int(self.state.epoch) != self.last_epoch:
+            self.last_epoch = int(self.state.epoch)
+            self.acc_meter.reset()

src/cross_rerank/utils.py

@@ -0,0 +1,137 @@
+import json
+import torch
+import torch.distributed as dist
+
+from typing import List, Union, Optional, Tuple, Mapping, Dict
+
+
+def save_json_to_file(objects: Union[List, dict], path: str, line_by_line: bool = False):
+    if line_by_line:
+        assert isinstance(objects, list), 'Only list can be saved in line by line format'
+
+    with open(path, 'w', encoding='utf-8') as writer:
+        if not line_by_line:
+            json.dump(objects, writer, ensure_ascii=False, indent=4, separators=(',', ':'))
+        else:
+            for obj in objects:
+                writer.write(json.dumps(obj, ensure_ascii=False, separators=(',', ':')))
+                writer.write('\n')
+
+
+def move_to_cuda(sample):
+    if len(sample) == 0:
+        return {}
+
+    def _move_to_cuda(maybe_tensor):
+        if torch.is_tensor(maybe_tensor):
+            return maybe_tensor.cuda(non_blocking=True)
+        elif isinstance(maybe_tensor, dict):
+            return {key: _move_to_cuda(value) for key, value in maybe_tensor.items()}
+        elif isinstance(maybe_tensor, list):
+            return [_move_to_cuda(x) for x in maybe_tensor]
+        elif isinstance(maybe_tensor, tuple):
+            return tuple([_move_to_cuda(x) for x in maybe_tensor])
+        elif isinstance(maybe_tensor, Mapping):
+            return type(maybe_tensor)({k: _move_to_cuda(v) for k, v in maybe_tensor.items()})
+        else:
+            return maybe_tensor
+
+    return _move_to_cuda(sample)
+
+
+def dist_gather_tensor(t: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
+    if t is None:
+        return None
+
+    t = t.contiguous()
+    all_tensors = [torch.empty_like(t) for _ in range(dist.get_world_size())]
+    dist.all_gather(all_tensors, t)
+
+    all_tensors[dist.get_rank()] = t
+    all_tensors = torch.cat(all_tensors, dim=0)
+    return all_tensors
+
+
+@torch.no_grad()
+def select_grouped_indices(scores: torch.Tensor,
+                           group_size: int,
+                           start: int = 0) -> torch.Tensor:
+    assert len(scores.shape) == 2
+    batch_size = scores.shape[0]
+    assert batch_size * group_size <= scores.shape[1]
+
+    indices = torch.arange(0, group_size, dtype=torch.long)
+    indices = indices.repeat(batch_size, 1)
+    indices += torch.arange(0, batch_size, dtype=torch.long).unsqueeze(-1) * group_size
+    indices += start
+
+    return indices.to(scores.device)
+
+
+def full_contrastive_scores_and_labels(
+        query: torch.Tensor,
+        key: torch.Tensor,
+        use_all_pairs: bool = True) -> Tuple[torch.Tensor, torch.Tensor]:
+    assert key.shape[0] % query.shape[0] == 0, '{} % {} > 0'.format(key.shape[0], query.shape[0])
+
+    train_n_passages = key.shape[0] // query.shape[0]
+    labels = torch.arange(0, query.shape[0], dtype=torch.long, device=query.device)
+    labels = labels * train_n_passages
+
+    # batch_size x (batch_size x n_psg)
+    qk = torch.mm(query, key.t())
+
+    if not use_all_pairs:
+        return qk, labels
+
+    # batch_size x dim
+    sliced_key = key.index_select(dim=0, index=labels)
+    assert query.shape[0] == sliced_key.shape[0]
+
+    # batch_size x batch_size
+    kq = torch.mm(sliced_key, query.t())
+    kq.fill_diagonal_(float('-inf'))
+
+    qq = torch.mm(query, query.t())
+    qq.fill_diagonal_(float('-inf'))
+
+    kk = torch.mm(sliced_key, sliced_key.t())
+    kk.fill_diagonal_(float('-inf'))
+
+    scores = torch.cat([qk, kq, qq, kk], dim=-1)
+
+    return scores, labels
+
+
+def slice_batch_dict(batch_dict: Dict[str, torch.Tensor], prefix: str) -> dict:
+    return {k[len(prefix):]: v for k, v in batch_dict.items() if k.startswith(prefix)}
+
+
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+
+    def __init__(self, name: str, round_digits: int = 3):
+        self.name = name
+        self.round_digits = round_digits
+        self.reset()
+
+    def reset(self):
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+    def __str__(self):
+        return '{}: {}'.format(self.name, round(self.avg, self.round_digits))
+
+
+if __name__ == '__main__':
+    query = torch.randn(4, 16)
+    key = torch.randn(4 * 3, 16)
+    scores, labels = full_contrastive_scores_and_labels(query, key)
+    print(scores.shape)
+    print(labels)

src/eval_cross.py

@@ -0,0 +1,512 @@
+import faiss
+import torch
+import logging
+import json
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+from typing import Optional
+from dataclasses import dataclass, field
+from transformers import HfArgumentParser
+from transformers import AutoTokenizer
+from bi.model import SharedBiEncoder
+from bi.preprocess import preprocess_question
+from cross_rerank.model import RerankerForInference
+#from src.process import process_query, process_text, concat_str
+import itertools
+from pyvi.ViTokenizer import tokenize
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Args:
+    encoder: str = field(
+        default="vinai/phobert-base-v2",
+        metadata={'help': 'The encoder name or path.'}
+    )
+    tokenizer: str = field(
+        default=None,
+        metadata={'help': 'The encoder name or path.'}
+    )
+    cross_checkpoint: str = field(
+        default="vinai/phobert-base-v2",
+        metadata={'help': 'The encoder name or path.'}
+    )
+    cross_tokenizer: str = field(
+        default=None,
+        metadata={'help': 'The encoder name or path.'}
+    )
+    sentence_pooling_method: str = field(
+        default="cls",
+        metadata={'help': 'Embedding method'}
+    )
+    fp16: bool = field(
+        default=False,
+        metadata={'help': 'Use fp16 in inference?'}
+    )
+    max_query_length: int = field(
+        default=32,
+        metadata={'help': 'Max query length.'}
+    )
+    max_passage_length: int = field(
+        default=256,
+        metadata={'help': 'Max passage length.'}
+    )
+    cross_max_length: int = field(
+        default=256,
+        metadata={'help': 'Max cross length.'}
+    )
+    cross_batch_size: int = field(
+        default=32,
+        metadata={'help': 'Inference batch size.'}
+    )
+    batch_size: int = field(
+        default=128,
+        metadata={'help': 'Inference batch size.'}
+    )
+    index_factory: str = field(
+        default="Flat",
+        metadata={'help': 'Faiss index factory.'}
+    )
+    k: int = field(
+        default=1000,
+        metadata={'help': 'How many neighbors to retrieve?'}
+    )
+    top_k: int = field(
+        default=1000,
+        metadata={'help': 'How many neighbors to rerank?'}
+    )
+    data_path: str = field(
+        default="/kaggle/input/zalo-data",
+        metadata={'help': 'Path to zalo data.'}
+    )
+    data_type: str = field(
+        default="test",
+        metadata={'help': 'Type data to test'}
+    )
+    corpus_file: str = field(
+        default="/kaggle/input/zalo-data",
+        metadata={'help': 'Path to zalo corpus.'}
+    )
+    
+    data_file: str = field(
+        default=None,
+        metadata={'help': 'Path to evaluated data.'}
+    )
+    
+    bi_data: bool = field(
+        default=False,
+        metadata={'help': 'Data for bi-encoder training'}
+    )
+    
+    save_embedding: bool = field(
+        default=False,
+        metadata={'help': 'Save embeddings in memmap at save_dir?'}
+    )
+    load_embedding: str = field(
+        default='',
+        metadata={'help': 'Path to saved embeddings.'}
+    )
+    
+    save_path: str = field(
+        default="embeddings.memmap",
+        metadata={'help': 'Path to save embeddings.'}
+    )
+
+def index(model: SharedBiEncoder, tokenizer:AutoTokenizer, corpus, batch_size: int = 16, max_length: int=512, index_factory: str = "Flat", save_path: str = None, save_embedding: bool = False, load_embedding: bool = False):
+    """
+    1. Encode the entire corpus into dense embeddings; 
+    2. Create faiss index; 
+    3. Optionally save embeddings.
+    """
+    if load_embedding != '':
+        test_tokens = tokenizer(['test'],
+                                padding=True,
+                                truncation=True,
+                                max_length=128,
+                                return_tensors="pt").to('cuda')
+        test = model.encoder.get_representation(test_tokens['input_ids'], test_tokens['attention_mask'])
+        test = test.cpu().numpy()
+        dtype = test.dtype
+        dim = test.shape[-1]
+
+        all_embeddings = np.memmap(
+            load_embedding,
+            mode="r",
+            dtype=dtype
+        ).reshape(-1, dim)
+    
+    else:
+        #df_corpus = pd.DataFrame()
+        #df_corpus['text'] = corpus
+        #pandarallel.initialize(progress_bar=True, use_memory_fs=False, nb_workers=12)
+        #df_corpus['processed_text'] = df_corpus['text'].parallel_apply(process_text)
+        #processed_corpus = df_corpus['processed_text'].tolist()
+        #model.to('cuda')
+        all_embeddings = []
+        for start_index in tqdm(range(0, len(corpus), batch_size), desc="Inference Embeddings",
+                                disable=len(corpus) < batch_size):
+            passages_batch = corpus[start_index:start_index + batch_size]
+            d_collated = tokenizer(
+                    passages_batch,
+                    padding=True,
+                    truncation=True,
+                    max_length=max_length,
+                    return_tensors="pt",
+                ).to('cuda')
+
+            with torch.no_grad():
+                corpus_embeddings = model.encoder.get_representation(d_collated['input_ids'], d_collated['attention_mask']) 
+            
+            corpus_embeddings = corpus_embeddings.cpu().numpy()
+            all_embeddings.append(corpus_embeddings)
+
+        all_embeddings = np.concatenate(all_embeddings, axis=0)
+        dim = all_embeddings.shape[-1]
+        
+        if save_embedding:
+            logger.info(f"saving embeddings at {save_path}...")
+            memmap = np.memmap(
+                save_path,
+                shape=all_embeddings.shape,
+                mode="w+",
+                dtype=all_embeddings.dtype
+            )
+
+            length = all_embeddings.shape[0]
+            # add in batch
+            save_batch_size = 10000
+            if length > save_batch_size:
+                for i in tqdm(range(0, length, save_batch_size), leave=False, desc="Saving Embeddings"):
+                    j = min(i + save_batch_size, length)
+                    memmap[i: j] = all_embeddings[i: j]
+            else:
+                memmap[:] = all_embeddings
+    # create faiss index
+    faiss_index = faiss.index_factory(dim, index_factory, faiss.METRIC_INNER_PRODUCT)
+
+    #if model.device == torch.device("cuda"):
+    if True:
+        co = faiss.GpuClonerOptions()
+        #co = faiss.GpuMultipleClonerOptions()
+        #co.useFloat16 = True
+        faiss_index = faiss.index_cpu_to_gpu(faiss.StandardGpuResources(), 0, faiss_index, co)
+        #faiss_index = faiss.index_cpu_to_all_gpus(faiss_index, co)
+
+    # NOTE: faiss only accepts float32
+    logger.info("Adding embeddings...")
+    all_embeddings = all_embeddings.astype(np.float32)
+    #print(all_embeddings[0])
+    faiss_index.train(all_embeddings)
+    faiss_index.add(all_embeddings)
+    return faiss_index
+
+
+def search(model: SharedBiEncoder, tokenizer:AutoTokenizer, questions, faiss_index: faiss.Index, k:int = 100, batch_size: int = 256, max_length: int=128):
+    """
+    1. Encode queries into dense embeddings;
+    2. Search through faiss index
+    """
+    #model.to('cuda')
+    q_embeddings = []
+    #questions = queries['tokenized_question'].tolist()
+    #questions = [process_query(x) for x in questions]
+    for start_index in tqdm(range(0, len(questions), batch_size), desc="Inference Embeddings",
+                            disable=len(questions) < batch_size):
+                    
+        q_collated = tokenizer(
+                    questions[start_index: start_index + batch_size],
+                    padding=True,
+                    truncation=True,
+                    max_length=128,
+                    return_tensors="pt",
+                ).to('cuda')
+
+        with torch.no_grad():
+            query_embeddings = model.encoder.get_representation(q_collated['input_ids'], q_collated['attention_mask'])
+        query_embeddings = query_embeddings.cpu().numpy()
+        q_embeddings.append(query_embeddings)
+    
+    q_embeddings = np.concatenate(q_embeddings, axis=0)
+    query_size = q_embeddings.shape[0]
+    all_scores = []
+    all_indices = []
+
+    for i in tqdm(range(0, query_size, batch_size), desc="Searching"):
+        j = min(i + batch_size, query_size)
+        q_embedding = q_embeddings[i: j]
+        score, indice = faiss_index.search(q_embedding.astype(np.float32), k=k)
+        all_scores.append(score)
+        all_indices.append(indice)
+    
+    all_scores = np.concatenate(all_scores, axis=0)
+    all_indices = np.concatenate(all_indices, axis=0)
+    return all_scores, all_indices
+
+def rerank(reranker: SharedBiEncoder, tokenizer:AutoTokenizer, questions, corpus, retrieved_ids, batch_size = 128, max_length = 256, top_k=30):
+    eos = tokenizer.eos_token
+    #questions = queries['tokenized_question'].tolist()
+    texts = []
+    for idx in range(len(questions)):
+        for j in range(top_k):
+            texts.append(questions[idx] + eos + eos + corpus[retrieved_ids[idx][j]])
+    reranked_ids = []
+    rerank_scores = []
+    
+    for start_index in tqdm(range(0, len(questions), batch_size), desc="Rerank",
+                            disable=len(questions) < batch_size):
+        batch_retrieved_ids = retrieved_ids[start_index: start_index+batch_size]
+        collated = tokenizer(
+                    texts[start_index*top_k: (start_index + batch_size)*top_k],
+                    padding=True,
+                    truncation=True,
+                    max_length=max_length,
+                    return_tensors="pt",
+                ).to('cuda')
+        reranked_scores = reranker(collated).logits
+        reranked_scores = reranked_scores.view(-1,top_k).to('cpu').tolist()
+        for m in range(len(reranked_scores)):
+            tuple_lst = [(batch_retrieved_ids[m][n], reranked_scores[m][n]) for n in range(top_k)]
+            tuple_lst.sort(key=lambda tup: tup[1], reverse=True)
+            reranked_ids.append([tup[0] for tup in tuple_lst]) 
+            rerank_scores.append([tup[1] for tup in tuple_lst])            
+            
+    return reranked_ids, rerank_scores
+        
+       
+    
+    
+def evaluate(preds, labels, cutoffs=[1,5,10,30,100]):
+    """
+    Evaluate MRR and Recall at cutoffs.
+    """
+    metrics = {}
+    
+    # MRR
+    mrrs = np.zeros(len(cutoffs))
+    for pred, label in zip(preds, labels):
+        jump = False
+        for i, x in enumerate(pred, 1):
+            if x in label:
+                for k, cutoff in enumerate(cutoffs):
+                    if i <= cutoff:
+                        mrrs[k] += 1 / i
+                jump = True
+            if jump:
+                break
+    mrrs /= len(preds)
+    for i, cutoff in enumerate(cutoffs):
+        mrr = mrrs[i]
+        metrics[f"MRR@{cutoff}"] = mrr
+
+    # Recall
+    recalls = np.zeros(len(cutoffs))
+    for pred, label in zip(preds, labels):
+        for k, cutoff in enumerate(cutoffs):
+            recall = np.intersect1d(label, pred[:cutoff])
+            recalls[k] += len(recall) / len(label)
+    recalls /= len(preds)
+    for i, cutoff in enumerate(cutoffs):
+        recall = recalls[i]
+        metrics[f"Recall@{cutoff}"] = recall
+
+    return metrics
+
+def calculate_score(ground_ids, retrieved_list):
+    all_count = 0
+    hit_count = 0
+    for i in range(len(ground_ids)):
+        all_check = True
+        hit_check = False
+        retrieved_ids = retrieved_list[i]
+        ans_ids = ground_ids[i]
+        for a_ids in ans_ids:
+            com = [a_id for a_id in a_ids if a_id in retrieved_ids]
+            if len(com) > 0:
+                hit_check = True
+            else:
+                all_check = False
+            
+        if hit_check:
+            hit_count += 1
+        if all_check:
+            all_count += 1
+                
+    all_acc = all_count/len(ground_ids)
+    hit_acc = hit_count/len(ground_ids)
+    return hit_acc, all_acc
+
+def check(ground_ids, retrieved_list, cutoffs=[1,5,10,30,100]):
+    metrics = {}
+    for cutoff in cutoffs:
+        retrieved_k = [x[:cutoff] for x in retrieved_list]
+        hit_acc, all_acc = calculate_score(ground_ids, retrieved_k)
+        metrics[f"All@{cutoff}"] = all_acc
+        metrics[f"Hit@{cutoff}"] = hit_acc
+    return metrics
+    
+def save_bi_data(tokenized_queries, ground_ids, indices, scores, file, org_questions=None):
+    rst = []
+    #tokenized_queries = test_data['tokenized_question'].tolist()
+    for i in range(len(tokenized_queries)):
+        scores_i = scores[i]
+        indices_i = indices[i]
+        ans_ids = ground_ids[i]
+        all_ans_id = [element for x in ans_ids for element in x]
+        neg_doc_ids = []
+        neg_scores = []
+        for count in range(len(indices_i)):
+            if indices_i[count] not in all_ans_id and indices_i[count] != -1:
+                neg_doc_ids.append(indices_i[count])
+                neg_scores.append(scores_i[count])
+                
+        for j in range(len(ans_ids)):
+            ans_id = ans_ids[j]
+            item = {}
+            if org_questions != None:
+                item['question'] = org_questions[i]
+            item['query'] = tokenized_queries[i]
+            item['positives'] = {}
+            item['negatives'] = {}
+            item['positives']['doc_id'] = []
+            item['positives']['score'] = []
+            item['negatives']['doc_id'] = neg_doc_ids
+            item['negatives']['score'] = neg_scores
+            for pos_id in ans_id:
+                item['positives']['doc_id'].append(pos_id)
+                try:
+                    idx = indices_i.index(pos_id)
+                    item['positives']['score'].append(scores_i[idx])
+                except:
+                    item['positives']['score'].append(scores_i[-1])
+                        
+            rst.append(item)
+            
+    with open(f'{file}.jsonl', 'w') as jsonl_file:
+        for item in rst:
+            json_line = json.dumps(item, ensure_ascii=False)
+            jsonl_file.write(json_line + '\n')
+                
+def main():
+    parser = HfArgumentParser([Args])
+    args: Args = parser.parse_args_into_dataclasses()[0]
+    print(args)
+    model = SharedBiEncoder(model_checkpoint=args.encoder,
+                            representation=args.sentence_pooling_method,
+                            fixed=True)
+    model.to('cuda')
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer if args.tokenizer else args.encoder)
+    reranker = RerankerForInference(model_checkpoint=args.cross_checkpoint)
+    reranker.to('cuda')
+    reranker_tokenizer = AutoTokenizer.from_pretrained(args.cross_tokenizer if args.cross_tokenizer else args.cross_checkpoint)
+    csv_file = True
+    if args.data_file:
+        if args.data_file.endswith("jsonl"):
+            test_data = []
+            with open(args.data_file, 'r') as jsonl_file:
+                for line in jsonl_file:
+                    temp = json.loads(line)
+                    test_data.append(temp)
+            csv_file=False
+        elif args.data_file.endswith("json"):
+            csv_file=False
+            with open(args.data_file, 'r') as json_file:
+                test_data = json.load(json_file) 
+        elif args.data_file.endswith("csv"):
+            test_data = pd.read_csv(args.data_file)
+            
+    elif args.data_type == 'eval':
+        test_data = pd.read_csv(args.data_path + "/tval.csv") 
+    elif args.data_type == 'train':
+        test_data = pd.read_csv(args.data_path + "/ttrain.csv")
+    elif args.data_type == 'all':
+        data1 = pd.read_csv(args.data_path + "/ttrain.csv")
+        data2 = pd.read_csv(args.data_path + "/ttest.csv")
+        data3 = pd.read_csv(args.data_path + "/tval.csv")
+        test_data = pd.concat([data1, data3, data2], ignore_index=True)
+          
+    else:
+        test_data = pd.read_csv(args.data_path + "/ttest.csv")
+    corpus_data = pd.read_csv(args.corpus_file)
+    #dcorpus = pd.DataFrame(corpus_data)
+    #pandarallel.initialize(progress_bar=True, use_memory_fs=False, nb_workers=12)
+    #dcorpus["full_text"] = dcorpus.parallel_apply(concat_str, axis=1)
+    corpus = corpus_data['tokenized_text'].tolist()
+    
+    if csv_file:
+        ans_ids = []
+        ground_ids = []
+        org_questions = test_data['question'].tolist()
+        questions = test_data['tokenized_question'].tolist()
+        for i in range(len(test_data)):
+            ans_ids.append(json.loads(test_data['best_ans_id'][i]))
+            ground_ids.append(json.loads(test_data['ans_id'][i]))
+        ground_truths = []
+        for sample in ans_ids:
+            temp = [corpus_data['law_id'][y] + "_" + str(corpus_data['article_id'][y]) for y in sample]
+            ground_truths.append(temp)
+    else:
+        ground_truths = []
+        ground_ids = []
+        org_questions = [sample['question'] for sample in test_data]
+        questions = [tokenize(preprocess_question(sample['question'], remove_end_phrase=False)) for sample in test_data]
+        for sample in test_data:
+            try:
+                temp = [it['law_id'] + "_" + it['article_id'] for it in sample['relevance_articles']]
+                tempp = [it['ans_id'] for it in sample['relevance_articles']]
+            except:
+                temp = [it['law_id'] + "_" + it['article_id'] for it in sample['relevant_articles']]
+                tempp = [it['ans_id'] for it in sample['relevant_articles']]
+            ground_truths.append(temp)
+            ground_ids.append(tempp)
+    
+    faiss_index = index(
+        model=model, 
+        tokenizer=tokenizer,
+        corpus=corpus, 
+        batch_size=args.batch_size,
+        max_length=args.max_passage_length,
+        index_factory=args.index_factory,
+        save_path=args.save_path,
+        save_embedding=args.save_embedding,
+        load_embedding=args.load_embedding
+    )
+    
+    scores, indices = search(
+        model=model, 
+        tokenizer=tokenizer,
+        questions=questions, 
+        faiss_index=faiss_index, 
+        k=args.k, 
+        batch_size=args.batch_size, 
+        max_length=args.max_query_length
+    )
+
+
+    retrieval_results, retrieval_ids = [], []
+    for indice in indices:
+        # filter invalid indices
+        indice = indice[indice != -1].tolist()
+        rst = []
+        for x in indice:
+            temp = corpus_data['law_id'][x] + "_" + str(corpus_data['article_id'][x])
+            if temp not in rst:
+                rst.append(temp)
+        retrieval_results.append(rst)
+        retrieval_ids.append(indice)
+        
+    rerank_ids, rerank_scores = rerank(reranker, reranker_tokenizer, questions, corpus, retrieval_ids, args.cross_batch_size, args.cross_max_length, args.top_k)
+
+    if args.bi_data:
+        save_bi_data(questions, ground_ids, rerank_ids, rerank_scores, args.data_type, org_questions)
+        
+    metrics = check(ground_ids, retrieval_ids)
+    print(metrics)
+    metrics = evaluate(retrieval_results, ground_truths)
+    print(metrics)
+    metrics = check(ground_ids, rerank_ids, cutoffs=[1,5,10,30])
+    print(metrics)
+
+if __name__ == "__main__":
+    main()

src/train_cross_encoder.py

@@ -0,0 +1,112 @@
+import logging
+
+import torch
+from typing import Dict
+from transformers.utils.logging import enable_explicit_format
+from transformers.trainer_callback import PrinterCallback
+from transformers import (
+    AutoTokenizer,
+    HfArgumentParser,
+    EvalPrediction,
+    set_seed,
+    PreTrainedTokenizerFast
+)
+
+from cross_rerank.logger_config import logger, LoggerCallback
+from cross_rerank.config import Arguments
+from cross_rerank.trainer import RerankerTrainer
+from cross_rerank.data_loader import CrossEncoderDataLoader
+from cross_rerank.collator import CrossEncoderCollator
+from cross_rerank.metrics import accuracy
+from cross_rerank.model import Reranker
+
+
+def _common_setup(args: Arguments):
+    if args.process_index > 0:
+        logger.setLevel(logging.WARNING)
+    enable_explicit_format()
+    set_seed(args.seed)
+
+
+def _compute_metrics(eval_pred: EvalPrediction) -> Dict:
+    preds = eval_pred.predictions
+    if isinstance(preds, tuple):
+        preds = preds[-1]
+    logits = torch.tensor(preds).float()
+    labels = torch.tensor(eval_pred.label_ids).long()
+    acc = accuracy(output=logits, target=labels)[0]
+
+    return {'acc': acc}
+
+
+def main():
+    parser = HfArgumentParser((Arguments,))
+    args: Arguments = parser.parse_args_into_dataclasses()[0]
+    _common_setup(args)
+    logger.info('Args={}'.format(str(args)))
+
+    try:
+        tokenizer: PreTrainedTokenizerFast = AutoTokenizer.from_pretrained(args.model_name_or_path)
+    except:
+        tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
+
+    model: Reranker = Reranker.from_pretrained(
+        all_args=args,
+        pretrained_model_name_or_path=args.model_name_or_path,
+        num_labels=1)
+
+    logger.info(model)
+    logger.info('Vocab size: {}'.format(len(tokenizer)))
+
+    data_collator = CrossEncoderCollator(
+        tokenizer=tokenizer,
+        pad_to_multiple_of=256 if args.fp16 else 256)
+
+    rerank_data_loader = CrossEncoderDataLoader(args=args, tokenizer=tokenizer)
+    train_dataset = rerank_data_loader.train_dataset
+    eval_dataset = rerank_data_loader.eval_dataset
+
+    trainer = RerankerTrainer(
+        model=model,
+        args=args,
+        train_dataset=train_dataset if args.do_train else None,
+        eval_dataset=eval_dataset if args.do_eval else None,
+        data_collator=data_collator,
+        compute_metrics=_compute_metrics,
+        tokenizer=tokenizer,
+    )
+    trainer.remove_callback(PrinterCallback)
+    trainer.add_callback(LoggerCallback)
+    rerank_data_loader.trainer = trainer
+    
+    if args.do_eval:
+        logger.info("*** Evaluate ***")
+        metrics = trainer.evaluate(metric_key_prefix="eval")
+        metrics["eval_samples"] = len(eval_dataset)
+
+        trainer.log_metrics("eval", metrics)
+        trainer.save_metrics("eval", metrics)
+
+    if args.do_train:
+        train_result = trainer.train(resume_from_checkpoint= args.resume_from_checkpoint)
+        trainer.save_model()
+
+        metrics = train_result.metrics
+        metrics["train_samples"] = len(train_dataset)
+
+        trainer.log_metrics("train", metrics)
+        trainer.save_metrics("train", metrics)
+
+    if args.do_eval:
+        logger.info("*** Evaluate ***")
+        metrics = trainer.evaluate(metric_key_prefix="eval")
+        metrics["eval_samples"] = len(eval_dataset)
+
+        trainer.log_metrics("eval", metrics)
+        trainer.save_metrics("eval", metrics)
+
+    return
+
+
+if __name__ == "__main__":
+    main()