upload

README.md

@@ -0,0 +1,5 @@
+# DistilBERT with word2vec token embeddings
+
+This model has a word2vec token embedding matrix with 256k entries. The word2vec was trained on 100GB data from C4, MSMARCO, News, Wikipedia, S2ORC, for 3 epochs.
+
+Then the model was trained on this dataset with MLM for 250k steps (batch size 64). The token embeddings were NOT updated.

config.json

@@ -0,0 +1,24 @@
+{
+  "_name_or_path": "train-w2v-model/c4_msmarco_news_s2orc_wiki/distilbert-256k/",
+  "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.17.0",
+  "vocab_size": 256000
+}

pytorch_model.bin

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

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": "c4_msmarco_news_s2orc_wiki/tokenizer-256k/special_tokens_map.json", "name_or_path": "train-w2v-model/c4_msmarco_news_s2orc_wiki/distilbert-256k/", "tokenizer_class": "PreTrainedTokenizerFast"}

train_script.py

@@ -0,0 +1,398 @@
+
+import argparse
+import logging
+import math
+import os
+from datetime import datetime
+import datasets
+import torch
+from torch.utils.data import DataLoader
+from tqdm.auto import tqdm
+import sys
+import transformers
+from accelerate import Accelerator, DistributedType
+from shutil import copyfile
+import wandb
+import numpy as np
+
+from transformers import (
+    MODEL_MAPPING,
+    AutoModelForMaskedLM,
+    AutoTokenizer,
+    DataCollatorForLanguageModeling,
+    SchedulerType,
+    get_scheduler
+)
+from transformers.utils.versions import require_version
+
+
+
+class TrainDataset(torch.utils.data.IterableDataset):
+    def __init__(self, filepath, tokenizer, max_length, batch_size, train_samples):
+        self.tokenizer = tokenizer
+        self.fIn = open(filepath)
+        self.max_length = max_length
+        self.batch_size = batch_size
+        self.train_samples = train_samples
+
+    def __iter__(self):
+        batch = []
+        for sent in self.fIn:
+            batch.append(sent.strip()[0:1000])
+
+            if len(batch) >= self.batch_size:
+                #Use multi process tokenization
+                encoded = self.tokenizer(batch, add_special_tokens=True, truncation=True, max_length=self.max_length, return_special_tokens_mask=True, padding=True)
+                #print(len(encoded['input_ids'][0]))
+                for idx in range(len(batch)):
+                    single_sample = {key: encoded[key][idx] for key in encoded}
+                    yield single_sample
+                
+                batch = []
+
+    def __len__(self):
+        return self.train_samples
+
+    
+
+
+
+## Dev dataset
+class DevDataset(torch.utils.data.Dataset):
+    def __init__(self, filepath, tokenizer, max_length):
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        with open(filepath) as fIn:
+            sentences = [sent.strip() for sent in fIn]
+
+        self.num_sentences = len(sentences)
+        self.tokenized = self.tokenizer(sentences, add_special_tokens=True, truncation=True, max_length=self.max_length, return_special_tokens_mask=True)
+
+    def __getitem__(self, idx):
+        return {key: self.tokenized[key][idx] for key in self.tokenized}
+
+    def __len__(self):
+        return self.num_sentences
+
+
+
+logger = logging.getLogger(__name__)
+require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt")
+MODEL_CONFIG_CLASSES = list(MODEL_MAPPING.keys())
+MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Finetune a transformers model on a Masked Language Modeling task")
+    parser.add_argument(
+        "--dataset_config_name",
+        type=str,
+        default=None,
+        help="The configuration name of the dataset to use (via the datasets library).",
+    )
+    parser.add_argument(
+        "--train_file", type=str, default=None, help="A text file data (1 text per line).."
+    )
+    parser.add_argument(
+        "--dev_file", type=str, default=None, help="A text file data (1 text per line)."
+    )
+    parser.add_argument(
+        "--model_name",
+        default="nicoladecao/msmarco-word2vec256000-distilbert-base-uncased",
+        type=str,
+        help="Path to pretrained model or model identifier from huggingface.co/models."
+    )
+    parser.add_argument(
+        "--per_device_batch_size",
+        type=int,
+        default=16,
+        help="Batch size (per device) for the training dataloader.",
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=5e-5,
+        help="Initial learning rate (after the potential warmup period) to use.",
+    )
+    parser.add_argument("--weight_decay", type=float, default=0.01, help="Weight decay to use.")
+    parser.add_argument("--num_train_epochs", type=int, default=1, help="Total number of training epochs to perform.")
+    parser.add_argument(
+        "--max_train_steps",
+        type=int,
+        help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
+    )
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help="Number of updates steps to accumulate before performing a backward/update pass.",
+    )
+    parser.add_argument(
+        "--lr_scheduler_type",
+        type=SchedulerType,
+        default="linear",
+        help="The scheduler type to use.",
+        choices=["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"],
+    )
+    parser.add_argument(
+        "--num_warmup_steps", type=int, default=1000, help="Number of steps for the warmup in the lr scheduler."
+    )
+    parser.add_argument(
+        "--model_type",
+        type=str,
+        default=None,
+        help="Model type to use if training from scratch.",
+        choices=MODEL_TYPES,
+    )
+    parser.add_argument(
+        "--max_seq_length",
+        type=int,
+        default=256,
+        help="The maximum total input sequence length after tokenization. Sequences longer than this will be truncated.",
+    )
+    parser.add_argument(
+        "--line_by_line",
+        type=bool,
+        default=True,
+        help="Whether distinct lines of text in the dataset are to be handled as distinct sequences.",
+    )
+    parser.add_argument(
+        "--overwrite_cache", type=bool, default=False, help="Overwrite the cached training and evaluation sets"
+    )
+    parser.add_argument(
+        "--mlm_probability", type=float, default=0.15, help="Ratio of tokens to mask for masked language modeling loss"
+    )
+    parser.add_argument("--mixed_precision", default="fp16")
+    parser.add_argument("--train_samples", required=True, type=int)
+    parser.add_argument("--eval_steps", default=10000, type=int)
+    parser.add_argument("--max_grad_norm", default=1.0, type=float)
+    parser.add_argument("--project", default="bert-word2vec")
+    parser.add_argument("--freeze_emb_layer", default=False, action='store_true')
+    parser.add_argument("--log_interval", default=1000, type=int)
+    parser.add_argument("--ckp_steps", default=50000, type=int)
+
+    args = parser.parse_args()
+
+
+    return args
+
+
+def main():
+    args = parse_args()
+
+    # Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
+    accelerator = Accelerator(mixed_precision=args.mixed_precision)
+    # Make one log on every process with the configuration for debugging.
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO,
+    )
+    logger.info(accelerator.state)
+
+    # Setup logging, we only want one process per machine to log things on the screen.
+    # accelerator.is_local_main_process is only True for one process per machine.
+    logger.setLevel(logging.INFO if accelerator.is_local_main_process else logging.ERROR)
+    if accelerator.is_local_main_process:
+        datasets.utils.logging.set_verbosity_warning()
+        transformers.utils.logging.set_verbosity_info()
+    else:
+        datasets.utils.logging.set_verbosity_error()
+        transformers.utils.logging.set_verbosity_error()
+
+
+    accelerator.wait_for_everyone()
+
+
+    #Load model
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
+    model = AutoModelForMaskedLM.from_pretrained(args.model_name)
+
+    #Freeze emb layer
+    if args.freeze_emb_layer:
+        model.distilbert.embeddings.word_embeddings.requires_grad_(False)
+   
+    # Logging & Co on main process
+    if accelerator.is_main_process:
+        exp_name = f'{args.model_name.replace("/", "-")}-{"freeze_emb" if args.freeze_emb_layer else "update_emb"}-{datetime.now().strftime("%Y-%m-%d_%H-%M-%S")}'
+        output_dir = os.path.join("output-mlm", exp_name)
+        wandb.init(project=args.project, name=exp_name, config=args)
+
+        os.makedirs(output_dir, exist_ok=False)
+
+        #Save tokenizer
+        tokenizer.save_pretrained(output_dir)
+
+        #Save train script
+        train_script_path = os.path.join(output_dir, '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))
+
+
+    total_batch_size = args.per_device_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+
+    train_dataset = TrainDataset(args.train_file, tokenizer, args.max_seq_length, batch_size=total_batch_size, train_samples=args.train_samples)
+    eval_dataset  = DevDataset(args.dev_file, tokenizer, args.max_seq_length)
+
+
+    # Data collator
+    # This one will take care of randomly masking the tokens.
+    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=args.mlm_probability)
+
+    # DataLoaders creation:
+    train_dataloader = DataLoader(train_dataset, collate_fn=data_collator, batch_size=args.per_device_batch_size)
+    eval_dataloader = DataLoader(eval_dataset, collate_fn=data_collator, batch_size=args.per_device_batch_size)
+
+    # Optimizer
+    # Split weights in two groups, one with weight decay and the other not.
+    no_decay = ["bias", "LayerNorm.weight"]
+    optimizer_grouped_parameters = [
+        {
+            "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
+            "weight_decay": args.weight_decay,
+        },
+        {
+            "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
+            "weight_decay": 0.0,
+        },
+    ]
+    optimizer = torch.optim.AdamW(optimizer_grouped_parameters, lr=args.learning_rate)
+
+    # Prepare everything with our `accelerator`.
+    model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(model, optimizer, train_dataloader, eval_dataloader)
+
+    # On TPU, the tie weights in our model have been disconnected, so we need to restore the ties.
+    if accelerator.distributed_type == DistributedType.TPU:
+        model.tie_weights()
+
+    # Note -> the training dataloader needs to be prepared before we grab his length below (cause its length will be
+    # shorter in multiprocess)
+
+    # Scheduler and math around the number of training steps.
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    if args.max_train_steps is None:
+        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+    else:
+        args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+    lr_scheduler = get_scheduler(
+        name=args.lr_scheduler_type,
+        optimizer=optimizer,
+        num_warmup_steps=args.num_warmup_steps,
+        num_training_steps=args.max_train_steps,
+    )
+
+
+    # Train!
+    logger.info("***** Running training *****")
+    logger.info(f"  Num examples = {args.train_samples}")
+    logger.info(f"  Num Epochs = {args.num_train_epochs}")
+    logger.info(f"  Instantaneous batch size per device = {args.per_device_batch_size}")
+    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
+    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
+    logger.info(f"  Total optimization steps = {args.max_train_steps}")
+    # Only show the progress bar once on each machine.
+    progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process, smoothing=0.05)
+    completed_steps = 0
+    train_loss_values = []
+
+    best_eval_loss = 999999
+    if accelerator.is_main_process:
+        best_ckp_dir = os.path.join(output_dir, "best")
+        tokenizer.save_pretrained(best_ckp_dir)
+
+    for epoch in range(args.num_train_epochs):
+        logger.info(f"Start epoch {epoch}")
+        model.train()
+        for step, batch in enumerate(train_dataloader):
+            outputs = model(**batch)
+            loss = outputs.loss
+            loss = loss / args.gradient_accumulation_steps
+           
+            if accelerator.is_main_process:
+                train_loss_values.append(loss.cpu().item())
+
+            accelerator.backward(loss)
+            accelerator.clip_grad_norm_(model.parameters(), args.max_grad_norm)
+            if step % args.gradient_accumulation_steps == 0:
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+                progress_bar.update(1)
+                completed_steps += 1
+
+                 ### Do logging
+                if accelerator.is_main_process:
+                    if completed_steps % args.log_interval == 0:
+                        wandb.log({"train/loss": np.mean(train_loss_values)}, step=completed_steps)
+                        train_loss_values = []
+
+           
+                if completed_steps % args.eval_steps == 0:
+                    model.eval()
+                    losses = []
+                    for step, batch in enumerate(eval_dataloader):
+                        with torch.no_grad():
+                            outputs = model(**batch)
+
+                        loss = outputs.loss
+                        losses.append(accelerator.gather(loss.repeat(args.per_device_batch_size)))
+
+                    losses = torch.cat(losses)
+                    losses = losses[: len(eval_dataset)]
+                    try:
+                        eval_loss = torch.mean(losses)
+                    except OverflowError:
+                        eval_loss = float("inf")
+
+                    logger.info(f"step {completed_steps}: perplexity: {eval_loss}")
+                    if accelerator.is_main_process:
+                        wandb.log({"eval/loss": eval_loss}, step=completed_steps)
+
+                    model.train()
+
+                    #Save model
+                    accelerator.wait_for_everyone()
+                    if accelerator.is_main_process:
+                        unwrapped_model = accelerator.unwrap_model(model)
+                        unwrapped_model.save_pretrained(output_dir, save_function=accelerator.save)
+                        with open(os.path.join(output_dir, "train_steps.log"), 'a') as fOut:
+                            fOut.write(f"{completed_steps}: {eval_loss}\n")
+
+                        #Save best model
+                        if eval_loss < best_eval_loss:
+                            best_eval_loss = eval_loss
+                            unwrapped_model.save_pretrained(best_ckp_dir, save_function=accelerator.save)
+                            with open(os.path.join(best_ckp_dir, "train_steps.log"), 'a') as fOut:
+                                fOut.write(f"{completed_steps}: {eval_loss}\n")
+                        
+                if accelerator.is_main_process and completed_steps % args.ckp_steps == 0:
+                    ckp_dir = os.path.join(output_dir, f"ckp-{int(completed_steps/1000)}k")
+                    unwrapped_model = accelerator.unwrap_model(model)
+                    unwrapped_model.save_pretrained(ckp_dir, save_function=accelerator.save)
+                    tokenizer.save_pretrained(ckp_dir)
+                    with open(os.path.join(ckp_dir, "train_steps.log"), 'a') as fOut:
+                        fOut.write(f"{completed_steps}: {eval_loss}\n")
+                   
+            
+                if completed_steps >= args.max_train_steps:
+                    break
+
+    if args.output_dir is not None:
+        accelerator.wait_for_everyone()
+        if accelerator.is_main_process:
+            unwrapped_model = accelerator.unwrap_model(model)
+            unwrapped_model.save_pretrained(output_dir, save_function=accelerator.save)
+            with open(os.path.join(output_dir, "train_steps.log"), 'a') as fOut:
+                fOut.write(f"{completed_steps}\n")
+        
+
+
+
+if __name__ == "__main__":
+    main()
+
+
+# Script was called via:
+#python train_mlm-iterable.py --train_file data/c4_msmarco_news_s2orc_wiki_train.txt --dev_file data/c4_msmarco_news_s2orc_wiki_dev.txt --train_samples 100000000 --model_name train-w2v-model/c4_msmarco_news_s2orc_wiki/distilbert-256k/ --freeze_emb_layer

