chore: import upstream snapshot with attribution

simlm/src/collators/__init__.py

@@ -0,0 +1,3 @@
+from .biencoder_collator import BiencoderCollator
+from .cross_encoder_collator import CrossEncoderCollator
+from .rlm_collator import DataCollatorForReplaceLM

simlm/src/collators/biencoder_collator.py

@@ -0,0 +1,56 @@
+import torch
+
+from dataclasses import dataclass
+from typing import List, Dict, Any
+from transformers import DataCollatorWithPadding, BatchEncoding
+
+
+def _unpack_doc_values(features: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    doc_examples = []
+    for f in features:
+        keys = list(f.keys())
+        lists_per_key = len(f[keys[0]])
+        for idx in range(lists_per_key):
+            doc_examples.append({k: f[k][idx] for k in keys})
+    return doc_examples
+
+
+@dataclass
+class BiencoderCollator(DataCollatorWithPadding):
+
+    def __call__(self, features: List[Dict[str, Any]]) -> BatchEncoding:
+        q_prefix, d_prefix = 'q_', 'd_'
+        query_examples = [{k[len(q_prefix):]: v for k, v in f.items() if k.startswith(q_prefix)} for f in features]
+        doc_examples = _unpack_doc_values(
+            [{k[len(d_prefix):]: v for k, v in f.items() if k.startswith(d_prefix)} for f in features])
+        assert len(doc_examples) % len(query_examples) == 0, \
+            '{} doc and {} queries'.format(len(doc_examples), len(query_examples))
+
+        # already truncated during tokenization
+        q_collated = self.tokenizer.pad(
+            query_examples,
+            padding=self.padding,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            return_tensors=self.return_tensors)
+        d_collated = self.tokenizer.pad(
+            doc_examples,
+            padding=self.padding,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            return_tensors=self.return_tensors)
+
+        # merge into a single BatchEncoding by adding prefix
+        for k in list(q_collated.keys()):
+            q_collated[q_prefix + k] = q_collated[k]
+            del q_collated[k]
+        for k in d_collated:
+            q_collated[d_prefix + k] = d_collated[k]
+
+        merged_batch_dict = q_collated
+        # dummy placeholder for field "labels", won't use it to compute loss
+        labels = torch.zeros(len(query_examples), dtype=torch.long)
+        merged_batch_dict['labels'] = labels
+
+        if 'kd_labels' in features[0]:
+            kd_labels = torch.stack([torch.tensor(f['kd_labels']) for f in features], dim=0).float()
+            merged_batch_dict['kd_labels'] = kd_labels
+        return merged_batch_dict

simlm/src/collators/collator_utils.py

@@ -0,0 +1,110 @@
+import torch
+import random
+import warnings
+
+from transformers import BertTokenizer, BertTokenizerFast, BatchEncoding
+from typing import List, Union, Tuple, Any, Dict
+
+
+def whole_word_mask(tokenizer: Union[BertTokenizer, BertTokenizerFast],
+                    input_tokens: List[str],
+                    mlm_prob: float,
+                    max_predictions=512) -> List[int]:
+    """
+    Get 0/1 labels for masked tokens with whole word mask proxy
+    """
+    if not isinstance(tokenizer, (BertTokenizer, BertTokenizerFast)):
+        warnings.warn(
+            "DataCollatorForWholeWordMask is only suitable for BertTokenizer-like tokenizers. "
+            "Please refer to the documentation for more information."
+        )
+
+    cand_indexes = []
+    for (i, token) in enumerate(input_tokens):
+        if token == "[CLS]" or token == "[SEP]":
+            continue
+
+        if len(cand_indexes) >= 1 and token.startswith("##"):
+            cand_indexes[-1].append(i)
+        else:
+            cand_indexes.append([i])
+
+    random.shuffle(cand_indexes)
+    num_to_predict = min(max_predictions, max(1, int(round(len(input_tokens) * mlm_prob))))
+    masked_lms = []
+    covered_indexes = set()
+    for index_set in cand_indexes:
+        if len(masked_lms) >= num_to_predict:
+            break
+        # If adding a whole-word mask would exceed the maximum number of
+        # predictions, then just skip this candidate.
+        if len(masked_lms) + len(index_set) > num_to_predict:
+            continue
+        is_any_index_covered = False
+        for index in index_set:
+            if index in covered_indexes:
+                is_any_index_covered = True
+                break
+        if is_any_index_covered:
+            continue
+        for index in index_set:
+            covered_indexes.add(index)
+            masked_lms.append(index)
+
+    if len(covered_indexes) != len(masked_lms):
+        raise ValueError("Length of covered_indexes is not equal to length of masked_lms.")
+    mask_labels = [1 if i in covered_indexes else 0 for i in range(len(input_tokens))]
+    return mask_labels
+
+
+def torch_mask_tokens(tokenizer: Union[BertTokenizer, BertTokenizerFast],
+                      inputs: torch.Tensor,
+                      mask_labels: torch.Tensor,
+                      all_use_mask_token: bool = False) -> Tuple[Any, Any]:
+    """
+    Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original. Set
+    'mask_labels' means we use whole word mask (wwm), we directly mask idxs according to it's ref.
+    """
+    if tokenizer.mask_token is None:
+        raise ValueError(
+            "This tokenizer does not have a mask token which is necessary for masked language modeling. Remove the --mlm flag if you want to use this tokenizer."
+        )
+    labels = inputs.clone()
+    masked_inputs = inputs.clone()
+    # We sample a few tokens in each sequence for masked-LM training (with probability args.mlm_probability defaults to 0.15 in Bert/RoBERTa)
+
+    probability_matrix = mask_labels.clone()
+
+    special_tokens_mask = [
+        tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True) for val in labels.tolist()
+    ]
+    probability_matrix.masked_fill_(torch.tensor(special_tokens_mask, dtype=torch.bool), value=0.0)
+    if tokenizer._pad_token is not None:
+        padding_mask = labels.eq(tokenizer.pad_token_id)
+        probability_matrix.masked_fill_(padding_mask, value=0.0)
+
+    masked_indices = probability_matrix.bool()
+    labels[~masked_indices] = -100  # We only compute loss on masked tokens
+
+    if all_use_mask_token:
+        masked_inputs[masked_indices] = tokenizer.convert_tokens_to_ids(tokenizer.mask_token)
+        return masked_inputs, labels
+
+    # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])
+    indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices
+    masked_inputs[indices_replaced] = tokenizer.convert_tokens_to_ids(tokenizer.mask_token)
+
+    # 10% of the time, we replace masked input tokens with random word
+    indices_random = torch.bernoulli(torch.full(labels.shape, 0.5)).bool() & masked_indices & ~indices_replaced
+    random_words = torch.randint(len(tokenizer), labels.shape, dtype=torch.long)
+    masked_inputs[indices_random] = random_words[indices_random]
+
+    # The rest of the time (10% of the time) we keep the masked input tokens unchanged
+    return masked_inputs, labels
+
+
+def merge_batch_dict(src_batch_dict: Union[Dict, BatchEncoding],
+                     tgt_batch_dict: Union[Dict, BatchEncoding],
+                     prefix: str = None):
+    for key in src_batch_dict:
+        tgt_batch_dict[(prefix or '') + key] = src_batch_dict[key].clone()

simlm/src/collators/cross_encoder_collator.py

@@ -0,0 +1,26 @@
+import torch
+
+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})
+
+        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)
+
+        collated_batch_dict['labels'] = torch.zeros(len(features), dtype=torch.long)
+        return collated_batch_dict

simlm/src/collators/rlm_collator.py

