{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# About this notebook\n\nThis is fully based on the work below and adding torch_xla to it to trian on TPU's\n- [Luke](https://arxiv.org/pdf/2010.01057v1.pdf)-base starter notebook https://www.kaggle.com/yasufuminakama/jigsaw4-luke-base-starter-train from nakama\n- Approach References\n    - https://www.kaggle.com/tanlikesmath/xlm-roberta-pytorch-xla-tpu#Training\n    - Thanks for sharing nakama ,@debarshichanda and @nbroad \n    \n Also just ran this in **debug** mode presently just to do a quick test . ","metadata":{}},{"cell_type":"code","source":"# for TPU\n!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n!python pytorch-xla-env-setup.py --apt-packages libomp5 libopenblas-dev","metadata":{"execution":{"iopub.status.busy":"2021-11-22T02:07:16.690972Z","iopub.execute_input":"2021-11-22T02:07:16.691363Z","iopub.status.idle":"2021-11-22T02:08:24.881131Z","shell.execute_reply.started":"2021-11-22T02:07:16.691266Z","shell.execute_reply":"2021-11-22T02:08:24.880009Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Directory settings","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# Directory settings\n# ====================================================\nimport os\n\nOUTPUT_DIR = './'\nif not os.path.exists(OUTPUT_DIR):\n    os.makedirs(OUTPUT_DIR)","metadata":{"execution":{"iopub.status.busy":"2021-11-22T02:08:24.885109Z","iopub.execute_input":"2021-11-22T02:08:24.885832Z","iopub.status.idle":"2021-11-22T02:08:24.893756Z","shell.execute_reply.started":"2021-11-22T02:08:24.885770Z","shell.execute_reply":"2021-11-22T02:08:24.892812Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# CFG","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# CFG\n# ====================================================\nclass CFG:\n    competition='Jigsaw4'\n    _wandb_kernel='gauravbrills'\n    debug=True\n    apex=True\n    print_freq=50\n    num_workers=8\n    model=\"studio-ousia/luke-base\"\n    scheduler='cosine' # ['linear', 'cosine']\n    batch_scheduler=True\n    num_cycles=0.5\n    num_warmup_steps=0\n    epochs=3\n    encoder_lr=1e-5\n    decoder_lr=1e-5\n    min_lr=1e-6\n    eps=1e-6\n    betas=(0.9, 0.999)\n    batch_size=64 #64\n    fc_dropout=0.\n    text=\"text\"\n    target=\"target\"\n    target_size=1\n    head=32\n    tail=32\n    max_len=head+tail\n    weight_decay=0.01\n    gradient_accumulation_steps=1\n    max_grad_norm=1000\n    margin=0.5\n    seed=2021\n    n_fold=5\n    trn_fold=[0, 1, 2, 3, 4]\n    train=True","metadata":{"execution":{"iopub.status.busy":"2021-11-22T02:08:24.895747Z","iopub.execute_input":"2021-11-22T02:08:24.896117Z","iopub.status.idle":"2021-11-22T02:08:24.906420Z","shell.execute_reply.started":"2021-11-22T02:08:24.896075Z","shell.execute_reply":"2021-11-22T02:08:24.905329Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Library","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# Library\n# ====================================================\nimport os\nimport gc\nimport re\nimport sys\nimport json\nimport time\nimport math\nimport string\nimport pickle\nimport random\nimport joblib\nimport itertools\nimport warnings\nwarnings.filterwarnings(\"ignore\")\n\nimport scipy as sp\nimport numpy as np\nimport pandas as pd\npd.set_option('display.max_rows', 500)\npd.set_option('display.max_columns', 500)\npd.set_option('display.width', 1000)\nfrom tqdm.auto import tqdm\nfrom sklearn.metrics import roc_auc_score\nfrom sklearn.model_selection import StratifiedKFold, GroupKFold, KFold\n\nimport torch\nimport torch.nn as nn\nfrom torch.nn import Parameter\nimport torch.nn.functional as F\nfrom torch.optim import Adam, SGD, AdamW\nfrom torch.utils.data import DataLoader, Dataset\n\nos.system('pip uninstall -q transformers -y')\nos.system('pip uninstall -q tokenizers -y')\nos.system('pip uninstall -q huggingface_hub -y')\n\nos.system('mkdir -p /tmp/pip/cache-tokenizers/')\nos.system('cp ../input/tokenizers-0103/tokenizers-0.10.3-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl /tmp/pip/cache-tokenizers/')\nos.system('pip install -q --no-index --find-links /tmp/pip/cache-tokenizers/ tokenizers')\n\nos.system('mkdir -p /tmp/pip/cache-huggingface-hub/')\nos.system('cp ../input/huggingface-hub-008/huggingface_hub-0.0.8-py3-none-any.whl /tmp/pip/cache-huggingface-hub/')\nos.system('pip install -q --no-index --find-links /tmp/pip/cache-huggingface-hub/ huggingface_hub')\n\nos.system('mkdir -p /tmp/pip/cache-transformers/')\nos.system('cp ../input/transformers-470/transformers-4.7.0-py3-none-any.whl /tmp/pip/cache-transformers/')\nos.system('pip install -q --no-index --find-links /tmp/pip/cache-transformers/ transformers')\n\nimport tokenizers\nimport transformers\nprint(f\"tokenizers.__version__: {tokenizers.__version__}\")\nprint(f\"transformers.__version__: {transformers.__version__}\")\nfrom transformers import LukeTokenizer, LukeModel, LukeConfig\nfrom transformers import get_linear_schedule_with_warmup, get_cosine_schedule_with_warmup\n# for TPU\nos.environ['XLA_USE_BF16']=\"1\"\nos.environ['XLA_TENSOR_ALLOCATOR_MAXSIZE'] = '100000000'\nimport torch_xla\nimport torch_xla.debug.metrics as met\nimport torch_xla.distributed.parallel_loader as pl\nimport torch_xla.utils.utils as xu\nimport torch_xla.core.xla_model as xm\nimport torch_xla.distributed.xla_multiprocessing as xmp\nimport torch_xla.test.test_utils as test_utils\n\n#device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n# for TPU\n\n#torch.set_default_tensor_type('torch.FloatTensor')\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")","metadata":{"execution":{"iopub.status.busy":"2021-11-22T02:08:24.909511Z","iopub.execute_input":"2021-11-22T02:08:24.910101Z","iopub.status.idle":"2021-11-22T02:09:04.909351Z","shell.execute_reply.started":"2021-11-22T02:08:24.910048Z","shell.execute_reply":"2021-11-22T02:09:04.908237Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import tensorflow as tf\n\ntry:\n    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\n    print(\"Running on TPU \", tpu.cluster_spec().as_dict()[\"worker\"])\nexcept ValueError:\n    tpu = None\nif tpu:\n    tf.config.experimental_connect_to_cluster(tpu)\n    tf.tpu.experimental.initialize_tpu_system(tpu)\n    strategy = tf.distribute.experimental.TPUStrategy(tpu)\nelse:\n    strategy = tf.distribute.get_strategy()","metadata":{"execution":{"iopub.status.busy":"2021-11-22T02:09:04.911036Z","iopub.execute_input":"2021-11-22T02:09:04.911303Z","iopub.status.idle":"2021-11-22T02:09:13.657003Z","shell.execute_reply.started":"2021-11-22T02:09:04.911267Z","shell.execute_reply":"2021-11-22T02:09:13.655903Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch_xla.version as xv\nprint('PYTORCH:', xv.__torch_gitrev__)\nprint('XLA:', xv.__xla_gitrev__)","metadata":{"execution":{"iopub.status.busy":"2021-11-22T02:09:13.659084Z","iopub.execute_input":"2021-11-22T02:09:13.659583Z","iopub.status.idle":"2021-11-22T02:09:13.667270Z","shell.execute_reply.started":"2021-11-22T02:09:13.659522Z","shell.execute_reply":"2021-11-22T02:09:13.666256Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Utils","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# Utils\n# ====================================================\ndef get_score(df):\n    score = len(df[df['less_toxic_pred'] < df['more_toxic_pred']]) / len(df)\n    return score\n\n\ndef get_logger(filename=OUTPUT_DIR+'train'):\n    from logging import getLogger, INFO, StreamHandler, FileHandler, Formatter\n    logger = getLogger(__name__)\n    logger.setLevel(INFO)\n    handler1 = StreamHandler()\n    handler1.setFormatter(Formatter(\"%(message)s\"))\n    handler2 = FileHandler(filename=f\"{filename}.log\")\n    handler2.setFormatter(Formatter(\"%(message)s\"))\n    logger.addHandler(handler1)\n    logger.addHandler(handler2)\n    return logger\n\nLOGGER = get_logger()\n\ndef seed_everything(seed=42):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n    \nseed_everything(seed=42)","metadata":{"execution":{"iopub.status.busy":"2021-11-22T02:09:13.669073Z","iopub.execute_input":"2021-11-22T02:09:13.669403Z","iopub.status.idle":"2021-11-22T02:09:13.715380Z","shell.execute_reply.started":"2021-11-22T02:09:13.669362Z","shell.execute_reply":"2021-11-22T02:09:13.714091Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Data Loading","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# Data Loading\n# ====================================================\ntrain = pd.read_csv('../input/jigsaw-toxic-severity-rating/validation_data.csv')\nif CFG.debug:\n    train = train.sample(n=1000, random_state=CFG.seed).reset_index(drop=True)\ntest = pd.read_csv('../input/jigsaw-toxic-severity-rating/comments_to_score.csv')\nsubmission = pd.read_csv('../input/jigsaw-toxic-severity-rating/sample_submission.csv')\nprint(train.shape)\nprint(test.shape, submission.shape)\ndisplay(train.head())\ndisplay(test.head())\ndisplay(submission.head())","metadata":{"execution":{"iopub.status.busy":"2021-11-22T02:09:13.718381Z","iopub.execute_input":"2021-11-22T02:09:13.718792Z","iopub.status.idle":"2021-11-22T02:09:14.530642Z","shell.execute_reply.started":"2021-11-22T02:09:13.718738Z","shell.execute_reply":"2021-11-22T02:09:14.530041Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# CV split","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# CV split\n# ====================================================\n#Fold = GroupKFold(n_splits=CFG.n_fold)\nFold = StratifiedKFold(n_splits=CFG.n_fold, shuffle=True, random_state=CFG.seed)\n\n#for n, (trn_index, val_index) in enumerate(Fold.split(train, train, train['worker'])):\nfor n, ( trn_index, val_index) in enumerate(Fold.split(X=train, y=train.worker)):\n    train.loc[val_index, 'fold'] = int(n)\ntrain['fold'] = train['fold'].astype(int)\ndisplay(train.groupby('fold').size())","metadata":{"execution":{"iopub.status.busy":"2021-11-22T02:09:14.531729Z","iopub.execute_input":"2021-11-22T02:09:14.532182Z","iopub.status.idle":"2021-11-22T02:09:14.564412Z","shell.execute_reply.started":"2021-11-22T02:09:14.532139Z","shell.execute_reply":"2021-11-22T02:09:14.563756Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# tokenizer","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# tokenizer\n# ====================================================\ntokenizer = LukeTokenizer.from_pretrained(CFG.model, lowercase=True)\ntokenizer.save_pretrained(OUTPUT_DIR+'tokenizer/')\nCFG.tokenizer = tokenizer","metadata":{"execution":{"iopub.status.busy":"2021-11-22T02:09:14.567111Z","iopub.execute_input":"2021-11-22T02:09:14.567790Z","iopub.status.idle":"2021-11-22T02:09:21.415292Z","shell.execute_reply.started":"2021-11-22T02:09:14.567752Z","shell.execute_reply":"2021-11-22T02:09:21.414538Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# Dataset\n# ====================================================\ndef prepare_input(text, cfg):\n    if cfg.tail == 0:\n        inputs = cfg.tokenizer.encode_plus(text, \n                                           return_tensors=None, \n                                           add_special_tokens=True, \n                                           max_length=cfg.max_len,\n                                           pad_to_max_length=True,\n                                           truncation=True)\n        for k, v in inputs.items():\n            inputs[k] = torch.tensor(v, dtype=torch.long)\n    else:\n        inputs = cfg.tokenizer.encode_plus(text,\n                                           return_tensors=None, \n                                           add_special_tokens=True, \n                                           truncation=True)\n        for k, v in inputs.items():\n            v_length = len(v)\n            if v_length > cfg.max_len:\n                v = np.hstack([v[:cfg.head], v[-cfg.tail:]])\n            if k == 'input_ids':\n                new_v = np.ones(cfg.max_len) * cfg.tokenizer.pad_token_id\n            else:\n                new_v = np.zeros(cfg.max_len)\n            new_v[:v_length] = v \n            inputs[k] = torch.tensor(new_v, dtype=torch.long)\n    return inputs\n\n\nclass TrainDataset(Dataset):\n    def __init__(self, cfg, df):\n        self.cfg = cfg\n        self.less_toxic = df['less_toxic'].fillna(\"none\").values\n        self.more_toxic = df['more_toxic'].fillna(\"none\").values\n\n    def __len__(self):\n        return len(self.less_toxic)\n\n    def __getitem__(self, item):\n        less_toxic_inputs = prepare_input(str(self.less_toxic[item]), self.cfg)\n        more_toxic_inputs = prepare_input(str(self.more_toxic[item]), self.cfg)\n        label = torch.tensor(1, dtype=torch.float)\n        return less_toxic_inputs, more_toxic_inputs, label\n\n\nclass TestDataset(Dataset):\n    def __init__(self, cfg, df):\n        self.cfg = cfg\n        self.text = df[cfg.text].fillna(\"none\").values\n\n    def __len__(self):\n        return len(self.text)\n\n    def __getitem__(self, item):\n        text = str(self.text[item])\n        inputs = prepare_input(text, self.cfg)\n        return inputs","metadata":{"execution":{"iopub.status.busy":"2021-11-22T02:09:21.416390Z","iopub.execute_input":"2021-11-22T02:09:21.417036Z","iopub.status.idle":"2021-11-22T02:09:21.433191Z","shell.execute_reply.started":"2021-11-22T02:09:21.417000Z","shell.execute_reply":"2021-11-22T02:09:21.432072Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# Model\n# ====================================================\nclass CustomModel(nn.Module):\n    def __init__(self, cfg, config_path=None, pretrained=False):\n        super().__init__()\n        self.cfg = cfg\n        if config_path is None:\n            self.config = LukeConfig.from_pretrained(cfg.model, output_hidden_states=True)\n        else:\n            self.config = torch.load(config_path)\n        if pretrained:\n            self.model = LukeModel.from_pretrained(cfg.model, config=self.config)\n        else:\n            self.model = LukeModel(self.config)\n        self.fc_dropout = nn.Dropout(cfg.fc_dropout)\n        self.fc = nn.Linear(self.config.hidden_size, cfg.target_size)\n        \n    def feature(self, inputs):\n        outputs = self.model(**inputs)\n        last_hidden_states = outputs[0]\n        feature = torch.mean(last_hidden_states, 1)\n        return feature\n\n    def forward(self, inputs):\n        feature = self.feature(inputs)\n        output = self.fc(self.fc_dropout(feature))\n        return output","metadata":{"execution":{"iopub.status.busy":"2021-11-22T02:09:21.434397Z","iopub.execute_input":"2021-11-22T02:09:21.434648Z","iopub.status.idle":"2021-11-22T02:09:21.451541Z","shell.execute_reply.started":"2021-11-22T02:09:21.434619Z","shell.execute_reply":"2021-11-22T02:09:21.450539Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Helper functions","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# Helper functions\n# ====================================================\nclass AverageMeter(object):\n    \"\"\"Computes and stores the average and current value\"\"\"\n    def __init__(self):\n        self.reset()\n\n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n\n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count\n\n\ndef asMinutes(s):\n    m = math.floor(s / 60)\n    s -= m * 60\n    return '%dm %ds' % (m, s)\n\n\ndef timeSince(since, percent):\n    now = time.time()\n    s = now - since\n    es = s / (percent)\n    rs = es - s\n    return '%s (remain %s)' % (asMinutes(s), asMinutes(rs))\n\n\ndef train_fn(fold, train_loader, model, criterion, optimizer, epoch, scheduler, device):\n    import torch_xla\n    model.train()\n    scaler = torch.cuda.amp.GradScaler(enabled=CFG.apex)\n    losses = AverageMeter()\n    start = end = time.time()\n    global_step = 0\n    for step, (less_toxic_inputs, more_toxic_inputs, labels) in enumerate(train_loader):\n        for k, v in less_toxic_inputs.items():\n            less_toxic_inputs[k] = v.to(device)\n        for k, v in more_toxic_inputs.items():\n            more_toxic_inputs[k] = v.to(device)\n        labels = labels.to(device)\n        batch_size = labels.size(0)\n        with torch.cuda.amp.autocast(enabled=CFG.apex):\n            less_toxic_y_preds = model(less_toxic_inputs)\n            more_toxic_y_preds = model(more_toxic_inputs)\n            loss = criterion(more_toxic_y_preds, less_toxic_y_preds, labels)\n        losses.update(loss.item(), batch_size)\n        if CFG.gradient_accumulation_steps > 1:\n            loss = loss / CFG.gradient_accumulation_steps\n        scaler.scale(loss).backward()\n        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), CFG.max_grad_norm)\n        if (step + 1) % CFG.gradient_accumulation_steps == 0:\n            scaler.step(optimizer)\n            scaler.update()\n            optimizer.zero_grad()\n            global_step += 1\n            if CFG.batch_scheduler:\n                scheduler.step()\n        end = time.time()\n        if step % CFG.print_freq == 0 or step == (len(train_loader)-1):\n            print('Epoch: [{0}][{1}/{2}] '\n                  'Elapsed {remain:s} '\n                  'Loss: {loss.val:.4f}({loss.avg:.4f}) '\n                  'Grad: {grad_norm:.4f}  '\n                  'LR: {lr:.8f}  '\n                  .format(epoch+1, step, len(train_loader), \n                          remain=timeSince(start, float(step+1)/len(train_loader)),\n                          loss=losses,\n                          grad_norm=grad_norm,\n                          lr=scheduler.get_lr()[0]))\n        #print({f\"[fold{fold}] loss\": losses.val,  f\"[fold{fold}] lr\": scheduler.get_lr()[0]})\n    return losses.avg\n\n\ndef inference_fn(test_loader, model, device):\n    preds = []\n    model.eval()\n    model.to(device)\n    tk0 = tqdm(test_loader, total=len(test_loader))\n    for inputs in tk0:\n        for k, v in inputs.items():\n            inputs[k] = v.to(device, dtype=torch.long)\n        with torch.no_grad():\n            y_preds = model(inputs)\n        preds.extend(y_preds.sigmoid().cpu().detach().numpy().tolist())\n    print(\"preds \",len(preds) )   \n    #predictions = np.concatenate(preds)\n    return preds","metadata":{"execution":{"iopub.status.busy":"2021-11-22T02:09:21.452993Z","iopub.execute_input":"2021-11-22T02:09:21.453373Z","iopub.status.idle":"2021-11-22T02:09:21.477201Z","shell.execute_reply.started":"2021-11-22T02:09:21.453343Z","shell.execute_reply":"2021-11-22T02:09:21.476048Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# ====================================================\n# train loop\n# ====================================================\ndef train_loop(folds, fold):\n    device = xm.xla_device()\n    \n    LOGGER.info(f\"========== fold: {fold} training ==========\")\n\n    # ====================================================\n    # loader\n    # ====================================================\n    \n    trn_idx = folds[folds['fold'] != fold].index\n    val_idx = folds[folds['fold'] == fold].index\n    \n    train_folds = folds.loc[trn_idx].reset_index(drop=True)\n    validation = folds.loc[val_idx].reset_index(drop=True)\n    \n    valid_folds = sorted(set(validation['less_toxic'].unique()) | set(validation['more_toxic'].unique()))\n    valid_folds = pd.DataFrame({'text': valid_folds}).reset_index()\n    \n    train_dataset = TrainDataset(CFG, train_folds)\n    valid_dataset = TestDataset(CFG, valid_folds)\n    train_sampler = torch.utils.data.distributed.DistributedSampler(\n              train_dataset,\n              num_replicas=xm.xrt_world_size(),\n              rank=xm.get_ordinal(),\n              shuffle=True)\n    train_loader = DataLoader(train_dataset,\n                              sampler= train_sampler,\n                              batch_size=CFG.batch_size,\n                              #shuffle=True,\n                              num_workers=0,#CFG.num_workers,\n                              pin_memory=True, drop_last=True)\n    gc.collect()\n    xm.master_print('parallel loader created... training now')\n    \n    valid_sampler = torch.utils.data.distributed.DistributedSampler(\n      valid_dataset,\n      num_replicas=xm.xrt_world_size(),\n      rank=xm.get_ordinal(),\n      shuffle=False)\n    valid_loader = DataLoader(valid_dataset,\n                              sampler= valid_sampler,\n                              batch_size=CFG.batch_size,\n                              #shuffle=False,\n                              num_workers=0,#CFG.num_workers, \n                              pin_memory=True, drop_last=False)\n    \n    \n    gc.collect()\n    # ====================================================\n    # model & optimizer\n    # ====================================================\n    model = CustomModel(CFG, config_path=None, pretrained=True)\n    xm.save(model.config, OUTPUT_DIR+'config.pth') # torch.save\n    model.to(device)\n    \n    def get_optimizer_params(model, encoder_lr, decoder_lr, weight_decay=0.0):\n        param_optimizer = list(model.named_parameters())\n        no_decay = [\"bias\", \"LayerNorm.bias\", \"LayerNorm.weight\"]\n        optimizer_parameters = [\n            {'params': [p for n, p in model.model.named_parameters() if not any(nd in n for nd in no_decay)],\n             'lr': encoder_lr, 'weight_decay': weight_decay},\n            {'params': [p for n, p in model.model.named_parameters() if any(nd in n for nd in no_decay)],\n             'lr': encoder_lr, 'weight_decay': 0.0},\n            {'params': [p for n, p in model.named_parameters() if \"model\" not in n],\n             'lr': decoder_lr, 'weight_decay': 0.0}\n        ]\n        return optimizer_parameters\n\n    optimizer_parameters = get_optimizer_params(model,\n                                                encoder_lr=CFG.encoder_lr, \n                                                decoder_lr=CFG.decoder_lr,\n                                                weight_decay=CFG.weight_decay)\n    optimizer = AdamW(optimizer_parameters, lr=CFG.encoder_lr, eps=CFG.eps, betas=CFG.betas)\n    \n    # ====================================================\n    # scheduler\n    # ====================================================\n    def get_scheduler(cfg, optimizer, num_train_steps):\n        if cfg.scheduler=='linear':\n            scheduler = get_linear_schedule_with_warmup(\n                optimizer, num_warmup_steps=cfg.num_warmup_steps, num_training_steps=num_train_steps\n            )\n        elif cfg.scheduler=='cosine':\n            scheduler = get_cosine_schedule_with_warmup(\n                optimizer, num_warmup_steps=cfg.num_warmup_steps, num_training_steps=num_train_steps, num_cycles=cfg.num_cycles\n            )\n        return scheduler\n    \n    num_train_steps = int(len(train_folds) / CFG.batch_size * CFG.epochs)\n    scheduler = get_scheduler(CFG, optimizer, num_train_steps)\n\n    # ====================================================\n    # loop\n    # ====================================================\n    criterion = nn.MarginRankingLoss(margin=CFG.margin)\n    \n    best_score = 0.\n\n    for epoch in range(CFG.epochs):\n\n        start_time = time.time()\n        LOGGER.info(\"epoch \",epoch)\n        para_loader = pl.ParallelLoader(train_loader, [device])\n\n        # train\n        avg_loss = train_fn(fold, para_loader.per_device_loader(device), model, \n                            criterion, optimizer, epoch, scheduler, device)\n        \n        del para_loader\n        para_loader = pl.ParallelLoader(valid_loader, [device])\n        \n        # eval\n        preds = inference_fn(para_loader.per_device_loader(device), model, device)\n        del para_loader\n        gc.collect()\n        \n        # scoring\n        valid_folds['pred'] = preds\n        if 'less_toxic_pred' in validation.columns:\n            validation = validation.drop(columns='less_toxic_pred')\n        if 'more_toxic_pred' in validation.columns:\n            validation = validation.drop(columns='more_toxic_pred')\n        rename_cols = {CFG.text: 'less_toxic', 'pred': 'less_toxic_pred'}\n        validation = validation.merge(valid_folds[[CFG.text, 'pred']].rename(columns=rename_cols), \n                                      on='less_toxic', how='left')\n        rename_cols = {CFG.text: 'more_toxic', 'pred': 'more_toxic_pred'}\n        validation = validation.merge(valid_folds[[CFG.text, 'pred']].rename(columns=rename_cols), \n                                      on='more_toxic', how='left')\n        score = get_score(validation)\n\n        elapsed = time.time() - start_time\n\n        LOGGER.info(f'Epoch {epoch+1} - avg_train_loss: {avg_loss:.4f}  time: {elapsed:.0f}s')\n        LOGGER.info(f'Epoch {epoch+1} - Score: {score:.4f}') \n        if score > best_score:\n            best_score = score\n            LOGGER.info(f'Epoch {epoch+1} - Save Best Score: {score:.4f} Model')\n            xm.save({'model': model.state_dict(), #torch.save\n                        'preds': preds},\n                        OUTPUT_DIR+f\"{CFG.model.replace('/', '-')}_fold{fold}_best.pth\")\n\n    preds = torch.load(OUTPUT_DIR+f\"{CFG.model.replace('/', '-')}_fold{fold}_best.pth\", #torch.load\n                       map_location=torch.device('cpu'))['preds']\n    valid_folds['pred'] = preds\n    if 'less_toxic_pred' in validation.columns:\n        validation = validation.drop(columns='less_toxic_pred')\n    if 'more_toxic_pred' in validation.columns:\n        validation = validation.drop(columns='more_toxic_pred')\n    rename_cols = {CFG.text: 'less_toxic', 'pred': 'less_toxic_pred'}\n    validation = validation.merge(valid_folds[[CFG.text, 'pred']].rename(columns=rename_cols), \n                                  on='less_toxic', how='left')\n    rename_cols = {CFG.text: 'more_toxic', 'pred': 'more_toxic_pred'}\n    validation = validation.merge(valid_folds[[CFG.text, 'pred']].rename(columns=rename_cols), \n                                  on='more_toxic', how='left')\n\n    #torch.cuda.empty_cache()\n    gc.collect()\n    \n    return validation","metadata":{"execution":{"iopub.status.busy":"2021-11-22T02:09:21.480219Z","iopub.execute_input":"2021-11-22T02:09:21.480671Z","iopub.status.idle":"2021-11-22T02:09:21.519007Z","shell.execute_reply.started":"2021-11-22T02:09:21.480624Z","shell.execute_reply":"2021-11-22T02:09:21.518219Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def _run():  \n    def get_result(oof_df):\n        score = get_score(oof_df)\n        LOGGER.info(f'Score: {score:<.4f}')\n    \n    if CFG.train:\n        # train \n        oof_df = pd.DataFrame()\n        for fold in range(CFG.n_fold):\n            if fold in CFG.trn_fold:\n                _oof_df = train_loop(train, fold) \n                oof_df = pd.concat([oof_df, _oof_df])\n                LOGGER.info(f\"========== fold: {fold} result ==========\")\n                get_result(_oof_df)\n        oof_df = oof_df.reset_index(drop=True)\n        # CV result\n        LOGGER.info(f\"========== CV ==========\")\n        get_result(oof_df)\n        # save result\n        oof_df.to_csv(OUTPUT_DIR+'oof_df.csv', index=False)\n    \n    #wandb.finish()  \n","metadata":{"execution":{"iopub.status.busy":"2021-11-22T02:21:29.528102Z","iopub.execute_input":"2021-11-22T02:21:29.529560Z","iopub.status.idle":"2021-11-22T02:21:29.537702Z","shell.execute_reply.started":"2021-11-22T02:21:29.529514Z","shell.execute_reply":"2021-11-22T02:21:29.537045Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import time\n\n# Start training processes\ndef _mp_fn(rank, flags):\n    torch.set_default_tensor_type('torch.FloatTensor')\n    a = _run()\n\nFLAGS={}\nstart_time = time.time()\nif __name__ == '__main__':\n    xmp.spawn(_mp_fn, args=(FLAGS,), nprocs=1, start_method='fork')","metadata":{"execution":{"iopub.status.busy":"2021-11-22T02:21:32.299880Z","iopub.execute_input":"2021-11-22T02:21:32.300201Z","iopub.status.idle":"2021-11-22T02:26:50.805757Z","shell.execute_reply.started":"2021-11-22T02:21:32.300170Z","shell.execute_reply":"2021-11-22T02:26:50.804464Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print('Time taken: ',time.time()-start_time)","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2021-11-22T02:28:54.211908Z","iopub.execute_input":"2021-11-22T02:28:54.213036Z","iopub.status.idle":"2021-11-22T02:28:54.221696Z","shell.execute_reply.started":"2021-11-22T02:28:54.212984Z","shell.execute_reply":"2021-11-22T02:28:54.220700Z"},"trusted":true},"execution_count":null,"outputs":[]}]}