upload

README.md

@@ -0,0 +1,4 @@
+# Model
+This model is based on [nicoladecao/msmarco-word2vec256000-distilbert-base-uncased](https://huggingface.co/nicoladecao/msmarco-word2vec256000-distilbert-base-uncased) with a 256k sized vocabulary initialized with word2vec.
+
+This model has been trained with MLM on the MS MARCO corpus collection for 785k steps. See train_mlm.py for the train script. It was run on 2x V100 GPUs.

config.json

@@ -0,0 +1,24 @@
+{
+  "_name_or_path": "nicoladecao/msmarco-word2vec256000-distilbert-base-uncased",
+  "activation": "gelu",
+  "architectures": [
+    "DistilBertForMaskedLM"
+  ],
+  "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.11.3",
+  "vocab_size": 256000
+}

optimizer.pt

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

pytorch_model.bin

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

rng_state.pth

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

scaler.pt

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

scheduler.pt

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

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"], "tokenizer_class": "PreTrainedTokenizerFast", "special_tokens_map_file": "/home/ukp-reimers/.cache/huggingface/transformers/fe09c361189d8238b9e387f10a088e93f70620bfe74b82036baff1fed512a153.dd8bd9bfd3664b530ea4e645105f557769387b3da9f79bdb55ed556bdd80611d", "name_or_path": "nicoladecao/msmarco-word2vec256000-distilbert-base-uncased"}

train_mlm.py

@@ -0,0 +1,120 @@
+"""
+This file runs Masked Language Model. You provide a training file. Each line is interpreted as a sentence / paragraph.
+Optionally, you can also provide a dev file.
+
+The fine-tuned model is stored in the output/model_name folder.
+
+python train_mlm.py model_name data/train_sentences.txt [data/dev_sentences.txt]
+"""
+
+from transformers import AutoModelForMaskedLM, AutoTokenizer
+from transformers import DataCollatorForLanguageModeling, DataCollatorForWholeWordMask
+from transformers import Trainer, TrainingArguments
+import sys
+import gzip
+from datetime import datetime
+import wandb
+
+wandb.init(project="bert-word2vec")
+
+model_name = "nicoladecao/msmarco-word2vec256000-distilbert-base-uncased"
+per_device_train_batch_size = 16
+save_steps = 5000
+eval_steps = 1000
+num_train_epochs = 3
+use_fp16 = True    #Set to True, if your GPU supports FP16 operations
+max_length = 250    #Max length for a text input
+do_whole_word_mask = True   #If set to true, whole words are masked
+mlm_prob = 15       #Probability that a word is replaced by a [MASK] token
+
+model = AutoModelForMaskedLM.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+## Freeze embedding layer
+model.distilbert.embeddings.requires_grad = False
+
+output_dir = "output/{}-{}".format(model_name.replace("/", "_"),  datetime.now().strftime("%Y-%m-%d_%H-%M-%S"))
+print("Save checkpoints to:", output_dir)
+
+
+##### Load our training datasets
+
+train_sentences = []
+train_path = 'data/train.txt'
+with gzip.open(train_path, 'rt', encoding='utf8') if train_path.endswith('.gz') else  open(train_path, 'r', encoding='utf8') as fIn:
+    for line in fIn:
+        line = line.strip()
+        if len(line) >= 10:
+            train_sentences.append(line)
+
+print("Train sentences:", len(train_sentences))
+
+dev_sentences = []
+
+dev_path = 'data/dev.txt'
+with gzip.open(dev_path, 'rt', encoding='utf8') if dev_path.endswith('.gz') else open(dev_path, 'r', encoding='utf8') as fIn:
+    for line in fIn:
+        line = line.strip()
+        if len(line) >= 10:
+            dev_sentences.append(line)
+
+print("Dev sentences:", len(dev_sentences))
+
+#A dataset wrapper, that tokenizes our data on-the-fly
+class TokenizedSentencesDataset:
+    def __init__(self, sentences, tokenizer, max_length, cache_tokenization=False):
+        self.tokenizer = tokenizer
+        self.sentences = sentences
+        self.max_length = max_length
+        self.cache_tokenization = cache_tokenization
+
+    def __getitem__(self, item):
+        if not self.cache_tokenization:
+            return self.tokenizer(self.sentences[item], add_special_tokens=True, truncation=True, max_length=self.max_length, return_special_tokens_mask=True)
+
+        if isinstance(self.sentences[item], str):
+            self.sentences[item] = self.tokenizer(self.sentences[item], add_special_tokens=True, truncation=True, max_length=self.max_length, return_special_tokens_mask=True)
+        return self.sentences[item]
+
+    def __len__(self):
+        return len(self.sentences)
+
+train_dataset = TokenizedSentencesDataset(train_sentences, tokenizer, max_length)
+dev_dataset = TokenizedSentencesDataset(dev_sentences, tokenizer, max_length, cache_tokenization=True) if len(dev_sentences) > 0 else None
+
+
+##### Training arguments
+
+if do_whole_word_mask:
+    data_collator = DataCollatorForWholeWordMask(tokenizer=tokenizer, mlm=True, mlm_probability=mlm_prob)
+else:
+    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=True, mlm_probability=mlm_prob)
+
+training_args = TrainingArguments(
+    output_dir=output_dir,
+    overwrite_output_dir=True,
+    num_train_epochs=num_train_epochs,
+    evaluation_strategy="steps" if dev_dataset is not None else "no",
+    per_device_train_batch_size=per_device_train_batch_size,
+    eval_steps=eval_steps,
+    save_steps=save_steps,
+    save_total_limit=1,
+    prediction_loss_only=True,
+    fp16=use_fp16
+)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    data_collator=data_collator,
+    train_dataset=train_dataset,
+    eval_dataset=dev_dataset
+)
+
+trainer.train()
+
+print("Save model to:", output_dir)
+model.save_pretrained(output_dir)
+tokenizer.save_pretrained(output_dir)
+
+print("Training done")

training_args.bin

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