@@ -0,0 +1,107 @@
+import copy
+
+from dataclasses import dataclass
+from typing import List, Dict, Optional, Any
+from transformers import BatchEncoding, BertTokenizerFast
+from transformers.data.data_collator import _torch_collate_batch
+from transformers.file_utils import PaddingStrategy
+
+from config import Arguments
+from .collator_utils import whole_word_mask, torch_mask_tokens, merge_batch_dict
+from logger_config import logger
+
+
+@dataclass
+class DataCollatorForReplaceLM:
+    tokenizer: BertTokenizerFast
+    pad_to_multiple_of: Optional[int] = None
+    args: Arguments = None
+
+    def __post_init__(self):
+        if self.tokenizer.mask_token is None:
+            raise ValueError(
+                "This tokenizer does not have a mask token which is necessary for masked language modeling. "
+                "You should pass `mlm=False` to train on causal language modeling instead."
+            )
+
+    def __call__(self, features: List[Dict]):
+        return self.torch_call(features)
+
+    def torch_call(self, examples: List[Dict[str, Any]]) -> BatchEncoding:
+        if 'title' in examples[0]:
+            text, text_pair = [ex['title'] for ex in examples], [ex['contents'] for ex in examples]
+        else:
+            text, text_pair = [ex['contents'] for ex in examples], None
+
+        batch_dict = self.tokenizer(text,
+                                    text_pair=text_pair,
+                                    max_length=self.args.rlm_max_length,
+                                    padding=PaddingStrategy.DO_NOT_PAD,
+                                    truncation=True)
+
+        encoder_mask_labels = []
+        decoder_mask_labels = []
+        extra_mlm_prob = self.args.rlm_decoder_mask_prob - self.args.rlm_encoder_mask_prob
+        # mlm_prob + (1 - mlm_prob) x = decoder_prob
+        # => x = (decoder_prob - mlm_prob) / (1 - mlm_prob)
+        # since we mask twice independently, we need to adjust extra_mlm_prob accordingly
+        extra_mlm_prob = extra_mlm_prob / (1 - self.args.rlm_encoder_mask_prob)
+
+        for input_ids in batch_dict['input_ids']:
+            ref_tokens = []
+            for token_id in input_ids:
+                token = self.tokenizer._convert_id_to_token(token_id)
+                ref_tokens.append(token)
+            encoder_mask_labels.append(whole_word_mask(self.tokenizer, ref_tokens,
+                                                       mlm_prob=self.args.rlm_encoder_mask_prob))
+
+            decoder_mask = encoder_mask_labels[-1][:]
+            # overlapping mask
+            if extra_mlm_prob > 1e-4:
+                decoder_mask = [max(m1, m2) for m1, m2 in zip(decoder_mask,
+                                whole_word_mask(self.tokenizer, ref_tokens, mlm_prob=extra_mlm_prob))]
+
+            assert len(decoder_mask) == len(encoder_mask_labels[-1])
+            decoder_mask_labels.append(decoder_mask)
+
+        encoder_batch_mask = _torch_collate_batch(encoder_mask_labels, self.tokenizer,
+                                                  pad_to_multiple_of=self.pad_to_multiple_of)
+
+        encoder_batch_dict = self.tokenizer.pad(batch_dict,
+                                                pad_to_multiple_of=self.pad_to_multiple_of,
+                                                return_tensors="pt")
+        encoder_inputs, encoder_labels = torch_mask_tokens(
+            self.tokenizer, encoder_batch_dict['input_ids'], encoder_batch_mask,
+            all_use_mask_token=self.args.all_use_mask_token)
+
+        clean_input_ids = encoder_batch_dict['input_ids'].clone()
+        encoder_batch_dict['input_ids'] = encoder_inputs
+        encoder_batch_dict['labels'] = encoder_labels
+
+        merged_batch_dict = BatchEncoding()
+        merge_batch_dict(encoder_batch_dict, merged_batch_dict, prefix='enc_')
+
+        decoder_batch_dict = copy.deepcopy(encoder_batch_dict)
+        if extra_mlm_prob > 1e-4:
+            decoder_batch_mask = _torch_collate_batch(decoder_mask_labels, self.tokenizer,
+                                                      pad_to_multiple_of=self.pad_to_multiple_of)
+            decoder_inputs, decoder_labels = torch_mask_tokens(
+                self.tokenizer, clean_input_ids, decoder_batch_mask,
+                all_use_mask_token=self.args.all_use_mask_token)
+
+            decoder_batch_dict['input_ids'] = decoder_inputs
+            decoder_batch_dict['labels'] = decoder_labels
+
+        merge_batch_dict(decoder_batch_dict, merged_batch_dict, prefix='dec_')
+
+        # simple integrity check
+        # logger.info('encoder mask cnt: {}, decoder mask cnt: {}, non-equal input_ids cnt: {}'.format(
+        #     (merged_batch_dict['enc_labels'] > 0).long().sum(),
+        #     (merged_batch_dict['dec_labels'] > 0).long().sum(),
+        #     (merged_batch_dict['dec_input_ids'] != merged_batch_dict['enc_input_ids']).long().sum()))
+
+        labels = clean_input_ids.clone()
+        for special_id in self.tokenizer.all_special_ids:
+            labels[labels == special_id] = -100
+        merged_batch_dict['labels'] = labels
+        return merged_batch_dict

simlm/src/config.py

@@ -0,0 +1,214 @@
+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"}
+    )
+
+    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_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."
+        },
+    )
+    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(10, self.logging_steps)
+
+        super(Arguments, self).__post_init__()
+
+        if self.output_dir:
+            os.makedirs(self.output_dir, exist_ok=True)
+
+        self.label_names = ['labels']

simlm/src/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')

simlm/src/inference/encode_main.py

@@ -0,0 +1,119 @@
+import os
+import tqdm
+import torch
+
+from contextlib import nullcontext
+from torch.utils.data import DataLoader
+from functools import partial
+from datasets import load_dataset
+from typing import Dict, List
+from transformers.file_utils import PaddingStrategy
+from transformers import (
+    AutoTokenizer,
+    PreTrainedTokenizerFast,
+    DataCollatorWithPadding,
+    HfArgumentParser,
+    BatchEncoding
+)
+
+from config import Arguments
+from logger_config import logger
+from utils import move_to_cuda
+from models import BiencoderModelForInference, BiencoderOutput
+
+parser = HfArgumentParser((Arguments,))
+args: Arguments = parser.parse_args_into_dataclasses()[0]
+
+
+def _psg_transform_func(tokenizer: PreTrainedTokenizerFast,
+                        examples: Dict[str, List]) -> BatchEncoding:
+    batch_dict = tokenizer(examples['title'],
+                           text_pair=examples['contents'],
+                           max_length=args.p_max_len,
+                           padding=PaddingStrategy.DO_NOT_PAD,
+                           truncation=True)
+    # for co-Condenser reproduction purpose only
+    if args.model_name_or_path.startswith('Luyu/'):
+        del batch_dict['token_type_ids']
+
+    return batch_dict
+
+
+@torch.no_grad()
+def _worker_encode_passages(gpu_idx: int):
+    def _get_out_path(shard_idx: int = 0) -> str:
+        return '{}/shard_{}_{}'.format(args.encode_save_dir, gpu_idx, shard_idx)
+
+    if os.path.exists(_get_out_path(0)):
+        logger.error('{} already exists, will skip encoding'.format(_get_out_path(0)))
+        return
+
+    dataset = load_dataset('json', data_files=args.encode_in_path)['train']
+    if args.dry_run:
+        dataset = dataset.select(range(4096))
+    dataset = dataset.shard(num_shards=torch.cuda.device_count(),
+                            index=gpu_idx,
+                            contiguous=True)
+
+    logger.info('GPU {} needs to process {} examples'.format(gpu_idx, len(dataset)))
+    torch.cuda.set_device(gpu_idx)
+
+    tokenizer: PreTrainedTokenizerFast = AutoTokenizer.from_pretrained(args.model_name_or_path)
+    model: BiencoderModelForInference = BiencoderModelForInference.build(args)
+    model.eval()
+    model.cuda()
+
+    dataset.set_transform(partial(_psg_transform_func, tokenizer))
+
+    data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8 if args.fp16 else None)
+    data_loader = DataLoader(
+        dataset,
+        batch_size=args.encode_batch_size,
+        shuffle=False,
+        drop_last=False,
+        num_workers=args.dataloader_num_workers,
+        collate_fn=data_collator,
+        pin_memory=True)
+
+    num_encoded_docs, encoded_embeds, cur_shard_idx = 0, [], 0
+    for batch_dict in tqdm.tqdm(data_loader, desc='passage encoding', mininterval=8):
+        batch_dict = move_to_cuda(batch_dict)
+
+        with torch.cuda.amp.autocast() if args.fp16 else nullcontext():
+            outputs: BiencoderOutput = model(query=None, passage=batch_dict)
+        encoded_embeds.append(outputs.p_reps.cpu())
+        num_encoded_docs += outputs.p_reps.shape[0]
+
+        if num_encoded_docs >= args.encode_shard_size:
+            out_path = _get_out_path(cur_shard_idx)
+            concat_embeds = torch.cat(encoded_embeds, dim=0)
+            logger.info('GPU {} save {} embeds to {}'.format(gpu_idx, concat_embeds.shape[0], out_path))
+            torch.save(concat_embeds, out_path)
+
+            cur_shard_idx += 1
+            num_encoded_docs = 0
+            encoded_embeds.clear()
+
+    if num_encoded_docs > 0:
+        out_path = _get_out_path(cur_shard_idx)
+        concat_embeds = torch.cat(encoded_embeds, dim=0)
+        logger.info('GPU {} save {} embeds to {}'.format(gpu_idx, concat_embeds.shape[0], out_path))
+        torch.save(concat_embeds, out_path)
+
+    logger.info('Done computing score for worker {}'.format(gpu_idx))
+
+
+def _batch_encode_passages():
+    logger.info('Args={}'.format(str(args)))
+    gpu_count = torch.cuda.device_count()
+    if gpu_count == 0:
+        logger.error('No gpu available')
+        return
+
+    logger.info('Use {} gpus'.format(gpu_count))
+    torch.multiprocessing.spawn(_worker_encode_passages, args=(), nprocs=gpu_count)
+    logger.info('Done batch encode passages')
+
+
+if __name__ == '__main__':
+    _batch_encode_passages()

simlm/src/inference/gen_teacher_scores.py