train_steps.log

@@ -0,0 +1,39 @@
+10000: 3.6185991764068604
+20000: 3.181567430496216
+30000: 3.019852638244629
+40000: 2.8929433822631836
+50000: 2.865853786468506
+60000: 2.8218629360198975
+70000: 2.7376461029052734
+90000: 2.698227882385254
+100000: 2.6650893688201904
+120000: 2.6339340209960938
+130000: 2.593796730041504
+160000: 2.570080280303955
+180000: 2.5539512634277344
+190000: 2.5419578552246094
+210000: 2.4972760677337646
+260000: 2.4895386695861816
+270000: 2.481090545654297
+290000: 2.4765520095825195
+300000: 2.463596820831299
+320000: 2.4584429264068604
+350000: 2.450732469558716
+360000: 2.443289279937744
+370000: 2.4305179119110107
+410000: 2.4060347080230713
+470000: 2.376832962036133
+510000: 2.3685810565948486
+550000: 2.3647472858428955
+600000: 2.3556222915649414
+670000: 2.3360767364501953
+690000: 2.327178955078125
+730000: 2.3191168308258057
+740000: 2.3143470287323
+830000: 2.3057608604431152
+840000: 2.2876601219177246
+980000: 2.253411293029785
+1080000: 2.241132974624634
+1230000: 2.234037160873413
+1320000: 2.2321970462799072
+1370000: 2.2040650844573975