Pretrain.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os, sys, random\n",
    "from pathlib import Path\n",
    "from functools import partial\n",
    "from datetime import datetime, timezone, timedelta\n",
    "from IPython.core.debugger import set_trace as bk\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "import torch.tensor as T\n",
    "import datasets\n",
    "from fastai.text.all import *\n",
    "from transformers import ElectraConfig, ElectraTokenizerFast, ElectraForMaskedLM, ElectraForPreTraining\n",
    "from hugdatafast import *\n",
    "from _utils.utils import *\n",
    "from _utils.would_like_to_pr import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Configuraton"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "c = MyConfig({\n",
    "    'device': 'cuda:0',\n",
    "    \n",
    "    'base_run_name': 'vanilla', # run_name = {base_run_name}_{seed}\n",
    "    'seed': 11081, # 11081 36 1188 76 1 4 4649 7 # None/False to randomly choose seed from [0,999999]\n",
    "\n",
    "    'adam_bias_correction': False,\n",
    "    'schedule': 'original_linear',\n",
    "    'sampling': 'fp32_gumbel',\n",
    "    'electra_mask_style': True,\n",
    "    'gen_smooth_label': False,\n",
    "    'disc_smooth_label': False,\n",
    "\n",
    "    'size': 'small',\n",
    "    'datas': ['openwebtext'],\n",
    "    \n",
    "    'logger': 'wandb',\n",
    "    'num_workers': 3,\n",
    "    'my_model': False, # only for my personal research\n",
    "})\n",
    "\n",
    "# only for my personal research\n",
    "hparam_update = {\n",
    "    \n",
    "}\n",
    "\n",
    "\"\"\" Vanilla ELECTRA settings\n",
    "'adam_bias_correction': False,\n",
    "'schedule': 'original_linear',\n",
    "'sampling': 'fp32_gumbel',\n",
    "'electra_mask_style': True,\n",
    "'gen_smooth_label': False,\n",
    "'disc_smooth_label': False,\n",
    "'size': 'small',\n",
    "'datas': ['openwebtext'],\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Check and Default\n",
    "assert c.sampling in ['fp32_gumbel', 'fp16_gumbel', 'multinomial']\n",
    "assert c.schedule in ['original_linear', 'separate_linear', 'one_cycle', 'adjusted_one_cycle']\n",
    "for data in c.datas: assert data in ['wikipedia', 'bookcorpus', 'openwebtext']\n",
    "assert c.logger in ['wandb', 'neptune', None, False]\n",
    "if not c.base_run_name: c.base_run_name = str(datetime.now(timezone(timedelta(hours=+8))))[6:-13].replace(' ','').replace(':','').replace('-','')\n",
    "if not c.seed: c.seed = random.randint(0, 999999)\n",
    "c.run_name = f'{c.base_run_name}_{c.seed}'\n",
    "if c.gen_smooth_label is True: c.gen_smooth_label = 0.1\n",
    "if c.disc_smooth_label is True: c.disc_smooth_label = 0.1\n",
    "\n",
    "# Setting of different sizes\n",
    "i = ['small', 'base', 'large'].index(c.size)\n",
    "c.mask_prob = [0.15, 0.15, 0.25][i]\n",
    "c.lr = [5e-4, 2e-4, 2e-4][i]\n",
    "c.bs = [128, 256, 2048][i]\n",
    "c.steps = [10**6, 766*1000, 400*1000][i]\n",
    "c.max_length = [128, 512, 512][i]\n",
    "generator_size_divisor = [4, 3, 4][i]\n",
    "disc_config = ElectraConfig.from_pretrained(f'google/electra-{c.size}-discriminator')\n",
    "gen_config = ElectraConfig.from_pretrained(f'google/electra-{c.size}-generator')\n",
    "# note that public electra-small model is actually small++ and don't scale down generator size \n",
    "gen_config.hidden_size = int(disc_config.hidden_size/generator_size_divisor)\n",
    "gen_config.num_attention_heads = disc_config.num_attention_heads//generator_size_divisor\n",
    "gen_config.intermediate_size = disc_config.intermediate_size//generator_size_divisor\n",
    "hf_tokenizer = ElectraTokenizerFast.from_pretrained(f\"google/electra-{c.size}-generator\")\n",
    "\n",
    "# logger\n",
    "if c.logger == 'neptune':\n",
    "  import neptune\n",
    "  from fastai.callback.neptune import NeptuneCallback\n",
    "  neptune.init(project_qualified_name='richard-wang/electra-pretrain')\n",
    "elif c.logger == 'wandb':\n",
    "  import wandb\n",
    "  from fastai.callback.wandb import  WandbCallback\n",
    "\n",
    "# Path to data\n",
    "Path('./datasets', exist_ok=True)\n",
    "Path('./checkpoints/pretrain').mkdir(exist_ok=True, parents=True)\n",
    "edl_cache_dir = Path(\"./datasets/electra_dataloader\")\n",
    "edl_cache_dir.mkdir(exist_ok=True)\n",
    "\n",
    "# Print info\n",
    "print(f\"process id: {os.getpid()}\")\n",
    "print(c)\n",
    "print(hparam_update)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if c.my_model: # only for use of my personal research \n",
    "  sys.path.insert(0, os.path.abspath(\"..\"))\n",
    "  from modeling.model import ModelForGenerator,ModelForDiscriminator\n",
    "  from hyperparameter import electra_hparam_from_hf\n",
    "  gen_hparam = electra_hparam_from_hf(gen_config, hf_tokenizer)\n",
    "  gen_hparam.update(hparam_update)\n",
    "  disc_hparam = electra_hparam_from_hf(disc_config, hf_tokenizer)\n",
    "  disc_hparam.update(hparam_update)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dsets = []\n",
    "ELECTRAProcessor = partial(ELECTRADataProcessor, hf_tokenizer=hf_tokenizer, max_length=c.max_length)\n",
    "\n",
    "# Wikipedia\n",
    "if 'wikipedia' in c.datas:\n",
    "  print('load/download wiki dataset')\n",
    "  wiki = datasets.load_dataset('wikipedia', '20200501.en', cache_dir='./datasets')['train']\n",
    "  print('load/create data from wiki dataset for ELECTRA')\n",
    "  e_wiki = ELECTRAProcessor(wiki).map(cache_file_name=f\"electra_wiki_{c.max_length}.arrow\", num_proc=1)\n",
    "  dsets.append(e_wiki)\n",
    "\n",
    "# OpenWebText\n",
    "if 'openwebtext' in c.datas:\n",
    "  print('load/download OpenWebText Corpus')\n",
    "  owt = datasets.load_dataset('openwebtext', cache_dir='./datasets')['train']\n",
    "  print('load/create data from OpenWebText Corpus for ELECTRA')\n",
    "  e_owt = ELECTRAProcessor(owt, apply_cleaning=False).map(cache_file_name=f\"electra_owt_{c.max_length}.arrow\", num_proc=1)\n",
    "  dsets.append(e_owt)\n",
    "\n",
    "assert len(dsets) == len(c.datas)\n",
    "\n",
    "merged_dsets = {'train': datasets.concatenate_datasets(dsets)}\n",
    "hf_dsets = HF_Datasets(merged_dsets, cols={'input_ids':TensorText,'sentA_length':noop},\n",
    "                       hf_toker=hf_tokenizer, n_inp=2)\n",
    "dls = hf_dsets.dataloaders(bs=c.bs, num_workers=c.num_workers, pin_memory=False,\n",
    "                           shuffle_train=True,\n",
    "                           srtkey_fc=False, \n",
    "                           cache_dir='./datasets/electra_dataloader', cache_name='dl_{split}.json')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Masked language model objective"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.1 MLM objective callback"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Modified from HuggingFace/transformers (https://github.com/huggingface/transformers/blob/0a3d0e02c5af20bfe9091038c4fd11fb79175546/src/transformers/data/data_collator.py#L102). \n",
    "It is a little bit faster cuz \n",
    "- intead of a[b] a on gpu b on cpu, tensors here are all in the same device\n",
    "- don't iterate the tensor when create special tokens mask\n",
    "And\n",
    "- doesn't require huggingface tokenizer\n",
    "- cost you only 550 µs for a (128,128) tensor on gpu, so dynamic masking is cheap   \n",
    "\"\"\"\n",
    "def mask_tokens(inputs, mask_token_index, vocab_size, special_token_indices, mlm_probability=0.15, replace_prob=0.1, orginal_prob=0.1, ignore_index=-100):\n",
    "  \"\"\" \n",
    "  Prepare masked tokens inputs/labels for masked language modeling: (1-replace_prob-orginal_prob)% MASK, replace_prob% random, orginal_prob% original within mlm_probability% of tokens in the sentence. \n",
    "  * ignore_index in nn.CrossEntropy is default to -100, so you don't need to specify ignore_index in loss\n",
    "  \"\"\"\n",
    "  \n",
    "  device = inputs.device\n",
    "  labels = inputs.clone()\n",
    "  \n",
    "  # Get positions to apply mlm (mask/replace/not changed). (mlm_probability)\n",
    "  probability_matrix = torch.full(labels.shape, mlm_probability, device=device)\n",
    "  special_tokens_mask = torch.full(inputs.shape, False, dtype=torch.bool, device=device)\n",
    "  for sp_id in special_token_indices:\n",
    "    special_tokens_mask = special_tokens_mask | (inputs==sp_id)\n",
    "  probability_matrix.masked_fill_(special_tokens_mask, value=0.0)\n",
    "  mlm_mask = torch.bernoulli(probability_matrix).bool()\n",
    "  labels[~mlm_mask] = ignore_index  # We only compute loss on mlm applied tokens\n",
    "\n",
    "  # mask  (mlm_probability * (1-replace_prob-orginal_prob))\n",
    "  mask_prob = 1 - replace_prob - orginal_prob\n",
    "  mask_token_mask = torch.bernoulli(torch.full(labels.shape, mask_prob, device=device)).bool() & mlm_mask\n",
    "  inputs[mask_token_mask] = mask_token_index\n",
    "\n",
    "  # replace with a random token (mlm_probability * replace_prob)\n",
    "  if int(replace_prob)!=0:\n",
    "    rep_prob = replace_prob/(replace_prob + orginal_prob)\n",
    "    replace_token_mask = torch.bernoulli(torch.full(labels.shape, rep_prob, device=device)).bool() & mlm_mask & ~mask_token_mask\n",
    "    random_words = torch.randint(vocab_size, labels.shape, dtype=torch.long, device=device)\n",
    "    inputs[replace_token_mask] = random_words[replace_token_mask]\n",
    "\n",
    "  # do nothing (mlm_probability * orginal_prob)\n",
    "  pass\n",
    "\n",
    "  return inputs, labels, mlm_mask\n",
    "\n",
    "class MaskedLMCallback(Callback):\n",
    "  @delegates(mask_tokens)\n",
    "  def __init__(self, mask_tok_id, special_tok_ids, vocab_size, ignore_index=-100, for_electra=False, **kwargs):\n",
    "    self.ignore_index = ignore_index\n",
    "    self.for_electra = for_electra\n",
    "    self.mask_tokens = partial(mask_tokens,\n",
    "                               mask_token_index=mask_tok_id,\n",
    "                               special_token_indices=special_tok_ids,\n",
    "                               vocab_size=vocab_size,\n",
    "                               ignore_index=-100,\n",
    "                               **kwargs)\n",
    "\n",
    "  def before_batch(self):\n",
    "    input_ids, sentA_lenths  = self.xb\n",
    "    masked_inputs, labels, is_mlm_applied = self.mask_tokens(input_ids)\n",
    "    if self.for_electra:\n",
    "      self.learn.xb, self.learn.yb = (masked_inputs, sentA_lenths, is_mlm_applied, labels), (labels,)\n",
    "    else:\n",
    "      self.learn.xb, self.learn.yb = (masked_inputs, sentA_lenths), (labels,)\n",
    "\n",
    "  @delegates(TfmdDL.show_batch)\n",
    "  def show_batch(self, dl, idx_show_ignored, verbose=True, **kwargs):\n",
    "    b = dl.one_batch()\n",
    "    input_ids, sentA_lenths  = b\n",
    "    masked_inputs, labels, is_mlm_applied = self.mask_tokens(input_ids.clone())\n",
    "    # check\n",
    "    assert torch.equal(is_mlm_applied, labels!=self.ignore_index)\n",
    "    assert torch.equal((~is_mlm_applied *masked_inputs + is_mlm_applied * labels), input_ids)\n",
    "    # change symbol to show the ignored position\n",
    "    labels[labels==self.ignore_index] = idx_show_ignored\n",
    "    # some notice to help understand the masking mechanism\n",
    "    if verbose: \n",
    "      print(\"We won't count loss from position where y is ignore index\")\n",
    "      print(\"Notice 1. Positions have label token in y will be either [Mask]/other token/orginal token in x\")\n",
    "      print(\"Notice 2. Special tokens (CLS, SEP) won't be masked.\")\n",
    "      print(\"Notice 3. Dynamic masking: every time you run gives you different results.\")\n",
    "    # show\n",
    "    tfm_b =(masked_inputs, sentA_lenths, is_mlm_applied, labels) if self.for_electra else (masked_inputs, sentA_lenths, labels)   \n",
    "    dl.show_batch(b=tfm_b, **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "mlm_cb = MaskedLMCallback(mask_tok_id=hf_tokenizer.mask_token_id, \n",
    "                          special_tok_ids=hf_tokenizer.all_special_ids, \n",
    "                          vocab_size=hf_tokenizer.vocab_size,\n",
    "                          mlm_probability=c.mask_prob,\n",
    "                          replace_prob=0.0 if c.electra_mask_style else 0.1, \n",
    "                          orginal_prob=0.15 if c.electra_mask_style else 0.1,\n",
    "                          for_electra=True)\n",
    "#mlm_cb.show_batch(dls[0], idx_show_ignored=hf_tokenizer.convert_tokens_to_ids(['#'])[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. ELECTRA (replaced token detection objective)\n",
    "see details in paper [ELECTRA: Pre-training Text Encoders as Discriminators Rather Than Generators](https://arxiv.org/abs/2003.10555)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ELECTRAModel(nn.Module):\n",
    "  \n",
    "  def __init__(self, generator, discriminator, hf_tokenizer):\n",
    "    super().__init__()\n",
    "    self.generator, self.discriminator = generator,discriminator\n",
    "    self.gumbel_dist = torch.distributions.gumbel.Gumbel(0.,1.)\n",
    "    self.hf_tokenizer = hf_tokenizer\n",
    "\n",
    "  def to(self, *args, **kwargs):\n",
    "    \"Also set dtype and device of contained gumbel distribution if needed\"\n",
    "    super().to(*args, **kwargs)\n",
    "    a_tensor = next(self.parameters())\n",
    "    device, dtype = a_tensor.device, a_tensor.dtype\n",
    "    if c.sampling=='fp32_gumbel': dtype = torch.float32\n",
    "    self.gumbel_dist = torch.distributions.gumbel.Gumbel(torch.tensor(0., device=device, dtype=dtype), torch.tensor(1., device=device, dtype=dtype))\n",
    "\n",
    "  def forward(self, masked_inputs, sentA_lenths, is_mlm_applied, labels):\n",
    "    \"\"\"\n",
    "    masked_inputs (Tensor[int]): (B, L)\n",
    "    sentA_lenths (Tensor[int]): (B, L)\n",
    "    is_mlm_applied (Tensor[boolean]): (B, L), True for positions chosen by mlm probability \n",
    "    labels (Tensor[int]): (B, L), -100 for positions where are not mlm applied\n",
    "    \"\"\"\n",
    "    attention_mask, token_type_ids = self._get_pad_mask_and_token_type(masked_inputs, sentA_lenths)\n",
    "    if c.my_model:\n",
    "      gen_logits = self.generator(masked_inputs, attention_mask, token_type_ids, is_mlm_applied)[0]\n",
    "      # already reduced before the mlm output layer, save more space and speed\n",
    "      mlm_gen_logits = gen_logits # ( #mlm_positions, vocab_size)\n",
    "    else:\n",
    "      gen_logits = self.generator(masked_inputs, attention_mask, token_type_ids)[0] # (B, L, vocab size)\n",
    "      # reduce size to save space and speed\n",
    "      mlm_gen_logits = gen_logits[is_mlm_applied, :] # ( #mlm_positions, vocab_size)\n",
    "    \n",
    "    with torch.no_grad():\n",
    "      # sampling\n",
    "      pred_toks = self.sample(mlm_gen_logits) # ( #mlm_positions, )\n",
    "      # produce inputs for discriminator\n",
    "      generated = masked_inputs.clone() # (B,L)\n",
    "      generated[is_mlm_applied] = pred_toks # (B,L)\n",
    "      # produce labels for discriminator\n",
    "      is_replaced = is_mlm_applied.clone() # (B,L)\n",
    "      is_replaced[is_mlm_applied] = (pred_toks != labels[is_mlm_applied]) # (B,L)\n",
    "\n",
    "    disc_logits = self.discriminator(generated, attention_mask, token_type_ids)[0] # (B, L)\n",
    "\n",
    "    return mlm_gen_logits, generated, disc_logits, is_replaced, attention_mask, is_mlm_applied\n",
    "\n",
    "  def _get_pad_mask_and_token_type(self, input_ids, sentA_lenths):\n",
    "    \"\"\"\n",
    "    Only cost you about 500 µs for (128, 128) on GPU, but so that your dataset won't need to save attention_mask and token_type_ids and won't be unnecessarily large, thus, prevent cpu processes loading batches from consuming lots of cpu memory and slow down the machine. \n",
    "    \"\"\"\n",
    "    attention_mask = input_ids != self.hf_tokenizer.pad_token_id\n",
    "    seq_len = input_ids.shape[1]\n",
    "    token_type_ids = torch.tensor([ ([0]*len + [1]*(seq_len-len)) for len in sentA_lenths.tolist()],  \n",
    "                                  device=input_ids.device)\n",
    "    return attention_mask, token_type_ids\n",
    "\n",
    "  def sample(self, logits):\n",
    "    \"Reimplement gumbel softmax cuz there is a bug in torch.nn.functional.gumbel_softmax when fp16 (https://github.com/pytorch/pytorch/issues/41663). Gumbel softmax is equal to what official ELECTRA code do, standard gumbel dist. = -ln(-ln(standard uniform dist.))\"\n",
    "    if c.sampling == 'fp32_gumbel':\n",
    "      return (logits.float() + self.gumbel_dist.sample(logits.shape)).argmax(dim=-1)\n",
    "    elif c.sampling == 'fp16_gumbel': # 5.06 ms\n",
    "      return (logits + self.gumbel_dist.sample(logits.shape)).argmax(dim=-1)\n",
    "    elif c.sampling == 'multinomial': # 2.X ms\n",
    "      return torch.multinomial(F.softmax(logits, dim=-1), 1).squeeze()\n",
    "\n",
    "class ELECTRALoss():\n",
    "  def __init__(self, loss_weights=(1.0, 50.0), gen_label_smooth=False, disc_label_smooth=False):\n",
    "    self.loss_weights = loss_weights\n",
    "    self.gen_loss_fc = LabelSmoothingCrossEntropyFlat(eps=gen_label_smooth) if gen_label_smooth else CrossEntropyLossFlat()\n",
    "    self.disc_loss_fc = nn.BCEWithLogitsLoss()\n",
    "    self.disc_label_smooth = disc_label_smooth\n",
    "    \n",
    "  def __call__(self, pred, targ_ids):\n",
    "    mlm_gen_logits, generated, disc_logits, is_replaced, non_pad, is_mlm_applied = pred\n",
    "    gen_loss = self.gen_loss_fc(mlm_gen_logits.float(), targ_ids[is_mlm_applied])\n",
    "    disc_logits = disc_logits.masked_select(non_pad) # -> 1d tensor\n",
    "    is_replaced = is_replaced.masked_select(non_pad) # -> 1d tensor\n",
    "    if self.disc_label_smooth:\n",
    "      is_replaced = is_replaced.float().masked_fill(~is_replaced, self.disc_label_smooth)\n",
    "    disc_loss = self.disc_loss_fc(disc_logits.float(), is_replaced.float())\n",
    "    return gen_loss * self.loss_weights[0] + disc_loss * self.loss_weights[1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 5. Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Seed & PyTorch benchmark\n",
    "torch.backends.cudnn.benchmark = True\n",
    "dls[0].rng = random.Random(c.seed) # for fastai dataloader\n",
    "random.seed(c.seed)\n",
    "np.random.seed(c.seed)\n",
    "torch.manual_seed(c.seed)\n",
    "\n",
    "# Generator and Discriminator\n",
    "if c.my_model:\n",
    "  generator = ModelForGenerator(gen_hparam)\n",
    "  discriminator = ModelForDiscriminator(disc_hparam)\n",
    "  discriminator.electra.embedding = generator.electra.embedding\n",
    "  # implicitly tie in/out embeddings of generator\n",
    "else:\n",
    "  generator = ElectraForMaskedLM(gen_config)\n",
    "  discriminator = ElectraForPreTraining(disc_config)\n",
    "  discriminator.electra.embeddings = generator.electra.embeddings\n",
    "  generator.generator_lm_head.weight = generator.electra.embeddings.word_embeddings.weight\n",
    "\n",
    "# ELECTRA training loop\n",
    "electra_model = ELECTRAModel(generator, discriminator, hf_tokenizer)\n",
    "electra_loss_func = ELECTRALoss(gen_label_smooth=c.gen_smooth_label, disc_label_smooth=c.disc_smooth_label)\n",
    "\n",
    "# Optimizer\n",
    "if c.adam_bias_correction: opt_func = partial(Adam, eps=1e-6, mom=0.9, sqr_mom=0.999, wd=0.01)\n",
    "else: opt_func = partial(Adam_no_bias_correction, eps=1e-6, mom=0.9, sqr_mom=0.999, wd=0.01)\n",
    "\n",
    "# Learning rate shedule\n",
    "if c.schedule.endswith('linear'):\n",
    "  lr_shed_func = linear_warmup_and_then_decay if c.schedule=='separate_linear' else linear_warmup_and_decay\n",
    "  lr_shedule = ParamScheduler({'lr': partial(linear_warmup_and_decay,\n",
    "                                             lr_max=c.lr,\n",
    "                                             warmup_steps=10000,\n",
    "                                             total_steps=c.steps,)})\n",
    "\n",
    "\n",
    "# Learner\n",
    "dls.to(torch.device(c.device))\n",
    "learn = Learner(dls, electra_model,\n",
    "                loss_func=electra_loss_func,\n",
    "                opt_func=opt_func ,\n",
    "                path='./checkpoints',\n",
    "                model_dir='pretrain',\n",
    "                cbs=[mlm_cb,\n",
    "                    RunSteps(c.steps, [0.0625, 0.125, 0.25, 0.5, 1.0], c.run_name+\"_{percent}\"),\n",
    "                     ],\n",
    "                )\n",
    "\n",
    "# logging\n",
    "if c.logger == 'neptune':\n",
    "  neptune.create_experiment(name=c.run_name, params={**c, **hparam_update})\n",
    "  learn.add_cb(NeptuneCallback(log_model_weights=False))\n",
    "elif c.logger == 'wandb':\n",
    "  wandb.init(name=c.run_name, project='electra_pretrain', config={**c, **hparam_update})\n",
    "  learn.add_cb(WandbCallback(log_preds=False, log_model=False))\n",
    "\n",
    "# Mixed precison and Gradient clip\n",
    "learn.to_native_fp16(init_scale=2.**11)\n",
    "learn.add_cb(GradientClipping(1.))\n",
    "\n",
    "# Print time and run name\n",
    "print(f\"{c.run_name} , starts at {datetime.now()}\")\n",
    "\n",
    "# Run\n",
    "if c.schedule == 'one_cycle': learn.fit_one_cycle(9999, lr_max=c.lr)\n",
    "elif c.schedule == 'adjusted_one_cycle': learn.fit_one_cycle(9999, lr_max=c.lr, div=1e5, pct_start=10000/c.steps)\n",
    "else: learn.fit(9999, cbs=[lr_shedule])"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7-final"
  },
  "orig_nbformat": 2,
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}