@@ -0,0 +1,174 @@
+import os
+import tqdm
+import torch
+
+from contextlib import nullcontext
+from torch.utils.data import DataLoader
+from functools import partial
+from datasets import Dataset, load_dataset
+from typing import Dict, List
+from transformers.file_utils import PaddingStrategy
+from transformers.modeling_outputs import SequenceClassifierOutput
+from transformers import (
+    AutoTokenizer,
+    PreTrainedTokenizerFast,
+    DataCollatorWithPadding,
+    HfArgumentParser,
+    BatchEncoding
+)
+
+from config import Arguments
+from logger_config import logger
+from utils import move_to_cuda
+from models import RerankerForInference
+from data_utils import load_corpus, load_queries, save_to_readable_format
+
+parser = HfArgumentParser((Arguments,))
+args: Arguments = parser.parse_args_into_dataclasses()[0]
+kd_gen_score_in_path = os.path.join(args.data_dir, '{}.jsonl'.format(args.kd_gen_score_split))
+kd_gen_score_out_path = os.path.join(args.data_dir, 'kd_{}.jsonl'.format(args.kd_gen_score_split))
+
+
+def _kd_gen_score_transform_func(tokenizer: PreTrainedTokenizerFast,
+                                 corpus: Dataset,
+                                 queries: Dict[str, str],
+                                 examples: Dict[str, List]) -> BatchEncoding:
+    input_docs: List[str] = []
+
+    # ATTENTION: this code should be consistent with CrossEncoderDataLoader
+    for doc_id in examples['doc_id']:
+        doc_id = int(doc_id)
+        prefix = ''
+        if corpus[doc_id].get('title', ''):
+            prefix = corpus[doc_id]['title'] + ': '
+        input_docs.append(prefix + corpus[doc_id]['contents'])
+
+    input_queries = [queries[query_id] for query_id in examples['query_id']]
+    batch_dict = tokenizer(input_queries,
+                           text_pair=input_docs,
+                           max_length=args.rerank_max_length,
+                           padding=PaddingStrategy.DO_NOT_PAD,
+                           truncation=True)
+
+    return batch_dict
+
+
+def _get_shard_path(worker_idx: int) -> str:
+    return '{}_shard_{}'.format(kd_gen_score_in_path, worker_idx)
+
+
+@torch.no_grad()
+def _worker_gen_teacher_score(gpu_idx: int):
+    dataset = load_dataset('json', data_files=kd_gen_score_in_path)['train']
+    if args.dry_run:
+        dataset = dataset.select(range(100))
+    dataset = dataset.shard(num_shards=torch.cuda.device_count(),
+                            index=gpu_idx,
+                            contiguous=True)
+
+    qid_pids = []
+    for ex in tqdm.tqdm(dataset, desc='get qid-pid pairs', mininterval=3):
+        for pos_doc_id in ex['positives']['doc_id']:
+            qid_pids.append((ex['query_id'], pos_doc_id))
+        for neg_doc_id in ex['negatives']['doc_id'][:args.kd_gen_score_n_neg]:
+            qid_pids.append((ex['query_id'], neg_doc_id))
+
+    dataset = Dataset.from_dict({'query_id': [t[0] for t in qid_pids],
+                                 'doc_id': [t[1] for t in qid_pids]})
+
+    query_ids, doc_ids = dataset['query_id'], dataset['doc_id']
+
+    logger.info('GPU {} needs to process {} examples'.format(gpu_idx, len(dataset)))
+    torch.cuda.set_device(gpu_idx)
+
+    tokenizer: PreTrainedTokenizerFast = AutoTokenizer.from_pretrained(args.model_name_or_path)
+    model: RerankerForInference = RerankerForInference.from_pretrained(args.model_name_or_path)
+    model.eval()
+    model.cuda()
+
+    corpus: Dataset = load_corpus(path=os.path.join(args.data_dir, 'passages.jsonl.gz'))
+    queries = load_queries(path='{}/{}_queries.tsv'.format(args.data_dir, args.kd_gen_score_split),
+                           task_type=args.task_type)
+    dataset.set_transform(partial(_kd_gen_score_transform_func, tokenizer, corpus, queries))
+
+    data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8 if args.fp16 else None)
+    data_loader = DataLoader(
+        dataset,
+        batch_size=args.kd_gen_score_batch_size,
+        shuffle=False,
+        drop_last=False,
+        num_workers=args.dataloader_num_workers,
+        collate_fn=data_collator,
+        pin_memory=True)
+
+    scores = []
+    for batch_dict in tqdm.tqdm(data_loader, desc='generate teacher score', mininterval=5):
+        batch_dict = move_to_cuda(batch_dict)
+
+        with torch.cuda.amp.autocast() if args.fp16 else nullcontext():
+            outputs: SequenceClassifierOutput = model(batch_dict)
+        scores.append(outputs.logits.squeeze(dim=-1).cpu())
+        assert len(scores[-1].shape) == 1
+
+    all_scores = torch.cat(scores, dim=-1)
+    assert all_scores.shape[0] == len(dataset), '{} != {}'
+    all_scores = all_scores.tolist()
+
+    with open(_get_shard_path(gpu_idx), 'w', encoding='utf-8') as writer:
+        for idx in range(len(query_ids)):
+            writer.write('{}\t{}\t{}\n'.format(query_ids[idx], doc_ids[idx], round(all_scores[idx], 5)))
+
+    logger.info('Done computing teacher score for worker {}'.format(gpu_idx))
+
+
+def _merge_teacher_scores(worker_cnt: int):
+    qid_to_pid_to_score = {}
+    for worker_idx in range(worker_cnt):
+        shard_path = _get_shard_path(worker_idx)
+        for line in tqdm.tqdm(open(shard_path, 'r', encoding='utf-8'),
+                              desc='Load shard {} score'.format(worker_idx), mininterval=3):
+            fs = line.strip().split('\t')
+            assert len(fs) == 3
+            qid, pid, score = fs
+            if qid not in qid_to_pid_to_score:
+                qid_to_pid_to_score[qid] = {}
+            qid_to_pid_to_score[qid][pid] = float(score)
+        os.remove(shard_path)
+
+    dataset = load_dataset('json', data_files=kd_gen_score_in_path)['train']
+    if args.dry_run:
+        dataset = dataset.select(range(100))
+
+    def _update_score(ex: Dict) -> Dict:
+        query_id = ex['query_id']
+        pid_to_score = qid_to_pid_to_score[query_id]
+        ex['negatives']['doc_id'] = [neg_doc_id for neg_doc_id in ex['negatives']['doc_id'] if neg_doc_id in pid_to_score]
+        ex['positives']['score'] = [pid_to_score[pos_doc_id] for pos_doc_id in ex['positives']['doc_id']]
+        ex['negatives']['score'] = [pid_to_score[neg_doc_id] for neg_doc_id in ex['negatives']['doc_id']]
+        return ex
+
+    dataset = dataset.map(_update_score, num_proc=4)
+    logger.info('Writing teacher score to {}'.format(kd_gen_score_out_path))
+    dataset.to_json(kd_gen_score_out_path, force_ascii=False, lines=True)
+
+
+def _batch_compute_teacher_score():
+    logger.info('Args={}'.format(str(args)))
+    gpu_count = torch.cuda.device_count()
+    if gpu_count == 0:
+        logger.error('No gpu available')
+        return
+
+    logger.info('Use {} gpus'.format(gpu_count))
+    torch.multiprocessing.spawn(_worker_gen_teacher_score, args=(), nprocs=gpu_count)
+    logger.info('Done batch generate teacher score')
+
+    _merge_teacher_scores(gpu_count)
+    logger.info('Done merge results')
+
+    corpus = load_corpus(path=os.path.join(args.data_dir, 'passages.jsonl.gz'))
+    save_to_readable_format(in_path=kd_gen_score_out_path, corpus=corpus)
+
+
+if __name__ == '__main__':
+    _batch_compute_teacher_score()

simlm/src/inference/rerank_main.py

@@ -0,0 +1,133 @@
+import os
+import tqdm
+import torch
+
+from contextlib import nullcontext
+from torch.utils.data import DataLoader
+from functools import partial
+from datasets import Dataset
+from typing import Dict, List
+from transformers.file_utils import PaddingStrategy
+from transformers.modeling_outputs import SequenceClassifierOutput
+from transformers import (
+    AutoTokenizer,
+    PreTrainedTokenizerFast,
+    DataCollatorWithPadding,
+    HfArgumentParser,
+    BatchEncoding
+)
+
+from config import Arguments
+from logger_config import logger
+from utils import move_to_cuda
+from models import RerankerForInference
+from data_utils import load_msmarco_predictions, load_corpus, load_queries, \
+    merge_rerank_predictions, get_rerank_shard_path
+
+parser = HfArgumentParser((Arguments,))
+args: Arguments = parser.parse_args_into_dataclasses()[0]
+
+
+def _rerank_transform_func(tokenizer: PreTrainedTokenizerFast,
+                           corpus: Dataset,
+                           queries: Dict[str, str],
+                           examples: Dict[str, List]) -> BatchEncoding:
+    input_docs: List[str] = []
+
+    # ATTENTION: this code should be consistent with RerankDataLoader
+    for doc_id in examples['doc_id']:
+        doc_id = int(doc_id)
+        prefix = ''
+        if corpus[doc_id].get('title', ''):
+            prefix = corpus[doc_id]['title'] + ': '
+        input_docs.append(prefix + corpus[doc_id]['contents'])
+
+    input_queries = [queries[query_id] for query_id in examples['query_id']]
+    batch_dict = tokenizer(input_queries,
+                           text_pair=input_docs,
+                           max_length=args.rerank_max_length,
+                           padding=PaddingStrategy.DO_NOT_PAD,
+                           truncation=True)
+
+    return batch_dict
+
+
+@torch.no_grad()
+def _worker_compute_reranker_score(gpu_idx: int):
+    preds = load_msmarco_predictions(args.rerank_in_path)
+    query_ids = sorted(list(preds.keys()))
+    qid_pid = []
+    for query_id in tqdm.tqdm(query_ids, desc='load qid-pid', mininterval=2):
+        qid_pid += [(scored_doc.qid, scored_doc.pid) for scored_doc in preds[query_id]
+                    if scored_doc.rank <= args.rerank_depth]
+
+    dataset = Dataset.from_dict({'query_id': [t[0] for t in qid_pid],
+                                 'doc_id': [t[1] for t in qid_pid]})
+    dataset = dataset.shard(num_shards=torch.cuda.device_count(),
+                            index=gpu_idx,
+                            contiguous=True)
+
+    logger.info('GPU {} needs to process {} examples'.format(gpu_idx, len(dataset)))
+    torch.cuda.set_device(gpu_idx)
+
+    query_ids, doc_ids = dataset['query_id'], dataset['doc_id']
+    assert len(dataset) == len(query_ids)
+
+    tokenizer: PreTrainedTokenizerFast = AutoTokenizer.from_pretrained(args.model_name_or_path)
+    model: RerankerForInference = RerankerForInference.from_pretrained(args.model_name_or_path)
+    model.eval()
+    model.cuda()
+
+    corpus: Dataset = load_corpus(path=os.path.join(args.data_dir, 'passages.jsonl.gz'))
+    queries = load_queries(path='{}/{}_queries.tsv'.format(args.data_dir, args.rerank_split),
+                           task_type=args.task_type)
+    dataset.set_transform(partial(_rerank_transform_func, tokenizer, corpus, queries))
+
+    data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8 if args.fp16 else None)
+    data_loader = DataLoader(
+        dataset,
+        batch_size=args.rerank_batch_size,
+        shuffle=False,
+        drop_last=False,
+        num_workers=args.dataloader_num_workers,
+        collate_fn=data_collator,
+        pin_memory=True)
+
+    scores = []
+    for batch_dict in tqdm.tqdm(data_loader, desc='passage rerank', mininterval=5):
+        batch_dict = move_to_cuda(batch_dict)
+
+        with torch.cuda.amp.autocast() if args.fp16 else nullcontext():
+            outputs: SequenceClassifierOutput = model(batch_dict)
+        scores.append(outputs.logits.squeeze(dim=-1).cpu())
+        assert len(scores[-1].shape) == 1
+
+    all_scores = torch.cat(scores, dim=-1)
+    assert all_scores.shape[0] == len(query_ids), '{} != {}'.format(all_scores.shape[0], len(query_ids))
+    all_scores = all_scores.tolist()
+
+    with open(get_rerank_shard_path(args, gpu_idx), 'w', encoding='utf-8') as writer:
+        for idx in range(len(query_ids)):
+            # dummy rank, since a query may be split across different workers
+            writer.write('{}\t{}\t{}\t{}\n'.format(query_ids[idx], doc_ids[idx], -1, round(all_scores[idx], 5)))
+
+    logger.info('Done computing rerank score for worker {}'.format(gpu_idx))
+
+
+def _batch_compute_reranker_score():
+    logger.info('Args={}'.format(str(args)))
+    gpu_count = torch.cuda.device_count()
+    if gpu_count == 0:
+        logger.error('No gpu available')
+        return
+
+    logger.info('Use {} gpus'.format(gpu_count))
+    torch.multiprocessing.spawn(_worker_compute_reranker_score, args=(), nprocs=gpu_count)
+    logger.info('Done batch compute rerank score')
+
+    merge_rerank_predictions(args, gpu_count)
+    logger.info('Done merge results')
+
+
+if __name__ == '__main__':
+    _batch_compute_reranker_score()

