{

 "cells": [

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "import sys \n",

    "sys.path.insert(0, '../')"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "from common.common import create_folder\n",

    "from common.pytorch import load_model\n",

    "import pytorch_pretrained_bert as Bert\n",

    "from model.utils import age_vocab\n",

    "from common.common import load_obj\n",

    "from dataLoader.MLM import MLMLoader\n",

    "from torch.utils.data import DataLoader\n",

    "import pandas as pd\n",

    "from model.MLM import BertForMaskedLM\n",

    "from model.optimiser import adam\n",

    "import sklearn.metrics as skm\n",

    "import numpy as np\n",

    "import torch\n",

    "import time\n",

    "import torch.nn as nn\n",

    "import os"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "class BertConfig(Bert.modeling.BertConfig):\n",

    "    def __init__(self, config):\n",

    "        super(BertConfig, self).__init__(\n",

    "            vocab_size_or_config_json_file=config.get('vocab_size'),\n",

    "            hidden_size=config['hidden_size'],\n",

    "            num_hidden_layers=config.get('num_hidden_layers'),\n",

    "            num_attention_heads=config.get('num_attention_heads'),\n",

    "            intermediate_size=config.get('intermediate_size'),\n",

    "            hidden_act=config.get('hidden_act'),\n",

    "            hidden_dropout_prob=config.get('hidden_dropout_prob'),\n",

    "            attention_probs_dropout_prob=config.get('attention_probs_dropout_prob'),\n",

    "            max_position_embeddings = config.get('max_position_embedding'),\n",

    "            initializer_range=config.get('initializer_range'),\n",

    "        )\n",

    "        self.seg_vocab_size = config.get('seg_vocab_size')\n",

    "        self.age_vocab_size = config.get('age_vocab_size')\n",

    "        \n",

    "class TrainConfig(object):\n",

    "    def __init__(self, config):\n",

    "        self.batch_size = config.get('batch_size')\n",

    "        self.use_cuda = config.get('use_cuda')\n",

    "        self.max_len_seq = config.get('max_len_seq')\n",

    "        self.train_loader_workers = config.get('train_loader_workers')\n",

    "        self.test_loader_workers = config.get('test_loader_workers')\n",

    "        self.device = config.get('device')\n",

    "        self.output_dir = config.get('output_dir')\n",

    "        self.output_name = config.get('output_name')\n",

    "        self.best_name = config.get('best_name')"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "file_config = {\n",

    "    'vocab':'',  # vocabulary idx2token, token2idx\n",

    "    'data': '',  # formated data \n",

    "    'model_path': '', # where to save model\n",

    "    'model_name': '', # model name\n",

    "    'file_name': '',  # log path\n",

    "}\n",

    "create_folder(file_config['model_path'])"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "global_params = {\n",

    "    'max_seq_len': 64,\n",

    "    'max_age': 110,\n",

    "    'month': 1,\n",

    "    'age_symbol': None,\n",

    "    'min_visit': 5,\n",

    "    'gradient_accumulation_steps': 1\n",

    "}\n",

    "\n",

    "optim_param = {\n",

    "    'lr': 3e-5,\n",

    "    'warmup_proportion': 0.1,\n",

    "    'weight_decay': 0.01\n",

    "}\n",

    "\n",

    "train_params = {\n",

    "    'batch_size': 256,\n",

    "    'use_cuda': True,\n",

    "    'max_len_seq': global_params['max_seq_len'],\n",

    "    'device': 'cuda:0'\n",

    "}"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "BertVocab = load_obj(file_config['vocab'])\n",

    "ageVocab, _ = age_vocab(max_age=global_params['max_age'], mon=global_params['month'], symbol=global_params['age_symbol'])"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "data = pd.read_parquet(file_config['data'])\n",

    "# remove patients with visits less than min visit\n",

    "data['length'] = data['caliber_id'].apply(lambda x: len([i for i in range(len(x)) if x[i] == 'SEP']))\n",

    "data = data[data['length'] >= global_params['min_visit']]\n",

    "data = data.reset_index(drop=True)"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "Dset = MLMLoader(data, BertVocab['token2idx'], ageVocab, max_len=train_params['max_len_seq'], code='caliber_id')\n",

    "trainload = DataLoader(dataset=Dset, batch_size=train_params['batch_size'], shuffle=True, num_workers=3)"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "model_config = {\n",

    "    'vocab_size': len(BertVocab['token2idx'].keys()), # number of disease + symbols for word embedding\n",

    "    'hidden_size': 288, # word embedding and seg embedding hidden size\n",

    "    'seg_vocab_size': 2, # number of vocab for seg embedding\n",

    "    'age_vocab_size': len(ageVocab.keys()), # number of vocab for age embedding\n",

    "    'max_position_embedding': train_params['max_len_seq'], # maximum number of tokens\n",

    "    'hidden_dropout_prob': 0.1, # dropout rate\n",

    "    'num_hidden_layers': 6, # number of multi-head attention layers required\n",

    "    'num_attention_heads': 12, # number of attention heads\n",

    "    'attention_probs_dropout_prob': 0.1, # multi-head attention dropout rate\n",

    "    'intermediate_size': 512, # the size of the \"intermediate\" layer in the transformer encoder\n",

    "    'hidden_act': 'gelu', # The non-linear activation function in the encoder and the pooler \"gelu\", 'relu', 'swish' are supported\n",

    "    'initializer_range': 0.02, # parameter weight initializer range\n",

    "}"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "conf = BertConfig(model_config)\n",

    "model = BertForMaskedLM(conf)"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "model = model.to(train_params['device'])\n",

    "optim = adam(params=list(model.named_parameters()), config=optim_param)"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "def cal_acc(label, pred):\n",

    "    logs = nn.LogSoftmax()\n",

    "    label=label.cpu().numpy()\n",

    "    ind = np.where(label!=-1)[0]\n",

    "    truepred = pred.detach().cpu().numpy()\n",

    "    truepred = truepred[ind]\n",

    "    truelabel = label[ind]\n",

    "    truepred = logs(torch.tensor(truepred))\n",

    "    outs = [np.argmax(pred_x) for pred_x in truepred.numpy()]\n",

    "    precision = skm.precision_score(truelabel, outs, average='micro')\n",

    "    return precision"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "def train(e, loader):\n",

    "    tr_loss = 0\n",

    "    temp_loss = 0\n",

    "    nb_tr_examples, nb_tr_steps = 0, 0\n",

    "    cnt= 0\n",

    "    start = time.time()\n",

    "\n",

    "    for step, batch in enumerate(loader):\n",

    "        cnt +=1\n",

    "        batch = tuple(t.to(train_params['device']) for t in batch)\n",

    "        age_ids, input_ids, posi_ids, segment_ids, attMask, masked_label = batch\n",

    "        loss, pred, label = model(input_ids, age_ids, segment_ids, posi_ids,attention_mask=attMask, masked_lm_labels=masked_label)\n",

    "        if global_params['gradient_accumulation_steps'] >1:\n",

    "            loss = loss/global_params['gradient_accumulation_steps']\n",

    "        loss.backward()\n",

    "        \n",

    "        temp_loss += loss.item()\n",

    "        tr_loss += loss.item()\n",

    "        \n",

    "        nb_tr_examples += input_ids.size(0)\n",

    "        nb_tr_steps += 1\n",

    "        \n",

    "        if step % 200==0:\n",

    "            print(\"epoch: {}\\t| cnt: {}\\t|Loss: {}\\t| precision: {:.4f}\\t| time: {:.2f}\".format(e, cnt, temp_loss/2000, cal_acc(label, pred), time.time()-start))\n",

    "            temp_loss = 0\n",

    "            start = time.time()\n",

    "            \n",

    "        if (step + 1) % global_params['gradient_accumulation_steps'] == 0:\n",

    "            optim.step()\n",

    "            optim.zero_grad()\n",

    "\n",

    "    print(\"** ** * Saving fine - tuned model ** ** * \")\n",

    "    model_to_save = model.module if hasattr(model, 'module') else model  # Only save the model it-self\n",

    "    create_folder(file_config['model_path'])\n",

    "    output_model_file = os.path.join(file_config['model_path'], file_config['model_name'])\n",

    "\n",

    "    torch.save(model_to_save.state_dict(), output_model_file)\n",

    "        \n",

    "    cost = time.time() - start\n",

    "    return tr_loss, cost"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "f = open(os.path.join(file_config['model_path'], file_config['file_name']), \"w\")\n",

    "f.write('{}\\t{}\\t{}\\n'.format('epoch', 'loss', 'time'))\n",

    "for e in range(50):\n",

    "    loss, time_cost = train(e, trainload)\n",

    "    loss = loss/data_len\n",

    "    f.write('{}\\t{}\\t{}\\n'.format(e, loss, time_cost))\n",

    "f.close()    "

   ]

  }

 ],

 "metadata": {

  "kernelspec": {

   "display_name": "Python 3",

   "language": "python",

   "name": "python3"

  },

  "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.4"

  }

 },

 "nbformat": 4,

 "nbformat_minor": 2

}