{"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":"code","source":"#TODO:[x] seperate datasets in one mention by adding one to the label of first token of each dataset.\n#TODO:[x] punish false positives by increasing the weight of class 0. \n#TODO:[x] have some unseen labels in the validation set. \n#TODO:[x] weight loss by the frequency of each label.\n#TODO:[x] lemmitization of clean text while keeping reference to original text\n#TODO:[] the frequency of each word in the known labels\n#TODO:[x] switch validation to inference","metadata":{"_uuid":"a5a1947b-fc1d-4f8e-8d74-cd6847479826","_cell_guid":"a8880c7f-d63a-4b1d-8085-907446b01861","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2021-06-03T02:32:08.631225Z","iopub.execute_input":"2021-06-03T02:32:08.631623Z","iopub.status.idle":"2021-06-03T02:32:08.636742Z","shell.execute_reply.started":"2021-06-03T02:32:08.631540Z","shell.execute_reply":"2021-06-03T02:32:08.635656Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport gc\nimport sys\n\nimport warnings\nwarnings.filterwarnings('ignore')","metadata":{"_uuid":"550ec9b3-57a8-43a4-883f-d0d8133151a6","_cell_guid":"eea197ec-1c64-4bc5-a069-dcd4f836d636","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2021-06-03T02:32:08.641055Z","iopub.execute_input":"2021-06-03T02:32:08.641458Z","iopub.status.idle":"2021-06-03T02:32:08.647317Z","shell.execute_reply.started":"2021-06-03T02:32:08.641430Z","shell.execute_reply":"2021-06-03T02:32:08.646549Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"COMPILE = 0\nVERSION = '2'\n\nPARENT_DIR = '/kaggle/input'\nDATA_DIR = os.path.join(PARENT_DIR, 'coleridge-preprocessed')\nINPUT_DIR = os.path.join(PARENT_DIR, 'coleridgeinitiative-show-us-the-data')\nMODELS_DIR = os.path.join(PARENT_DIR, 'coleridge-models-v4')\nGLOVE_DIR = os.path.join(PARENT_DIR, \"glove-global-vectors-for-word-representation\")\n# LABELS_DIR = os.path.join(PARENT_DIR, \"coleridge-datagov\")\nOUT_DIR = '/kaggle/working'\n\nTRAIN_DATA = None\nLABEL_TOKENS = None","metadata":{"_uuid":"c49680a5-8275-4c58-b573-8bf5d3092aae","_cell_guid":"6de786cf-b682-4aa4-89ed-243f145682c6","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2021-06-03T02:32:08.650678Z","iopub.execute_input":"2021-06-03T02:32:08.651069Z","iopub.status.idle":"2021-06-03T02:32:08.657009Z","shell.execute_reply.started":"2021-06-03T02:32:08.651035Z","shell.execute_reply":"2021-06-03T02:32:08.655997Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\n\n!pip install pytorch-crf --no-index --find-links=file://{PARENT_DIR}/torchcrf/torchcrf\n    \nimport re\nimport json\nimport math\nimport time\nimport random\nfrom typing import Optional\nfrom collections import Counter\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport matplotlib.pyplot as plt\n\nimport torch\nprint(torch.__version__)\nimport torch.nn as nn\nfrom torch import Tensor\nimport torch.nn.functional as F\nfrom torch.nn.utils import clip_grad_norm_\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\n\nimport spacy\nfrom torchcrf import CRF\nimport torchtext\nprint(torchtext.__version__)\nfrom torchtext.data import Field\nfrom gensim.parsing.preprocessing import remove_stopwords\n\nfrom nltk.stem import PorterStemmer\nfrom sklearn.model_selection import train_test_split\nfrom coleridge_utility import *\n\nfrom multiprocessing import Pool\nfrom IPython.display import display\nfrom tqdm.notebook import tqdm\ntqdm.pandas()","metadata":{"_uuid":"6311b2dc-723e-4ad6-b394-66abea289dff","_cell_guid":"09c944d1-13e0-400a-a565-bb0819e771ba","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2021-06-03T02:32:08.681380Z","iopub.execute_input":"2021-06-03T02:32:08.681671Z","iopub.status.idle":"2021-06-03T02:32:21.975977Z","shell.execute_reply.started":"2021-06-03T02:32:08.681647Z","shell.execute_reply":"2021-06-03T02:32:21.975001Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def seed_everything(seed):\n    random.seed(seed)\n    np.random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.\n    torch.backends.cudnn.benchmark = False\n    torch.backends.cudnn.deterministic = True\n    \nRANDOM_SEED = 50\nseed_everything(RANDOM_SEED)","metadata":{"_uuid":"1bf3896d-c0f8-4c3e-8305-7f344b1412ef","_cell_guid":"dfec0cc1-8212-4753-83a4-3f113f873b50","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2021-06-03T02:32:21.977809Z","iopub.execute_input":"2021-06-03T02:32:21.978168Z","iopub.status.idle":"2021-06-03T02:32:21.988661Z","shell.execute_reply.started":"2021-06-03T02:32:21.978132Z","shell.execute_reply":"2021-06-03T02:32:21.987800Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Helper Functions","metadata":{"_uuid":"75794f0b-a94b-4713-bb26-04f8ccbcf045","_cell_guid":"5cf00b96-5af2-4971-bc1c-1911f0e3d4b8","trusted":true}},{"cell_type":"code","source":"gpu_info = !nvidia-smi\ngpu_info = '\\n'.join(gpu_info)\nif gpu_info.find('failed') >= 0:\n  print('Select the Runtime > \"Change runtime type\" menu to enable a GPU accelerator, ')\n  print('and then re-execute this cell.')\nelse:\n  print(gpu_info)","metadata":{"_uuid":"47c176c1-3ebe-4ef7-ab6e-40828c520930","_cell_guid":"8e802261-7e9b-407d-90e7-35db6677d0f4","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2021-06-03T02:32:21.989949Z","iopub.execute_input":"2021-06-03T02:32:21.990292Z","iopub.status.idle":"2021-06-03T02:32:22.046123Z","shell.execute_reply.started":"2021-06-03T02:32:21.990253Z","shell.execute_reply":"2021-06-03T02:32:22.045367Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def create_scheduler(optim):\n    sched = None\n    if SCHED_TYPE == 'plateau':\n        sched = ReduceLROnPlateau(\n            optim, 'max',\n            patience=1,\n            factor=0.5,\n            threshold=5e-6,\n            min_lr=1e-6)\n    return sched\n\n\ndef create_optimizer(estimator, lr):\n    return torch.optim.AdamW(estimator.parameters(), lr=lr)\n\n\ndef create_model(model_type, **params):\n    return model_type(**params).to(DEVICE)\n\n\ndef save_model(estimator, optim, sched=None, val_score:float=0, epoch=0):\n    checkpoint = {'model_params': estimator.params,\n                  'model_state_dict': estimator.state_dict(),\n                  'optimizer_state_dict': optim.state_dict(),\n                  'learning_rate': optim.param_groups[0]['lr'],\n                  'val_score': val_score}\n    if sched is not None:\n        checkpoint = {**checkpoint,\n                      'sheduler_state_dict': sched.state_dict(),\n                      'epoch': epoch}\n    file_dir = os.path.join(OUT_DIR, VERSION)\n    if not os.path.exists(file_dir):\n        os.makedirs(file_dir)\n    filepath = os.path.join(file_dir, MODEL_FILENAME)\n    torch.save(checkpoint, filepath)\n\n    \ndef load_model(model_type, for_training=False):\n    local_path = os.path.join(OUT_DIR, VERSION, MODEL_FILENAME)\n    if os.path.exists(local_path):\n        filepath = local_path\n    else:\n        filepath = os.path.join(MODELS_DIR, VERSION, MODEL_FILENAME)\n\n    print(f'Loading model from {filepath}')\n    checkpoint = torch.load(filepath, map_location=DEVICE)\n    print('Model score:', checkpoint['val_score'])\n    \n    checkpoint['model_params']['device'] = DEVICE\n    estimator = create_model(model_type, **checkpoint['model_params'])\n    estimator.load_state_dict(checkpoint['model_state_dict'])\n    estimator.to(DEVICE)\n    \n    if for_training:\n        optim = create_optimizer(estimator, lr=checkpoint['learning_rate'])\n        optim.load_state_dict(checkpoint['optimizer_state_dict'])\n        if 'sheduler_state_dict' in checkpoint:\n            sched = create_scheduler(optim)\n            sched.load_state_dict(checkpoint['sheduler_state_dict'])\n        return estimator, optim, sched, checkpoint\n    return estimator\n\ndef save_vocab(vocab, filename):\n    file_dir = os.path.join(OUT_DIR, VERSION)\n    if not os.path.exists(file_dir):\n        os.makedirs(file_dir)\n    filepath = os.path.join(file_dir, filename)\n    torch.save(vocab, filepath)\n    \ndef load_vocab(filename):\n    local_path = os.path.join(OUT_DIR, VERSION, filename)\n    if os.path.exists(local_path):\n        filepath = local_path\n    else:\n        filepath = os.path.join(MODELS_DIR, VERSION, filename)\n    print(f'Loading vocab from {filepath}')\n    return torch.load(filepath, map_location=DEVICE)","metadata":{"_uuid":"82eab60e-27d9-4e50-bbab-7d2034e74095","_cell_guid":"eb1b773d-47d3-4cc6-ae5e-30c92245b99a","execution":{"iopub.status.busy":"2021-06-03T02:32:22.048352Z","iopub.execute_input":"2021-06-03T02:32:22.048732Z","iopub.status.idle":"2021-06-03T02:32:22.064970Z","shell.execute_reply.started":"2021-06-03T02:32:22.048692Z","shell.execute_reply":"2021-06-03T02:32:22.064136Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def jaccard(a, b):\n    if a == b:\n        return 1\n    a = set(a.lower().split())\n    b = set(b.lower().split())\n    c = a.intersection(b)\n    return float(len(c)) / (len(a) + len(b) - len(c))\n\ndef get_precision_recall(tp, fp, fn):\n    precision = tp / (tp + fp or 1)\n    recall = tp / (tp + fn or 1)\n    return precision, recall\n\ndef fbeta_score(precision, recall, beta):\n    if precision == recall == 0:\n        return 0\n    fbeta = (1 + beta * beta) * ((precision * recall) / ((beta * beta * precision) + recall))\n    return fbeta\n\ndef coleridge_initiative_jaccard(output_str, target_str, verbose=True):\n    gts = target_str.split('|')\n    pds = sorted(output_str.split('|'))\n    if verbose:\n        print(\"Ground truth : \" , target_str, gts)\n        print(\"Prediction : \", output_str, pds)\n    assert(len(gts) and len(pds)), (len(gts), len(pds))\n    if '' in gts:\n        assert(len(gts) == 1), (gts, target_str)\n        if pds != gts:\n            return 0, len(pds), 0\n    if '' in pds:\n        assert(len(pds) == 1), (pds, output_str)\n        if pds != gts:\n            return 0, 0, len(gts)\n\n    #### Counting True Positives (TP) and False Positives (FP)\n    tpc = 0\n    fpc = 0\n    \n    for pd in pds:\n        score = -1\n        _gt = None\n        for gt in gts:\n            js = jaccard(pd, gt)\n            if js > score:\n                score = js\n                _gt = gt\n        if score >= 0.5:\n            tpc += 1\n            gts.remove(_gt)\n        else:\n            fpc += 1\n    \n    fnc = len(gts)\n    return tpc, fpc, fnc\n\ndef score_coleridge_initiative(output, target, beta=0.5, verbose=False):\n    \n    '''\n    This function will calculate the FBeta score for Coleridge Initiative competition \n    if given appropriate arguments\n    \n    Arguments - \n    output - Your submission dataframe that has both ground truth and prediction columns.\n    gt_col - This is the column name of ground truth column.\n    pred_col - This is the column name of predictions column.\n    beta - Beta value to calculate FBeta score.\n    \n    Returns - \n    This function will return the FBeta (beta=0.5) score.\n    \n    ## Set verbose = True to print logs    \n    '''\n    \n    ### Jaccard Similarity\n    evaluation = np.array([coleridge_initiative_jaccard(p, t, verbose=False) \n                           for p, t in zip(output, target)], dtype=np.int32)\n    tp, fp, fn = np.sum(evaluation, axis=0)\n    precision, recall = get_precision_recall(tp, fp, fn)\n    fbeta = fbeta_score(precision, recall, 0.5)\n    \n    if verbose:\n        print(\"True Positives (TP) : \", tp)\n        print(\"False Positives (FP) : \", fp)\n        print(\"False Negatives (FN) : \", fn)\n        print(\"Precision : \", precision)\n        print(\"Recall : \", recall)\n        print(\"FBeta Score : \", fbeta)\n\n    return fbeta\n\ndef tokenized_to_text(text, label):\n    assert(label.shape == text.shape), (label.shape, text.shape)\n    if isinstance(label, torch.Tensor):\n        label = label.cpu()\n    nl_id = 1\n    label[text == PAD] = LABEL_PAD\n    eos_row = np.repeat(['\\n'], text.shape[1])\n    text = np.vstack([text, eos_row])\n    eos_row = np.repeat([nl_id], label.shape[1])\n    label = np.vstack([label, eos_row])\n\n    out = text.T[label.T > 0]\n    label = label.T[label.T > 0]\n\n    def get_out(idx):\n        txt = out[idx]\n        lab = label[idx]\n        assert(lab > 0)\n        if lab == TAGS_TO_IDX['B']:\n            return f'|{txt}'\n        return txt\n\n    f_itos = np.vectorize(get_out)\n    out = f_itos(np.arange(len(out)))\n    out = ' '.join(out)\n    out = np.array(out.split('\\n'))[:-1] \n    out = [clean_predictions(o) for o in out]\n    \n    return out\n    \ndef extract_text(texts, numericals):\n    out = np.array([tokenized_to_text(texts[i], numericals[i]) for i in range(len(texts)-1)]).flatten()\n    if len(texts):\n        out = np.r_[out, tokenized_to_text(texts[-1], numericals[-1])]\n    return out\n\ndef merge_article_labels(ids, labels, remove_similars=False, remove_subsets=False):\n    out = {idx: list() for idx in set(ids)}\n    for idx, lab in zip(ids, labels):\n        if len(lab):\n            out[idx].append(lab)\n            \n    for k, v in out.items():\n        v = sorted(set('|'.join(v).split('|')))\n        if remove_similars:\n            v = remove_similar_labels(v)\n        if remove_subsets:\n            v = remove_subset_labels(v)\n        v = '|'.join(v)\n        v = clean_predictions(v)\n        out[k] = v\n        \n    return pd.DataFrame(out.items(), columns=['Id', 'PredictionString'])\n        \ndef remove_subset_labels(labels):\n    if len(labels) <= 1:\n            return labels\n    unq_labels = list(labels)\n    tok_labels = [lab.split() for lab in labels]\n\n    for i1 in range(len(labels)):\n        if len(unq_labels) == 1:\n            break\n        lab1 = labels[i1]\n        tok1 = tok_labels[i1]\n        for i2 in range(len(labels)):\n            lab2 = labels[i2]\n            tok2 = tok_labels[i2]\n            if lab1 is lab2:\n                continue\n            if all(t in tok2 for t in tok1):\n                unq_labels.remove(lab1)\n                break\n    assert(len(unq_labels)), (labels, tok_labels, unq_labels)\n    return unq_labels\n                    \ndef remove_similar_labels(labels):\n    if len(labels) <= 1:\n        return labels\n    unq_labels = list(labels)\n    for i1 in range(len(labels)):\n        l1 = labels[i1]\n        for i2 in range(i1, len(labels)):\n            l2 = labels[i2]\n            if l1 is l2:\n                continue\n            if jaccard(l1, l2) >= 0.5:\n                if l1.count(' ') > l2.count(' '):\n                    if l2 in unq_labels:\n                        unq_labels.remove(l2)\n                elif l1.count(' ') < l2.count(' '):\n                    if l1 in unq_labels:\n                        unq_labels.remove(l1)\n                elif len(l1) > len(l2):\n                    if l2 in unq_labels:\n                        unq_labels.remove(l2)\n                else:\n                    if l1 in unq_labels:\n                        unq_labels.remove(l1)\n    assert(len(unq_labels)), (labels, unq_labels)\n    return unq_labels\n\ndef clean_predictions(labels):\n    out = labels.replace(' |', '|').replace('| ', '|')\n    while '||' in out:\n        out = out.replace('||', '|')\n    out = out.strip(' |')\n    assert('<' not in out), out\n    return out\n\ndef print_random_samples(ids, outputs, targets, n=10):\n    n = min(n, len(outputs))\n    idx = np.random.choice(range(len(outputs)), n, replace=False)\n    scores = [coleridge_initiative_jaccard(o, t, verbose=False) for o, t in zip(outputs, targets)]\n    sample = pd.DataFrame({'target': [f'[{x}]' for x in targets[idx]],\n                           'output': [f'[{x}]' for x in outputs[idx]],\n                           'tp-fp-fn': [scores[i] for i in idx]},\n                          index=ids[idx])\n\n    sample_styler = sample.style.set_properties(**{'text-align': 'left', 'width': '500px'})\n    sample_styler = sample_styler.set_table_styles([dict(selector='th', props=[('text-align', 'left')])])\n    display(sample_styler)","metadata":{"_uuid":"c68b070e-ded0-454a-818d-926632a48572","_cell_guid":"8d5e98db-ee89-4a0b-82de-752376ffbfbf","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2021-06-03T02:32:22.068353Z","iopub.execute_input":"2021-06-03T02:32:22.068926Z","iopub.status.idle":"2021-06-03T02:32:22.106947Z","shell.execute_reply.started":"2021-06-03T02:32:22.068886Z","shell.execute_reply":"2021-06-03T02:32:22.106179Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def clean_label(txt):\n    return re.sub('[^A-Za-z0-9]+', ' ', txt.lower()).strip()\n\ndef find_sublist(sub, bigger):\n    sub, bigger = tuple(sub), tuple(bigger)\n    if not bigger or not sub:\n        return []\n    \n    first, rest = sub[0], sub[1:]\n    pos = 0\n    out = []\n    try:\n        while True:\n            pos = bigger.index(first, pos) + 1\n            if not rest or bigger[pos:pos+len(rest)] == rest:\n                out.append(pos - 1)\n    except ValueError:\n        return out\n\ndef is_sublist(sub, bigger):\n    return len(find_sublist(sub, bigger)) > 0\n\ndef get_row_chunks2(row, is_test):\n    assert(MAX_LEN > 0), MAX_LEN\n    assert(OVERLAP < MAX_LEN), (OVERLAP, MAX_LEN)\n\n    full_chunks, empty_chunks = [], []\n    doc_feats = {feat: eval_iterable_field(row[feat]) for feat in USED_FEATURES}\n    tok_texts = doc_feats['tokenized_text'] \n    if not is_test:\n        labels = eval_iterable_field(row['label'])\n#         labels = [lab for lab in labels if lab in UNIQUE_LABELS]\n#         labels = remove_subset_labels(labels)\n#         assert(len(labels)), (labels, row['label'])\n        labels = sorted(labels, key=lambda x: len(x))\n        tok_labels = [tokenize(label) for label in labels]\n        found_labels = []\n        \n    def get_chunk_labels(chunk_txt):\n        if is_test:\n            return None, None\n        _found_labels = []\n        chunk_label = [TAGS_TO_IDX['O']] * len(chunk_txt)\n        for i, tok_label in enumerate(tok_labels):\n            positions = find_sublist(tok_label, chunk_txt)\n            if len(positions):\n                _found_labels.append(labels[i])\n            for pos in positions:\n                epos = pos + len(tok_label)\n                if chunk_label[pos] != TAGS_TO_IDX['I']:\n                    chunk_label[pos] = TAGS_TO_IDX['B']\n                    chunk_label[pos+1: epos] = [TAGS_TO_IDX['I']] * (epos - pos - 1)\n        found_labels.extend(_found_labels)\n\n        assert(len(chunk_label) == len(chunk_txt)), (len(chunk_label) == len(chunk_txt))\n        return chunk_label, _found_labels\n\n    for i in range(len(tok_texts)):\n        sect_feats = {feat: doc_feats[feat][i] for feat in doc_feats}\n        tok_text = sect_feats['tokenized_text']\n        \n        sent_pos = [0] + find_sublist([SEP], tok_text)\n        if tok_text[-1] != SEP:\n            sent_pos += [len(tok_text)]\n        prev_feats = None\n        for s, e in zip(sent_pos[:-1], sent_pos[1:]):\n            if tok_text[s] == SEP:\n                s += 1\n            sent_feats = {feat: sect_feats[feat][s: e] for feat in sect_feats}\n            if prev_feats is not None:\n                sent_feats = {feat: prev_feats[feat]+sent_feats[feat] for feat in sent_feats}\n                prev_feats = None\n            tok_sent = sent_feats['tokenized_text']\n            assert(SEP not in tok_sent), (s, e, tok_sent, tok_text)\n            if len(tok_sent) < MIN_LEN:\n                prev_feats = sent_feats\n                if e != sent_pos[-1] or len(tok_sent) < 4:\n                    continue\n\n            sidx = 0\n            while True:\n                chunk_feats = {feat: sent_feats[feat][sidx: sidx+MAX_LEN] for feat in sent_feats}\n                tok_chunk = chunk_feats['tokenized_text']\n                if 'pos' in chunk_feats:\n#                     chunk_feats['pos'] = [p if tok_chunk[pi] not in LABEL_TOKENS else POS_TAGS['IMP']\n#                                           for pi, p in enumerate(chunk_feats['pos'])]\n                    chunk_feats['pos'] = np.array(chunk_feats['pos'], dtype=np.int8)\n                if 'is_upper' in chunk_feats:\n                    chunk_feats['is_upper'] = np.array(chunk_feats['is_upper'], dtype=np.int8)\n                if 'is_keyword' in chunk_feats:\n                    chunk_feats['is_keyword'] = np.array(chunk_feats['is_keyword'], dtype=np.int8)\n                if 'is_stopword' in chunk_feats:\n                    chunk_feats['is_stopword'] = np.array(chunk_feats['is_stopword'], dtype=np.int8)\n                bio_chunk, label_chunk = get_chunk_labels(tok_chunk)\n                assert(all(len(feat) == len(tok_chunk) for _, feat in chunk_feats.items())), tuple((k, len(feat)) for k, feat in chunk_feats.items())\n                chunk_data = (row.name, *chunk_feats.values()) + ((bio_chunk, label_chunk) if not is_test else ())\n                if label_chunk is None or len(label_chunk):\n                    full_chunks.append(chunk_data)\n                elif 'is_upper' in chunk_feats and sum(chunk_feats['is_upper']) > 2:\n                    empty_chunks.append(chunk_data)\n                elif 'is_keyword' in chunk_feats and any(chunk_feats['is_keyword']):\n                    empty_chunks.append(chunk_data)\n#                 else:\n#                     empty_chunks.append(chunk_data)\n                if sidx + MAX_LEN >= len(tok_sent):\n                    break\n                sidx = sidx + MAX_LEN - OVERLAP\n#     if not is_test:\n#         assert(set(found_labels) == set(labels)), (set(found_labels), set(labels), row.name)\n    random.shuffle(empty_chunks)\n    return full_chunks + empty_chunks#[:len(full_chunks)]\n\ndef get_row_chunks(row, is_valid, is_test):\n    assert(MAX_LEN > 0), MAX_LEN\n    assert(OVERLAP < MAX_LEN), (OVERLAP, MAX_LEN)\n\n    full_chunks, empty_chunks = [], []\n    doc_feats = {feat: eval_iterable_field(row[feat]) for feat in USED_FEATURES}\n    tok_texts = doc_feats['tokenized_text'] \n    if not is_test:\n        labels = eval_iterable_field(row['label'])\n#         labels = [lab for lab in labels if lab in UNIQUE_LABELS]\n#         labels = remove_subset_labels(labels)\n#         assert(len(labels)), (labels, row['label'])\n        labels = sorted(labels, key=lambda x: len(x))\n        tok_labels = [tokenize(label) for label in labels]\n        found_labels = []\n\n    def get_sentence_positions(chunk_toks):\n        sent_pos = [0] + find_sublist([SEP], chunk_toks)\n        if chunk_toks[-1] != SEP:\n            sent_pos += [len(chunk_toks)]\n        _sent_positions = np.full(len(chunk_toks), -1)\n        for i, (s, e) in enumerate(zip(sent_pos[:-1], sent_pos[1:])):\n            if chunk_toks[s] == SEP and s != 0:\n                s += 1\n            if e != len(chunk_toks) and chunk_toks[e] == SEP:\n                e += 1\n            _sent_positions[s: e] = i\n        assert(-1 not in _sent_positions), (_sent_positions, chunk_toks)\n        return np.array(_sent_positions, dtype=np.uint8)\n    \n    def get_chunk_labels(chunk_txt):\n        if is_test:\n            return None, None\n        _found_labels = []\n        chunk_label = [TAGS_TO_IDX['O']] * len(chunk_txt)\n        for i, tok_label in enumerate(tok_labels):\n            positions = find_sublist(tok_label, chunk_txt)\n            if len(positions):\n                _found_labels.append(labels[i])\n            for pos in positions:\n                epos = pos + len(tok_label)\n                chunk_label[pos] = TAGS_TO_IDX['B']\n                chunk_label[pos+1: epos] = [TAGS_TO_IDX['I']] * (epos - pos - 1)\n        found_labels.extend(_found_labels)\n\n        assert(len(chunk_label) == len(chunk_txt)), (len(chunk_label) == len(chunk_txt))\n        return chunk_label, _found_labels\n\n    for i in range(len(tok_texts)):\n        sect_feats = {feat: doc_feats[feat][i] for feat in doc_feats}\n        tok_text = sect_feats['tokenized_text']\n        if len(tok_text) < MIN_LEN:\n            continue\n\n        sidx = 0\n        while True:\n            chunk_feats = {feat: sect_feats[feat][sidx: sidx+MAX_LEN] for feat in sect_feats}\n            tok_chunk = chunk_feats['tokenized_text']\n            if 'pos' in chunk_feats:\n#                     chunk_feats['pos'] = [p if tok_chunk[pi] not in LABEL_TOKENS else POS_TAGS['IMP']\n#                                           for pi, p in enumerate(chunk_feats['pos'])]\n                chunk_feats['pos'] = np.array(chunk_feats['pos'], dtype=np.int8)\n            if 'is_upper' in chunk_feats:\n                chunk_feats['is_upper'] = np.array(chunk_feats['is_upper'], dtype=np.int8)\n            if 'is_keyword' in chunk_feats:\n                chunk_feats['is_keyword'] = np.array(chunk_feats['is_keyword'], dtype=np.int8)\n            if 'is_stopword' in chunk_feats:\n                chunk_feats['is_stopword'] = np.array(chunk_feats['is_stopword'], dtype=np.int8)\n#             chunk_feats['sents'] = get_sentence_positions(tok_chunk)\n            bio_chunk, label_chunk = get_chunk_labels(tok_chunk)\n            assert(all(len(feat) == len(tok_chunk) for _, feat in chunk_feats.items())), tuple((k, len(feat)) for k, feat in chunk_feats.items())\n            chunk_data = (row.name, *chunk_feats.values()) + ((bio_chunk, label_chunk) if not is_test else ())\n            if is_valid or is_test or len(label_chunk):\n                full_chunks.append(chunk_data)\n            elif 'is_upper' in chunk_feats and sum(chunk_feats['is_upper']) > tok_chunk.count(SEP) + 2:\n                empty_chunks.append(chunk_data)\n            elif 'is_keyword' in chunk_feats and any(chunk_feats['is_keyword']):\n                empty_chunks.append(chunk_data)\n#                 else:\n#                     empty_chunks.append(chunk_data)\n            if sidx + MAX_LEN >= len(tok_text):\n                break\n            sidx = sidx + MAX_LEN - OVERLAP\n#     if not is_test:\n#         assert(set(found_labels) == set(labels)), (set(found_labels), set(labels), row.name)\n    random.shuffle(empty_chunks)\n    return full_chunks + empty_chunks#[:len(full_chunks)]\n\ndef get_row_chunks_subprocess(arg):\n    return get_row_chunks(*arg)\n    \ndef prepare_data_chunks(dt, is_valid, is_test, chunk_size=4000):\n    print('Getting data chunks...')\n    args = tuple((row, is_valid, is_test) for _, row in dt.iterrows())\n    n_chunks = int(np.ceil(len(args) / chunk_size))\n    for c in tqdm(range(n_chunks)):\n        with Pool() as pool:\n            _args = args[c*chunk_size: (c+1)*chunk_size]\n            chunks = [_chunk for _chunks in tqdm(pool.imap_unordered(get_row_chunks_subprocess, _args),\n                                                 total=len(_args), ncols=LEN_TBAR)\n                      for _chunk in _chunks]\n        cols = ['Id', *USED_FEATURES] + (['bio_label', 'label'] if not is_test else [])\n        chunks = pd.DataFrame(chunks, columns=cols).astype({'Id': 'category'})\n        for col in cols[1:]:\n            chunks[col] = chunks[col].apply(tuple)        \n        chunks.to_csv(os.path.join(OUT_DIR, f'data_chunk_{c+1}.csv'), index=False)\n        del chunks\n        gc.collect()\n        \n    print()\n    df = None    \n    for c in tqdm(range(n_chunks), desc='Merging chunks'):\n        df_filepath = os.path.join(OUT_DIR, f'data_chunk_{c+1}.csv')\n        _df = pd.read_csv(df_filepath)\n        if df is not None:\n            df = pd.concat([df, _df], ignore_index=True)\n        else:\n            df = _df\n        del _df\n        gc.collect()\n        os.remove(df_filepath)\n    for col in tqdm(cols[1:], desc='Evaluating fields'):\n        if col in ['pos', 'is_upper', 'is_keyword', 'is_stopword', 'sent_pos', 'bio_label']:\n            df[col] = df[col].apply(lambda x: np.array(eval_iterable_field(x), dtype=np.uint8))\n        else:\n            df[col] = df[col].apply(eval_iterable_field)\n    return df\n\ndef group_similar_labels(labels):\n    labs = sorted(list(set(labels)), key=lambda x: -len(x.split()))\n    labs_toks = [lab.split() for lab in labs]\n    all_labs = list(zip(labs, labs_toks))\n    similar_uniques = []\n    while len(all_labs):\n        lab, tok = all_labs.pop(0)\n        similars = [lab]\n        similars_toks = [tok]\n        i = 0\n        while i < len(all_labs):\n            lab2, tok2 = all_labs[i]\n            if any(jaccard(lab1, lab2) >= 0.4 for lab1 in similars) or any(all(t in tok1 for t in tok2) for tok1 in similars_toks):\n                lab, tok = all_labs.pop(i)\n                similars.append(lab)\n                similars_toks.append(tok)\n            else:\n                i += 1\n        similar_uniques.append(similars)\n    return similar_uniques\n\ndef get_label_tokens(labels):\n    labels = [tokenize(remove_stopwords(clean_label(lab))) for lab in labels]\n#     similar_labs = group_similar_labels(labels)\n    toks = set(np.hstack(labels))\n#     toks = Counter(toks)\n#     count = sum(toks.values())\n#     toks = {k: v/count for k, v in toks.items()}\n    return toks","metadata":{"_uuid":"6171af98-c5db-4196-b126-c5962edcddae","_cell_guid":"a8f71aac-1818-490f-94c3-308993b6e76a","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2021-06-03T02:32:22.110720Z","iopub.execute_input":"2021-06-03T02:32:22.111035Z","iopub.status.idle":"2021-06-03T02:32:22.164861Z","shell.execute_reply.started":"2021-06-03T02:32:22.111008Z","shell.execute_reply":"2021-06-03T02:32:22.164014Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Coleridge(torch.utils.data.Dataset):\n    def __init__(self, samples, is_test):\n        super().__init__()\n        self.is_test = is_test\n        self.len = len(samples)\n        if not is_test:\n            self.full_samples = samples[samples.label.apply(lambda x: len(x) > 0)].values.tolist()\n            self.empty_samples = samples[samples.label.apply(lambda x: len(x) == 0)].values.tolist()\n            self.len = len(self.full_samples) * 4\n            self.f_idx = 0\n            self.e_idx = 0\n        else:\n            self.samples = samples.values.tolist()\n\n    def __len__(self):\n        return self.len\n\n    def __getitem__(self, idx):\n        if not self.is_test:\n            if self.f_idx >= len(self.full_samples):\n                self.f_idx = 0\n            if self.e_idx >= len(self.empty_samples):\n                self.e_idx = 0\n            if np.random.rand() < 0.25:\n                self.f_idx += 1\n                return self.full_samples[self.f_idx - 1]\n            else:\n                self.e_idx += 1\n                return self.empty_samples[self.e_idx - 1]\n        return self.samples[idx]\n\ndef collate(batch):\n    (idx, \n     texts, \n#      stems, \n#      pos, \n     uppers,\n     keywords,\n     stopwords,\n#      sents, \n     labels, _) = zip(*batch)\n#     inp = ENGLISH.process(texts, device=DEVICE).long()  # pad and numericalize\n    texts = np.array(ENGLISH.pad(texts)).T  # pad\n#     inp = ENGLISH.process(stems, device=DEVICE).long()  # pad and numericalize\n#     pos = LABEL.process(pos, device=DEVICE).long()\n    uppers = LABEL.process(uppers, device=DEVICE).long()\n    keywords = LABEL.process(keywords, device=DEVICE).long()\n    stopwords = LABEL.process(stopwords, device=DEVICE).long()\n#     sents = LABEL.process(sents, device=DEVICE).long()\n    labels = LABEL.process(labels, device=DEVICE).long()  # pad\n    return (idx, \n            texts, \n#             inp, \n#             pos, \n            uppers, \n            keywords, \n            stopwords,\n#             sents,\n            labels)\n\ndef collate_test(batch):\n    (idx, \n     texts, \n#      stems, \n#      pos, \n     uppers,\n     keywords,\n     stopwords,\n#      sents\n    ) = zip(*batch)\n#     inp = ENGLISH.process(texts, device=DEVICE).long()  # pad and numericalize\n    texts = np.array(ENGLISH.pad(texts)).T\n#     inp = ENGLISH.process(stems, device=DEVICE).long()  # pad and numericalize\n#     pos = LABEL.process(pos, device=DEVICE).long()\n    uppers = LABEL.process(uppers, device=DEVICE).long()\n    keywords = LABEL.process(keywords, device=DEVICE).long()\n    stopwords = LABEL.process(stopwords, device=DEVICE).long()\n#     sents = LABEL.process(sents, device=DEVICE).long()\n    return (idx, \n            texts, \n#             inp, \n#             pos, \n            uppers, \n            keywords, \n            stopwords,\n#             sents\n           )\n\ndef create_dataloader(dt, used_feats, is_valid=False, is_test=False):\n    if isinstance(dt, list):\n        dt = prepare_dataset(dt, used_feats, chunk_size=4000, is_test=is_test)\n    ds = prepare_data_chunks(dt, is_valid=is_valid, is_test=is_test, chunk_size=4000)\n    dataset = Coleridge(ds, is_test=(is_test or is_valid))\n    dataloader = torch.utils.data.DataLoader(\n        dataset=dataset,\n        batch_size=TEST_BATCH_SIZE if (is_test or is_valid) else BATCH_SIZE,\n        collate_fn=collate_test if is_test else collate,\n        shuffle=False if is_test else True)\n    \n    return dt, ds, dataset, dataloader\n\ndef train_valid_split(df, val_size, verbose=True):\n    seed_everything(RANDOM_SEED)\n    unique_labels = sorted(set(np.hstack(df.label.apply(eval))))\n#     unique_labels = [lab for lab in unique_labels if lab in UNIQUE_LABELS]\n    print('Number of unique labels:', len(unique_labels))\n    similar_uniques = group_similar_labels(unique_labels)\n    trn_count = int(len(df) * (1 - val_size))\n    trn_df = df.copy()\n    random.shuffle(similar_uniques)\n    similar_uniques = iter(similar_uniques)\n    while len(trn_df) > trn_count:\n        try:\n            labs = next(similar_uniques)\n        except StopIteration:\n            break\n        idx = []\n        for lab in labs:\n            idx.extend(list(trn_df[trn_df.label.str.contains(f\"'{lab}'\")].index))  \n        idx = set(idx)\n        if len(trn_df) - len(idx) < trn_count * 0.99:\n            continue\n        trn_df.drop(index=idx, inplace=True)\n    val_df = df.drop(index=trn_df.index)\n    \n    if verbose:\n        trn_labs = set(lab for lab_list in trn_df.label for lab in eval(lab_list))\n        val_labs = set(lab for lab_list in val_df.label for lab in eval(lab_list))\n        val_uniques = list(val_labs ^ val_labs.intersection(trn_labs))\n        sz = len(val_uniques)\n        n_splits = 5\n        pad_sz = int(np.ceil(sz / n_splits) * n_splits - sz)\n        val_uniques += [''] * pad_sz\n        split_sz = len(val_uniques) // n_splits\n        out = pd.DataFrame({f'val_uniques_{ni}': val_uniques[ni*split_sz:(ni+1)*split_sz] for ni in range(n_splits)})\n        styler = out.head(100).style.set_properties(**{'text-align': 'left', 'width': '300px'})\n        styler = styler.set_table_styles([dict(selector='th', props=[('text-align', 'left')])])\n        display(styler)\n    return trn_df, val_df","metadata":{"_uuid":"959d5384-5a51-49f6-a3be-5407b9503d6c","_cell_guid":"bbb6cb59-afc1-449d-aaf7-1fd11da15aa7","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2021-06-03T02:32:22.166292Z","iopub.execute_input":"2021-06-03T02:32:22.166686Z","iopub.status.idle":"2021-06-03T02:32:22.193189Z","shell.execute_reply.started":"2021-06-03T02:32:22.166649Z","shell.execute_reply":"2021-06-03T02:32:22.192193Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Transformer Model","metadata":{"_uuid":"a34bb383-cdf6-4dfd-ac77-acc5774bba07","_cell_guid":"6c19d848-6507-43f4-b04b-b3d6b75d3ee8","trusted":true}},{"cell_type":"code","source":"class PositionalEncoding(nn.Module):\n    def __init__(self, d_model, max_len):\n        super().__init__()\n        pe = torch.zeros(max_len, d_model)\n        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)\n        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (\n            -math.log(10000.0) / d_model))\n        pe[:, 0::2] = torch.sin(position * div_term)\n        pe[:, 1::2] = torch.cos(position * div_term)\n        self.register_buffer('pe', pe)\n\n    def forward(self, pos):\n        return self.pe[pos, :]\n\nclass TransformerModel(nn.Module):\n    def __init__(self, **params):\n        print(params)\n        super().__init__()\n        self.params = params.copy()\n        self.d_model = params['d_model']\n        nhead = params['nhead']\n        max_len = params['max_len']\n        dim_feedforward = params['dim_feedforward']\n        n_enc_layers = params['num_encoder_layers']\n        dropout_rate = params['dropout']\n        activation = params['activation']\n#         word_vocab_size = len(ENGLISH.vocab)\n        target_size = len(TAGS_TO_IDX)\n#         assert(ENGLISH.vocab.stoi[PAD] == LABEL.pad_token + 1), (ENGLISH.vocab.stoi[PAD], LABEL.pad_token)\n        self.pad_idx = LABEL.pad_token + 1\n        \n        self.norm = nn.LayerNorm(self.d_model, elementwise_affine=False)\n#         self.norm_2 = nn.LayerNorm(dim_feedforward*2, elementwise_affine=False)\n        encoder_layer = nn.TransformerEncoderLayer(d_model=self.d_model, \n                                                   nhead=nhead,\n                                                   dim_feedforward=dim_feedforward, \n                                                   dropout=dropout_rate, \n                                                   activation=activation)\n        self.encoder = nn.TransformerEncoder(encoder_layer=encoder_layer,\n                                             num_layers=n_enc_layers,\n                                             norm=self.norm)\n        self.lstm = nn.GRU(input_size=self.d_model,\n                           hidden_size=self.d_model//2,\n                           num_layers=n_enc_layers,\n                           dropout=dropout_rate,\n                           bidirectional=True)\n#         self.lstm_2 = nn.LSTM(input_size=self.d_model,\n#                               hidden_size=dim_feedforward,\n#                               num_layers=1,\n#                               dropout=dropout_rate,\n#                               bidirectional=True)\n\n        self.linear = nn.Linear(self.d_model, target_size)\n#         self.linear = nn.Linear(dim_feedforward*2, target_size)\n        self.crf = CRF(target_size)\n        \n        self.position_embeddings = PositionalEncoding(self.d_model, max_len)\n#         self.word_embeddings = \\\n#             nn.Embedding(num_embeddings=word_vocab_size,\n#                          embedding_dim=self.d_model,\n#                          padding_idx=self.pad_idx,\n#                          #  scale_grad_by_freq=True\n#                         )\n#         self.pos_embeddings = \\\n#             nn.Embedding(num_embeddings=len(POS_TAGS)+1,\n#                          embedding_dim=self.d_model,\n#                          padding_idx=self.pad_idx,\n#                          #  scale_grad_by_freq=True\n#                         )\n        self.upper_embeddings = \\\n            nn.Embedding(num_embeddings=4, # +1 for pad token.\n                         embedding_dim=self.d_model,\n                         padding_idx=self.pad_idx,\n#                          scale_grad_by_freq=True\n                         )\n        self.keyword_embeddings = \\\n            nn.Embedding(num_embeddings=3, # +1 for pad token.\n                         embedding_dim=self.d_model,\n                         padding_idx=self.pad_idx,\n#                          scale_grad_by_freq=True\n                         )\n        self.stopword_embeddings = \\\n            nn.Embedding(num_embeddings=3, # +1 for pad token.\n                         embedding_dim=self.d_model,\n                         padding_idx=self.pad_idx,\n#                          scale_grad_by_freq=True\n                         )\n        self.target_embeddings = \\\n            nn.Embedding(num_embeddings=target_size + 1, # +1 for pad token.\n                         embedding_dim=self.d_model,\n                         padding_idx=self.pad_idx,\n#                          scale_grad_by_freq=True\n                         )\n        \n        self.dropout = nn.Dropout(dropout_rate)\n        self.device = params['device']        \n#         self.init_weights()\n\n    def get_pad_mask(self, arr):\n        pad_mask = (arr.transpose(1, 0) == self.pad_idx)\n        return pad_mask.to(self.device)\n\n    def get_positions(self, tens):\n        pos = torch.arange(0, tens.shape[0], device=self.device).unsqueeze(0)\n        pos = pos.repeat(tens.shape[1], 1)\n        emb_pos = self.position_embeddings(pos).transpose(0, 1)        \n        return emb_pos\n    \n    def init_weights(self):\n        initrange = 0.1\n#         self.word_embeddings.weight.data = ENGLISH.vocab.vectors.to(DEVICE)\n#         self.pos_embeddings.weight.data.uniform_(-initrange, initrange)\n        self.upper_embeddings.weight.data.uniform_(-initrange, initrange)\n        self.keyword_embeddings.weight.data.uniform_(-initrange, initrange)\n        self.stopword_embeddings.weight.data.uniform_(-initrange, initrange)\n        self.target_embeddings.weight.data.uniform_(-initrange, initrange)\n        self.linear.bias.data.zero_()\n        self.linear.weight.data.uniform_(-initrange, initrange)\n        self.crf.reset_parameters()\n        \n    def forward(self, \n#                 text,\n#                 pos,\n                uppers,\n                keywords,\n                stopwords,\n               ):\n#         pos = pos + 1\n        uppers = uppers + 1\n        keywords = keywords + 1\n        stopwords = stopwords + 1\n        \n        emb_src = (\n#             self.word_embeddings(text)\n#             + self.pos_embeddings(pos)\n            self.upper_embeddings(uppers)\n            + self.keyword_embeddings(keywords)\n            + self.stopword_embeddings(stopwords)\n        ) * np.sqrt(self.d_model)\n        emb_src = self.norm(emb_src)\n        output, _ = self.lstm(emb_src)\n        output = self.linear(self.norm(output)).squeeze(-1)\n        return output\n\n    def forward2(self, \n#                 text,\n#                 pos,\n                uppers,\n                keywords,\n                stopwords,\n#                 sents\n               ):\n#         pos = pos + 1\n        uppers = uppers + 1\n        keywords = keywords + 1\n        stopwords = stopwords + 1\n#         sents = sents + 1\n        \n        emb_tok_positions = self.get_positions(sents)\n#         emb_sent_positions = self.position_embeddings(sents)\n        emb_src = (\n#             self.word_embeddings(text)\n#             + self.pos_embeddings(pos)\n            self.upper_embeddings(uppers)\n            + self.keyword_embeddings(keywords)\n            + self.stopword_embeddings(stopwords)\n        ) * np.sqrt(self.d_model)\n        emb_src = self.norm(emb_src \n                            + emb_src_positions \n#                             + emb_sent_positions\n                           )\n        output = self.encoder(src=emb_src,\n                              src_key_padding_mask=self.get_pad_mask(sents))\n        output = self.linear(self.norm(output)).squeeze(-1)\n        return output","metadata":{"_uuid":"1df04d26-e9ec-4884-9142-4a4b207a5697","_cell_guid":"e1ea3e9f-201e-4ddf-a709-6d84c457ebdb","execution":{"iopub.status.busy":"2021-06-03T02:32:22.195016Z","iopub.execute_input":"2021-06-03T02:32:22.195564Z","iopub.status.idle":"2021-06-03T02:32:22.221209Z","shell.execute_reply.started":"2021-06-03T02:32:22.195528Z","shell.execute_reply":"2021-06-03T02:32:22.220360Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_epoch(estimator, trn_iterator, optim, criterion,\n                device=\"cpu\"):\n    estimator.train()\n    \n    loss_sum = 0\n    count = 0\n    outputs = []\n    labels = []\n    ids = []\n    \n    tbar = tqdm(trn_iterator, ncols=LEN_TBAR)\n    for batch in tbar:\n        (idx, \n         text, \n#          inp, \n#          pos, \n         uppers, \n         keywords,\n         stopwords,\n#          sents,\n         label) = batch  # (S, N)\n        mask = (label != LABEL_PAD)\n        \n        optim.zero_grad()\n        output = estimator(\n#             inp,\n#             pos,\n            uppers,\n            keywords,\n            stopwords,\n#             sents\n        )\n        _output = LABEL.process(estimator.crf.decode(output, mask=mask))\n        \n        ids.append(idx)\n        labels.append(tokenized_to_text(text, label))\n        outputs.append(tokenized_to_text(text, _output))\n        \n#         print(list(_output))\n        #print(_output[:, 0].tolist())\n        loss = estimator.crf(output, label, mask=mask) * -1\n#         loss_2 = criterion(output.reshape(-1, output.shape[2]), label.reshape(-1))\n        loss.backward()\n        clip_grad_norm_(estimator.parameters(), max_norm=1)\n        optim.step()\n\n        loss_sum += loss.item()\n        count += 1\n\n        tbar.set_description(\n            'trn_loss - {:.6f}'.format(loss_sum / count))\n\n    loss = loss_sum / count\n\n    ids = np.hstack(ids)\n    labels = np.hstack(labels)\n    outputs = np.hstack(outputs)\n    \n    labels = merge_article_labels(ids, labels, remove_similars=False, remove_subsets=False)\n    outputs = merge_article_labels(ids, outputs, remove_similars=False, remove_subsets=False)\n\n    score = score_coleridge_initiative(outputs.PredictionString, labels.PredictionString, beta=0.5, verbose=True)\n    print_random_samples(labels.Id, outputs.PredictionString, labels.PredictionString, 20)\n    \n    return loss, score\n\n\ndef val_epoch(estimator, val_iterator, criterion, device=\"cpu\"):\n    estimator.eval()\n\n    loss_sum = 0\n    count = 0\n    outputs = []\n    labels = []\n    ids = []\n    \n    tbar = tqdm(val_iterator, ncols=LEN_TBAR)\n    for batch in tbar:\n        (idx, \n         text, \n#          inp, \n#          pos, \n         uppers,\n         keywords,\n         stopwords,\n#          sents,\n         label) = batch\n        mask = (label != LABEL_PAD)\n        \n        with torch.no_grad():\n            output = estimator(\n#                 inp,\n#                 pos,\n                uppers,\n                keywords,\n                stopwords,\n#                 sents\n            )\n            _output = LABEL.process(estimator.crf.decode(output, mask=mask))\n            \n        ids.append(idx)\n        labels.append(tokenized_to_text(text, label))\n        outputs.append(tokenized_to_text(text, _output)) \n        \n        loss = estimator.crf(output, label, mask=mask) * -1\n#         loss_2 = criterion(output.reshape(-1, output.shape[2]), label.reshape(-1))\n        loss_sum += loss.item()\n        count += 1\n        \n        tbar.set_description(\n            'val_loss - {:.6f}'.format(loss_sum / count))\n\n    loss = loss_sum / count\n    ids = np.hstack(ids)\n    labels = np.hstack(labels)\n    outputs = np.hstack(outputs)\n    \n    labels = merge_article_labels(ids, labels, remove_similars=False, remove_subsets=False)\n    outputs = merge_article_labels(ids, outputs, remove_similars=False, remove_subsets=False)\n\n    score = score_coleridge_initiative(outputs.PredictionString, labels.PredictionString, \n                                       beta=0.5, verbose=True)\n    print_random_samples(labels.Id, outputs.PredictionString, labels.PredictionString, 20)\n    \n    return loss, score\n\n\ndef train_transformer(estimator_type, estimator_params, epochs, batch_size,\n                      device, early_stopping, eps, nworkers, cont, debug,\n                      debug_size, data_size, valid_size):\n    print('DEBUG MODE:', bool(debug))\n    seed_everything(RANDOM_SEED)\n    \n    train = TRAIN_DATA\n    print('# of documents:', train.shape[0])\n    \n    data_size = int(len(train) * (debug_size if debug else data_size))\n    train = train.sample(data_size)\n    print('# of documents used:', train.shape[0])\n    \n    train, valid = train_valid_split(train, valid_size)\n#     train, valid = train_test_split(train, test_size=valid_size)\n    print('# of train documents:', train.shape[0])\n    print('# of valid documents:', valid.shape[0])\n    \n    used_feats = (*USED_FEATURES, 'label')\n    _, trn_data, trn_dataset, trn_dataloader = create_dataloader(dt=train, used_feats=used_feats)\n    _, _, val_dataset, val_dataloader = create_dataloader(dt=valid, used_feats=used_feats, is_valid=True)\n    print('# of train samples:', len(trn_dataset))\n    print('# of valid samples:', len(val_dataset))\n    del trn_dataset, val_dataset\n    gc.collect()\n        \n    last_score = 0\n    last_epoch = 0\n    if cont and not debug:\n        estimator, optim, sched, cp = \\\n            load_model(MODEL, for_training=True)\n        last_score = cp['val_score']\n        last_epoch = cp['epoch']\n        print('Previous Validation score:', last_score)\n    else:\n        # Build and save vocabularies.\n#         ENGLISH.build_vocab(trn_data.tokenized_text, \n# #                             [list(LABEL_TOKENS.keys())], \n#                             specials=[PAD, UNK, SEP], \n#                             vectors=torchtext.vocab.Vectors(GLOVE_DIR + f'/glove.6B.{D_MODEL}d.txt'))\n#         ENGLISH.pad_token, ENGLISH.unk_token = PAD, UNK\n#         save_vocab(ENGLISH, VOCAB_FILENAME)\n        \n        estimator = create_model(estimator_type, **estimator_params)\n        estimator.init_weights()\n        optim = create_optimizer(estimator, lr=LEARNING_RATE)\n        sched = create_scheduler(optim)\n        \n#     print('Vocabulary size:', len(ENGLISH.vocab))\n    del trn_data\n    gc.collect()\n        \n    criterion = nn.CrossEntropyLoss(ignore_index=LABEL_PAD).to(device)\n    over_fit = 0\n    seed_everything(RANDOM_SEED)\n    for epoch in range(last_epoch, epochs):      \n        if SCHED_TYPE == 'warmup':\n            for g in optim.param_groups:\n                g['lr'] = ((estimator.d_model ** -0.5) * \n                           min((epoch + 1) ** -0.5, \n                               (epoch + 1) * (WARMUP_STEPS ** -1.5)))\n        print(f'\\nEpoch {epoch + 1} - Learning rate:',\n              optim.param_groups[0]['lr'])\n        trn_loss, trn_score = train_epoch(estimator, trn_dataloader, optim, criterion, device)\n        print(\"  TRAINING -- loss: {:.6f} - score: {:.6f}\"\\\n              .format(trn_loss, trn_score))\n\n        val_loss, val_score = val_epoch(estimator, val_dataloader, criterion, device)\n        color = '\\033[91m' if val_score > last_score and epoch != 0 else ''\n        print(color + \"  VALIDATION -- loss: {:.6f} - score: {:.6f} - best: {:.6f}\"\\\n              .format(val_loss, val_score, last_score) + '\\033[0m')\n\n        if SCHED_TYPE == 'plateau':\n            sched.step(val_score)\n\n        if val_score > last_score + eps:\n            last_score = val_score\n            over_fit = 0\n            save_model(estimator, optim, sched, val_score, epoch)\n        else:\n            over_fit += 1\n\n        if early_stopping > 0 and over_fit >= early_stopping:\n            print(\"early stop epoch \", epoch + 1)\n            break\n\n    return estimator","metadata":{"_uuid":"bdd12204-03b0-4bd3-8182-5cfb594375c5","_cell_guid":"7b923b68-8b7f-4a8f-86ec-7ff258f4c7dc","execution":{"iopub.status.busy":"2021-06-03T02:32:22.224259Z","iopub.execute_input":"2021-06-03T02:32:22.224548Z","iopub.status.idle":"2021-06-03T02:32:22.253715Z","shell.execute_reply.started":"2021-06-03T02:32:22.224522Z","shell.execute_reply":"2021-06-03T02:32:22.252752Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"MODEL = TransformerModel  # @param [\"TransformerModel\"] {type:\"raw\"}\n\nMODEL_FILENAME = f'{MODEL.__name__}_best.pth'  # @param {type:\"string\"}\n\nVOCAB_FILENAME = 'vocab.pkl'\n\n\n# @markdown # Model Settings\n\nMAX_LEN = 256  # @param {type:\"integer\"}\n\nMIN_LEN = 10  # @param {type:\"integer\"}\n\nOVERLAP = 32  # @param {type:\"integer\"}\n\nD_MODEL = 128  # @param {type:\"integer\"}\n\nNHEAD = 4  # @param {type:\"integer\"}\n\nN_ENC_LAYERS = 2 # @param {type:\"integer\"}\n\nDIM_FEEDFORWARD = 64  # @param {type:\"integer\"}\n\nDROPOUT = 0.1  # @param {type:\"slider\", min:0, max:1, step:0.05}\n\nACTIVATION = \"gelu\"  # @param [\"relu\", \"gelu\"]\n\nLEARNING_RATE = 0.0001  # @param {type:\"number\"}\n\n\n# @markdown # Training Settings\n\nEPOCHS = 50  # @param {type:\"integer\"}\n\nBATCH_SIZE = 512  # @param {type:\"integer\"}\n\nTEST_BATCH_SIZE = 1028  # @param {type:\"integer\"}\n\nVALID_SIZE = 0.1  # @param {type:\"slider\", min:0, max:1, step:0.05}\n\nEARLY_STOPPING = 13  # @param {type:\"integer\"}\n\nNUM_WORKERS = 0  # @param {type:\"integer\"}\n\nSCHED_TYPE = \"NONE\"  # @param [\"plateau\", \"warmup\", \"NONE\"]\n\nWARMUP_STEPS = int(EPOCHS * 0.3)  # @param {type:\"integer\"}","metadata":{"_uuid":"d07b0fcf-7836-4491-bff1-c551e791dd67","_cell_guid":"000331f2-9acb-4cfb-a4fe-b4a14466d509","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2021-06-03T02:32:22.254838Z","iopub.execute_input":"2021-06-03T02:32:22.255292Z","iopub.status.idle":"2021-06-03T02:32:22.268659Z","shell.execute_reply.started":"2021-06-03T02:32:22.255255Z","shell.execute_reply":"2021-06-03T02:32:22.267754Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"RETRAIN = 1  # @param {type:\"boolean\"}\n\nCONTINUE = 0  # @param {type:\"boolean\"}\n\nDEBUG = 0  # @param {type:\"boolean\"}\n\nDATA_SIZE = 1  # @param {type:\"slider\", min:0.01, max:1, step:0.01}\n\nDEBUG_SIZE = 0.4 if torch.cuda.is_available() else 0.1\n\n\nPAD = '<pad>'\nUNK = '<unk>'\nLABEL_PAD = -1\nTAGS_TO_IDX = {'O': 0, 'B': 1, 'I': 2}\nUSED_FEATURES = ['tokenized_text', \n#                  'pos', \n                 'is_upper', \n                 'is_keyword', \n                 'is_stopword']\n\n# ENGLISH = Field(tokenize=tokenize, lower=True, pad_token=None, unk_token=None)\nENGLISH = Field(use_vocab=False, pad_token=PAD, lower=True)\nLABEL = Field(use_vocab=False, pad_token=LABEL_PAD)#, dtype=torch.float)\n\n# if LABEL_TOKENS is None:\n#     LABEL_TOKENS = get_label_tokens(UNIQUE_LABELS)\n    \nDEVICE = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")","metadata":{"_uuid":"f4366314-55c4-447d-b349-563246de5a2c","_cell_guid":"f7a04861-a3e3-40d7-a873-83794a3823ec","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2021-06-03T02:32:22.270140Z","iopub.execute_input":"2021-06-03T02:32:22.270591Z","iopub.status.idle":"2021-06-03T02:32:22.332889Z","shell.execute_reply.started":"2021-06-03T02:32:22.270551Z","shell.execute_reply":"2021-06-03T02:32:22.332112Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if __name__ == '__main__':\n    print('Using Device -', DEVICE)\n    print('Using Model -', MODEL.__name__)\n    model_params = {\n        'max_len': MAX_LEN,\n        'd_model': D_MODEL,\n        'nhead': NHEAD,\n        'num_encoder_layers': N_ENC_LAYERS,\n        'dim_feedforward': DIM_FEEDFORWARD,\n        'dropout': DROPOUT,\n        'activation': ACTIVATION,\n        'device': DEVICE\n    }\n\n    train_params = {\n        'estimator_type': MODEL,\n        'estimator_params': model_params,\n        'epochs': EPOCHS,\n        'batch_size': BATCH_SIZE,\n        'early_stopping': EARLY_STOPPING,\n        'nworkers': NUM_WORKERS,\n        'device': DEVICE,\n        'eps': 1e-6,\n        'cont': CONTINUE,\n        'debug': DEBUG,\n        'data_size': DATA_SIZE,\n        'debug_size': DEBUG_SIZE,\n        'valid_size': VALID_SIZE\n    }\n        \n    if RETRAIN and not COMPILE:\n        if not DEBUG:\n            files = os.listdir(OUT_DIR)\n            for f in os.listdir(MODELS_DIR):\n                if f not in files:     \n                    from distutils.dir_util import copy_tree\n                    copy_tree(MODELS_DIR, OUT_DIR)\n                    break\n        if not CONTINUE or DEBUG:\n            VERSION = 1\n            while os.path.exists(os.path.join(MODELS_DIR, str(VERSION))):\n                VERSION = int(VERSION) + 1\n            VERSION = str(VERSION)\n            \n        if TRAIN_DATA is None:\n            print(f\"Loading preprocessed data from {DATA_DIR}...\")\n            TRAIN_DATA = pd.read_csv(DATA_DIR + '/train_preprocessed.csv', \n                                     index_col='Id')\n#         if CONTINUE:\n#             ENGLISH = load_vocab(VOCAB_FILENAME)\n        print('Version:', VERSION)\n        print('Parameters:', train_params)\n        model = train_transformer(**train_params)","metadata":{"_uuid":"42cfd570-2ff2-4c88-b30b-c3448aeced8a","_cell_guid":"668d5774-c57d-436b-b4f9-a7af68212f04","execution":{"iopub.status.busy":"2021-06-03T02:32:22.334156Z","iopub.execute_input":"2021-06-03T02:32:22.334521Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Test","metadata":{"_uuid":"15d1f391-5e84-4c91-95e9-25d4b3ee0da8","_cell_guid":"e5581b5a-1357-499a-a42a-2ecd03508a66","trusted":true}},{"cell_type":"code","source":"def get_predictions(dt_loader):\n    preds = []\n    ids = []\n\n    estimator = load_model(MODEL, for_training=False)\n    estimator.eval()\n    \n    tbar = tqdm(dt_loader, ncols=LEN_TBAR)\n    for batch in tbar:\n        (idx, \n        text, \n        inp, \n#         pos, \n        uppers, \n        keywords, \n        stopwords) = batch\n        mask = (inp != ENGLISH.vocab.stoi[PAD])\n        with torch.no_grad():\n            output = estimator(inp, \n#                                pos, \n                               uppers, \n                               keywords, \n                               stopwords)\n            _output = LABEL.process(estimator.crf.decode(output, mask=mask))\n        ids.append(idx)\n        preds.append(tokenized_to_text(text, _output))\n        \n    return np.hstack(ids), np.hstack(preds)\n\ndef remove_known_labels_subprocess(label):\n    out = []\n    for lab in label.split('|'):\n        if all(jaccard(lab.strip(), known_lab) < 0.5 for known_lab in KNOWN_LABELS):\n            out.append(lab)\n    return '|'.join(out)   \n\ndef remove_singles_subprocess(label):\n    out = []\n    for lab in label.split('|'):\n        if len(set(lab.split())) > 1 or lab in UNIQUE_LABELS:\n            out.append(lab)\n    return '|'.join(out)     \n    \ndef postprocess_predictions(preds, func):\n    with Pool() as pool:\n        return pool.map(func, preds)    \n    \ndef remove_similars_2d(labs1, labs2):\n    arr1 = labs1.split('|')\n    arr2 = labs2.split('|')\n    out = list(arr1)\n    for lab1 in arr1:\n        if any(jaccard(lab1, lab2) >= 0.5 for lab2 in arr2):\n            out.remove(lab1)\n    return '|'.join(out)\n\ndef remove_subsets_2d(labs1, labs2):\n    arr1 = labs1.split('|')\n    arr2 = labs2.split('|')\n    tok_arr1 = [lab.split() for lab in arr1]\n    tok_arr2 = [lab.split() for lab in arr2]\n    out = list(arr1)\n    for i, lab1 in enumerate(arr1):\n        tok_lab1 = tok_arr1[i]\n        if any(all(tok1 in tok_lab2 for tok1 in tok_lab1) for tok_lab2 in tok_arr2):\n            out.remove(lab1)\n    return '|'.join(out)","metadata":{"_uuid":"4d43bc0a-1132-4f12-ab4c-b6cb9595bb90","_cell_guid":"e01b35f9-f595-466e-91ce-e9148c6c04f6","collapsed":false,"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print('TESTING...')\nis_local = False\n\n# post-processing:\nremove_knowns = False\nremove_singles = True\nget_fuzzy_predictions = True\n\nENGLISH = load_vocab(VOCAB_FILENAME)\n\nnrows = 100\ntest_dir = os.path.join(INPUT_DIR, 'test' if not is_local else 'train')\ntest_paths = [os.path.join(test_dir, fname) for fname in os.listdir(test_dir)]\nused_feats = USED_FEATURES + (['label'] if get_fuzzy_predictions else [])\nif is_local:\n    random.shuffle(test_paths)\n    test_paths = test_paths[:nrows]\n    \ntst_dt, _, _, tst_dataloader = create_dataloader(dt=test_paths, used_feats=used_feats, is_test=True)\nids, predictions = get_predictions(tst_dataloader)\n\ndel tst_dataloader\ndel ENGLISH\ngc.collect()\n\n# Prepare submission file:\nsubmission = merge_article_labels(ids, predictions, remove_similars=False, remove_subsets=True)\nif remove_singles:\n    submission.PredictionString = postprocess_predictions(submission.PredictionString, remove_singles_subprocess)\nif remove_knowns:\n    submission.PredictionString = postprocess_predictions(submission.PredictionString, remove_known_labels_subprocess)\nif get_fuzzy_predictions:\n    fuzzy_predictions = tst_dt.label.apply(lambda x: '|'.join(remove_subset_labels(eval(x))))\n    submission.PredictionString = submission.apply(lambda x: remove_subsets_2d(x.PredictionString, fuzzy_predictions.loc[x.Id]), axis=1)\n    submission.PredictionString = submission.apply(lambda x: remove_similars_2d(x.PredictionString, fuzzy_predictions.loc[x.Id]), axis=1)\n    submission.PredictionString = submission.apply(lambda x: clean_predictions('|'.join([x.PredictionString, fuzzy_predictions.loc[x.Id]])), axis=1)\n\nif is_local:\n    labels = pd.read_csv(DATA_DIR + '/train_preprocessed.csv', index_col='Id', usecols=['Id', 'label'])\n    labels = labels.label.apply(lambda x: '|'.join(remove_subset_labels(eval(x))))\n    submission = submission.merge(labels, on='Id')\n    score = score_coleridge_initiative(submission.PredictionString, submission.label, beta=0.5, verbose=True)\n\nstyler = submission.head().style.set_properties(**{'text-align': 'left', 'width': '300px'})\nstyler = styler.set_table_styles([dict(selector='th', props=[('text-align', 'left')])])\ndisplay(styler) \n\nsubmission.to_csv('submission.csv', index=False)","metadata":{"_uuid":"eae8bfd0-6fa5-4af6-af9b-0fa9b9ba3690","_cell_guid":"4ad0ad1e-5bd1-4eab-95b6-5783037dc8fb","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Fuzzy Prediction","metadata":{}},{"cell_type":"code","source":"# print('TESTING...')\n\n# is_local = False\n# nrows = 9000\n\n# test_dir = os.path.join(INPUT_DIR, 'train' if is_local else 'test')\n# test_paths = [os.path.join(test_dir, fname) for fname in os.listdir(test_dir)]\n# if is_local:\n#     test_paths = test_paths[:nrows]\n    \n# submission = prepare_dataset(test_paths, ['label'], chunk_size=4000, is_test=True).reset_index()\n\n# submission.label = submission.label.apply(lambda x: clean_predictions('|'.join(eval(x))))\n# submission.columns = ['Id', 'PredictionString']\n    \n# styler = submission.head(10).style.set_properties(**{'text-align': 'left', 'width': '300px'})\n# styler = styler.set_table_styles([dict(selector='th', props=[('text-align', 'left')])])\n# display(styler) \n\n# submission.to_csv('submission.csv', index=False)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Draft","metadata":{"_uuid":"09c1e223-7aa5-49ae-936b-372903edb6f0","_cell_guid":"de56b694-892a-4089-8ece-9815413d02b4","trusted":true}},{"cell_type":"code","source":"# if TRAIN_DATA is None:\n#     print(f\"Loading preprocessed data from {DATA_DIR}...\")\n#     TRAIN_DATA = pd.read_csv(DATA_DIR + '/train_preprocessed.csv', index_col='Id')","metadata":{"_uuid":"3dd553ac-c0a4-48c7-ab35-2a561c8ea520","_cell_guid":"9382c3a5-6261-444c-8a7b-68958313994b","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# eval(TRAIN_DATA.loc['ac096f58-1866-45fe-882a-9cc1823ce90a'].text)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# sorted(UNIQUE_LABELS, key=lambda x: len(x.split()))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# trn_data, trn_dataset, trn_dataloader = create_dataloader(dt=TRAIN_DATA)#.iloc[:100])\n# ENGLISH.build_vocab(trn_data.text, specials=[PAD, UNK, SEP])\n# ENGLISH.pad_token = PAD\n# ENGLISH.unk_token = UNK","metadata":{"_uuid":"e775f60a-020b-4cea-8c60-374053b8d42a","_cell_guid":"3d4c23df-e181-4be2-a79b-358a8dbb96cc","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# trn_data['len'] = trn_data.text.apply(lambda x: len(x))\n# trn_data.sort_values(by='len')[trn_data.len == 5].iloc[:50]","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# eval(TRAIN_DATA.loc['9e631b3f-8ac1-4911-bdb2-eab3b685f76d'].text)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# def get_text(filepath):\n#     with open(filepath, 'r') as f:\n#         dt = json.load(f)\n#         for sect in dt:\n#             title = sect.get('section_title').strip()\n#             sect_text = sect.get('text').strip()\n#             print('\\nSECTION:', title, '\\n',  sect_text)\n# t=get_text('../input/coleridgeinitiative-show-us-the-data/train/ac096f58-1866-45fe-882a-9cc1823ce90a.json')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# with open('../input/coleridge-datagov/dataset_names.txt') as f:\n#     df = f.readlines()\n#     df = [x.strip(' \\n') for x in df if len(x) > 4]\n#     df = pd.DataFrame(df, columns=['title'])\n#     print(df)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# add_labels_path = LABELS_DIR + '/data_set_800_with6000popular.csv'\n# df = [clean_label(lab) for lab in pd.read_csv(add_labels_path).title.values.tolist()]# if len(lab.split()) > 1]\n# sorted(df, key=len)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# ENGLISH = load_vocab(VOCAB_FILENAME)\n# model = load_model(MODEL, for_training=False)","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}