simlm/src/inference/search_main.py

@@ -0,0 +1,193 @@
+import json
+import os
+import glob
+import tqdm
+import torch
+
+from contextlib import nullcontext
+from torch.utils.data import DataLoader
+from functools import partial
+from collections import defaultdict
+from datasets import Dataset
+from typing import Dict, List, Tuple
+from transformers.file_utils import PaddingStrategy
+from transformers import (
+    AutoTokenizer,
+    PreTrainedTokenizerFast,
+    DataCollatorWithPadding,
+    HfArgumentParser,
+    BatchEncoding
+)
+
+from config import Arguments
+from logger_config import logger
+from utils import move_to_cuda, save_json_to_file
+from metrics import compute_mrr, trec_eval, ScoredDoc
+from data_utils import load_queries, load_qrels, load_msmarco_predictions, save_preds_to_msmarco_format
+from models import BiencoderModelForInference, BiencoderOutput
+
+parser = HfArgumentParser((Arguments,))
+args: Arguments = parser.parse_args_into_dataclasses()[0]
+assert os.path.exists(args.encode_save_dir)
+
+
+def _get_all_shards_path() -> List[str]:
+    path_list = glob.glob('{}/shard_*_*'.format(args.encode_save_dir))
+    assert len(path_list) > 0
+
+    def _parse_worker_idx_shard_idx(p: str) -> Tuple:
+        worker_idx, shard_idx = [int(f) for f in os.path.basename(p).split('_')[-2:]]
+        return worker_idx, shard_idx
+
+    path_list = sorted(path_list, key=lambda path: _parse_worker_idx_shard_idx(path))
+    logger.info('Embeddings path list: {}'.format(path_list))
+    return path_list
+
+
+def _get_topk_result_save_path(worker_idx: int) -> str:
+    return '{}/top{}_{}_{}.txt'.format(args.search_out_dir, args.search_topk, args.search_split, worker_idx)
+
+
+def _query_transform_func(tokenizer: PreTrainedTokenizerFast,
+                          examples: Dict[str, List]) -> BatchEncoding:
+    batch_dict = tokenizer(examples['query'],
+                           max_length=args.q_max_len,
+                           padding=PaddingStrategy.DO_NOT_PAD,
+                           truncation=True)
+
+    return batch_dict
+
+
+@torch.no_grad()
+def _worker_encode_queries(gpu_idx: int) -> Tuple:
+    # fail fast if shard does not exist
+    _get_all_shards_path()
+
+    query_id_to_text = load_queries(path=os.path.join(args.data_dir, '{}_queries.tsv'.format(args.search_split)),
+                                    task_type=args.task_type)
+    query_ids = sorted(list(query_id_to_text.keys()))
+    queries = [query_id_to_text[query_id] for query_id in query_ids]
+    dataset = Dataset.from_dict({'query_id': query_ids,
+                                 'query': queries})
+    dataset = dataset.shard(num_shards=torch.cuda.device_count(),
+                            index=gpu_idx,
+                            contiguous=True)
+
+    # only keep data for current shard
+    query_ids = dataset['query_id']
+    query_id_to_text = {qid: query_id_to_text[qid] for qid in query_ids}
+
+    logger.info('GPU {} needs to process {} examples'.format(gpu_idx, len(dataset)))
+    torch.cuda.set_device(gpu_idx)
+
+    tokenizer: PreTrainedTokenizerFast = AutoTokenizer.from_pretrained(args.model_name_or_path)
+    model: BiencoderModelForInference = BiencoderModelForInference.build(args)
+    model.eval()
+    model.cuda()
+
+    dataset.set_transform(partial(_query_transform_func, tokenizer))
+
+    data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8)
+    data_loader = DataLoader(
+        dataset,
+        batch_size=512,
+        shuffle=False,
+        drop_last=False,
+        num_workers=args.dataloader_num_workers,
+        collate_fn=data_collator,
+        pin_memory=True)
+
+    encoded_embeds = []
+    for batch_dict in tqdm.tqdm(data_loader, desc='query encoding', mininterval=5):
+        batch_dict = move_to_cuda(batch_dict)
+
+        with torch.cuda.amp.autocast() if args.fp16 else nullcontext():
+            outputs: BiencoderOutput = model(query=batch_dict, passage=None)
+        encoded_embeds.append(outputs.q_reps)
+
+    query_embeds = torch.cat(encoded_embeds, dim=0)
+    logger.info('Done query encoding for worker {}'.format(gpu_idx))
+
+    return query_embeds, query_ids, query_id_to_text
+
+
+@torch.no_grad()
+def _worker_batch_search(gpu_idx: int):
+    embeds_path_list = _get_all_shards_path()
+
+    query_embeds, query_ids, query_id_to_text = _worker_encode_queries(gpu_idx)
+    assert query_embeds.shape[0] == len(query_ids), '{} != {}'.format(query_embeds.shape[0], len(query_ids))
+
+    query_id_to_topk = defaultdict(list)
+    psg_idx_offset = 0
+    for shard_idx, shard_path in enumerate(embeds_path_list):
+        shard_psg_embed = torch.load(shard_path, map_location=lambda storage, loc: storage).to(query_embeds.device)
+        logger.info('Load {} passage embeddings from {}'.format(shard_psg_embed.shape[0], shard_path))
+
+        for start in tqdm.tqdm(range(0, len(query_ids), args.search_batch_size),
+                               desc="search shard {}".format(shard_idx),
+                               mininterval=5):
+            batch_query_embed = query_embeds[start:(start + args.search_batch_size)]
+            batch_query_ids = query_ids[start:(start + args.search_batch_size)]
+            batch_score = torch.mm(batch_query_embed, shard_psg_embed.t())
+            batch_sorted_score, batch_sorted_indices = torch.topk(batch_score, k=args.search_topk, dim=-1, largest=True)
+            for batch_idx, query_id in enumerate(batch_query_ids):
+                cur_scores = batch_sorted_score[batch_idx].cpu().tolist()
+                cur_indices = [idx + psg_idx_offset for idx in batch_sorted_indices[batch_idx].cpu().tolist()]
+                query_id_to_topk[query_id] += list(zip(cur_scores, cur_indices))
+                query_id_to_topk[query_id] = sorted(query_id_to_topk[query_id], key=lambda t: (-t[0], t[1]))
+                query_id_to_topk[query_id] = query_id_to_topk[query_id][:args.search_topk]
+
+        psg_idx_offset += shard_psg_embed.shape[0]
+
+    out_path = _get_topk_result_save_path(worker_idx=gpu_idx)
+    with open(out_path, 'w', encoding='utf-8') as writer:
+        for query_id in query_id_to_text:
+            for rank, (score, doc_id) in enumerate(query_id_to_topk[query_id]):
+                writer.write('{}\t{}\t{}\t{}\n'.format(query_id, doc_id, rank + 1, round(score, 4)))
+
+    logger.info('Write scores to {} done'.format(out_path))
+
+
+def _compute_and_save_metrics(worker_cnt: int):
+    preds: Dict[str, List[ScoredDoc]] = {}
+    for worker_idx in range(worker_cnt):
+        path = _get_topk_result_save_path(worker_idx)
+        preds.update(load_msmarco_predictions(path))
+    out_path = os.path.join(args.search_out_dir, '{}.msmarco.txt'.format(args.search_split))
+    save_preds_to_msmarco_format(preds, out_path)
+    logger.info('Merge done: save {} predictions to {}'.format(len(preds), out_path))
+
+    path_qrels = os.path.join(args.data_dir, '{}_qrels.txt'.format(args.search_split))
+    if os.path.exists(path_qrels):
+        qrels = load_qrels(path=path_qrels)
+        all_metrics = trec_eval(qrels=qrels, predictions=preds)
+        all_metrics['mrr'] = compute_mrr(qrels=qrels, predictions=preds)
+
+        logger.info('{} trec metrics = {}'.format(args.search_split, json.dumps(all_metrics, ensure_ascii=False, indent=4)))
+        save_json_to_file(all_metrics, os.path.join(args.search_out_dir, 'metrics_{}.json'.format(args.search_split)))
+    else:
+        logger.warning('No qrels found for {}'.format(args.search_split))
+
+    # do some cleanup
+    for worker_idx in range(worker_cnt):
+        path = _get_topk_result_save_path(worker_idx)
+        os.remove(path)
+
+
+def _batch_search_queries():
+    logger.info('Args={}'.format(str(args)))
+    gpu_count = torch.cuda.device_count()
+    if gpu_count == 0:
+        logger.error('No gpu available')
+        return
+
+    logger.info('Use {} gpus'.format(gpu_count))
+    torch.multiprocessing.spawn(_worker_batch_search, args=(), nprocs=gpu_count)
+    logger.info('Done batch search queries')
+
+    _compute_and_save_metrics(gpu_count)
+
+
+if __name__ == '__main__':
+    _batch_search_queries()

