upload

README.md

@@ -0,0 +1,137 @@
+---
+pipeline_tag: sentence-similarity
+tags:
+- sentence-transformers
+- feature-extraction
+- sentence-similarity
+- transformers
+---
+
+# dense_encoder-msmarco-distilbert-word2vec256k-MLM_210k
+
+This model is based on [vocab-transformers/msmarco-distilbert-word2vec256k-MLM_445k](https://huggingface.co/vocab-transformers/msmarco-distilbert-word2vec256k-MLM_445k) with a 256k sized vocabulary initialized with word2vec that has been trained with MLM for 445k steps with frozen embedding matrix on the MSMARCO corpus.
+
+It has been trained on MS MARCO using [MarginMSELoss](https://github.com/UKPLab/sentence-transformers/blob/master/examples/training/ms_marco/train_bi-encoder_margin-mse.py) with a frozen embedding matrix. See the train_script.py in this repository.
+
+
+Performance:
+- MS MARCO dev: (running) (MRR@10)
+- TREC-DL 2019: 66.72 (nDCG@10)
+- TREC-DL 2020: 69.14 (nDCG@10)
+
+
+
+## Usage (Sentence-Transformers)
+
+
+This is a [sentence-transformers](https://www.SBERT.net) model: It maps sentences & paragraphs to a 768 dimensional dense vector space and can be used for tasks like clustering or semantic search.
+
+
+Using this model becomes easy when you have [sentence-transformers](https://www.SBERT.net) installed:
+
+```
+pip install -U sentence-transformers
+```
+
+Then you can use the model like this:
+
+```python
+from sentence_transformers import SentenceTransformer
+sentences = ["This is an example sentence", "Each sentence is converted"]
+
+model = SentenceTransformer('{MODEL_NAME}')
+embeddings = model.encode(sentences)
+print(embeddings)
+```
+
+
+
+## Usage (HuggingFace Transformers)
+Without [sentence-transformers](https://www.SBERT.net), you can use the model like this: First, you pass your input through the transformer model, then you have to apply the right pooling-operation on-top of the contextualized word embeddings.
+
+```python
+from transformers import AutoTokenizer, AutoModel
+import torch
+
+
+#Mean Pooling - Take attention mask into account for correct averaging
+def mean_pooling(model_output, attention_mask):
+    token_embeddings = model_output[0] #First element of model_output contains all token embeddings
+    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
+
+
+# Sentences we want sentence embeddings for
+sentences = ['This is an example sentence', 'Each sentence is converted']
+
+# Load model from HuggingFace Hub
+tokenizer = AutoTokenizer.from_pretrained('{MODEL_NAME}')
+model = AutoModel.from_pretrained('{MODEL_NAME}')
+
+# Tokenize sentences
+encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
+
+# Compute token embeddings
+with torch.no_grad():
+    model_output = model(**encoded_input)
+
+# Perform pooling. In this case, mean pooling.
+sentence_embeddings = mean_pooling(model_output, encoded_input['attention_mask'])
+
+print("Sentence embeddings:")
+print(sentence_embeddings)
+```
+
+
+
+## Evaluation Results
+
+<!--- Describe how your model was evaluated -->
+
+For an automated evaluation of this model, see the *Sentence Embeddings Benchmark*: [https://seb.sbert.net](https://seb.sbert.net?model_name={MODEL_NAME})
+
+
+## Training
+The model was trained with the parameters:
+
+**DataLoader**:
+
+`torch.utils.data.dataloader.DataLoader` of length 7858 with parameters:
+```
+{'batch_size': 64, 'sampler': 'torch.utils.data.sampler.RandomSampler', 'batch_sampler': 'torch.utils.data.sampler.BatchSampler'}
+```
+
+**Loss**:
+
+`sentence_transformers.losses.MarginMSELoss.MarginMSELoss` 
+
+Parameters of the fit()-Method:
+```
+{
+    "epochs": 30,
+    "evaluation_steps": 0,
+    "evaluator": "NoneType",
+    "max_grad_norm": 1,
+    "optimizer_class": "<class 'transformers.optimization.AdamW'>",
+    "optimizer_params": {
+        "lr": 2e-05
+    },
+    "scheduler": "WarmupLinear",
+    "steps_per_epoch": null,
+    "warmup_steps": 1000,
+    "weight_decay": 0.01
+}
+```
+
+
+## Full Model Architecture
+```
+SentenceTransformer(
+  (0): Transformer({'max_seq_length': 250, 'do_lower_case': False}) with Transformer model: DistilBertModel 
+  (1): Pooling({'word_embedding_dimension': 768, 'pooling_mode_cls_token': False, 'pooling_mode_mean_tokens': True, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False})
+)
+```
+
+## Citing & Authors
+
+<!--- Describe where people can find more information -->

config.json

@@ -0,0 +1,24 @@
+{
+  "_name_or_path": "/home/word2vec-bert/models/distilbert-word2vec256k-MLM_445k/",
+  "activation": "gelu",
+  "architectures": [
+    "DistilBertModel"
+  ],
+  "attention_dropout": 0.1,
+  "dim": 768,
+  "dropout": 0.1,
+  "hidden_dim": 3072,
+  "initializer_range": 0.02,
+  "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.16.2",
+  "vocab_size": 256000
+}

config_sentence_transformers.json

@@ -0,0 +1,7 @@
+{
+  "__version__": {
+    "sentence_transformers": "2.2.0",
+    "transformers": "4.16.2",
+    "pytorch": "1.10.2"
+  }
+}

modules.json

@@ -0,0 +1,14 @@
+[
+  {
+    "idx": 0,
+    "name": "0",
+    "path": "",
+    "type": "sentence_transformers.models.Transformer"
+  },
+  {
+    "idx": 1,
+    "name": "1",
+    "path": "1_Pooling",
+    "type": "sentence_transformers.models.Pooling"
+  }
+]

pytorch_model.bin

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

sentence_bert_config.json

@@ -0,0 +1,4 @@
+{
+  "max_seq_length": 250,
+  "do_lower_case": false
+}

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 @@
+{"do_lower_case": true, "unk_token": "[UNK]", "sep_token": "[SEP]", "pad_token": "[PAD]", "cls_token": "[CLS]", "mask_token": "[MASK]", "tokenize_chinese_chars": true, "strip_accents": null, "model_max_length": 512, "model_input_names": ["input_ids", "attention_mask"], "special_tokens_map_file": "/home/ukp-reimers/.cache/huggingface/transformers/fe09c361189d8238b9e387f10a088e93f70620bfe74b82036baff1fed512a153.dd8bd9bfd3664b530ea4e645105f557769387b3da9f79bdb55ed556bdd80611d", "name_or_path": "/home/word2vec-bert/models/distilbert-word2vec256k-MLM_445k/", "tokenizer_class": "PreTrainedTokenizerFast"}

train_script.py

@@ -0,0 +1,234 @@
+import sys
+import json
+from torch.utils.data import DataLoader
+from sentence_transformers import SentenceTransformer, LoggingHandler, util, models, evaluation, losses, InputExample
+import logging
+from datetime import datetime
+import gzip
+import os
+import tarfile
+import tqdm
+from torch.utils.data import Dataset
+import random
+from shutil import copyfile
+import pickle
+import argparse
+
+#### Just some code to print debug information to stdout
+logging.basicConfig(format='%(asctime)s - %(message)s',
+                    datefmt='%Y-%m-%d %H:%M:%S',
+                    level=logging.INFO,
+                    handlers=[LoggingHandler()])
+#### /print debug information to stdout
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--train_batch_size", default=64, type=int)
+parser.add_argument("--max_seq_length", default=250, type=int)
+parser.add_argument("--model_name", default="nicoladecao/msmarco-word2vec256000-distilbert-base-uncased")
+parser.add_argument("--max_passages", default=0, type=int)
+parser.add_argument("--epochs", default=30, type=int)
+parser.add_argument("--pooling", default="mean")
+parser.add_argument("--negs_to_use", default=None, help="From which systems should negatives be used? Multiple systems seperated by comma. None = all")
+parser.add_argument("--warmup_steps", default=1000, type=int)
+parser.add_argument("--lr", default=2e-5, type=float)
+parser.add_argument("--num_negs_per_system", default=5, type=int)
+parser.add_argument("--use_all_queries", default=False, action="store_true")
+args = parser.parse_args()
+
+logging.info(str(args))
+
+
+
+# The  model we want to fine-tune
+train_batch_size = args.train_batch_size          #Increasing the train batch size improves the model performance, but requires more GPU memory
+model_name = args.model_name
+max_passages = args.max_passages
+max_seq_length = args.max_seq_length            #Max length for passages. Increasing it, requires more GPU memory
+
+num_negs_per_system = args.num_negs_per_system  # We used different systems to mine hard negatives. Number of hard negatives to add from each system
+num_epochs = args.epochs         # Number of epochs we want to train
+
+# Load our embedding model
+
+logging.info("Create new SBERT model")
+word_embedding_model = models.Transformer(model_name, max_seq_length=max_seq_length)
+pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension(), args.pooling)
+model = SentenceTransformer(modules=[word_embedding_model, pooling_model])
+
+#Freeze embedding layer
+word_embedding_model.auto_model.embeddings.requires_grad = False
+
+model_save_path = f'output/train_bi-encoder-margin_mse-word2vec-{model_name.replace("/", "-")}-batch_size_{train_batch_size}-{datetime.now().strftime("%Y-%m-%d_%H-%M-%S")}'
+
+
+# Write self to path
+os.makedirs(model_save_path, exist_ok=True)
+
+train_script_path = os.path.join(model_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))
+
+
+### Now we read the MS Marco dataset
+data_folder = 'msmarco-data'
+
+#### Read the corpus files, that contain all the passages. Store them in the corpus dict
+corpus = {}         #dict in the format: passage_id -> passage. Stores all existent passages
+collection_filepath = os.path.join(data_folder, 'collection.tsv')
+if not os.path.exists(collection_filepath):
+    tar_filepath = os.path.join(data_folder, 'collection.tar.gz')
+    if not os.path.exists(tar_filepath):
+        logging.info("Download collection.tar.gz")
+        util.http_get('https://msmarco.blob.core.windows.net/msmarcoranking/collection.tar.gz', tar_filepath)
+
+    with tarfile.open(tar_filepath, "r:gz") as tar:
+        tar.extractall(path=data_folder)
+
+logging.info("Read corpus: collection.tsv")
+with open(collection_filepath, 'r', encoding='utf8') as fIn:
+    for line in fIn:
+        pid, passage = line.strip().split("\t")
+        pid = int(pid)
+        corpus[pid] = passage
+
+
+### Read the train queries, store in queries dict
+queries = {}        #dict in the format: query_id -> query. Stores all training queries
+queries_filepath = os.path.join(data_folder, 'queries.train.tsv')
+if not os.path.exists(queries_filepath):
+    tar_filepath = os.path.join(data_folder, 'queries.tar.gz')
+    if not os.path.exists(tar_filepath):
+        logging.info("Download queries.tar.gz")
+        util.http_get('https://msmarco.blob.core.windows.net/msmarcoranking/queries.tar.gz', tar_filepath)
+
+    with tarfile.open(tar_filepath, "r:gz") as tar:
+        tar.extractall(path=data_folder)
+
+
+with open(queries_filepath, 'r', encoding='utf8') as fIn:
+    for line in fIn:
+        qid, query = line.strip().split("\t")
+        qid = int(qid)
+        queries[qid] = query
+
+
+# Load a dict (qid, pid) -> ce_score that maps query-ids (qid) and paragraph-ids (pid)
+# to the CrossEncoder score computed by the cross-encoder/ms-marco-MiniLM-L-6-v2 model
+ce_scores_file = os.path.join(data_folder, 'cross-encoder-ms-marco-MiniLM-L-6-v2-scores.pkl.gz')
+if not os.path.exists(ce_scores_file):
+    logging.info("Download cross-encoder scores file")
+    util.http_get('https://huggingface.co/datasets/sentence-transformers/msmarco-hard-negatives/resolve/main/cross-encoder-ms-marco-MiniLM-L-6-v2-scores.pkl.gz', ce_scores_file)
+
+logging.info("Load CrossEncoder scores dict")
+with gzip.open(ce_scores_file, 'rb') as fIn:
+    ce_scores = pickle.load(fIn)
+
+# As training data we use hard-negatives that have been mined using various systems
+hard_negatives_filepath = os.path.join(data_folder, 'msmarco-hard-negatives.jsonl.gz')
+if not os.path.exists(hard_negatives_filepath):
+    logging.info("Download cross-encoder scores file")
+    util.http_get('https://huggingface.co/datasets/sentence-transformers/msmarco-hard-negatives/resolve/main/msmarco-hard-negatives.jsonl.gz', hard_negatives_filepath)
+
+
+logging.info("Read hard negatives train file")
+train_queries = {}
+negs_to_use = None
+with gzip.open(hard_negatives_filepath, 'rt') as fIn:
+    for line in tqdm.tqdm(fIn):
+        if max_passages > 0 and len(train_queries) >= max_passages:
+            break
+        data = json.loads(line)
+
+        #Get the positive passage ids
+        pos_pids = data['pos']
+
+        #Get the hard negatives
+        neg_pids = set()
+        if negs_to_use is None:
+            if args.negs_to_use is not None:    #Use specific system for negatives
+                negs_to_use = args.negs_to_use.split(",")
+            else:   #Use all systems
+                negs_to_use = list(data['neg'].keys())
+            logging.info("Using negatives from the following systems: {}".format(", ".join(negs_to_use)))
+
+        for system_name in negs_to_use:
+            if system_name not in data['neg']:
+                continue
+
+            system_negs = data['neg'][system_name]
+            negs_added = 0
+            for pid in system_negs:
+                if pid not in neg_pids:
+                    neg_pids.add(pid)
+                    negs_added += 1
+                    if negs_added >= num_negs_per_system:
+                        break
+
+        if args.use_all_queries or (len(pos_pids) > 0 and len(neg_pids) > 0):
+            train_queries[data['qid']] = {'qid': data['qid'], 'query': queries[data['qid']], 'pos': pos_pids, 'neg': neg_pids}
+
+logging.info("Train queries: {}".format(len(train_queries)))
+
+# We create a custom MSMARCO dataset that returns triplets (query, positive, negative)
+# on-the-fly based on the information from the mined-hard-negatives jsonl file.
+class MSMARCODataset(Dataset):
+    def __init__(self, queries, corpus, ce_scores):
+        self.queries = queries
+        self.queries_ids = list(queries.keys())
+        self.corpus = corpus
+        self.ce_scores = ce_scores
+
+        for qid in self.queries:
+            self.queries[qid]['pos'] = list(self.queries[qid]['pos'])
+            self.queries[qid]['neg'] = list(self.queries[qid]['neg'])
+            random.shuffle(self.queries[qid]['neg'])
+
+    def __getitem__(self, item):
+        query = self.queries[self.queries_ids[item]]
+        query_text = query['query']
+        qid = query['qid']
+
+        if len(query['pos']) > 0:
+            pos_id = query['pos'].pop(0)    #Pop positive and add at end
+            pos_text = self.corpus[pos_id]
+            query['pos'].append(pos_id)
+        else:   #We only have negatives, use two negs
+            pos_id = query['neg'].pop(0)    #Pop negative and add at end
+            pos_text = self.corpus[pos_id]
+            query['neg'].append(pos_id)
+
+        #Get a negative passage
+        neg_id = query['neg'].pop(0)    #Pop negative and add at end
+        neg_text = self.corpus[neg_id]
+        query['neg'].append(neg_id)
+
+        pos_score = self.ce_scores[qid][pos_id]
+        neg_score = self.ce_scores[qid][neg_id]
+
+        return InputExample(texts=[query_text, pos_text, neg_text], label=pos_score-neg_score)
+
+    def __len__(self):
+        return len(self.queries)
+
+# For training the SentenceTransformer model, we need a dataset, a dataloader, and a loss used for training.
+train_dataset = MSMARCODataset(queries=train_queries, corpus=corpus, ce_scores=ce_scores)
+train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=train_batch_size, drop_last=True)
+train_loss = losses.MarginMSELoss(model=model)
+
+# Train the model
+model.fit(train_objectives=[(train_dataloader, train_loss)],
+          epochs=num_epochs,
+          warmup_steps=args.warmup_steps,
+          use_amp=True,
+          checkpoint_path=model_save_path,
+          checkpoint_save_steps=10000,
+          optimizer_params = {'lr': args.lr},
+          )
+
+# Train latest model
+model.save(model_save_path)
+
+# Script was called via:
+#python train_bi-encoder_margin-mse_word2vec.py --model /home/word2vec-bert/models/distilbert-word2vec256k-MLM_445k/