upload

README.md

@@ -0,0 +1,14 @@
+#cross_encoder-msmarco-word2vec256k
+
+This CrossEncoder was trained with MarginMSE loss from the [nicoladecao/msmarco-word2vec256000-distilbert-base-uncased](https://hf.co/nicoladecao/msmarco-word2vec256000-distilbert-base-uncased) checkpoint. **Word embedding matrix has been frozen during training**.
+
+You can load the model with [sentence-transformers](https://sbert.net):
+```python
+from sentence_transformers import CrossEncoder
+from torch import nn
+model = CrossEncoder(model_name, default_activation_function=nn.Identity())
+```
+
+Performance on TREC Deep Learning (nDCG@10):
+- TREC-DL 19:  72.49
+- TREC-DL 20:  72.71

config.json

@@ -0,0 +1,30 @@
+{
+  "_name_or_path": "nicoladecao/msmarco-word2vec256000-distilbert-base-uncased",
+  "activation": "gelu",
+  "architectures": [
+    "DistilBertForSequenceClassification"
+  ],
+  "attention_dropout": 0.1,
+  "dim": 768,
+  "dropout": 0.1,
+  "hidden_dim": 3072,
+  "id2label": {
+    "0": "LABEL_0"
+  },
+  "initializer_range": 0.02,
+  "label2id": {
+    "LABEL_0": 0
+  },
+  "max_position_embeddings": 512,
+  "model_type": "distilbert",
+  "n_heads": 12,
+  "n_layers": 6,
+  "pad_token_id": 0,
+  "qa_dropout": 0.1,
+  "seq_classif_dropout": 0.2,
+  "sinusoidal_pos_embds": false,
+  "tie_weights_": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.11.3",
+  "vocab_size": 256000
+}

pytorch_model.bin

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

special_tokens_map.json

@@ -0,0 +1 @@
+{"unk_token": "[UNK]", "sep_token": "[SEP]", "pad_token": "[PAD]", "cls_token": "[CLS]", "mask_token": "[MASK]"}

tokenizer_config.json

@@ -0,0 +1 @@
+{"model_max_length": 512, "unk_token": "[UNK]", "cls_token": "[CLS]", "sep_token": "[SEP]", "pad_token": "[PAD]", "mask_token": "[MASK]", "model_input_names": ["input_ids", "attention_mask"], "special_tokens_map_file": "/root/.cache/huggingface/transformers/fe09c361189d8238b9e387f10a088e93f70620bfe74b82036baff1fed512a153.dd8bd9bfd3664b530ea4e645105f557769387b3da9f79bdb55ed556bdd80611d", "name_or_path": "nicoladecao/msmarco-word2vec256000-distilbert-base-uncased", "tokenizer_class": "PreTrainedTokenizerFast"}

train_script.py

@@ -0,0 +1,233 @@
+import gzip
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoConfig, AdamW
+import sys
+import torch
+import transformers
+from torch.utils.data import Dataset, DataLoader
+from torch.cuda.amp import autocast
+import tqdm
+from datetime import datetime
+from shutil import copyfile
+import os
+####################################
+
+import gzip
+from collections import defaultdict
+import logging
+import tqdm
+import numpy as np
+import sys
+import pytrec_eval
+from sentence_transformers import SentenceTransformer, util, CrossEncoder
+import torch
+
+
+######### Evaluation
+queries_filepath = '/home/msmarco/data/trec2019/msmarco-test2019-queries.tsv.gz'
+queries_eval = {}
+with gzip.open(queries_filepath, 'rt', encoding='utf8') as fIn:
+    for line in fIn:
+        qid, query = line.strip().split("\t")[0:2]
+        queries_eval[qid] = query
+
+rel = defaultdict(lambda: defaultdict(int))
+
+with open('/home/msmarco/data/trec2019/2019qrels-pass.txt') as fIn:
+    for line in fIn:
+        qid, _, pid, score = line.strip().split()
+        score = int(score)
+        if score > 0:
+            rel[qid][pid] = score
+
+relevant_qid = []
+for qid in queries_eval:
+    if len(rel[qid]) > 0:
+        relevant_qid.append(qid)
+
+# Read top 1k
+passage_cand = {}
+
+with gzip.open('/home/msmarco/data/trec2019/msmarco-passagetest2019-top1000.tsv.gz', 'rt', encoding='utf8') as fIn:
+    for line in fIn:
+        qid, pid, query, passage = line.strip().split("\t")
+        if qid not in passage_cand:
+            passage_cand[qid] = []
+
+        passage_cand[qid].append([pid, passage])
+
+
+
+def eval_modal(model_path):
+    run = {}
+    model = CrossEncoder(model_path, max_length=512)
+
+    for qid in relevant_qid:
+        query = queries_eval[qid]
+
+        cand = passage_cand[qid]
+        pids = [c[0] for c in cand]
+        corpus_sentences = [c[1] for c in cand]
+
+        ## CrossEncoder
+        cross_inp = [[query, sent] for sent in corpus_sentences]
+        if model.config.num_labels > 1:
+            cross_scores = model.predict(cross_inp, apply_softmax=True)[:, 1].tolist()
+        else:
+            cross_scores = model.predict(cross_inp, activation_fct=torch.nn.Identity()).tolist()
+
+        cross_scores_sparse = {}
+        for idx, pid in enumerate(pids):
+            cross_scores_sparse[pid] = cross_scores[idx]
+
+        sparse_scores = cross_scores_sparse
+        run[qid] = {}
+        for pid in sparse_scores:
+            run[qid][pid] = float(sparse_scores[pid])
+
+    evaluator = pytrec_eval.RelevanceEvaluator(rel, {'ndcg_cut.10'})
+    scores = evaluator.evaluate(run)
+    scores_mean = np.mean([ele["ndcg_cut_10"] for ele in scores.values()])
+
+    print("NDCG@10: {:.2f}".format(scores_mean * 100))
+    return scores_mean
+
+################################
+
+model_name = sys.argv[1]
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+config = AutoConfig.from_pretrained(model_name)
+config.num_labels = 1
+model = AutoModelForSequenceClassification.from_pretrained(model_name, config=config)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+## Freeze embedding layer
+model.distilbert.embeddings.word_embeddings.requires_grad_(False)
+
+
+
+
+
+#######################
+
+queries = {}
+corpus = {}
+
+output_save_path = 'output-ce-emb_frozen/{}-{}'.format(model_name.replace("/", "-"), datetime.now().strftime("%Y-%m-%d_%H-%M-%S"))
+output_save_path_latest = output_save_path+"-latest"
+tokenizer.save_pretrained(output_save_path)
+tokenizer.save_pretrained(output_save_path_latest)
+
+
+# Write self to path
+train_script_path = os.path.join(output_save_path, 'train_script.py')
+copyfile(__file__, train_script_path)
+with open(train_script_path, 'a') as fOut:
+    fOut.write("\n\n# Script was called via:\n#python " + " ".join(sys.argv))
+
+
+####
+train_script_path = os.path.join(output_save_path_latest, 'train_script.py')
+copyfile(__file__, train_script_path)
+with open(train_script_path, 'a') as fOut:
+    fOut.write("\n\n# Script was called via:\n#python " + " ".join(sys.argv))
+
+
+
+#### Read train file
+with gzip.open('/home/msmarco/data/collection.tsv.gz', 'rt') as fIn:
+    for line in fIn:
+        pid, passage = line.strip().split("\t")
+        corpus[pid] = passage
+
+with open('/home/msmarco/data/queries.train.tsv', 'r') as fIn:
+    for line in fIn:
+        qid, query = line.strip().split("\t")
+        queries[qid] = query
+
+
+############## Train Dataset
+class MSEDataset(Dataset):
+    def __init__(self, filepath):
+        super().__init__()
+
+        self.examples = []
+        with open(filepath) as fIn:
+            for line in fIn:
+                pos_score, neg_score, qid, pid1, pid2 = line.strip().split("\t")
+                self.examples.append([qid, pid1, pid2, float(pos_score)-float(neg_score)])
+
+    def __len__(self):
+        return len(self.examples)
+
+    def __getitem__(self, item):
+        return self.examples[item]
+
+train_batch_size = 32
+train_dataset = MSEDataset('/home/msmarco/data/bert_cat_ensemble_msmarcopassage_train_scores_ids.tsv')
+train_dataloader = DataLoader(train_dataset, drop_last=True, shuffle=True, batch_size=16)
+
+
+############## Optimizer
+
+weight_decay = 0.01
+max_grad_norm = 1
+param_optimizer = list(model.named_parameters())
+
+no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
+optimizer_grouped_parameters = [
+    {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': weight_decay},
+    {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
+]
+
+optimizer = AdamW(optimizer_grouped_parameters, lr=1e-5)
+scheduler = transformers.get_linear_schedule_with_warmup(optimizer, num_warmup_steps=1000, num_training_steps=len(train_dataloader))
+scaler = torch.cuda.amp.GradScaler()
+
+loss_fct = torch.nn.MSELoss()
+### Start training
+model.to(device)
+
+auto_save = 10000
+best_ndcg_score = 0
+for step_idx, batch in tqdm.tqdm(enumerate(train_dataloader), total=len(train_dataloader)):
+    batch_queries = [queries[qid] for qid in batch[0]]
+    pos = [corpus[cid] for cid in batch[1]]
+    neg = [corpus[cid] for cid in batch[2]]
+    scores = batch[3].float().to(device)    #torch.tensor(batch[3], dtype=torch.float, device=device)
+
+    with autocast():
+        inp_pos = tokenizer(batch_queries, pos, max_length=512, padding=True, truncation='longest_first', return_tensors='pt').to(device)
+        pred_pos = model(**inp_pos).logits.squeeze()
+
+        inp_neg = tokenizer(batch_queries, neg, max_length=512, padding=True, truncation='longest_first', return_tensors='pt').to(device)
+        pred_neg = model(**inp_neg).logits.squeeze()
+
+        pred_diff = pred_pos - pred_neg
+        loss_value = loss_fct(pred_diff, scores)
+
+    
+    scaler.scale(loss_value).backward()
+    scaler.unscale_(optimizer)
+    torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
+    scaler.step(optimizer)
+    scaler.update()
+
+    optimizer.zero_grad()
+    scheduler.step()
+
+    if (step_idx+1) % auto_save == 0:
+        print("Step:", step_idx+1)
+        model.save_pretrained(output_save_path_latest)
+        ndcg_score = eval_modal(output_save_path_latest)
+
+        if ndcg_score >= best_ndcg_score:
+            best_ndcg_score = ndcg_score
+            print("Save to:", output_save_path)
+            model.save_pretrained(output_save_path)
+
+model.save_pretrained(output_save_path)
+
+
+# Script was called via:
+#python train_ce_emb_frozen.py nicoladecao/msmarco-word2vec256000-distilbert-base-uncased