simlm/src/loaders/__init__.py

@@ -0,0 +1,3 @@
+from .biencoder_dataloader import RetrievalDataLoader
+from .cross_encoder_dataloader import CrossEncoderDataLoader
+from .rlm_dataloader import ReplaceLMDataloader

simlm/src/loaders/biencoder_dataloader.py

@@ -0,0 +1,112 @@
+import os
+import random
+
+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 RetrievalDataLoader:
+
+    def __init__(self, args: Arguments, tokenizer: PreTrainedTokenizerFast):
+        self.args = args
+        self.negative_size = args.train_n_passages - 1
+        assert self.negative_size > 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.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: List[int] = group_doc_ids(
+            examples=examples,
+            negative_size=self.negative_size,
+            offset=current_epoch + self.args.seed,
+            use_first_positive=self.args.use_first_positive
+        )
+        assert len(input_doc_ids) == len(examples['query']) * self.args.train_n_passages
+
+        input_docs: List[str] = [self.corpus[doc_id]['contents'] for doc_id in input_doc_ids]
+        input_titles: List[str] = [self.corpus[doc_id]['title'] for doc_id in input_doc_ids]
+
+        query_batch_dict = self.tokenizer(examples['query'],
+                                          max_length=self.args.q_max_len,
+                                          padding=PaddingStrategy.DO_NOT_PAD,
+                                          truncation=True)
+        doc_batch_dict = self.tokenizer(input_titles,
+                                        text_pair=input_docs,
+                                        max_length=self.args.p_max_len,
+                                        padding=PaddingStrategy.DO_NOT_PAD,
+                                        truncation=True)
+
+        merged_dict = {'q_{}'.format(k): v for k, v in query_batch_dict.items()}
+        step_size = self.args.train_n_passages
+        for k, v in doc_batch_dict.items():
+            k = 'd_{}'.format(k)
+            merged_dict[k] = []
+            for idx in range(0, len(v), step_size):
+                merged_dict[k].append(v[idx:(idx + step_size)])
+
+        if self.args.do_kd_biencoder:
+            qid_to_doc_id_to_score = {}
+
+            def _update_qid_pid_score(q_id: str, ex: Dict):
+                assert len(ex['doc_id']) == len(ex['score'])
+                if q_id not in qid_to_doc_id_to_score:
+                    qid_to_doc_id_to_score[q_id] = {}
+                for doc_id, score in zip(ex['doc_id'], ex['score']):
+                    qid_to_doc_id_to_score[q_id][int(doc_id)] = score
+
+            for idx, query_id in enumerate(examples['query_id']):
+                _update_qid_pid_score(query_id, examples['positives'][idx])
+                _update_qid_pid_score(query_id, examples['negatives'][idx])
+
+            merged_dict['kd_labels'] = []
+            for idx in range(0, len(input_doc_ids), step_size):
+                qid = examples['query_id'][idx // step_size]
+                cur_kd_labels = [qid_to_doc_id_to_score[qid][doc_id] for doc_id in input_doc_ids[idx:idx + step_size]]
+                merged_dict['kd_labels'].append(cur_kd_labels)
+            assert len(merged_dict['kd_labels']) == len(examples['query_id']), \
+                '{} != {}'.format(len(merged_dict['kd_labels']), len(examples['query_id']))
+
+        # Custom formatting function must return a dict
+        return merged_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:
+                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)), 3):
+                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:
+                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

simlm/src/loaders/cross_encoder_dataloader.py

@@ -0,0 +1,91 @@
+import os.path
+import random
+
+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: PreTrainedTokenizerFast):
+        self.args = args
+        self.negative_size = args.train_n_passages - 1
+        assert self.negative_size > 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.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
+        )
+        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(prefix + self.corpus[doc_id]['contents'])
+            input_queries.append(examples['query'][idx // self.args.train_n_passages])
+
+        batch_dict = self.tokenizer(input_queries,
+                                    text_pair=input_docs,
+                                    max_length=self.args.rerank_max_length,
+                                    padding=PaddingStrategy.DO_NOT_PAD,
+                                    truncation=True)
+
+        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:
+                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)), 3):
+                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:
+                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

simlm/src/loaders/loader_utils.py

@@ -0,0 +1,43 @@
+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) -> 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 p_idx == 0 or ex_pos['score'][p_idx] >= ex_pos['score'][0]
+                               or ex_pos['score'][p_idx] > max(examples['negatives'][idx]['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:
+        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

simlm/src/loaders/rlm_dataloader.py

@@ -0,0 +1,39 @@
+import random
+
+from typing import Tuple
+from transformers import PreTrainedTokenizerFast
+from datasets import Dataset, load_dataset
+
+from config import Arguments
+from logger_config import logger
+
+
+def split_dataset(dataset: Dataset,
+                  num_eval_examples: int,
+                  max_train_samples: int = None) -> Tuple[Dataset, Dataset]:
+    indices = list(range(len(dataset)))
+    random.Random(123).shuffle(indices)
+    eval_dataset = dataset.select(indices[:num_eval_examples])
+    train_dataset = dataset.select(indices[num_eval_examples:])
+
+    if max_train_samples is not None:
+        train_dataset = train_dataset.select(range(max_train_samples))
+
+    # Log a few random samples from the training set:
+    for index in random.sample(range(len(train_dataset)), 3):
+        logger.info(f"Sample {index} of the training set: {train_dataset[index]}.")
+
+    return train_dataset, eval_dataset
+
+
+class ReplaceLMDataloader:
+
+    def __init__(self, args: Arguments, tokenizer: PreTrainedTokenizerFast):
+        self.args = args
+        self.tokenizer = tokenizer
+        data_files = args.train_file.strip().split(',')
+        self.corpus: Dataset = load_dataset('json', data_files=data_files)['train']
+        self.train_dataset, self.eval_dataset = split_dataset(
+            self.corpus,
+            num_eval_examples=args.rlm_num_eval_samples,
+            max_train_samples=args.max_train_samples)

simlm/src/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)

simlm/src/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)

simlm/src/models/__init__.py

@@ -0,0 +1,3 @@
+from .biencoder_model import BiencoderModel, BiencoderModelForInference, BiencoderOutput
+from .cross_encoder_model import Reranker, RerankerForInference
+from .rlm import ReplaceLM, ReplaceLMOutput

simlm/src/models/biencoder_model.py

