{"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":"# 🤗 Bert for Question Answering Baseline: Training\n\nThis code is adapted from my work in the Tweet Sentiment Extraction Competition.\n\nIt tackles the task as a Question Answering one, where the question is implicit and can be understood as : \"Which datasets are mentionned ?\"\n\n\nThe approach is quite naïve and has a lot of flaws. Feel free to ask any question in the comments.\n\nInference Kernel : https://www.kaggle.com/theoviel/bert-for-question-answering-baseline-inference","metadata":{}},{"cell_type":"markdown","source":"# Initialization","metadata":{}},{"cell_type":"markdown","source":"## Imports","metadata":{}},{"cell_type":"code","source":"import re\nimport os\nimport gc\nimport glob\nimport json\nimport time\nimport torch\nimport random\nimport datetime\nimport tokenizers\nimport numpy as np\nimport transformers\nimport pandas as pd\nimport torch.nn as nn\nimport seaborn as sns\nimport matplotlib.pyplot as plt\n\nfrom tokenizers import *\nfrom transformers import *\nfrom functools import partial\nfrom tqdm.notebook import tqdm\nfrom torch.nn import functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom sklearn.model_selection import GroupKFold","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-06-04T11:17:37.104804Z","iopub.execute_input":"2021-06-04T11:17:37.105471Z","iopub.status.idle":"2021-06-04T11:17:47.353816Z","shell.execute_reply.started":"2021-06-04T11:17:37.105328Z","shell.execute_reply":"2021-06-04T11:17:47.352778Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Params","metadata":{}},{"cell_type":"code","source":"SEED = 2020\n\nDATA_PATH = \"../input/coleridgeinitiative-show-us-the-data/\"\nDATA_PATH_TRAIN = DATA_PATH + 'train/'\nDATA_PATH_TEST = DATA_PATH + 'test/'\n\nNUM_WORKERS = 4\n\nVOCABS = {\n    \"bert-base-uncased\": \"../input/vocabs/bert-base-uncased-vocab.txt\",\n}\n\nMODEL_PATHS = {\n    'bert-base-uncased': '../input/bertconfigs/uncased_L-12_H-768_A-12/uncased_L-12_H-768_A-12/',\n    'bert-large-uncased-whole-word-masking-finetuned-squad': '../input/bertconfigs/wwm_uncased_L-24_H-1024_A-16/wwm_uncased_L-24_H-1024_A-16/',\n    'albert-large-v2': '../input/albert-configs/albert-large-v2/albert-large-v2/',\n    'albert-base-v2': '../input/albert-configs/albert-base-v2/albert-base-v2/',\n    'distilbert': '../input/albert-configs/distilbert/distilbert/',\n}","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-06-04T11:17:47.355599Z","iopub.execute_input":"2021-06-04T11:17:47.356086Z","iopub.status.idle":"2021-06-04T11:17:47.361521Z","shell.execute_reply.started":"2021-06-04T11:17:47.356053Z","shell.execute_reply":"2021-06-04T11:17:47.360607Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Data Preparation\nInstead of having labels at article level, I refine the definition to section level. This allows for less sparsity. ","metadata":{}},{"cell_type":"code","source":"def load_text(id_, root=\"\"):\n    with open(os.path.join(root, id_ + \".json\")) as f:\n        text = json.load(f)\n    return text\n\n\ndef clean_text(txt):\n    return re.sub('[^A-Za-z0-9]+', ' ', str(txt).lower()).strip()\n\n\ndef create_data():\n    new_df = []\n\n    for idx in tqdm(range(len(df))):\n        article = load_text(df['Id'][idx], DATA_PATH_TRAIN)\n        id_, pub_title, dataset_title, dataset_label, cleaned_label = df.iloc[idx]\n        \n        \n        for i, section in enumerate(article):\n            text = section['text']\n            title = section['section_title']\n            \n            cleaned_text = clean_text(section['text'])\n            \n            found = cleaned_label in cleaned_text\n    \n            dic = {\n                \"id\": [id_], \n                \"section_id\": [i],\n                \"pub_title\": [pub_title], \n                \"dataset_title\": [dataset_title], \n                \"dataset_label\": [dataset_label], \n                \"cleaned_label\": [cleaned_label],\n                \"text\": [text],\n                \"cleaned_text\": [cleaned_text],\n                \"label_found\": [found],\n            }\n            new_df.append(pd.DataFrame.from_dict(dic))\n            \n    return pd.concat(new_df).reset_index(drop=True)","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-06-04T11:17:47.362639Z","iopub.execute_input":"2021-06-04T11:17:47.363064Z","iopub.status.idle":"2021-06-04T11:17:47.376625Z","shell.execute_reply.started":"2021-06-04T11:17:47.363035Z","shell.execute_reply":"2021-06-04T11:17:47.375329Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# df = pd.read_csv(DATA_PATH + 'train.csv')\n\n# new_df = create_data()  # Quite slow, could be sped-up with multi-processing.","metadata":{"execution":{"iopub.status.busy":"2021-06-04T11:17:47.378014Z","iopub.execute_input":"2021-06-04T11:17:47.378333Z","iopub.status.idle":"2021-06-04T11:31:18.307002Z","shell.execute_reply.started":"2021-06-04T11:17:47.378305Z","shell.execute_reply":"2021-06-04T11:31:18.305784Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# sns.countplot(x=new_df['label_found'])\n# plt.show()","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-06-04T11:31:18.320107Z","iopub.execute_input":"2021-06-04T11:31:18.320535Z","iopub.status.idle":"2021-06-04T11:31:18.572742Z","shell.execute_reply.started":"2021-06-04T11:31:18.320504Z","shell.execute_reply":"2021-06-04T11:31:18.571647Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"A lot of sections have no label found, I won't use them for training.","metadata":{}},{"cell_type":"code","source":"# df = new_df[new_df['label_found']].reset_index(drop=True)\n\n# df['length'] = df['cleaned_text'].apply(lambda x: len(x.split()))\n# df = df[df['length'] < 3000]  # remove too long texts\n\n# df.to_csv(\"df_train.csv\", index=False)  # saving, just in case","metadata":{"execution":{"iopub.status.busy":"2021-06-04T11:31:18.574488Z","iopub.execute_input":"2021-06-04T11:31:18.574781Z","iopub.status.idle":"2021-06-04T11:31:32.372984Z","shell.execute_reply.started":"2021-06-04T11:31:18.574752Z","shell.execute_reply":"2021-06-04T11:31:32.371914Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training a Model","metadata":{}},{"cell_type":"markdown","source":"# Utils","metadata":{"_kg_hide-input":true}},{"cell_type":"code","source":"def seed_everything(seed):\n    \"\"\"\n    Seeds basic parameters for reproductibility of results.\n\n    Args:\n        seed (int): Number of the seed.\n    \"\"\"\n\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    torch.backends.cudnn.benchmark = False\n\n\ndef save_model_weights(model, filename, verbose=1, cp_folder=\"\"):\n    \"\"\"\n    Saves the weights of a PyTorch model.\n\n    Args:\n        model (torch model): Model to save the weights of.\n        filename (str): Name of the checkpoint.\n        verbose (int, optional): Whether to display infos. Defaults to 1.\n        cp_folder (str, optional): Folder to save to. Defaults to \"\".\n    \"\"\"\n\n    if verbose:\n        print(f\"\\n -> Saving weights to {os.path.join(cp_folder, filename)}\\n\")\n    torch.save(model.state_dict(), os.path.join(cp_folder, filename))\n\n\ndef count_parameters(model, all=False):\n    \"\"\"\n    Count the parameters of a model.\n\n    Args:\n        model (torch model): Model to count the parameters of.\n        all (bool, optional):  Whether to count not trainable parameters. Defaults to False.\n\n    Returns:\n        int: Number of parameters.\n    \"\"\"\n\n    if all:\n        return sum(p.numel() for p in model.parameters())\n    else:\n        return sum(p.numel() for p in model.parameters() if p.requires_grad)\n","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-06-04T11:31:32.37435Z","iopub.execute_input":"2021-06-04T11:31:32.374967Z","iopub.status.idle":"2021-06-04T11:31:32.386613Z","shell.execute_reply.started":"2021-06-04T11:31:32.374923Z","shell.execute_reply":"2021-06-04T11:31:32.385732Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Data","metadata":{}},{"cell_type":"markdown","source":"## Tokenizer","metadata":{}},{"cell_type":"code","source":"class EncodedText:\n    def __init__(self, ids, offsets):\n        self.ids = ids\n        self.offsets = offsets\n\n\ndef create_tokenizer_and_tokens(config):\n    if \"roberta\" in config.selected_model:\n        raise NotImplementedError\n        \n    elif \"albert\" in config.selected_model:\n        raise NotImplementedError\n        \n    else:\n        tokenizer = BertWordPieceTokenizer(\n            MODEL_PATHS[config.selected_model] + 'vocab.txt',\n            lowercase=config.lowercase,\n        )\n\n        tokens = {\n            'cls': tokenizer.token_to_id('[CLS]'),\n            'sep': tokenizer.token_to_id('[SEP]'),\n            'pad': tokenizer.token_to_id('[PAD]'),\n        }\n    \n    return tokenizer, tokens","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-06-04T11:31:32.387978Z","iopub.execute_input":"2021-06-04T11:31:32.388669Z","iopub.status.idle":"2021-06-04T11:31:32.417855Z","shell.execute_reply.started":"2021-06-04T11:31:32.388613Z","shell.execute_reply":"2021-06-04T11:31:32.416287Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Locating labels","metadata":{}},{"cell_type":"code","source":"def locate_label_string(text, label):\n    \"\"\"\n    Finds the label in the text\n    \"\"\"\n    len_label = len(label) - 1\n\n    candidates_idx = [i for i, e in enumerate(text) if e == label[1]]\n    for idx in candidates_idx:\n        if \" \" + text[idx: idx + len_label] == label:\n            idx_start = idx\n            idx_end = idx + len_label\n            break\n\n    assert (\n        text[idx_start:idx_end] == label[1:]\n    ), f'\"{text[idx_start: idx_end]}\" instead of \"{label}\" in \"{text}\"'\n\n    char_targets = np.zeros(len(text))\n    char_targets[idx_start:idx_end] = 1\n\n    return idx_start, idx_end, char_targets\n\n\ndef locate_label_tokens(offsets, char_targets):\n    \"\"\"\n    Finds the tokens corresponding to the found labels\n    \"\"\"\n    target_idx = []\n    for idx, (offset1, offset2) in enumerate(offsets):\n        if sum(char_targets[offset1:offset2]) > 0:\n            target_idx.append(idx)\n\n    if not len(target_idx):\n        for idx, (offset1, offset2) in enumerate(offsets):\n            if sum(char_targets[offset1:offset2]) > 0:\n                target_idx.append(idx)\n\n    return target_idx[0], target_idx[-1]","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-06-04T11:31:32.419598Z","iopub.execute_input":"2021-06-04T11:31:32.419928Z","iopub.status.idle":"2021-06-04T11:31:32.435646Z","shell.execute_reply.started":"2021-06-04T11:31:32.419895Z","shell.execute_reply":"2021-06-04T11:31:32.434471Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Process sample","metadata":{}},{"cell_type":"code","source":"def process_data(\n    text,\n    label,\n    tokenizer,\n    tokens,\n    max_len=100,\n    model_name=\"bert\",\n):\n    \"\"\"\n    Prepares the data for the question answering task.\n    Adapted from Abishek's work on the Tweet Sentiment extraction competition, \n    check his work for more details !\n    \"\"\"\n    target_start, target_end = 0, 0\n    text = \" \" + \" \".join(str(text).split())\n    label = \" \" + \" \".join(str(label).split())\n\n    if label != \" \":\n        idx_start, idx_end, char_targets = locate_label_string(\n            text, label\n        )\n\n    tokenized = tokenizer.encode(text)\n    input_ids_text = tokenized.ids[1:-1]\n\n    # print(input_ids_text, len(input_ids_text))\n\n    offsets = tokenized.offsets[1:-1]\n\n    if label != \" \":\n        target_start, target_end = locate_label_tokens(offsets, char_targets)\n\n    if target_end >= max_len - 2:  # target is too far in the sentence, we crop its beginning.\n        n_tok_to_crop = target_start - max_len // 2\n        new_str_start = offsets[n_tok_to_crop][0]\n\n        input_ids_text = input_ids_text[n_tok_to_crop:]\n\n        offsets = [tuple(t) for t in np.array(offsets[n_tok_to_crop:]) - new_str_start]\n        text = text[new_str_start:]\n\n        target_start -= n_tok_to_crop\n        target_end -= n_tok_to_crop\n\n    input_ids = (\n        [tokens[\"cls\"]]\n        + input_ids_text[:max_len - 2]\n        + [tokens[\"sep\"]]\n    )\n\n    if \"roberta\" in model_name:\n        token_type_ids = [0] * len(input_ids)\n    else:\n        token_type_ids = [1] * len(input_ids)\n\n    text_offsets = [(0, 0)] + offsets[:max_len - 2] + [(0, 0)]\n\n    target_start += 1\n    target_end += 1\n\n    # target_end = min(target_end, max_len - 1)\n\n    assert len(input_ids) == len(token_type_ids) and len(input_ids) == len(text_offsets), (len(input_ids), len(text_offsets))  # noqa\n\n    padding_length = max_len - len(input_ids)\n    if padding_length > 0:\n        input_ids = input_ids + ([tokens[\"pad\"]] * padding_length)\n        token_type_ids = token_type_ids + ([0] * padding_length)\n        text_offsets = text_offsets + ([(0, 0)] * padding_length)\n\n    return {\n        \"ids\": input_ids,\n        \"token_type_ids\": token_type_ids,\n        \"targets_start\": target_start,\n        \"targets_end\": target_end,\n        \"text\": text,\n        \"label\": label,\n        \"offsets\": text_offsets,\n    }","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-06-04T11:31:32.438078Z","iopub.execute_input":"2021-06-04T11:31:32.43938Z","iopub.status.idle":"2021-06-04T11:31:32.457337Z","shell.execute_reply.started":"2021-06-04T11:31:32.439342Z","shell.execute_reply":"2021-06-04T11:31:32.456348Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Dataset","metadata":{}},{"cell_type":"code","source":"class SectionDataset(Dataset):\n    def __init__(\n        self,\n        df,\n        tokenizer,\n        tokens,\n        max_len=512,\n        model_name=\"bert\",\n    ):\n        self.tokens = tokens\n        self.tokenizer = tokenizer\n        self.max_len = max_len\n        self.model_name = model_name\n\n        self.texts = df[\"cleaned_text\"].values\n        self.labels = df[\"cleaned_label\"].values\n\n    def __len__(self):\n        return len(self.texts)\n\n    def __getitem__(self, idx):\n        data = process_data(\n            self.texts[idx],\n            self.labels[idx],\n            self.tokenizer,\n            self.tokens,\n            max_len=self.max_len,\n            model_name=self.model_name,\n        )\n\n        return {\n            \"ids\": torch.tensor(data[\"ids\"], dtype=torch.long),\n            \"token_type_ids\": torch.tensor(data[\"token_type_ids\"], dtype=torch.long),\n            \"target_start\": torch.tensor(data[\"targets_start\"], dtype=torch.long),\n            \"target_end\": torch.tensor(data[\"targets_end\"], dtype=torch.long),\n            \"text\": data[\"text\"],\n            \"label\": data[\"label\"],\n            \"offsets\": torch.tensor(data[\"offsets\"], dtype=torch.long),\n        }","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-06-04T11:31:32.459357Z","iopub.execute_input":"2021-06-04T11:31:32.460096Z","iopub.status.idle":"2021-06-04T11:31:32.473848Z","shell.execute_reply.started":"2021-06-04T11:31:32.460049Z","shell.execute_reply":"2021-06-04T11:31:32.472839Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model","metadata":{}},{"cell_type":"code","source":"TRANSFORMERS = {\n    \"roberta-base\": (RobertaModel, \"roberta-base\"),\n    'albert-base-v2': (AlbertModel, 'albert-base-v2'),\n    'albert-large-v2': (AlbertModel, 'albert-large-v2'),\n    'albert-xlarge-v2': (AlbertModel, 'albert-xlarge-v2'),\n    'albert-xxlarge-v2': (AlbertModel, 'albert-xxlarge-v2'),\n    \"bert-base-uncased\": (BertModel, \"bert-base-uncased\"),\n    \"bert-base-cased\": (BertModel, \"bert-base-cased\"),\n    \"bert-large-uncased-whole-word-masking\": (BertModel, \"bert-large-uncased-whole-word-masking\"),\n    \"distilbert-base-uncased-distilled-squad\": (\n        DistilBertModel,\n        \"distilbert-base-uncased-distilled-squad\"\n    )\n}\n\n\nclass QATransformer(nn.Module):\n    def __init__(self, model):\n        super().__init__()\n        self.name = model\n\n        self.pad_idx = 1 if \"roberta\" in self.name else 0\n\n        model_class, pretrained_weights = TRANSFORMERS[model]\n\n        self.transformer = model_class.from_pretrained(\n            pretrained_weights, output_hidden_states=True\n        )\n\n        self.nb_features = self.transformer.pooler.dense.out_features\n\n        self.logits = nn.Sequential(\n            nn.Linear(self.nb_features, self.nb_features),\n            nn.Tanh(),\n            nn.Linear(self.nb_features, 2),\n        )\n\n    def forward(self, tokens, token_type_ids):\n        \"\"\"\n        Usual torch forward function\n\n        Arguments:\n            tokens {torch tensor} -- Sentence tokens\n            token_type_ids {torch tensor} -- Sentence tokens ids\n        \"\"\"\n\n        hidden_states = self.transformer(\n            tokens,\n            attention_mask=(tokens != self.pad_idx).long(),\n            token_type_ids=token_type_ids,\n        )[-1]\n\n        features = hidden_states[-1]\n        logits = self.logits(features)\n\n        start_logits, end_logits = logits[:, :, 0], logits[:, :, 1]\n\n        return start_logits, end_logits","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-06-04T11:31:32.476077Z","iopub.execute_input":"2021-06-04T11:31:32.476938Z","iopub.status.idle":"2021-06-04T11:31:32.490847Z","shell.execute_reply.started":"2021-06-04T11:31:32.476886Z","shell.execute_reply":"2021-06-04T11:31:32.489917Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Metric","metadata":{}},{"cell_type":"code","source":"def jaccard_similarity(s1, s2): # could be wrong, see CPMP's thread\n    l1 = s1.split(\" \")\n    l2 = s2.split(\" \")\n    intersection = len(list(set(l1).intersection(l2)))\n    union = (len(l1) + len(l2)) - intersection\n    return float(intersection) / union\n\n\ndef jaccard_similarity(str1, str2): # the right one\n    a = set(str1.lower().split()) \n    b = set(str2.lower().split())\n    c = a.intersection(b)\n    return float(len(c)) / (len(a) + len(b) - len(c))\n\n\ndef compute_score(y_true, y_pred, beta=0.5):\n    \"\"\"\n    From https://www.kaggle.com/tungmphung/coleridge-initiative-local-score-computation/\n    \"\"\"\n    TP, FP, FN = 0, 0, 0\n\n    for truth, pred in zip(y_true, y_pred):\n        true_datasets = truth.split('|')\n        # Predicted strings for each publication are sorted alphabetically\n        # and processed in that order.\n        pred_datasets = sorted(pred.split('|'))\n\n        for true_dataset in true_datasets:\n            if len(pred_datasets):\n                match_scores = [jaccard_similarity(true_dataset, pred_dataset)\n                                for pred_dataset in pred_datasets]\n                # The prediction with the highest score for a given ground truth\n                # is matched with that ground truth.\n                match_index = np.argmax(match_scores)\n\n                if match_scores[match_index] >= 0.5:\n                    # Any matched predictions where the Jaccard score meets or\n                    # exceeds the threshold of 0.5 are counted as true positives (TP),\n                    TP += 1\n                else:\n                    # the remainder as false positives (FP).\n                    FP += 1\n\n                del(pred_datasets[match_index])\n            else:\n                # Any ground truths with no nearest predictions are counted as\n                # false negatives (FN).\n                FN += 1\n        # Any unmatched predictions are counted as false positives (FP).\n        FP += len(pred_datasets)\n\n    precision = TP / (TP + FP)\n    recall = TP / (TP + FN)\n    f_score = (1 + beta**2)*(precision*recall)/((beta**2)*precision + recall)\n\n    return f_score","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-06-04T11:31:32.492767Z","iopub.execute_input":"2021-06-04T11:31:32.493608Z","iopub.status.idle":"2021-06-04T11:31:32.510502Z","shell.execute_reply.started":"2021-06-04T11:31:32.493559Z","shell.execute_reply":"2021-06-04T11:31:32.509396Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Preds from probas","metadata":{}},{"cell_type":"code","source":"def get_string_from_idx(text, idx_start, idx_end, offsets):\n    \"\"\"\n    Uses the offsets to retrieve the predicted string based on the start and end indices\n    \"\"\"\n    if idx_end < idx_start:\n        idx_end = idx_start\n\n    predicted_string = \"\"\n    for i in range(idx_start, idx_end + 1):\n        predicted_string += text[offsets[i][0]: offsets[i][1]]\n        if i + 1 < len(offsets) and offsets[i][1] < offsets[i + 1][0]:\n            predicted_string += \" \"\n\n    return predicted_string\n\ndef get_pred_from_logits(data, start_logits, end_logits, from_proba=False):\n    if not from_proba:\n        start_logits = torch.softmax(start_logits, dim=1).cpu().detach().numpy()\n        end_logits = torch.softmax(end_logits, dim=1).cpu().detach().numpy()\n\n    offsets = data[\"offsets\"].cpu().numpy()\n\n    preds = []\n    for i in range(len(start_logits)):\n        start_idx = np.argmax(start_logits[i])\n        end_idx = np.argmax(end_logits[i])\n        preds.append(get_string_from_idx(data[\"text\"][i], start_idx, end_idx, offsets[i]))\n\n    return preds","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-06-04T11:31:32.512241Z","iopub.execute_input":"2021-06-04T11:31:32.513073Z","iopub.status.idle":"2021-06-04T11:31:32.534279Z","shell.execute_reply.started":"2021-06-04T11:31:32.513022Z","shell.execute_reply":"2021-06-04T11:31:32.532744Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Loss","metadata":{}},{"cell_type":"code","source":"def ce_loss(\n    pred, truth, smoothing=False, trg_pad_idx=-1, eps=0.1\n):\n    \"\"\"\n    Computes the cross entropy loss with label smoothing\n\n    Args:\n        pred (torch tensor): Prediction\n        truth (torch tensor): Target\n        smoothing (bool, optional): Whether to use smoothing. Defaults to False.\n        trg_pad_idx (int, optional): Indices to ignore in the loss. Defaults to -1.\n        eps (float, optional): Smoothing coefficient. Defaults to 0.1.\n\n    Returns:\n        [type]: [description]\n    \"\"\"\n    truth = truth.contiguous().view(-1)\n\n    one_hot = torch.zeros_like(pred).scatter(1, truth.view(-1, 1), 1)\n\n    if smoothing:\n        n_class = pred.size(1)\n        one_hot = one_hot * (1 - eps) + (1 - one_hot) * eps / (n_class - 1)\n\n    loss = -one_hot * F.log_softmax(pred, dim=1)\n\n    if trg_pad_idx >= 0:\n        loss = loss.sum(dim=1)\n        non_pad_mask = truth.ne(trg_pad_idx)\n        loss = loss.masked_select(non_pad_mask)\n\n    return loss.sum()\n\n\ndef qa_loss_fn(start_logits, end_logits, start_positions, end_positions, config):\n    \"\"\"\n    Loss function for the question answering task.\n    It is the sum of the cross entropy for the start and end logits\n\n    Args:\n        start_logits (torch tensor): Start logits\n        end_logits (torch tensor): End logits\n        start_positions (torch tensor): Start ground truth\n        end_positions (torch tensor): End ground truth\n        config (dict): Dictionary of parameters for the CE Loss.\n\n    Returns:\n        torch tensor: Loss value\n    \"\"\"\n    bs = start_logits.size(0)\n\n    start_loss = ce_loss(\n        start_logits,\n        start_positions,\n        smoothing=config[\"smoothing\"],\n        eps=config[\"eps\"],\n    )\n\n    end_loss = ce_loss(\n        end_logits,\n        end_positions,\n        smoothing=config[\"smoothing\"],\n        eps=config[\"eps\"],\n    )\n\n    total_loss = start_loss + end_loss\n\n    return total_loss / bs\n","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-06-04T11:31:32.537607Z","iopub.execute_input":"2021-06-04T11:31:32.538335Z","iopub.status.idle":"2021-06-04T11:31:32.555299Z","shell.execute_reply.started":"2021-06-04T11:31:32.538287Z","shell.execute_reply":"2021-06-04T11:31:32.553767Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Trim tensors\n\nHelps for speedup","metadata":{}},{"cell_type":"code","source":"def trim_tensors(tokens, input_ids, model_name='bert', min_len=10):\n    \"\"\"\n    Trim tensors so that within a batch, padding is shortened.\n    This speeds up training for RNNs and Transformers\n\n    Arguments:\n        tokens {torch tensor} -- Text tokens\n\n    Keyword Arguments:\n        min_len {int} -- Minimum length to trim to (default: {10})\n\n    Returns:\n        torch tensor -- trimmed tokens\n    \"\"\"\n    pad_token = 1 if \"roberta\" in model_name else 0\n    max_len = max(torch.max(torch.sum((tokens != pad_token), 1)), min_len)\n    return tokens[:, :max_len], input_ids[:, :max_len]","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-06-04T11:31:32.556636Z","iopub.execute_input":"2021-06-04T11:31:32.556938Z","iopub.status.idle":"2021-06-04T11:31:32.573082Z","shell.execute_reply.started":"2021-06-04T11:31:32.55691Z","shell.execute_reply":"2021-06-04T11:31:32.571962Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Fit function","metadata":{}},{"cell_type":"code","source":"def fit(\n    model,\n    train_dataset,\n    val_dataset,\n    loss_config,\n    epochs=5,\n    batch_size=8,\n    weight_decay=0,\n    warmup_prop=0.0,\n    lr=5e-4,\n    cfg=dict()\n):\n\n    train_loader = DataLoader(\n        train_dataset, batch_size=batch_size, num_workers=NUM_WORKERS, shuffle=True,\n    )\n    val_loader = DataLoader(\n        val_dataset, batch_size=batch_size, shuffle=False, num_workers=NUM_WORKERS\n    )\n\n    opt_params = []\n    no_decay = [\"bias\", \"LayerNorm.bias\", \"LayerNorm.weight\"]\n    for n, p in model.named_parameters():\n        wd = 0 if any(nd in n for nd in no_decay) else weight_decay\n        opt_params.append(\n            {\"params\": [p], \"weight_decay\": wd, \"lr\": lr}\n        )\n\n    optimizer = AdamW(opt_params, lr=lr, betas=(0.5, 0.999))\n\n    n_steps = epochs * len(train_loader)\n    num_warmup_steps = int(warmup_prop * n_steps)\n    scheduler = get_linear_schedule_with_warmup(\n        optimizer, num_warmup_steps, n_steps\n    )\n\n    total_steps = 0\n    for epoch in range(epochs):\n        model.train()\n        start_time = time.time()\n\n        optimizer.zero_grad()\n        avg_loss = 0\n\n        for step, data in enumerate(train_loader):\n            total_steps += 1\n\n            ids, token_type_ids = trim_tensors(\n                data[\"ids\"], data[\"token_type_ids\"], model.name\n            )\n\n            start_logits, end_logits = model(ids.cuda(), token_type_ids.cuda())\n\n            loss = qa_loss_fn(\n                start_logits,\n                end_logits,\n                data[\"target_start\"].cuda(),\n                data[\"target_end\"].cuda(),\n                config=loss_config,\n            )\n\n            avg_loss += loss.item() / len(train_loader)\n            loss.backward()\n\n            nn.utils.clip_grad_norm_(model.parameters(), 10.0)\n\n            optimizer.step()\n            scheduler.step()\n            model.zero_grad()\n\n        model.eval()\n        avg_val_loss = 0.\n        preds, truths = [], []\n\n        with torch.no_grad():\n\n            for data in val_loader:\n                ids, token_type_ids = trim_tensors(\n                    data[\"ids\"], data[\"token_type_ids\"], model.name\n                )\n\n                start_logits, end_logits = model(ids.cuda(), token_type_ids.cuda())\n\n                loss = qa_loss_fn(\n                    start_logits.detach(),\n                    end_logits.detach(),\n                    data[\"target_start\"].cuda(),\n                    data[\"target_end\"].cuda(),\n                    config=loss_config,\n                )\n\n                avg_val_loss += loss.item() / len(val_loader)\n\n                preds += get_pred_from_logits(data, start_logits, end_logits)\n                truths += data['label']\n\n        score = compute_score(truths, preds)\n\n        dt = time.time() - start_time\n        lr = scheduler.get_last_lr()[0]\n        print(f\"Epoch {epoch + 1}/{epochs} \\t lr={lr:.1e} \\t t={dt:.0f}s \\t\", end=\"\")\n        print(\n            f\"loss={avg_loss:.3f} \\t val_loss={avg_val_loss:.3f} \\t val_score={score:.4f}\"\n        )\n        save_model_weights(model, f\"{cfg['selected_model']}_{cfg['fold']}_{epoch+1}.pt\", cp_folder=\"\")\n\n    del loss, data, avg_val_loss, avg_loss, train_loader, val_loader\n    torch.cuda.empty_cache()\n    gc.collect()\n\n    return preds\n","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-06-04T11:31:32.574884Z","iopub.execute_input":"2021-06-04T11:31:32.575326Z","iopub.status.idle":"2021-06-04T11:31:32.597286Z","shell.execute_reply.started":"2021-06-04T11:31:32.575291Z","shell.execute_reply":"2021-06-04T11:31:32.595991Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## $k$-fold","metadata":{}},{"cell_type":"code","source":"def k_fold(config, df, save=True):\n    tokenizer, tokens = create_tokenizer_and_tokens(config)\n    \n    time = re.sub(' ', '_', str(datetime.datetime.now())[:16])\n    score = 0\n    \n    gkf = GroupKFold(n_splits=config.k)\n    folds = list(gkf.split(X=df, groups=df['dataset_title']))\n    \n    pred_oof = [''] * len(df)\n    \n    for fold, (train_idx, val_idx) in enumerate(folds):\n        if fold in config.selected_folds:\n            print(f\"\\n-------------   Fold {fold + 1} / {len(folds)}  -------------\\n\")\n            seed_everything(config.seed + fold)\n\n            model = QATransformer(config.selected_model).cuda()\n            model.zero_grad()\n\n            train_dataset = SectionDataset(\n                df.iloc[train_idx], \n                tokenizer, \n                tokens, \n                max_len=config.max_len, \n                model_name=config.selected_model\n            )\n\n            val_dataset = SectionDataset(\n                df.iloc[val_idx], \n                tokenizer, \n                tokens,\n                max_len=config.max_len, \n                model_name=config.selected_model\n            )\n\n            n_parameters = count_parameters(model)\n\n            print(f\"    -> {len(train_dataset)} training texts\")\n            print(f\"    -> {len(val_dataset)} validation texts\")\n            print(f\"    -> {n_parameters} trainable parameters\\n\")  \n\n            preds = fit(\n                model, \n                train_dataset, \n                val_dataset, \n                config.loss_config,\n                epochs=config.epochs, \n                batch_size=config.batch_size, \n                weight_decay=config.weight_decay, \n                lr=config.lr, \n                warmup_prop=config.warmup_prop,\n                cfg={'selected_model': config.selected_model,\n                    'fold': fold+1}\n            )\n\n            for i, idx in enumerate(val_idx):\n                pred_oof[idx] = preds[i]\n\n            if save:\n                save_model_weights(model, f'{config.selected_model}_{fold + 1}.pt', cp_folder=\"\")\n\n            del model, train_dataset, val_dataset\n            torch.cuda.empty_cache()\n            gc.collect()","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-06-04T11:31:32.59923Z","iopub.execute_input":"2021-06-04T11:31:32.599583Z","iopub.status.idle":"2021-06-04T11:31:32.614301Z","shell.execute_reply.started":"2021-06-04T11:31:32.599549Z","shell.execute_reply":"2021-06-04T11:31:32.613337Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Main","metadata":{}},{"cell_type":"code","source":"class Config:\n    # General\n    k = 5\n    seed = 2021\n    selected_folds = [0] # [0, 1, 2, 3, 4]\n\n    # Texts\n    max_len = 256\n    \n    # Architecture\n    selected_model = \"bert-base-uncased\"\n    lowercase = True\n    \n    # Loss function\n    loss_config = {\n        \"smoothing\": False,\n        \"eps\": 0.1,\n    }\n    \n    # Training\n    batch_size = 16\n    batch_size_val = batch_size * 2\n    weight_decay = 1.\n    \n    epochs = 1 # 3\n    lr = 5e-5\n    warmup_prop = 0.1","metadata":{"execution":{"iopub.status.busy":"2021-06-04T11:31:32.615465Z","iopub.execute_input":"2021-06-04T11:31:32.615727Z","iopub.status.idle":"2021-06-04T11:31:32.629532Z","shell.execute_reply.started":"2021-06-04T11:31:32.615701Z","shell.execute_reply":"2021-06-04T11:31:32.628466Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df = pd.read_csv(\"../input/bert-question-answering-training/df_train.csv\")\n\nk_fold(\n    Config,\n    df,\n    save=True,\n)","metadata":{"execution":{"iopub.status.busy":"2021-06-04T11:31:32.631218Z","iopub.execute_input":"2021-06-04T11:31:32.631669Z","iopub.status.idle":"2021-06-04T11:31:51.750938Z","shell.execute_reply.started":"2021-06-04T11:31:32.631624Z","shell.execute_reply":"2021-06-04T11:31:51.748874Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Thanks for reading !\n\nHopefully this work comes helpful in beating the public LB baselines...\n\nDon't forget to upvote :)\n","metadata":{}}]}