@@ -0,0 +1,219 @@
+import os
+import copy
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from dataclasses import dataclass
+from typing import Optional, Dict, Tuple
+from torch import Tensor
+from transformers import (
+    AutoModel,
+    PreTrainedModel,
+)
+from transformers.modeling_outputs import ModelOutput
+
+from config import Arguments
+from logger_config import logger
+from utils import dist_gather_tensor, select_grouped_indices, full_contrastive_scores_and_labels
+
+
+@dataclass
+class BiencoderOutput(ModelOutput):
+    q_reps: Optional[Tensor] = None
+    p_reps: Optional[Tensor] = None
+    loss: Optional[Tensor] = None
+    labels: Optional[Tensor] = None
+    scores: Optional[Tensor] = None
+
+
+class BiencoderModel(nn.Module):
+    def __init__(self, args: Arguments,
+                 lm_q: PreTrainedModel,
+                 lm_p: PreTrainedModel):
+        super().__init__()
+        self.lm_q = lm_q
+        self.lm_p = lm_p
+        self.cross_entropy = nn.CrossEntropyLoss(reduction='mean')
+        self.kl_loss_fn = torch.nn.KLDivLoss(reduction="batchmean", log_target=True)
+        self.args = args
+        self.pooler = nn.Linear(self.lm_q.config.hidden_size, args.out_dimension) if args.add_pooler else nn.Identity()
+
+        from trainers import BiencoderTrainer
+        self.trainer: Optional[BiencoderTrainer] = None
+
+    def forward(self, query: Dict[str, Tensor] = None,
+                passage: Dict[str, Tensor] = None):
+        assert self.args.process_index >= 0
+
+        scores, labels, q_reps, p_reps, all_scores, all_labels = self._compute_scores(query, passage)
+
+        start = self.args.process_index * q_reps.shape[0]
+        group_indices = select_grouped_indices(scores=scores,
+                                               group_size=self.args.train_n_passages,
+                                               start=start * self.args.train_n_passages)
+
+        if not self.args.do_kd_biencoder:
+            # training biencoder from scratch
+            if self.args.use_scaled_loss:
+                loss = self.cross_entropy(all_scores, all_labels)
+                loss *= self.args.world_size if self.args.loss_scale <= 0 else self.args.loss_scale
+            else:
+                loss = self.cross_entropy(scores, labels)
+        else:
+            # training biencoder with kd
+            # batch_size x train_n_passage
+            group_scores = torch.gather(input=scores, dim=1, index=group_indices)
+            assert group_scores.shape[1] == self.args.train_n_passages
+            group_log_scores = torch.log_softmax(group_scores, dim=-1)
+            kd_log_target = torch.log_softmax(query['kd_labels'], dim=-1)
+
+            kd_loss = self.kl_loss_fn(input=group_log_scores, target=kd_log_target)
+
+            # (optionally) mask out hard negatives
+            if self.training and self.args.kd_mask_hn:
+                scores = torch.scatter(input=scores, dim=1, index=group_indices[:, 1:], value=float('-inf'))
+            if self.args.use_scaled_loss:
+                ce_loss = self.cross_entropy(all_scores, all_labels)
+                ce_loss *= self.args.world_size if self.args.loss_scale <= 0 else self.args.loss_scale
+            else:
+                ce_loss = self.cross_entropy(scores, labels)
+
+            loss = self.args.kd_cont_loss_weight * ce_loss + kd_loss
+
+        total_n_psg = self.args.world_size * q_reps.shape[0] * self.args.train_n_passages
+
+        return BiencoderOutput(loss=loss, q_reps=q_reps, p_reps=p_reps,
+                               labels=labels.contiguous(),
+                               scores=scores[:, :total_n_psg].contiguous())
+
+    def _compute_scores(self, query: Dict[str, Tensor] = None,
+                        passage: Dict[str, Tensor] = None) -> Tuple:
+        q_reps = self._encode(self.lm_q, query)
+        p_reps = self._encode(self.lm_p, passage)
+
+        all_q_reps = dist_gather_tensor(q_reps)
+        all_p_reps = dist_gather_tensor(p_reps)
+        assert all_p_reps.shape[0] == self.args.world_size * q_reps.shape[0] * self.args.train_n_passages
+
+        all_scores, all_labels = full_contrastive_scores_and_labels(
+            query=all_q_reps, key=all_p_reps,
+            use_all_pairs=self.args.full_contrastive_loss)
+
+        if self.args.l2_normalize:
+            if self.args.t_warmup:
+                scale = 1 / self.args.t * min(1.0, self.trainer.state.global_step / self.args.warmup_steps)
+                scale = max(1.0, scale)
+            else:
+                scale = 1 / self.args.t
+            all_scores = all_scores * scale
+
+        start = self.args.process_index * q_reps.shape[0]
+        local_query_indices = torch.arange(start, start + q_reps.shape[0], dtype=torch.long).to(q_reps.device)
+        # batch_size x (world_size x batch_size x train_n_passage)
+        scores = all_scores.index_select(dim=0, index=local_query_indices)
+        labels = all_labels.index_select(dim=0, index=local_query_indices)
+
+        return scores, labels, q_reps, p_reps, all_scores, all_labels
+
+    def _encode(self, encoder: PreTrainedModel, input_dict: dict) -> Optional[torch.Tensor]:
+        if not input_dict:
+            return None
+        outputs = encoder(**{k: v for k, v in input_dict.items() if k not in ['kd_labels']}, return_dict=True)
+        hidden_state = outputs.last_hidden_state
+        embeds = hidden_state[:, 0]
+        embeds = self.pooler(embeds)
+        if self.args.l2_normalize:
+            embeds = F.normalize(embeds, dim=-1)
+        return embeds.contiguous()
+
+    @classmethod
+    def build(cls, args: Arguments, **hf_kwargs):
+        # load local
+        if os.path.isdir(args.model_name_or_path):
+            if not args.share_encoder:
+                _qry_model_path = os.path.join(args.model_name_or_path, 'query_model')
+                _psg_model_path = os.path.join(args.model_name_or_path, 'passage_model')
+                if not os.path.exists(_qry_model_path):
+                    _qry_model_path = args.model_name_or_path
+                    _psg_model_path = args.model_name_or_path
+                logger.info(f'loading query model weight from {_qry_model_path}')
+                lm_q = AutoModel.from_pretrained(_qry_model_path, **hf_kwargs)
+                logger.info(f'loading passage model weight from {_psg_model_path}')
+                lm_p = AutoModel.from_pretrained(_psg_model_path, **hf_kwargs)
+            else:
+                logger.info(f'loading shared model weight from {args.model_name_or_path}')
+                lm_q = AutoModel.from_pretrained(args.model_name_or_path, **hf_kwargs)
+                lm_p = lm_q
+        # load pre-trained
+        else:
+            lm_q = AutoModel.from_pretrained(args.model_name_or_path, **hf_kwargs)
+            lm_p = copy.deepcopy(lm_q) if not args.share_encoder else lm_q
+
+        model = cls(args=args, lm_q=lm_q, lm_p=lm_p)
+        return model
+
+    def save(self, output_dir: str):
+        if not self.args.share_encoder:
+            os.makedirs(os.path.join(output_dir, 'query_model'), exist_ok=True)
+            os.makedirs(os.path.join(output_dir, 'passage_model'), exist_ok=True)
+            self.lm_q.save_pretrained(os.path.join(output_dir, 'query_model'))
+            self.lm_p.save_pretrained(os.path.join(output_dir, 'passage_model'))
+        else:
+            self.lm_q.save_pretrained(output_dir)
+        if self.args.add_pooler:
+            torch.save(self.pooler.state_dict(), os.path.join(output_dir, 'pooler.pt'))
+
+
+class BiencoderModelForInference(BiencoderModel):
+    def __init__(self, args: Arguments,
+                 lm_q: PreTrainedModel,
+                 lm_p: PreTrainedModel):
+        nn.Module.__init__(self)
+        self.args = args
+        self.lm_q = lm_q
+        self.lm_p = lm_p
+        self.pooler = nn.Linear(self.lm_q.config.hidden_size, args.out_dimension) if args.add_pooler else nn.Identity()
+
+    @torch.no_grad()
+    def forward(self, query: Dict[str, Tensor] = None,
+                passage: Dict[str, Tensor] = None):
+        q_reps = self._encode(self.lm_q, query)
+        p_reps = self._encode(self.lm_p, passage)
+        return BiencoderOutput(q_reps=q_reps, p_reps=p_reps)
+
+    @classmethod
+    def build(cls, args: Arguments, **hf_kwargs):
+        model_name_or_path = args.model_name_or_path
+
+        # load local
+        if os.path.isdir(model_name_or_path):
+            _qry_model_path = os.path.join(model_name_or_path, 'query_model')
+            _psg_model_path = os.path.join(model_name_or_path, 'passage_model')
+            if os.path.exists(_qry_model_path):
+                logger.info(f'found separate weight for query/passage encoders')
+                logger.info(f'loading query model weight from {_qry_model_path}')
+                lm_q = AutoModel.from_pretrained(_qry_model_path, **hf_kwargs)
+                logger.info(f'loading passage model weight from {_psg_model_path}')
+                lm_p = AutoModel.from_pretrained(_psg_model_path, **hf_kwargs)
+            else:
+                logger.info(f'try loading tied weight')
+                logger.info(f'loading model weight from {model_name_or_path}')
+                lm_q = AutoModel.from_pretrained(model_name_or_path, **hf_kwargs)
+                lm_p = lm_q
+        else:
+            logger.info(f'try loading tied weight {model_name_or_path}')
+            lm_q = AutoModel.from_pretrained(model_name_or_path, **hf_kwargs)
+            lm_p = lm_q
+
+        model = cls(args=args, lm_q=lm_q, lm_p=lm_p)
+
+        pooler_path = os.path.join(args.model_name_or_path, 'pooler.pt')
+        if os.path.exists(pooler_path):
+            logger.info('loading pooler weights from local files')
+            state_dict = torch.load(pooler_path, map_location="cpu")
+            model.pooler.load_state_dict(state_dict)
+        else:
+            assert not args.add_pooler
+            logger.info('No pooler will be loaded')
+        return model

simlm/src/models/cross_encoder_model.py

@@ -0,0 +1,96 @@
+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
+
+
+class Reranker(nn.Module):
+    def __init__(self, hf_model: PreTrainedModel, args: Arguments):
+        super().__init__()
+        self.hf_model = hf_model
+        self.args = args
+
+        self.cross_entropy = nn.CrossEntropyLoss(reduction='mean')
+        self.kl_loss_fn = torch.nn.KLDivLoss(reduction="batchmean", log_target=True)
+
+    def forward(self, batch: Dict[str, torch.Tensor]) -> SequenceClassifierOutput:
+        input_batch_dict = {k: v for k, v in batch.items() if k != 'labels'}
+
+        if self.args.rerank_forward_factor > 1:
+            assert torch.sum(batch['labels']).long().item() == 0
+            assert all(len(v.shape) == 2 for v in input_batch_dict.values())
+
+            is_train = self.hf_model.training
+            self.hf_model.eval()
+
+            with torch.no_grad():
+                outputs: SequenceClassifierOutput = self.hf_model(**input_batch_dict, return_dict=True)
+                outputs.logits = outputs.logits.view(-1, self.args.train_n_passages)
+                # make sure the target passage is not masked out
+                outputs.logits[:, 0].fill_(float('inf'))
+
+                k = self.args.train_n_passages // self.args.rerank_forward_factor
+                _, topk_indices = torch.topk(outputs.logits, k=k, dim=-1, largest=True)
+                topk_indices += self.args.train_n_passages * torch.arange(0, topk_indices.shape[0],
+                                                                          dtype=torch.long,
+                                                                          device=topk_indices.device).unsqueeze(-1)
+                topk_indices = topk_indices.view(-1)
+
+                input_batch_dict = {k: v.index_select(dim=0, index=topk_indices) for k, v in input_batch_dict.items()}
+
+            self.hf_model.train(is_train)
+
+        n_psg_per_query = self.args.train_n_passages // self.args.rerank_forward_factor
+
+        if self.args.rerank_use_rdrop and self.training:
+            input_batch_dict = {k: torch.cat([v, v], dim=0) for k, v in input_batch_dict.items()}
+
+        outputs: SequenceClassifierOutput = self.hf_model(**input_batch_dict, return_dict=True)
+
+        if self.args.rerank_use_rdrop and self.training:
+            logits = outputs.logits.view(2, -1, n_psg_per_query)
+            outputs.logits = logits[0, :, :].contiguous()
+            log_prob = torch.log_softmax(logits, dim=2)
+            log_prob1, log_prob2 = log_prob[0, :, :], log_prob[1, :, :]
+            rdrop_loss = 0.5 * (self.kl_loss_fn(log_prob1, log_prob2) + self.kl_loss_fn(log_prob2, log_prob1))
+            ce_loss = 0.5 * (self.cross_entropy(log_prob1, batch['labels'])
+                             + self.cross_entropy(log_prob2, batch['labels']))
+
+            outputs.loss = rdrop_loss + ce_loss
+        else:
+            outputs.logits = outputs.logits.view(-1, n_psg_per_query)
+            loss = self.cross_entropy(outputs.logits, batch['labels'])
+            outputs.loss = loss
+
+        return outputs
+
+    @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, hf_model: Optional[PreTrainedModel] = None):
+        super().__init__()
+        self.hf_model = hf_model
+        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)

simlm/src/models/rlm.py

@@ -0,0 +1,136 @@
+import copy
+import os
+import torch
+import torch.nn as nn
+
+from contextlib import nullcontext
+from torch import Tensor
+from torch.distributions import Categorical
+from typing import Dict, Optional, Tuple
+from dataclasses import dataclass
+from transformers import AutoModelForMaskedLM, ElectraModel
+from transformers.modeling_outputs import MaskedLMOutput, ModelOutput
+from transformers.models.bert import BertForMaskedLM
+
+from logger_config import logger
+from config import Arguments
+from utils import slice_batch_dict
+
+
+@dataclass
+class ReplaceLMOutput(ModelOutput):
+    loss: Optional[Tensor] = None
+    encoder_mlm_loss: Optional[Tensor] = None
+    decoder_mlm_loss: Optional[Tensor] = None
+    g_mlm_loss: Optional[Tensor] = None
+    replace_ratio: Optional[Tensor] = None
+
+
+class ReplaceLM(nn.Module):
+    def __init__(self, args: Arguments,
+                 bert: BertForMaskedLM):
+        super(ReplaceLM, self).__init__()
+        self.encoder = bert
+        self.decoder = copy.deepcopy(self.encoder.bert.encoder.layer[-args.rlm_decoder_layers:])
+        self.cross_entropy = nn.CrossEntropyLoss(reduction='mean')
+
+        self.generator: ElectraModel = AutoModelForMaskedLM.from_pretrained(args.rlm_generator_model_name)
+
+        if args.rlm_freeze_generator:
+            self.generator.eval()
+            self.generator.requires_grad_(False)
+
+        self.args = args
+
+        from trainers.rlm_trainer import ReplaceLMTrainer
+        self.trainer: Optional[ReplaceLMTrainer] = None
+
+    def forward(self, model_input: Dict[str, torch.Tensor]) -> ReplaceLMOutput:
+        enc_prefix, dec_prefix = 'enc_', 'dec_'
+        encoder_inputs = slice_batch_dict(model_input, enc_prefix)
+        decoder_inputs = slice_batch_dict(model_input, dec_prefix)
+        labels = model_input['labels']
+
+        enc_sampled_input_ids, g_mlm_loss = self._replace_tokens(encoder_inputs)
+        if self.args.rlm_freeze_generator:
+            g_mlm_loss = torch.tensor(0, dtype=torch.float, device=g_mlm_loss.device)
+        dec_sampled_input_ids, _ = self._replace_tokens(decoder_inputs, no_grad=True)
+
+        encoder_inputs['input_ids'] = enc_sampled_input_ids
+        decoder_inputs['input_ids'] = dec_sampled_input_ids
+        # use the un-masked version of labels
+        encoder_inputs['labels'] = labels
+        decoder_inputs['labels'] = labels
+
+        is_replaced = (encoder_inputs['input_ids'] != labels) & (labels >= 0)
+        replace_cnt = is_replaced.long().sum().item()
+        total_cnt = (encoder_inputs['attention_mask'] == 1).long().sum().item()
+        replace_ratio = torch.tensor(replace_cnt / total_cnt, device=g_mlm_loss.device)
+
+        encoder_out: MaskedLMOutput = self.encoder(
+            **encoder_inputs,
+            output_hidden_states=True,
+            return_dict=True)
+
+        # batch_size x 1 x hidden_dim
+        cls_hidden = encoder_out.hidden_states[-1][:, :1]
+        # batch_size x seq_length x embed_dim
+        dec_inputs_embeds = self.encoder.bert.embeddings(decoder_inputs['input_ids'])
+        hiddens = torch.cat([cls_hidden, dec_inputs_embeds[:, 1:]], dim=1)
+
+        attention_mask = self.encoder.get_extended_attention_mask(
+            encoder_inputs['attention_mask'],
+            encoder_inputs['attention_mask'].shape,
+            encoder_inputs['attention_mask'].device
+        )
+
+        for layer in self.decoder:
+            layer_out = layer(hiddens, attention_mask)
+            hiddens = layer_out[0]
+
+        decoder_mlm_loss = self.mlm_loss(hiddens, labels)
+
+        loss = decoder_mlm_loss + encoder_out.loss + g_mlm_loss * self.args.rlm_generator_mlm_weight
+
+        return ReplaceLMOutput(loss=loss,
+                               encoder_mlm_loss=encoder_out.loss.detach(),
+                               decoder_mlm_loss=decoder_mlm_loss.detach(),
+                               g_mlm_loss=g_mlm_loss.detach(),
+                               replace_ratio=replace_ratio)
+
+    def _replace_tokens(self, batch_dict: Dict[str, torch.Tensor],
+                        no_grad: bool = False) -> Tuple:
+        with torch.no_grad() if self.args.rlm_freeze_generator or no_grad else nullcontext():
+            outputs: MaskedLMOutput = self.generator(
+                **batch_dict,
+                return_dict=True)
+
+        with torch.no_grad():
+            sampled_input_ids = Categorical(logits=outputs.logits).sample()
+            is_mask = (batch_dict['labels'] >= 0).long()
+            sampled_input_ids = batch_dict['input_ids'] * (1 - is_mask) + sampled_input_ids * is_mask
+
+        return sampled_input_ids.long(), outputs.loss
+
+    def mlm_loss(self, hiddens: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+        pred_scores = self.encoder.cls(hiddens)
+        mlm_loss = self.cross_entropy(
+            pred_scores.view(-1, self.encoder.config.vocab_size),
+            labels.view(-1))
+        return mlm_loss
+
+    @classmethod
+    def from_pretrained(cls, all_args: Arguments,
+                        model_name_or_path: str, *args, **kwargs):
+        hf_model = AutoModelForMaskedLM.from_pretrained(model_name_or_path, *args, **kwargs)
+        model = cls(all_args, hf_model)
+        decoder_save_path = os.path.join(model_name_or_path, 'decoder.pt')
+        if os.path.exists(decoder_save_path):
+            logger.info('loading extra weights from local files')
+            state_dict = torch.load(decoder_save_path, map_location="cpu")
+            model.decoder.load_state_dict(state_dict)
+        return model
+
+    def save_pretrained(self, output_dir: str):
+        self.encoder.save_pretrained(output_dir)
+        torch.save(self.decoder.state_dict(), os.path.join(output_dir, 'decoder.pt'))

simlm/src/train_biencoder.py

@@ -0,0 +1,101 @@
+import logging
+
+import torch
+from typing import Dict
+from functools import partial
+from transformers.utils.logging import enable_explicit_format
+from transformers.trainer_callback import PrinterCallback
+from transformers import (
+    AutoTokenizer,
+    HfArgumentParser,
+    EvalPrediction,
+    Trainer,
+    set_seed,
+    PreTrainedTokenizerFast
+)
+
+from logger_config import logger, LoggerCallback
+from config import Arguments
+from trainers import BiencoderTrainer
+from loaders import RetrievalDataLoader
+from collators import BiencoderCollator
+from metrics import accuracy, batch_mrr
+from models import BiencoderModel
+
+
+def _common_setup(args: Arguments):
+    if args.process_index > 0:
+        logger.setLevel(logging.WARNING)
+    enable_explicit_format()
+    set_seed(args.seed)
+
+
+def _compute_metrics(args: Arguments, eval_pred: EvalPrediction) -> Dict[str, float]:
+    # field consistent with BiencoderOutput
+    preds = eval_pred.predictions
+    scores = torch.tensor(preds[-1]).float()
+    labels = torch.arange(0, scores.shape[0], dtype=torch.long) * args.train_n_passages
+    labels = labels % scores.shape[1]
+
+    topk_metrics = accuracy(output=scores, target=labels, topk=(1, 3))
+    mrr = batch_mrr(output=scores, target=labels)
+
+    return {'mrr': mrr, 'acc1': topk_metrics[0], 'acc3': topk_metrics[1]}
+
+
+def main():
+    parser = HfArgumentParser((Arguments,))
+    args: Arguments = parser.parse_args_into_dataclasses()[0]
+    _common_setup(args)
+    logger.info('Args={}'.format(str(args)))
+
+    tokenizer: PreTrainedTokenizerFast = AutoTokenizer.from_pretrained(args.model_name_or_path)
+    model: BiencoderModel = BiencoderModel.build(args=args)
+    logger.info(model)
+    logger.info('Vocab size: {}'.format(len(tokenizer)))
+
+    data_collator = BiencoderCollator(
+        tokenizer=tokenizer,
+        pad_to_multiple_of=8 if args.fp16 else None)
+
+    retrieval_data_loader = RetrievalDataLoader(args=args, tokenizer=tokenizer)
+    train_dataset = retrieval_data_loader.train_dataset
+    eval_dataset = retrieval_data_loader.eval_dataset
+
+    trainer: Trainer = BiencoderTrainer(
+        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=partial(_compute_metrics, args),
+        tokenizer=tokenizer,
+    )
+    trainer.remove_callback(PrinterCallback)
+    trainer.add_callback(LoggerCallback)
+    retrieval_data_loader.trainer = trainer
+    model.trainer = trainer
+
+    if args.do_train:
+        train_result = trainer.train()
+        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()

simlm/src/train_cross_encoder.py

@@ -0,0 +1,102 @@
+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,
+    Trainer,
+    set_seed,
+    PreTrainedTokenizerFast
+)
+
+from logger_config import logger, LoggerCallback
+from config import Arguments
+from trainers.reranker_trainer import RerankerTrainer
+from loaders import CrossEncoderDataLoader
+from collators import CrossEncoderCollator
+from metrics import accuracy
+from models 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)))
+
+    tokenizer: PreTrainedTokenizerFast = 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=8 if args.fp16 else None)
+
+    rerank_data_loader = CrossEncoderDataLoader(args=args, tokenizer=tokenizer)
+    train_dataset = rerank_data_loader.train_dataset
+    eval_dataset = rerank_data_loader.eval_dataset
+
+    trainer: 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_train:
+        train_result = trainer.train()
+        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()

simlm/src/train_rlm.py

@@ -0,0 +1,101 @@
+import logging
+import numpy as np
+
+from typing import Dict
+from transformers.utils.logging import enable_explicit_format
+from transformers.trainer_callback import PrinterCallback
+from transformers import (
+    AutoTokenizer,
+    HfArgumentParser,
+    set_seed,
+    PreTrainedTokenizerFast,
+    EvalPrediction,
+)
+
+from logger_config import logger, LoggerCallback
+from config import Arguments
+from loaders import ReplaceLMDataloader
+from collators import DataCollatorForReplaceLM
+from trainers import ReplaceLMTrainer
+from models import ReplaceLM
+
+
+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[str, float]:
+    preds = eval_pred.predictions
+
+    avg_enc_mlm_loss = float(np.mean(preds[0]))
+    avg_dec_mlm_loss = float(np.mean(preds[1]))
+    avg_g_mlm_loss = float(np.mean(preds[2]))
+    avg_replace_ratio = float(np.mean(preds[3]))
+
+    return {'avg_enc_mlm_loss': round(avg_enc_mlm_loss, 4),
+            'avg_dec_mlm_loss': round(avg_dec_mlm_loss, 4),
+            'avg_g_mlm_loss': round(avg_g_mlm_loss, 4),
+            'avg_replace_ratio': round(avg_replace_ratio, 4)}
+
+
+def main():
+    parser = HfArgumentParser((Arguments,))
+    args: Arguments = parser.parse_args_into_dataclasses()[0]
+    _common_setup(args)
+    logger.info('Args={}'.format(str(args)))
+
+    tokenizer: PreTrainedTokenizerFast = AutoTokenizer.from_pretrained(args.model_name_or_path)
+    model: ReplaceLM = ReplaceLM.from_pretrained(
+        all_args=args, model_name_or_path=args.model_name_or_path)
+    logger.info(model)
+    logger.info('Vocab size: {}'.format(len(tokenizer)))
+
+    dataloader = ReplaceLMDataloader(args=args, tokenizer=tokenizer)
+    train_dataset, eval_dataset = dataloader.train_dataset, dataloader.eval_dataset
+
+    data_collator = DataCollatorForReplaceLM(
+        tokenizer,
+        pad_to_multiple_of=8 if args.fp16 else None,
+        args=args,
+    )
+
+    trainer: ReplaceLMTrainer = ReplaceLMTrainer(
+        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)
+
+    model.trainer = trainer
+
+    if args.do_train:
+        train_result = trainer.train()
+        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()
+        metrics["eval_samples"] = len(eval_dataset)
+
+        trainer.log_metrics("eval", metrics)
+        trainer.save_metrics("eval", metrics)
+
+    return
+
+
+if __name__ == "__main__":
+    main()

simlm/src/trainers/__init__.py

@@ -0,0 +1,3 @@
+from .biencoder_trainer import BiencoderTrainer
+from .reranker_trainer import RerankerTrainer
+from .rlm_trainer import ReplaceLMTrainer

simlm/src/trainers/biencoder_trainer.py

@@ -0,0 +1,74 @@
+import os
+import torch
+
+from typing import Optional, Dict, Tuple
+from transformers.trainer import Trainer
+
+from logger_config import logger
+from metrics import accuracy, batch_mrr
+from models import BiencoderOutput, BiencoderModel
+from utils import AverageMeter
+
+
+def _unpack_qp(inputs: Dict[str, torch.Tensor]) -> Tuple:
+    q_prefix, d_prefix, kd_labels_key = 'q_', 'd_', 'kd_labels'
+    query_batch_dict = {k[len(q_prefix):]: v for k, v in inputs.items() if k.startswith(q_prefix)}
+    doc_batch_dict = {k[len(d_prefix):]: v for k, v in inputs.items() if k.startswith(d_prefix)}
+
+    if kd_labels_key in inputs:
+        assert len(query_batch_dict) > 0
+        query_batch_dict[kd_labels_key] = inputs[kd_labels_key]
+
+    if not query_batch_dict:
+        query_batch_dict = None
+    if not doc_batch_dict:
+        doc_batch_dict = None
+
+    return query_batch_dict, doc_batch_dict
+
+
+class BiencoderTrainer(Trainer):
+    def __init__(self, *pargs, **kwargs):
+        super(BiencoderTrainer, self).__init__(*pargs, **kwargs)
+        self.model: BiencoderModel
+
+        self.acc1_meter = AverageMeter('Acc@1', round_digits=2)
+        self.acc3_meter = AverageMeter('Acc@3', round_digits=2)
+        self.mrr_meter = AverageMeter('mrr', 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(output_dir)
+        if self.tokenizer is not None:
+            self.tokenizer.save_pretrained(output_dir)
+
+    def compute_loss(self, model, inputs, return_outputs=False):
+        query, passage = _unpack_qp(inputs)
+        outputs: BiencoderOutput = model(query=query, passage=passage)
+        loss = outputs.loss
+
+        if self.model.training:
+            step_acc1, step_acc3 = accuracy(output=outputs.scores.detach(), target=outputs.labels, topk=(1, 3))
+            step_mrr = batch_mrr(output=outputs.scores.detach(), target=outputs.labels)
+
+            self.acc1_meter.update(step_acc1)
+            self.acc3_meter.update(step_acc3)
+            self.mrr_meter.update(step_mrr)
+
+            if self.state.global_step > 0 and self.state.global_step % self.args.logging_steps == 0:
+                log_info = ', '.join(map(str, [self.mrr_meter, self.acc1_meter, self.acc3_meter]))
+                logger.info('step: {}, {}'.format(self.state.global_step, log_info))
+
+            self._reset_meters_if_needed()
+
+        return (loss, outputs) if return_outputs else loss
+
+    def _reset_meters_if_needed(self):
+        if int(self.state.epoch) != self.last_epoch:
+            self.last_epoch = int(self.state.epoch)
+            self.acc1_meter.reset()
+            self.acc3_meter.reset()
+            self.mrr_meter.reset()

simlm/src/trainers/reranker_trainer.py

@@ -0,0 +1,48 @@
+import os
+
+from typing import Optional, Union
+from transformers.trainer import Trainer
+from transformers.modeling_outputs import SequenceClassifierOutput
+
+from logger_config import logger
+from metrics import accuracy
+from utils import AverageMeter
+
+
+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):
+        outputs: SequenceClassifierOutput = model(inputs)
+        loss = outputs.loss
+
+        if self.model.training:
+            labels = inputs['labels']
+            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 _reset_meters_if_needed(self):
+        if int(self.state.epoch) != self.last_epoch:
+            self.last_epoch = int(self.state.epoch)
+            self.acc_meter.reset()

simlm/src/trainers/rlm_trainer.py

@@ -0,0 +1,53 @@
+import os
+
+from typing import Optional
+from transformers.trainer import Trainer
+
+from logger_config import logger
+from models import ReplaceLM, ReplaceLMOutput
+from utils import AverageMeter
+
+
+class ReplaceLMTrainer(Trainer):
+    def __init__(self, *pargs, **kwargs):
+        super(ReplaceLMTrainer, self).__init__(*pargs, **kwargs)
+        self.model: ReplaceLM
+
+        self.enc_mlm_loss = AverageMeter('enc_mlm_loss', round_digits=3)
+        self.dec_mlm_loss = AverageMeter('dec_mlm_loss', round_digits=3)
+        self.g_mlm_loss = AverageMeter('g_mlm_loss', round_digits=3)
+        self.replace_ratio = AverageMeter('replace_ratio', round_digits=3)
+        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:
+            self.tokenizer.save_pretrained(output_dir)
+
+    def compute_loss(self, model, inputs, return_outputs=False):
+        outputs: ReplaceLMOutput = model(model_input=inputs)
+        loss = outputs.loss
+
+        if self.model.training:
+            self.enc_mlm_loss.update(outputs.encoder_mlm_loss.item())
+            self.dec_mlm_loss.update(outputs.decoder_mlm_loss.item())
+            self.g_mlm_loss.update(outputs.g_mlm_loss.item())
+            self.replace_ratio.update(outputs.replace_ratio.item())
+            if self.state.global_step > 0 and self.state.global_step % self.args.logging_steps == 0:
+                log_info = ', '.join(map(str, [self.enc_mlm_loss, self.dec_mlm_loss, self.g_mlm_loss, self.replace_ratio]))
+                logger.info('step: {}, {}'.format(self.state.global_step, log_info))
+
+            self._reset_meters_if_needed()
+
+        return (loss, outputs) if return_outputs else loss
+
+    def _reset_meters_if_needed(self):
+        if int(self.state.epoch) != self.last_epoch:
+            self.last_epoch = int(self.state.epoch)
+            self.enc_mlm_loss.reset()
+            self.dec_mlm_loss.reset()
+            self.g_mlm_loss.reset()
+            self.replace_ratio.reset()

simlm/src/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)