inference_colwise_gittables_drop_seed1_contra_multi_testmmr copy 2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[nltk_data] Downloading package punkt to /home/zhihao/nltk_data...\n",
      "[nltk_data]   Package punkt is already up-to-date!\n",
      "[nltk_data] Downloading package stopwords to /home/zhihao/nltk_data...\n",
      "[nltk_data]   Package stopwords is already up-to-date!\n"
     ]
    }
   ],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "import argparse\n",
    "import json\n",
    "import math\n",
    "import os\n",
    "import random\n",
    "from time import time\n",
    "import mlflow\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from sklearn.metrics import confusion_matrix, f1_score\n",
    "from collections import defaultdict\n",
    "\n",
    "# import pytrec_eval\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.nn import CrossEntropyLoss, BCEWithLogitsLoss\n",
    "from torch.utils.data import DataLoader, RandomSampler\n",
    "from transformers import BertTokenizer, BertForSequenceClassification, BertConfig, AutoTokenizer\n",
    "from transformers import AdamW, get_linear_schedule_with_warmup\n",
    "from accelerate import Accelerator\n",
    "from copy import deepcopy\n",
    "\n",
    "torch.backends.cuda.matmul.allow_tf32 = True\n",
    "\n",
    "from watchog.dataset import (\n",
    "    # collate_fn,\n",
    "    TURLColTypeTablewiseDataset,\n",
    "    TURLRelExtTablewiseDataset,\n",
    "    SatoCVTablewiseDataset,\n",
    "    ColPoplTablewiseDataset\n",
    ")\n",
    "\n",
    "from watchog.dataset import TableDataset, SupCLTableDataset, SemtableCVTablewiseDataset, GittablesColwiseDataset, GittablesTablewiseDataset\n",
    "from watchog.model import BertMultiPairPooler, BertForMultiOutputClassification, BertForMultiOutputClassificationColPopl, Verifier\n",
    "from watchog.model import SupCLforTable, UnsupCLforTable, lm_mp\n",
    "from watchog.utils import load_checkpoint, f1_score_multilabel, collate_fn, get_col_pred, ColPoplEvaluator\n",
    "from watchog.utils import task_num_class_dict\n",
    "from accelerate import DistributedDataParallelKwargs\n",
    "import wandb\n",
    "\n",
    "from argparse import Namespace\n",
    "import torch\n",
    "import random\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import os\n",
    "import pickle\n",
    "import json\n",
    "import re\n",
    "import transformers\n",
    "from torch.utils import data\n",
    "from torch.nn.utils import rnn\n",
    "from transformers import AutoTokenizer\n",
    "\n",
    "from typing import List\n",
    "from functools import reduce\n",
    "import operator\n",
    "\n",
    "from itertools import chain\n",
    "import copy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Detected kernel version 5.4.0, which is below the recommended minimum of 5.5.0; this can cause the process to hang. It is recommended to upgrade the kernel to the minimum version or higher.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "args={\"wandb\": false, \"model\": \"Watchog\", \"unlabeled_train_only\": false, \"context_encoding_type\": \"v0\", \"pool_version\": \"v0.2\", \"random_sample\": false, \"comment\": \"debug\", \"shortcut_name\": \"bert-base-uncased\", \"max_length\": 128, \"adaptive_max_length\": false, \"max_num_col\": 8, \"batch_size\": 3, \"epoch\": 1, \"random_seed\": 4649, \"train_n_seed_cols\": -1, \"num_classes\": 101, \"multi_gpu\": false, \"fp16\": false, \"warmup\": 0.0, \"lr\": 5e-05, \"task\": \"gt-semtab22-dbpedia-all0\", \"colpair\": false, \"metadata\": false, \"from_scratch\": false, \"cl_tag\": \"wikitables/simclr/bert_None_10_32_256_5e-05_sample_row4,sample_row4_tfidf_entity_column_0.05_0_last.pt\", \"dropout_prob\": 0.5, \"eval_test\": true, \"small_tag\": \"semi1\", \"data_path\": \"/data/zhihao/TU/\", \"pretrained_ckpt_path\": \"/data/zhihao/TU/Watchog/model/\"}\n",
      "gt-semtab22-dbpedia-all0/wikitables-simclr-bert_None_10_32_256_5e-05_sample_row4,sample_row4_tfidf_entity_column_0.05_0_last.pt_bert-base-uncased-poolsemi1-max_colsv0.2-rand8-bsFalse-ml3-ne128-do10.5\n",
      "Namespace(augment_op='sample_row4,sample_row4', batch_size=32, data_path='/data/zhihao/TU/TURL/', fp16=True, gpus='0', lm='bert', logdir='/data/zhihao/TU/Watchog/model/', lr=5e-05, max_len=256, mode='simclr', model='Watchog', n_epochs=10, pretrain_data='wikitables', pretrained_model_path='', projector=768, run_id=0, sample_meth='tfidf_entity', save_model=10, single_column=False, size=None, table_order='column', temperature=0.05)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "A parameter name that contains `beta` will be renamed internally to `bias`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `gamma` will be renamed internally to `weight`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `beta` will be renamed internally to `bias`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `gamma` will be renamed internally to `weight`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `beta` will be renamed internally to `bias`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `gamma` will be renamed internally to `weight`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `beta` will be renamed internally to `bias`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `gamma` will be renamed internally to `weight`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `beta` will be renamed internally to `bias`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `gamma` will be renamed internally to `weight`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `beta` will be renamed internally to `bias`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `gamma` will be renamed internally to `weight`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `beta` will be renamed internally to `bias`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `gamma` will be renamed internally to `weight`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `beta` will be renamed internally to `bias`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `gamma` will be renamed internally to `weight`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `beta` will be renamed internally to `bias`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `gamma` will be renamed internally to `weight`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `beta` will be renamed internally to `bias`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `gamma` will be renamed internally to `weight`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `beta` will be renamed internally to `bias`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `gamma` will be renamed internally to `weight`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `beta` will be renamed internally to `bias`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `gamma` will be renamed internally to `weight`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `beta` will be renamed internally to `bias`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `gamma` will be renamed internally to `weight`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `beta` will be renamed internally to `bias`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `gamma` will be renamed internally to `weight`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `beta` will be renamed internally to `bias`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `gamma` will be renamed internally to `weight`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `beta` will be renamed internally to `bias`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `gamma` will be renamed internally to `weight`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `beta` will be renamed internally to `bias`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `gamma` will be renamed internally to `weight`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `beta` will be renamed internally to `bias`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `gamma` will be renamed internally to `weight`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `beta` will be renamed internally to `bias`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `gamma` will be renamed internally to `weight`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `beta` will be renamed internally to `bias`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `gamma` will be renamed internally to `weight`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `beta` will be renamed internally to `bias`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `gamma` will be renamed internally to `weight`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `beta` will be renamed internally to `bias`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `gamma` will be renamed internally to `weight`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `beta` will be renamed internally to `bias`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `gamma` will be renamed internally to `weight`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `beta` will be renamed internally to `bias`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `gamma` will be renamed internally to `weight`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `beta` will be renamed internally to `bias`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `gamma` will be renamed internally to `weight`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `beta` will be renamed internally to `bias`. Please use a different name to suppress this warning.\n",
      "A parameter name that contains `gamma` will be renamed internally to `weight`. Please use a different name to suppress this warning.\n"
     ]
    }
   ],
   "source": [
    "    device = torch.device(1)\n",
    "\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument(\"--wandb\", type=bool, default=False)\n",
    "    parser.add_argument(\"--model\", type=str, default=\"Watchog\")\n",
    "    parser.add_argument(\"--unlabeled_train_only\", type=bool, default=False)\n",
    "    parser.add_argument(\"--context_encoding_type\", type=str, default=\"v0\")\n",
    "    parser.add_argument(\"--pool_version\", type=str, default=\"v0.2\")\n",
    "    parser.add_argument(\"--random_sample\", type=bool, default=False)\n",
    "    parser.add_argument(\"--comment\", type=str, default=\"debug\", help=\"to distinguish the runs\")\n",
    "    parser.add_argument(\n",
    "        \"--shortcut_name\",\n",
    "        default=\"bert-base-uncased\",\n",
    "        type=str,\n",
    "        help=\"Huggingface model shortcut name \",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--max_length\",\n",
    "        default=128,\n",
    "        type=int,\n",
    "        help=\n",
    "        \"The maximum total input sequence length after tokenization. Sequences longer \"\n",
    "        \"than this will be truncated, sequences shorter will be padded.\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--adaptive_max_length\",\n",
    "        default=False,\n",
    "        type=bool,\n",
    "    )    \n",
    "    parser.add_argument(\n",
    "        \"--max_num_col\",\n",
    "        default=8,\n",
    "        type=int,\n",
    "    )   \n",
    "\n",
    "    parser.add_argument(\n",
    "        \"--batch_size\",\n",
    "        default=3,\n",
    "        type=int,\n",
    "        help=\"Batch size\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--epoch\",\n",
    "        default=1,\n",
    "        type=int,\n",
    "        help=\"Number of epochs for training\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--random_seed\",\n",
    "        default=4649,\n",
    "        type=int,\n",
    "        help=\"Random seed\",\n",
    "    )\n",
    "    \n",
    "    parser.add_argument(\n",
    "        \"--train_n_seed_cols\",\n",
    "        default=-1,\n",
    "        type=int,\n",
    "        help=\"number of seeding columns in training\",\n",
    "    )\n",
    "\n",
    "    parser.add_argument(\n",
    "        \"--num_classes\",\n",
    "        default=78,\n",
    "        type=int,\n",
    "        help=\"Number of classes\",\n",
    "    )\n",
    "    parser.add_argument(\"--multi_gpu\",\n",
    "                        action=\"store_true\",\n",
    "                        default=False,\n",
    "                        help=\"Use multiple GPU\")\n",
    "    parser.add_argument(\"--fp16\",\n",
    "                        action=\"store_true\",\n",
    "                        default=False,\n",
    "                        help=\"Use FP16\")\n",
    "    parser.add_argument(\"--warmup\",\n",
    "                        type=float,\n",
    "                        default=0.,\n",
    "                        help=\"Warmup ratio\")\n",
    "    parser.add_argument(\"--lr\", type=float, default=5e-5, help=\"Learning rate\")\n",
    "    parser.add_argument(\"--task\",\n",
    "                        type=str,\n",
    "                        default='gt-semtab22-dbpedia-all0',\n",
    "                        choices=[\n",
    "                            \"sato0\", \"sato1\", \"sato2\", \"sato3\", \"sato4\",\n",
    "                            \"msato0\", \"msato1\", \"msato2\", \"msato3\", \"msato4\",\n",
    "                            \"gt-dbpedia0\", \"gt-dbpedia1\", \"gt-dbpedia2\", \"gt-dbpedia3\", \"gt-dbpedia4\",\n",
    "                            \"gt-dbpedia-all0\", \"gt-dbpedia-all1\", \"gt-dbpedia-all2\", \"gt-dbpedia-all3\", \"gt-dbpedia-all4\",\n",
    "                            \"gt-schema-all0\", \"gt-schema-all1\", \"gt-schema-all2\", \"gt-schema-all3\", \"gt-schema-all4\",\n",
    "                            \"gt-semtab22-dbpedia\", \"gt-semtab22-dbpedia0\", \"gt-semtab22-dbpedia1\", \"gt-semtab22-dbpedia2\", \"gt-semtab22-dbpedia3\", \"gt-semtab22-dbpedia4\",\n",
    "                            \"gt-semtab22-dbpedia-all\", \"gt-semtab22-dbpedia-all0\", \"gt-semtab22-dbpedia-all1\", \"gt-semtab22-dbpedia-all2\", \"gt-semtab22-dbpedia-all3\", \"gt-semtab22-dbpedia-all4\",\n",
    "                            \"gt-semtab22-schema-class-all\", \"gt-semtab22-schema-property-all\",\n",
    "                            \"turl\", \"turl-re\", \"col-popl-1\", \"col-popl-2\", \"col-popl-3\", \"row-popl\",\n",
    "                            \"col-popl-turl-0\", \"col-popl-turl-1\", \"col-popl-turl-2\",\n",
    "                            \"col-popl-turl-mdonly-0\", \"col-popl-turl-mdonly-1\", \"col-popl-turl-mdonly-2\"\n",
    "                        ],\n",
    "                        help=\"Task names}\")\n",
    "    parser.add_argument(\"--colpair\",\n",
    "                        action=\"store_true\",\n",
    "                        help=\"Use column pair embedding\")\n",
    "    parser.add_argument(\"--metadata\",\n",
    "                        action=\"store_true\",\n",
    "                        help=\"Use column header metadata\")\n",
    "    parser.add_argument(\"--from_scratch\",\n",
    "                        action=\"store_true\",\n",
    "                        help=\"Training from scratch\")\n",
    "    parser.add_argument(\"--cl_tag\",\n",
    "                        type=str,\n",
    "                        default=\"wikitables/simclr/bert_None_10_32_256_5e-05_sample_row4,sample_row4_tfidf_entity_column_0.05_0_last.pt\",\n",
    "                        help=\"path to the pre-trained file\")\n",
    "    parser.add_argument(\"--dropout_prob\",\n",
    "                        type=float,\n",
    "                        default=0.5)\n",
    "    parser.add_argument(\"--eval_test\",\n",
    "                        action=\"store_true\",\n",
    "                        help=\"evaluate on testset and do not save the model file\")\n",
    "    parser.add_argument(\"--small_tag\",\n",
    "                        type=str,\n",
    "                        default=\"semi1\",\n",
    "                        help=\"e.g., by_table_t5_v1\")\n",
    "    parser.add_argument(\"--data_path\",\n",
    "                        type=str,\n",
    "                        default=\"/data/zhihao/TU/\")\n",
    "    parser.add_argument(\"--pretrained_ckpt_path\",\n",
    "                        type=str,\n",
    "                        default=\"/data/zhihao/TU/Watchog/model/\")    \n",
    "\n",
    "    args = parser.parse_args([])\n",
    "    task = args.task\n",
    "    if args.small_tag != \"\":\n",
    "        args.eval_test = True\n",
    "    \n",
    "    args.num_classes = task_num_class_dict[task]\n",
    "    if args.colpair:\n",
    "        assert \"turl-re\" == task, \"colpair can be only used for Relation Extraction\"\n",
    "    if args.metadata:\n",
    "        assert \"turl-re\" == task or \"turl\" == task, \"metadata can be only used for TURL datasets\"\n",
    "    if \"col-popl\":\n",
    "        # metrics = {\n",
    "        #     \"accuracy\": CategoricalAccuracy(tie_break=True),\n",
    "        # }\n",
    "        if args.train_n_seed_cols != -1:\n",
    "            if \"col-popl\" in task:\n",
    "                assert args.train_n_seed_cols == int(task[-1]),  \"# of seed columns must match\"\n",
    "\n",
    "    print(\"args={}\".format(json.dumps(vars(args))))\n",
    "\n",
    "    max_length = args.max_length\n",
    "    batch_size = args.batch_size\n",
    "    num_train_epochs = args.epoch\n",
    "\n",
    "    shortcut_name = args.shortcut_name\n",
    "\n",
    "    if args.colpair and args.metadata:\n",
    "        taskname = \"{}-colpair-metadata\".format(task)\n",
    "    elif args.colpair:\n",
    "        taskname = \"{}-colpair\".format(task)\n",
    "    elif args.metadata:\n",
    "        taskname = \"{}-metadata\".format(task)\n",
    "    elif args.train_n_seed_cols == -1 and 'popl' in task:\n",
    "        taskname = \"{}-mix\".format(task)\n",
    "    else:\n",
    "        taskname = \"\".join(task)\n",
    "    cv = int(task[-1])\n",
    "\n",
    "    if args.from_scratch:\n",
    "        if \"gt\" in task:\n",
    "            tag_name = \"{}/{}-{}-{}-pool{}-max_cols{}-rand{}-bs{}-ml{}-ne{}-do{}{}\".format(\n",
    "                taskname,  \"{}-fromscratch\".format(shortcut_name), args.small_tag, args.comment, args.pool_version, args.max_num_col, args.random_sample,\n",
    "                batch_size, max_length, num_train_epochs, args.dropout_prob, \n",
    "                '-rs{}'.format(args.random_seed) if args.random_seed != 4649 else '')\n",
    "        else:\n",
    "            tag_name = \"{}/{}-{}-{}-bs{}-ml{}-ne{}-do{}{}\".format(\n",
    "                taskname,  \"{}-fromscratch\".format(shortcut_name), args.small_tag, args.comment, \n",
    "                batch_size, max_length, num_train_epochs, args.dropout_prob, \n",
    "                '-rs{}'.format(args.random_seed) if args.random_seed != 4649 else '')\n",
    "        \n",
    "    else:\n",
    "        if \"gt\" in task:\n",
    "            tag_name = \"{}/{}_{}-pool{}-max_cols{}-rand{}-bs{}-ml{}-ne{}-do{}{}\".format(\n",
    "                taskname, args.cl_tag.replace('/', '-'),  shortcut_name, args.small_tag, args.pool_version, args.max_num_col, args.random_sample,\n",
    "                batch_size, max_length, num_train_epochs, args.dropout_prob,\n",
    "                '-rs{}'.format(args.random_seed) if args.random_seed != 4649 else '')\n",
    "        else:\n",
    "            tag_name = \"{}/{}_{}-{}-bs{}-ml{}-ne{}-do{}{}\".format(\n",
    "                taskname, args.cl_tag.replace('/', '-'),  shortcut_name, args.small_tag,\n",
    "                batch_size, max_length, num_train_epochs, args.dropout_prob,\n",
    "                '-rs{}'.format(args.random_seed) if args.random_seed != 4649 else '')\n",
    "\n",
    "    # if args.eval_test:\n",
    "    #     if args.small_tag != '':\n",
    "    #         tag_name = tag_name.replace('outputs', 'small_outputs')\n",
    "    #         tag_name += '-' + args.small_tag\n",
    "    print(tag_name)\n",
    "    file_path = os.path.join(args.data_path, \"Watchog\", \"outputs\", tag_name)\n",
    "\n",
    "    dirpath = os.path.dirname(file_path)\n",
    "    if not os.path.exists(dirpath):\n",
    "        print(\"{} not exists. Created\".format(dirpath))\n",
    "        os.makedirs(dirpath)\n",
    "    \n",
    "    if args.fp16:\n",
    "        torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = True\n",
    "        \n",
    "      \n",
    "        \n",
    "    # accelerator = Accelerator(mixed_precision=\"no\" if not args.fp16 else \"fp16\")   \n",
    "    ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)\n",
    "    accelerator = Accelerator(mixed_precision=\"no\" if not args.fp16 else \"fp16\", kwargs_handlers=[ddp_kwargs])\n",
    "\n",
    "\n",
    "    ckpt_path = os.path.join(args.pretrained_ckpt_path, args.cl_tag)\n",
    "    # ckpt_path = '/efs/checkpoints/{}.pt'.format(args.cl_tag)\n",
    "    ckpt = torch.load(ckpt_path, map_location=device)\n",
    "    ckpt_hp = ckpt['hp']\n",
    "    print(ckpt_hp)\n",
    " \n",
    "    setattr(ckpt_hp, 'batch_size', args.batch_size)\n",
    "    setattr(ckpt_hp, 'hidden_dropout_prob', args.dropout_prob)\n",
    "    setattr(ckpt_hp, 'shortcut_name', args.shortcut_name)\n",
    "    setattr(ckpt_hp, 'num_labels', args.num_classes)\n",
    "    \n",
    "    \n",
    "    \n",
    "    tokenizer = BertTokenizer.from_pretrained(shortcut_name)\n",
    "    padder = collate_fn(tokenizer.pad_token_id)\n",
    "    if task == \"turl-re\" and args.colpair:\n",
    "        model = BertForMultiOutputClassification(ckpt_hp, device=device, lm=ckpt['hp'].lm, col_pair='Pair')\n",
    "    elif \"col-popl\" in task:\n",
    "        model = BertForMultiOutputClassificationColPopl(ckpt_hp, device=device, lm=ckpt['hp'].lm, n_seed_cols=int(task[i][-1]), cls_for_md=\"md\" in task)\n",
    "    else:\n",
    "        model = BertForMultiOutputClassification(ckpt_hp, device=device, lm=ckpt['hp'].lm, version=\"v0\", use_attention_mask=True)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# best_state_dict = torch.load(\"/data/zhihao/TU/Watchog/outputs/gt-semtab22-dbpedia-all0/bert-base-uncased-fromscratch-semi1-AttnMask-max-unlabeled@8-poolv0-unlabeled8-randFalse-bs16-ml128-ne50-do0.1_best_f1_micro.pt\", map_location=device)\n",
    "best_state_dict = torch.load(\"/data/zhihao/TU/Watchog/outputs/gt-semtab22-dbpedia-all0/bert-base-uncased-fromscratch-semi1-Repeat@3-AttnMask-max-unlabeled@8-seed@2-poolv0-unlabeled8-randTrue-bs16-ml128-ne50-do0.1_best_last_0.pt\", map_location=device)\n",
    "model.load_state_dict(best_state_dict, strict=False)\n",
    "model = model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from watchog.model import VerifierMulti\n",
    "# verifier = Verifier(norm=\"batch_norm\", num_layers=None).to(device)\n",
    "verifier = VerifierMulti(norm=\"batch_norm\", num_layers=None).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "veri_state_dict = torch.load(\"/data/zhihao/TU/Watchog/outputs/gt-semtab22-dbpedia-all0/bert-base-uncased-fromscratch-semi1-Random1-Multi-Iter@0-Reg@0-mode@ffn-context@None-data@rand_seed2_mmr0.5_dp3-lr@0.0005-warmup@0.0-dp@0.0-norm@batch_norm-pos@0.5-poolv0-unlabeled8-randFalse-bs64-ml128-ne200-do0.0_verifier_multi_last.pt\", map_location=device)\n",
    "verifier.load_state_dict(veri_state_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def collate_fn(pad_token_id, data_only=True):\n",
    "    '''padder for input batch'''\n",
    "\n",
    "    def padder(samples):    \n",
    "        data = torch.nn.utils.rnn.pad_sequence(\n",
    "            [sample[\"data\"] for sample in samples], padding_value=pad_token_id)\n",
    "        if not data_only:\n",
    "            label = torch.nn.utils.rnn.pad_sequence(\n",
    "                [sample[\"label\"] for sample in samples], padding_value=-1)\n",
    "        else:\n",
    "            label = torch.cat([sample[\"label\"] for sample in samples])\n",
    "        batch = {\"data\": data, \"label\": label}\n",
    "        if \"idx\" in samples[0]:\n",
    "            batch[\"idx\"] = [sample[\"idx\"] for sample in samples]\n",
    "        if \"cls_indexes\" in samples[0]:\n",
    "            cls_indexes = torch.nn.utils.rnn.pad_sequence(\n",
    "                [sample[\"cls_indexes\"] for sample in samples], padding_value=0)\n",
    "            batch[\"cls_indexes\"] = cls_indexes\n",
    "        if \"target_col_mask\" in samples[0]:\n",
    "            target_col_mask = torch.nn.utils.rnn.pad_sequence(\n",
    "                [sample[\"target_col_mask\"] for sample in samples], padding_value=-1)\n",
    "            batch[\"target_col_mask\"] = target_col_mask\n",
    "        if \"table_embedding\" in samples[0]:\n",
    "            table_embeddings = [sample[\"table_embedding\"] for sample in samples]\n",
    "            batch[\"table_embedding\"] = torch.stack(table_embeddings, dim=0)\n",
    "        return batch\n",
    "        \n",
    "    return padder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers import SentenceTransformer, util\n",
    "\n",
    "def maximal_marginal_relevance(query_embedding, answer_embeddings, top_k=3, lambda_param=0.7):\n",
    "    \"\"\"\n",
    "    Perform Maximal Marginal Relevance (MMR) using PyTorch.\n",
    "    \n",
    "    Args:\n",
    "    - query_embedding (Tensor): Embedding of the query, shape (d,).\n",
    "    - answer_embeddings (Tensor): Embeddings of the candidate answers, shape (n, d).\n",
    "    - top_k (int): Number of diverse answers to retrieve.\n",
    "    - lambda_param (float): Trade-off parameter between relevance and diversity.\n",
    "    \n",
    "    Returns:\n",
    "    - selected_indices (list): Indices of the selected answers.\n",
    "    \"\"\"\n",
    "    # Initialize selected and remaining indices\n",
    "    selected_indices = []\n",
    "    remaining_indices = list(range(answer_embeddings.size(0)))\n",
    "\n",
    "    # Compute relevance scores (similarity with the query)\n",
    "    query_similarities = util.cos_sim(query_embedding, answer_embeddings)[0]\n",
    "\n",
    "    for _ in range(top_k):\n",
    "        if not remaining_indices:\n",
    "            break\n",
    "\n",
    "        mmr_scores = []\n",
    "        for idx in remaining_indices:\n",
    "            # Relevance score\n",
    "            relevance_score = query_similarities[idx]\n",
    "\n",
    "            # Diversity score (max similarity to already selected answers)\n",
    "            if selected_indices:\n",
    "                diversity_score = max(util.cos_sim(answer_embeddings[idx], answer_embeddings[selected_indices])[0])\n",
    "            else:\n",
    "                diversity_score = torch.tensor(0.0)\n",
    "\n",
    "            # Compute MMR score\n",
    "            mmr_score = lambda_param * relevance_score - (1 - lambda_param) * diversity_score\n",
    "            mmr_scores.append(mmr_score)\n",
    "\n",
    "        # Select the answer with the highest MMR score\n",
    "        mmr_scores = torch.stack(mmr_scores)\n",
    "        best_idx = remaining_indices[torch.argmax(mmr_scores)]\n",
    "        selected_indices.append(best_idx)\n",
    "        remaining_indices.remove(best_idx)\n",
    "\n",
    "    return selected_indices\n",
    "\n",
    "\n",
    "from sklearn.cluster import KMeans\n",
    "class GittablesTablewiseIterateMMRDataset(data.Dataset):\n",
    "\n",
    "    def __init__(\n",
    "            self,\n",
    "            cv: int,\n",
    "            split: str,\n",
    "            src: str,  # train or test\n",
    "            tokenizer: AutoTokenizer,\n",
    "            max_length: int = 256,\n",
    "            gt_only: bool = False,\n",
    "            device: torch.device = None,\n",
    "            base_dirpath: str = \"/data/zhihao/TU/GitTables/semtab_gittables/2022\",\n",
    "            base_tag: str = '', # blank, comma\n",
    "            small_tag: str = \"\",\n",
    "            train_ratio: float = 1.0,\n",
    "            max_unlabeled=8,\n",
    "            adaptive_max_length=True,\n",
    "            random_sample=False, # TODO\n",
    "            train_only=False,\n",
    "            lbd=0.5,\n",
    "            model=None): # TODO\n",
    "        if device is None:\n",
    "            device = torch.device('cpu')\n",
    "        basename = small_tag+ \"_cv_{}.csv\"\n",
    "    \n",
    "        if split in [\"train\", \"valid\"]:\n",
    "            df_list = []\n",
    "            for i in range(5):\n",
    "                if i == cv:\n",
    "                    continue\n",
    "                filepath = os.path.join(base_dirpath, basename.format(i))\n",
    "                df_list.append(pd.read_csv(filepath))\n",
    "                print(split, i)\n",
    "            df = pd.concat(df_list, axis=0)\n",
    "        else:\n",
    "            # test\n",
    "            filepath = os.path.join(base_dirpath, basename.format(cv))\n",
    "            df = pd.read_csv(filepath)\n",
    "            print(split)\n",
    "\n",
    "\n",
    "        if gt_only:\n",
    "            df = df[df[\"class_id\"] > -1]\n",
    "        if train_only and split != \"train\":\n",
    "            df = df[df[\"class_id\"] > -1]\n",
    "\n",
    "        \n",
    "        data_list = []\n",
    "        \n",
    "        df['class_id'] = df['class_id'].astype(int)\n",
    "        df.drop(df[(df['data'].isna()) & (df['class_id'] == -1)].index, inplace=True)\n",
    "        df['col_idx'] = df['col_idx'].astype(int)\n",
    "        df['data'] = df['data'].astype(str)\n",
    "        \n",
    "        num_tables = len(df.groupby(\"table_id\"))\n",
    "        valid_index = int(num_tables * 0.8)\n",
    "        num_train = int(train_ratio * num_tables * 0.8)        \n",
    "        \n",
    "        # df.drop(df[(df['data'] == '') & (df['class_id'] == -1)].index, inplace=True)\n",
    "        total_num_cols = 0\n",
    "        self.correctness_checklist = []\n",
    "        self.max_rest_num = 0\n",
    "        self.max_col_num = 0\n",
    "        for i, (index, group_df) in enumerate(df.groupby(\"table_id\")):\n",
    "            if (split == \"train\") and ((i >= num_train) or (i >= valid_index)):\n",
    "                break\n",
    "            if split == \"valid\" and i < valid_index:\n",
    "                continue\n",
    "            #     break\n",
    "            group_df.sort_values(by=['col_idx'], inplace=True)\n",
    "            group_df = group_df.reset_index(drop=True)\n",
    "            # scores = bm25.get_scores(test_dataset_gt[idx][\"data\"].cpu().tolist())\n",
    "            # semantic_scores = torch.tensor(group_df[\"data\"].apply(lambda x: max(bm25.get_scores(tokenizer.encode(\n",
    "            #         tokenizer.cls_token + \" \" + x, add_special_tokens=False, max_length=max_length, truncation=True)))))\n",
    "            \n",
    "            # labeled_columns_index = group_df[group_df['class_id'] > -1].index.to_list()\n",
    "            # assert len(labeled_columns_index) < 8\n",
    "            # if set(labeled_columns_index).issubset(set(semantic_scores.topk(min(len(group_df), max_unlabeled)).indices.tolist())):\n",
    "            #     self.correctness_checklist.append(1)\n",
    "            # else:\n",
    "            #     self.correctness_checklist.append(0)\n",
    "            if len(group_df) > max_unlabeled:\n",
    "                # col_embs = sb_model.encode(group_df[\"data\"].to_list())\n",
    "                \n",
    "                with torch.no_grad():\n",
    "                    model = model.to(device)\n",
    "                    model.eval()\n",
    "                    token_list = group_df[\"data\"].apply(lambda x: tokenizer.encode(\n",
    "                            tokenizer.cls_token + \" \" + x, add_special_tokens=False, max_length=max_length, truncation=True)).tolist()\n",
    "                    token_list = torch.nn.utils.rnn.pad_sequence(\n",
    "                    [torch.LongTensor(tokens) for tokens in token_list], padding_value=0).T.to(device)\n",
    "                    cls_indexes = torch.nonzero(\n",
    "                                token_list == tokenizer.cls_token_id)\n",
    "                    logits, embs = model(token_list, cls_indexes=cls_indexes.to(device), get_enc=True)\n",
    "                    col_embs = embs.cpu()\n",
    "                    if len(embs.size()) == 1:\n",
    "                        col_embs = col_embs.unsqueeze(0)\n",
    "                # col_embs = torch.tensor(sb_model.encode(group_df[\"data\"].to_list()))   \n",
    "                        \n",
    "            labeled_columns = group_df[group_df['class_id'] > -1]\n",
    "            labeled_columns.sort_values(by=['col_idx'], inplace=True)\n",
    "            for index, target_column in labeled_columns.iterrows():\n",
    "                target_col_idx = target_column[\"col_idx\"]\n",
    "                \n",
    "                if len(group_df) > max_unlabeled:\n",
    "                    other_columns = group_df.drop(index)\n",
    "                    selected_index =  maximal_marginal_relevance(col_embs[index], col_embs[torch.arange(len(col_embs)) != index], top_k=max_unlabeled-1, lambda_param=lbd)\n",
    "                    assert len(set(selected_index)) == len(selected_index)\n",
    "                    other_columns = other_columns.iloc[selected_index]\n",
    "                    assert len(other_columns) == max_unlabeled - 1\n",
    "                    # try:\n",
    "                        \n",
    "                    # except:\n",
    "                    #     print(\"here\")\n",
    "                else:\n",
    "                    other_columns = group_df.drop(index)\n",
    "                target_table = pd.concat([target_column.to_frame().T, other_columns], ignore_index=True)\n",
    "                \n",
    "                # if len(group_df) >= 8:\n",
    "                #     assert len(target_table) == max_unlabeled\n",
    "                # else:\n",
    "                #     assert len(target_table) == len(group_df)\n",
    "    \n",
    "                \n",
    "                target_cls = target_column[\"class_id\"]\n",
    "                target_table.sort_values(by=['col_idx'], inplace=True)\n",
    "                col_idx_list = target_table[\"col_idx\"].tolist()\n",
    "                # target_table.sort_values(by=['col_idx'], inplace=True)\n",
    "\n",
    "                if max_length <= 128 and adaptive_max_length:\n",
    "                    cur_maxlen = min(max_length, 512 // len(target_table) - 1)\n",
    "                else:\n",
    "                    cur_maxlen = max_length\n",
    "                    \n",
    "                token_ids_list = target_table[\"data\"].apply(lambda x: tokenizer.encode(\n",
    "                    tokenizer.cls_token + \" \" + x, add_special_tokens=False, max_length=cur_maxlen, truncation=True)).tolist(\n",
    "                    )\n",
    "                token_ids = torch.LongTensor(reduce(operator.add,\n",
    "                                                    token_ids_list)).to(device)\n",
    "\n",
    "                target_col_mask = []\n",
    "                cls_index_value = 0\n",
    "                context_id = 1\n",
    "                meet_target = False\n",
    "                for idx, col_i in enumerate(col_idx_list):\n",
    "                    if col_i == target_col_idx:\n",
    "                        target_col_mask += [0] * len(token_ids_list[idx])\n",
    "                        meet_target = True\n",
    "                    else:\n",
    "                        target_col_mask += [context_id] * len(token_ids_list[idx])\n",
    "                        context_id += 1\n",
    "                    if not meet_target:\n",
    "                        cls_index_value += len(token_ids_list[idx])\n",
    "                cls_index_list = [cls_index_value] \n",
    "                for cls_index in cls_index_list:\n",
    "                    assert token_ids[\n",
    "                        cls_index] == tokenizer.cls_token_id, \"cls_indexes validation\"\n",
    "                cls_indexes = torch.LongTensor(cls_index_list).to(device)\n",
    "                class_ids = torch.LongTensor(\n",
    "                    [target_cls]).to(device)\n",
    "                target_col_mask = torch.LongTensor(target_col_mask).to(device)\n",
    "                data_list.append(\n",
    "                    [index,\n",
    "                    len(target_table), token_ids, class_ids, cls_indexes, target_col_mask])                \n",
    "        print(split, len(data_list))\n",
    "        self.table_df = pd.DataFrame(data_list,\n",
    "                                     columns=[\n",
    "                                         \"table_id\", \"num_col\", \"data_tensor\",\n",
    "                                         \"label_tensor\", \"cls_indexes\", \"target_col_mask\"\n",
    "                                     ])\n",
    "        \"\"\"\n",
    "        # NOTE: msato contains a small portion of single-col tables. keep it to be consistent.  \n",
    "        if multicol_only:\n",
    "            # Check\n",
    "            num_all_tables = len(self.table_df)\n",
    "            self.table_df = self.table_df[self.table_df[\"num_col\"] > 1]\n",
    "            assert len(self.table_df) == num_all_tables\n",
    "        \"\"\"\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.table_df)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return {\n",
    "            \"data\": self.table_df.iloc[idx][\"data_tensor\"],\n",
    "            \"label\": self.table_df.iloc[idx][\"label_tensor\"],\n",
    "            \"cls_indexes\": self.table_df.iloc[idx][\"cls_indexes\"],\n",
    "            \"target_col_mask\": self.table_df.iloc[idx][\"target_col_mask\"],\n",
    "        }\n",
    "        #\"idx\": torch.LongTensor([idx])}\n",
    "        #\"cls_indexes\": self.table_df.iloc[idx][\"cls_indexes\"]}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GittablesTablewiseIterateDataset(data.Dataset):\n",
    "\n",
    "    def __init__(\n",
    "            self,\n",
    "            cv: int,\n",
    "            split: str,\n",
    "            src: str,  # train or test\n",
    "            tokenizer: AutoTokenizer,\n",
    "            max_length: int = 256,\n",
    "            gt_only: bool = False,\n",
    "            device: torch.device = None,\n",
    "            base_dirpath: str = \"/data/zhihao/TU/GitTables/semtab_gittables/2022\",\n",
    "            base_tag: str = '', # blank, comma\n",
    "            small_tag: str = \"\",\n",
    "            train_ratio: float = 1.0,\n",
    "            max_unlabeled=8,\n",
    "            adaptive_max_length=True,\n",
    "            random_sample=False, # TODO\n",
    "            train_only=False, \n",
    "            seed=None): # TODO\n",
    "        if device is None:\n",
    "            device = torch.device('cpu')\n",
    "        basename = small_tag+ \"_cv_{}.csv\"\n",
    "    \n",
    "        if split in [\"train\", \"valid\"]:\n",
    "            df_list = []\n",
    "            for i in range(5):\n",
    "                if i == cv:\n",
    "                    continue\n",
    "                filepath = os.path.join(base_dirpath, basename.format(i))\n",
    "                df_list.append(pd.read_csv(filepath))\n",
    "                print(split, i)\n",
    "            df = pd.concat(df_list, axis=0)\n",
    "        else:\n",
    "            # test\n",
    "            filepath = os.path.join(base_dirpath, basename.format(cv))\n",
    "            df = pd.read_csv(filepath)\n",
    "            print(split)\n",
    "\n",
    "\n",
    "        if gt_only:\n",
    "            df = df[df[\"class_id\"] > -1]\n",
    "        if train_only and split != \"train\":\n",
    "            df = df[df[\"class_id\"] > -1]\n",
    "\n",
    "        \n",
    "        data_list = []\n",
    "        \n",
    "        df['class_id'] = df['class_id'].astype(int)\n",
    "        df.drop(df[(df['data'].isna()) & (df['class_id'] == -1)].index, inplace=True)\n",
    "        df['col_idx'] = df['col_idx'].astype(int)\n",
    "        df['data'] = df['data'].astype(str)\n",
    "        \n",
    "        num_tables = len(df.groupby(\"table_id\"))\n",
    "        valid_index = int(num_tables * 0.8)\n",
    "        num_train = int(train_ratio * num_tables * 0.8)        \n",
    "        \n",
    "        # df.drop(df[(df['data'] == '') & (df['class_id'] == -1)].index, inplace=True)\n",
    "        total_num_cols = 0\n",
    "        for i, (index, group_df) in enumerate(df.groupby(\"table_id\")):\n",
    "            if (split == \"train\") and ((i >= num_train) or (i >= valid_index)):\n",
    "                break\n",
    "            if split == \"valid\" and i < valid_index:\n",
    "                continue\n",
    "            #     break\n",
    "            num_cols = len(group_df)\n",
    "            group_df.sort_values(by=['col_idx'], inplace=True)\n",
    "            labeled_columns = group_df[group_df['class_id'] > -1]\n",
    "            unlabeled_columns = group_df[group_df['class_id'] == -1]\n",
    "            num_unlabeled = min(max(max_unlabeled-len(labeled_columns), 0), len(unlabeled_columns)) if max_unlabeled > 0 else len(unlabeled_columns)\n",
    "            unlabeled_columns = unlabeled_columns.sample(num_unlabeled, random_state=seed) if random_sample else unlabeled_columns[0:num_unlabeled]\n",
    "            # group_df = pd.concat([group_df[group_df['class_id'] > -1], unlabeled_columns.sample(min(10-len(labeled_columns), len(unlabeled_columns)))])\n",
    "            # group_df = pd.concat([group_df[group_df['class_id'] > -1], unlabeled_columns[0:min(max(10-len(labeled_columns), 0), len(unlabeled_columns))]])\n",
    "            group_df = pd.concat([group_df[group_df['class_id'] > -1], unlabeled_columns]) \n",
    "            group_df.sort_values(by=['col_idx'], inplace=True) # TODO \n",
    "\n",
    "            if max_length <= 128 and adaptive_max_length:\n",
    "                cur_maxlen = min(max_length, 512 // len(group_df) - 1)\n",
    "            else:\n",
    "                cur_maxlen = max_length\n",
    "                \n",
    "            token_ids_list = group_df[\"data\"].apply(lambda x: tokenizer.encode(\n",
    "                tokenizer.cls_token + \" \" + x, add_special_tokens=False, max_length=cur_maxlen, truncation=True)).tolist(\n",
    "                )\n",
    "            token_ids = torch.LongTensor(reduce(operator.add,\n",
    "                                                token_ids_list)).to(device)\n",
    "            for col_i in range(len(token_ids_list)):\n",
    "                if group_df[\"class_id\"].values[col_i] == -1:\n",
    "                    continue\n",
    "                target_col_mask = []\n",
    "                cls_index_value = 0\n",
    "                context_id = 1\n",
    "                for col_j in range(len(token_ids_list)):\n",
    "                    if col_j == col_i:\n",
    "                        target_col_mask += [0] * len(token_ids_list[col_j])\n",
    "                    else:\n",
    "                        target_col_mask += [context_id] * len(token_ids_list[col_j])\n",
    "                        context_id += 1\n",
    "                    if col_j < col_i:\n",
    "                        cls_index_value += len(token_ids_list[col_j])\n",
    "                cls_index_list = [cls_index_value] \n",
    "                for cls_index in cls_index_list:\n",
    "                    assert token_ids[\n",
    "                        cls_index] == tokenizer.cls_token_id, \"cls_indexes validation\"\n",
    "                cls_indexes = torch.LongTensor(cls_index_list).to(device)\n",
    "                class_ids = torch.LongTensor(\n",
    "                    [group_df[\"class_id\"].values[col_i]]).to(device)\n",
    "                target_col_mask = torch.LongTensor(target_col_mask).to(device)\n",
    "                col_idx = group_df[\"col_idx\"].values\n",
    "                data_list.append(\n",
    "                    [index,\n",
    "                    len(group_df), token_ids, class_ids, cls_indexes, target_col_mask, num_cols, col_idx])                \n",
    "        print(split, len(data_list))\n",
    "        self.table_df = pd.DataFrame(data_list,\n",
    "                                     columns=[\n",
    "                                         \"table_id\", \"num_col\", \"data_tensor\",\n",
    "                                         \"label_tensor\", \"cls_indexes\", \"target_col_mask\", \"initial_num_col\", \"col_idx\"\n",
    "                                     ])\n",
    "        \"\"\"\n",
    "        # NOTE: msato contains a small portion of single-col tables. keep it to be consistent.  \n",
    "        if multicol_only:\n",
    "            # Check\n",
    "            num_all_tables = len(self.table_df)\n",
    "            self.table_df = self.table_df[self.table_df[\"num_col\"] > 1]\n",
    "            assert len(self.table_df) == num_all_tables\n",
    "        \"\"\"\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.table_df)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return {\n",
    "            \"data\": self.table_df.iloc[idx][\"data_tensor\"],\n",
    "            \"label\": self.table_df.iloc[idx][\"label_tensor\"],\n",
    "            \"cls_indexes\": self.table_df.iloc[idx][\"cls_indexes\"],\n",
    "            \"target_col_mask\": self.table_df.iloc[idx][\"target_col_mask\"],\n",
    "            \"initial_num_col\": self.table_df.iloc[idx][\"initial_num_col\"],\n",
    "            \"col_idx\": self.table_df.iloc[idx][\"col_idx\"]\n",
    "        }\n",
    "        #\"idx\": torch.LongTensor([idx])}\n",
    "        #\"cls_indexes\": self.table_df.iloc[idx][\"cls_indexes\"]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GittablesTablewiseIterateNearbyDataset(data.Dataset):\n",
    "\n",
    "    def __init__(\n",
    "            self,\n",
    "            cv: int,\n",
    "            split: str,\n",
    "            src: str,  # train or test\n",
    "            tokenizer: AutoTokenizer,\n",
    "            max_length: int = 256,\n",
    "            gt_only: bool = False,\n",
    "            device: torch.device = None,\n",
    "            base_dirpath: str = \"/data/zhihao/TU/GitTables/semtab_gittables/2022\",\n",
    "            base_tag: str = '', # blank, comma\n",
    "            small_tag: str = \"\",\n",
    "            train_ratio: float = 1.0,\n",
    "            max_unlabeled=8,\n",
    "            adaptive_max_length=True,\n",
    "            random_sample=False, # TODO\n",
    "            train_only=False): # TODO\n",
    "        if device is None:\n",
    "            device = torch.device('cpu')\n",
    "        basename = small_tag+ \"_cv_{}.csv\"\n",
    "    \n",
    "        if split in [\"train\", \"valid\"]:\n",
    "            df_list = []\n",
    "            for i in range(5):\n",
    "                if i == cv:\n",
    "                    continue\n",
    "                filepath = os.path.join(base_dirpath, basename.format(i))\n",
    "                df_list.append(pd.read_csv(filepath))\n",
    "                print(split, i)\n",
    "            df = pd.concat(df_list, axis=0)\n",
    "        else:\n",
    "            # test\n",
    "            filepath = os.path.join(base_dirpath, basename.format(cv))\n",
    "            df = pd.read_csv(filepath)\n",
    "            print(split)\n",
    "\n",
    "\n",
    "        if gt_only:\n",
    "            df = df[df[\"class_id\"] > -1]\n",
    "        if train_only and split != \"train\":\n",
    "            df = df[df[\"class_id\"] > -1]\n",
    "\n",
    "        \n",
    "        data_list = []\n",
    "        \n",
    "        df['class_id'] = df['class_id'].astype(int)\n",
    "        df.drop(df[(df['data'].isna()) & (df['class_id'] == -1)].index, inplace=True)\n",
    "        df['col_idx'] = df['col_idx'].astype(int)\n",
    "        df['data'] = df['data'].astype(str)\n",
    "        \n",
    "        num_tables = len(df.groupby(\"table_id\"))\n",
    "        valid_index = int(num_tables * 0.8)\n",
    "        num_train = int(train_ratio * num_tables * 0.8)        \n",
    "        \n",
    "        # df.drop(df[(df['data'] == '') & (df['class_id'] == -1)].index, inplace=True)\n",
    "        total_num_cols = 0\n",
    "        for i, (index, group_df) in enumerate(df.groupby(\"table_id\")):\n",
    "            if (split == \"train\") and ((i >= num_train) or (i >= valid_index)):\n",
    "                break\n",
    "            if split == \"valid\" and i < valid_index:\n",
    "                continue\n",
    "            #     break\n",
    "            labeled_columns = group_df[group_df['class_id'] > -1]\n",
    "            for index, target_column in labeled_columns.iterrows():\n",
    "                target_col_idx = target_column[\"col_idx\"]\n",
    "                \n",
    "                other_columns = group_df.drop(index)\n",
    "                \n",
    "                other_columns['distance'] = abs(other_columns['col_idx'] - target_col_idx)\n",
    "                other_columns.sort_values(by='distance', inplace=True)\n",
    "                nearby_columns = other_columns.head(max_unlabeled-1)\n",
    "                target_table = pd.concat([target_column.to_frame().T, nearby_columns], ignore_index=True)\n",
    "                \n",
    "                \n",
    "                \n",
    "                # other_labeled = other_columns[other_columns['class_id'] > -1].head(max_unlabeled-1)\n",
    "                # other_unlabeled = other_columns[other_columns['class_id'] == -1]\n",
    "                # other_unlabeled['distance'] = abs(other_unlabeled['col_idx'] - target_col_idx)\n",
    "                # other_unlabeled.sort_values(by='distance', inplace=True)\n",
    "                # nearby_columns = other_unlabeled.head(max_unlabeled-len(other_labeled)-1)\n",
    "                # target_table = pd.concat([target_column.to_frame().T, other_labeled, nearby_columns], ignore_index=True)\n",
    "\n",
    "                # other_columns['distance'] = abs(other_columns['col_idx'] - target_col_idx)\n",
    "                # other_columns.sort_values(by='distance', inplace=True)\n",
    "                # nearby_columns = other_columns.head(max_unlabeled-1)\n",
    "                # target_table = pd.concat([target_column.to_frame().T, nearby_columns], ignore_index=True)\n",
    "                \n",
    "                \n",
    "                target_cls = target_column[\"class_id\"]\n",
    "                target_table.sort_values(by=['col_idx'], inplace=True)\n",
    "                col_idx_list = target_table[\"col_idx\"].tolist()\n",
    "\n",
    "                if max_length <= 128 and adaptive_max_length:\n",
    "                    cur_maxlen = min(max_length, 512 // len(target_table) - 1)\n",
    "                else:\n",
    "                    cur_maxlen = max_length\n",
    "                    \n",
    "                token_ids_list = target_table[\"data\"].apply(lambda x: tokenizer.encode(\n",
    "                    tokenizer.cls_token + \" \" + x, add_special_tokens=False, max_length=cur_maxlen, truncation=True)).tolist(\n",
    "                    )\n",
    "                token_ids = torch.LongTensor(reduce(operator.add,\n",
    "                                                    token_ids_list)).to(device)\n",
    "\n",
    "                target_col_mask = []\n",
    "                cls_index_value = 0\n",
    "                context_id = 1\n",
    "                meet_target = False\n",
    "                for idx, col_i in enumerate(col_idx_list):\n",
    "                    if col_i == target_col_idx:\n",
    "                        target_col_mask += [0] * len(token_ids_list[idx])\n",
    "                        meet_target = True\n",
    "                    else:\n",
    "                        target_col_mask += [context_id] * len(token_ids_list[idx])\n",
    "                        context_id += 1\n",
    "                    if not meet_target:\n",
    "                        cls_index_value += len(token_ids_list[idx])\n",
    "                cls_index_list = [cls_index_value] \n",
    "                for cls_index in cls_index_list:\n",
    "                    assert token_ids[\n",
    "                        cls_index] == tokenizer.cls_token_id, \"cls_indexes validation\"\n",
    "                cls_indexes = torch.LongTensor(cls_index_list).to(device)\n",
    "                class_ids = torch.LongTensor(\n",
    "                    [target_cls]).to(device)\n",
    "                target_col_mask = torch.LongTensor(target_col_mask).to(device)\n",
    "                data_list.append(\n",
    "                    [index,\n",
    "                    len(target_table), token_ids, class_ids, cls_indexes, target_col_mask])                \n",
    "        print(split, len(data_list))\n",
    "        self.table_df = pd.DataFrame(data_list,\n",
    "                                     columns=[\n",
    "                                         \"table_id\", \"num_col\", \"data_tensor\",\n",
    "                                         \"label_tensor\", \"cls_indexes\", \"target_col_mask\"\n",
    "                                     ])\n",
    "        \"\"\"\n",
    "        # NOTE: msato contains a small portion of single-col tables. keep it to be consistent.  \n",
    "        if multicol_only:\n",
    "            # Check\n",
    "            num_all_tables = len(self.table_df)\n",
    "            self.table_df = self.table_df[self.table_df[\"num_col\"] > 1]\n",
    "            assert len(self.table_df) == num_all_tables\n",
    "        \"\"\"\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.table_df)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return {\n",
    "            \"data\": self.table_df.iloc[idx][\"data_tensor\"],\n",
    "            \"label\": self.table_df.iloc[idx][\"label_tensor\"],\n",
    "            \"cls_indexes\": self.table_df.iloc[idx][\"cls_indexes\"],\n",
    "            \"target_col_mask\": self.table_df.iloc[idx][\"target_col_mask\"],\n",
    "        }\n",
    "        #\"idx\": torch.LongTensor([idx])}\n",
    "        #\"cls_indexes\": self.table_df.iloc[idx][\"cls_indexes\"]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from sentence_transformers import SentenceTransformer, util\n",
    "\n",
    "def maximal_marginal_relevance(query_embedding, answer_embeddings, top_k=3, lambda_param=0.7):\n",
    "    \"\"\"\n",
    "    Perform Maximal Marginal Relevance (MMR) using PyTorch.\n",
    "    \n",
    "    Args:\n",
    "    - query_embedding (Tensor): Embedding of the query, shape (d,).\n",
    "    - answer_embeddings (Tensor): Embeddings of the candidate answers, shape (n, d).\n",
    "    - top_k (int): Number of diverse answers to retrieve.\n",
    "    - lambda_param (float): Trade-off parameter between relevance and diversity.\n",
    "    \n",
    "    Returns:\n",
    "    - selected_indices (list): Indices of the selected answers.\n",
    "    \"\"\"\n",
    "    # Initialize selected and remaining indices\n",
    "    selected_indices = []\n",
    "    remaining_indices = list(range(answer_embeddings.size(0)))\n",
    "\n",
    "    # Compute relevance scores (similarity with the query)\n",
    "    query_similarities = util.cos_sim(query_embedding, answer_embeddings)[0]\n",
    "\n",
    "    for _ in range(top_k):\n",
    "        if not remaining_indices:\n",
    "            break\n",
    "\n",
    "        mmr_scores = []\n",
    "        for idx in remaining_indices:\n",
    "            # Relevance score\n",
    "            relevance_score = query_similarities[idx]\n",
    "\n",
    "            # Diversity score (max similarity to already selected answers)\n",
    "            if selected_indices:\n",
    "                diversity_score = max(util.cos_sim(answer_embeddings[idx], answer_embeddings[selected_indices])[0])\n",
    "            else:\n",
    "                diversity_score = torch.tensor(0.0)\n",
    "\n",
    "            # Compute MMR score\n",
    "            mmr_score = lambda_param * relevance_score - (1 - lambda_param) * diversity_score\n",
    "            mmr_scores.append(mmr_score)\n",
    "\n",
    "        # Select the answer with the highest MMR score\n",
    "        mmr_scores = torch.stack(mmr_scores)\n",
    "        best_idx = remaining_indices[torch.argmax(mmr_scores)]\n",
    "        selected_indices.append(best_idx)\n",
    "        remaining_indices.remove(best_idx)\n",
    "\n",
    "    return selected_indices\n",
    "\n",
    "\n",
    "from sklearn.cluster import KMeans\n",
    "class GittablesTablewiseIterateMMRRetrieveDataset(data.Dataset):\n",
    "\n",
    "    def __init__(\n",
    "            self,\n",
    "            cv: int,\n",
    "            split: str,\n",
    "            src: str,  # train or test\n",
    "            tokenizer: AutoTokenizer,\n",
    "            max_length: int = 256,\n",
    "            gt_only: bool = False,\n",
    "            device: torch.device = None,\n",
    "            base_dirpath: str = \"/data/zhihao/TU/GitTables/semtab_gittables/2022\",\n",
    "            base_tag: str = '', # blank, comma\n",
    "            small_tag: str = \"\",\n",
    "            train_ratio: float = 1.0,\n",
    "            max_unlabeled=8,\n",
    "            adaptive_max_length=True,\n",
    "            random_sample=False, # TODO\n",
    "            train_only=False,\n",
    "            training_corpus_embs=None,\n",
    "            df_train_groups=None, \n",
    "            K=1,\n",
    "            lbd=0.5,\n",
    "            model=None,\n",
    "            retrieve_model=\"doduo\"): # TODO\n",
    "        if device is None:\n",
    "            device = torch.device('cpu')\n",
    "        basename = small_tag+ \"_cv_{}.csv\"\n",
    "        self.extra_col_num = []\n",
    "        self.init_col_num = []\n",
    "        if split in [\"train\", \"valid\"]:\n",
    "            df_list = []\n",
    "            for i in range(5):\n",
    "                if i == cv:\n",
    "                    continue\n",
    "                filepath = os.path.join(base_dirpath, basename.format(i))\n",
    "                df_list.append(pd.read_csv(filepath))\n",
    "                print(split, i)\n",
    "            df = pd.concat(df_list, axis=0)\n",
    "        else:\n",
    "            # test\n",
    "            filepath = os.path.join(base_dirpath, basename.format(cv))\n",
    "            df = pd.read_csv(filepath)\n",
    "            print(split)\n",
    "\n",
    "\n",
    "        if gt_only:\n",
    "            df = df[df[\"class_id\"] > -1]\n",
    "        if train_only and split != \"train\":\n",
    "            df = df[df[\"class_id\"] > -1]\n",
    "\n",
    "        \n",
    "        data_list = []\n",
    "        \n",
    "        df['class_id'] = df['class_id'].astype(int)\n",
    "        df.drop(df[(df['data'].isna()) & (df['class_id'] == -1)].index, inplace=True)\n",
    "        df['col_idx'] = df['col_idx'].astype(int)\n",
    "        df['data'] = df['data'].astype(str)\n",
    "        \n",
    "        num_tables = len(df.groupby(\"table_id\"))\n",
    "        valid_index = int(num_tables * 0.8)\n",
    "        num_train = int(train_ratio * num_tables * 0.8)        \n",
    "        \n",
    "        \n",
    "\n",
    "\n",
    "                \n",
    "        total_num_cols = 0\n",
    "        self.correctness_checklist = []\n",
    "        for i, (index, group_df) in enumerate(df.groupby(\"table_id\")):\n",
    "            if (split == \"train\") and ((i >= num_train) or (i >= valid_index)):\n",
    "                break\n",
    "            if split == \"valid\" and i < valid_index:\n",
    "                continue\n",
    "            #     break\n",
    "            group_df = group_df.sort_values(by=['col_idx'])\n",
    "            group_df = group_df.reset_index(drop=True)\n",
    "\n",
    "            if retrieve_model == \"doduo\":\n",
    "                with torch.no_grad():\n",
    "                    model.eval()\n",
    "                    token_list = group_df[\"data\"].apply(lambda x: tokenizer.encode(\n",
    "                            tokenizer.cls_token + \" \" + x, add_special_tokens=False, max_length=max_length, truncation=True)).tolist()\n",
    "                    token_list = torch.nn.utils.rnn.pad_sequence(\n",
    "                    [torch.LongTensor(tokens) for tokens in token_list], padding_value=0).T.to(device)\n",
    "                    cls_indexes = torch.nonzero(\n",
    "                                token_list == tokenizer.cls_token_id)\n",
    "                    logits, embs = model(token_list, cls_indexes=cls_indexes, get_enc=True)\n",
    "                    col_embs = embs.cpu()\n",
    "                    if len(embs.size()) == 1:\n",
    "                        col_embs = col_embs.unsqueeze(0)\n",
    "            elif retrieve_model == \"sb\":\n",
    "                col_embs = torch.tensor(sb_model.encode(group_df[\"data\"].to_list()))\n",
    "            else:\n",
    "                raise ValueError(\"retrieve_model should be either doduo or sb\")\n",
    "                        \n",
    "            labeled_columns = group_df[group_df['class_id'] > -1]\n",
    "            labeled_columns.sort_values(by=['col_idx'], inplace=True)\n",
    "            for index, target_column in labeled_columns.iterrows():\n",
    "                target_col_idx = target_column[\"col_idx\"]\n",
    "                \n",
    "                if len(group_df) < max_unlabeled:\n",
    "                    # if len(group_df) > max_unlabeled//2:\n",
    "                    #     target_table = group_df\n",
    "                    # else:\n",
    "                    col_emb_target = col_embs[index]\n",
    "                    dist = torch.nn.functional.cosine_similarity(col_emb_target.reshape(1, -1), training_corpus_embs)\n",
    "                    closest_tables = dist.argsort(descending=True)[:K].tolist()\n",
    "                    extra_df = []\n",
    "                    start_col_idx = len(group_df)\n",
    "                    for extra_i in closest_tables:\n",
    "                        extra_df_i = deepcopy(df_train_groups[extra_i]).iloc[1:] \n",
    "                        extra_df_i[\"col_idx\"] = extra_df_i[\"col_idx\"] + start_col_idx\n",
    "                        extra_df.append(extra_df_i) \n",
    "                        start_col_idx += len(extra_df_i)\n",
    "                        \n",
    "                    self.init_col_num.append(len(group_df))\n",
    "                    self.extra_col_num.append(start_col_idx - len(group_df))\n",
    "                    \n",
    "                    \n",
    "                    extra_df = pd.concat([target_column.to_frame().T]+extra_df, ignore_index=True)\n",
    "                    extra_df = extra_df.reset_index(drop=True)\n",
    "                    if len(extra_df) + len(group_df) < max_unlabeled+1:\n",
    "                        target_table = pd.concat([group_df, extra_df.iloc[1:]], ignore_index=True)\n",
    "                    else:\n",
    "                        with torch.no_grad():\n",
    "                            model.eval()\n",
    "                            token_list = extra_df[\"data\"].apply(lambda x: tokenizer.encode(\n",
    "                                    tokenizer.cls_token + \" \" + x, add_special_tokens=False, max_length=max_length, truncation=True)).tolist()\n",
    "                            token_list = torch.nn.utils.rnn.pad_sequence(\n",
    "                            [torch.LongTensor(tokens) for tokens in token_list], padding_value=0).T.to(device)\n",
    "                            cls_indexes = torch.nonzero(\n",
    "                                        token_list == tokenizer.cls_token_id)\n",
    "                            logits, embs_i = model(token_list, cls_indexes=cls_indexes, get_enc=True)\n",
    "                            col_embs_i = embs_i.cpu()\n",
    "                            if len(embs_i.size()) == 1:\n",
    "                                col_embs_i = col_embs_i.unsqueeze(0)\n",
    "                                \n",
    "                        other_columns = extra_df.drop(0)\n",
    "                        selected_index =  maximal_marginal_relevance(col_embs_i[0], col_embs_i[torch.arange(len(col_embs_i)) != 0], top_k=max_unlabeled-len(group_df), lambda_param=lbd)\n",
    "                        assert len(set(selected_index)) == len(selected_index)\n",
    "                        other_columns = other_columns.iloc[selected_index]\n",
    "                        assert len(other_columns) == max_unlabeled - len(group_df)  \n",
    "                        target_table = pd.concat([group_df, other_columns], ignore_index=True)     \n",
    "                # else:\n",
    "                #     self.extra_col_num.append(0)\n",
    "                #     other_columns = group_df.drop(index)\n",
    "                #     selected_index =  maximal_marginal_relevance(col_embs[index], col_embs[torch.arange(len(col_embs)) != index], top_k=max_unlabeled-1, lambda_param=lbd)\n",
    "                #     assert len(set(selected_index)) == len(selected_index)\n",
    "                #     other_columns = other_columns.iloc[selected_index]\n",
    "                #     assert len(other_columns) == max_unlabeled - 1\n",
    "                #     target_table = pd.concat([target_column.to_frame().T, other_columns], ignore_index=True)\n",
    "\n",
    "    \n",
    "                \n",
    "                    target_cls = target_column[\"class_id\"]\n",
    "                    target_table.sort_values(by=['col_idx'], inplace=True)\n",
    "                    col_idx_list = target_table[\"col_idx\"].tolist()\n",
    "                    # target_table.sort_values(by=['col_idx'], inplace=True)\n",
    "\n",
    "                    if max_length <= 128 and adaptive_max_length:\n",
    "                        cur_maxlen = min(max_length, 512 // len(target_table) - 1)\n",
    "                    else:\n",
    "                        cur_maxlen = max_length\n",
    "                        \n",
    "                    token_ids_list = target_table[\"data\"].apply(lambda x: tokenizer.encode(\n",
    "                        tokenizer.cls_token + \" \" + x, add_special_tokens=False, max_length=cur_maxlen, truncation=True)).tolist(\n",
    "                        )\n",
    "                    token_ids = torch.LongTensor(reduce(operator.add,\n",
    "                                                        token_ids_list)).to(device)\n",
    "\n",
    "                    target_col_mask = []\n",
    "                    cls_index_value = 0\n",
    "                    context_id = 1\n",
    "                    meet_target = False\n",
    "                    for idx, col_i in enumerate(col_idx_list):\n",
    "                        if col_i == target_col_idx:\n",
    "                            target_col_mask += [0] * len(token_ids_list[idx])\n",
    "                            meet_target = True\n",
    "                        else:\n",
    "                            target_col_mask += [context_id] * len(token_ids_list[idx])\n",
    "                            context_id += 1\n",
    "                        if not meet_target:\n",
    "                            cls_index_value += len(token_ids_list[idx])\n",
    "                    cls_index_list = [cls_index_value] \n",
    "                    for cls_index in cls_index_list:\n",
    "                        assert token_ids[\n",
    "                            cls_index] == tokenizer.cls_token_id, \"cls_indexes validation\"\n",
    "                    cls_indexes = torch.LongTensor(cls_index_list).to(device)\n",
    "                    class_ids = torch.LongTensor(\n",
    "                        [target_cls]).to(device)\n",
    "                    target_col_mask = torch.LongTensor(target_col_mask).to(device)\n",
    "                    data_list.append(\n",
    "                        [index,\n",
    "                        len(target_table), token_ids, class_ids, cls_indexes, target_col_mask])                \n",
    "        print(split, len(data_list))\n",
    "        self.table_df = pd.DataFrame(data_list,\n",
    "                                     columns=[\n",
    "                                         \"table_id\", \"num_col\", \"data_tensor\",\n",
    "                                         \"label_tensor\", \"cls_indexes\", \"target_col_mask\"\n",
    "                                     ])\n",
    "        \"\"\"\n",
    "        # NOTE: msato contains a small portion of single-col tables. keep it to be consistent.  \n",
    "        if multicol_only:\n",
    "            # Check\n",
    "            num_all_tables = len(self.table_df)\n",
    "            self.table_df = self.table_df[self.table_df[\"num_col\"] > 1]\n",
    "            assert len(self.table_df) == num_all_tables\n",
    "        \"\"\"\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.table_df)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return {\n",
    "            \"data\": self.table_df.iloc[idx][\"data_tensor\"],\n",
    "            \"label\": self.table_df.iloc[idx][\"label_tensor\"],\n",
    "            \"cls_indexes\": self.table_df.iloc[idx][\"cls_indexes\"],\n",
    "            \"target_col_mask\": self.table_df.iloc[idx][\"target_col_mask\"],\n",
    "        }\n",
    "        #\"idx\": torch.LongTensor([idx])}\n",
    "        #\"cls_indexes\": self.table_df.iloc[idx][\"cls_indexes\"]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test\n",
      "test 1085\n"
     ]
    }
   ],
   "source": [
    "src = None\n",
    "test_dataset_iter =  GittablesTablewiseIterateDataset(cv=cv,\n",
    "                            split=\"test\", src=src,\n",
    "                            tokenizer=tokenizer,\n",
    "                            max_length=max_length,\n",
    "                            gt_only='all' not in task,\n",
    "                            device=device,\n",
    "                            base_dirpath=os.path.join(args.data_path, \"GitTables/semtab_gittables/2022\"),\n",
    "                            small_tag=\"semi1\",\n",
    "                            random_sample=True,\n",
    "                            seed=2)\n",
    "padder = collate_fn(tokenizer.pad_token_id)\n",
    "test_dataloader_iter = DataLoader(test_dataset_iter,\n",
    "                                batch_size=1,\n",
    "                            #   collate_fn=collate_fn)\n",
    "                            collate_fn=padder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_train = []\n",
    "for i in range(5):\n",
    "    if i == cv:\n",
    "        continue\n",
    "    filepath = os.path.join(\"/data/zhihao/TU/GitTables/semtab_gittables/2022\",\"semi_cv_{}.csv\".format(i))\n",
    "    df_train.append(pd.read_csv(filepath))\n",
    "df_train = pd.concat(df_train)\n",
    "df_train_labeled = df_train[df_train[\"class_id\"] > -1]\n",
    "model.eval()\n",
    "model = model.to(device)\n",
    "training_corpus_embs_doduo = []\n",
    "training_labels = []\n",
    "with torch.no_grad():\n",
    "    for label, x in zip(df_train_labeled[\"class_id\"], df_train_labeled[\"data\"]):\n",
    "        training_labels.append(label)\n",
    "        cls_indexes = torch.LongTensor([[0, 0]]).to(device)\n",
    "        x = torch.LongTensor( tokenizer.encode(\n",
    "                    tokenizer.cls_token + \" \" + x, add_special_tokens=False, max_length=max_length, truncation=True)).to(device).reshape(1, -1)\n",
    "        logits, embs = model(x, cls_indexes=cls_indexes, get_enc=True)\n",
    "        training_corpus_embs_doduo.append(embs.detach().cpu().reshape(1, -1))\n",
    "    training_corpus_embs_doduo = torch.cat(training_corpus_embs_doduo, dim=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# verifier iterate test\n",
    "import itertools\n",
    "\n",
    "def is_sublist(A, B):\n",
    "    it = iter(B)\n",
    "    return all(x in it for x in A)\n",
    "def get_permutation(x):\n",
    "    new = []\n",
    "    x = x.tolist()\n",
    "    if len(x) == 1:\n",
    "        x = x[0]\n",
    "    for k in x:\n",
    "        if k not in new:\n",
    "            new.append(k)\n",
    "    return new\n",
    "\n",
    "num_cols = []\n",
    "for batch_idx, batch in enumerate(test_dataloader_iter):\n",
    "    target_col_mask = batch[\"target_col_mask\"].T\n",
    "    max_score = -float(\"inf\")\n",
    "\n",
    "\n",
    "    init_permutation_i = get_permutation(target_col_mask)\n",
    "    num_cols.append(len(init_permutation_i))\n",
    "num_cols = torch.tensor(num_cols)\n",
    "num_cols = num_cols[num_cols<8]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test\n",
      "test 1085\n"
     ]
    }
   ],
   "source": [
    "lbd = 0.3\n",
    "test_dataset_mmr = GittablesTablewiseIterateMMRDataset(cv=cv,\n",
    "                            split=\"test\", src=src,\n",
    "                            tokenizer=tokenizer,\n",
    "                            max_length=max_length,\n",
    "                            gt_only='all' not in task,\n",
    "                            device=device,\n",
    "                            base_dirpath=os.path.join(args.data_path, \"GitTables/semtab_gittables/2022\"),\n",
    "                            small_tag=\"semi1\",\n",
    "                            lbd=lbd,\n",
    "                            model=model)\n",
    "padder = collate_fn(tokenizer.pad_token_id)\n",
    "test_dataloader_mmr = DataLoader(test_dataset_mmr,\n",
    "                                batch_size=1,\n",
    "                            #   collate_fn=collate_fn)\n",
    "                            collate_fn=padder)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "**********************************MMR=0.3, DP=1***************************\n",
      "ts_micro_f1=0.5456, ts_macro_f1=0.2782\n",
      "<8 ts_micro_f1=0.5761, ts_macro_f1=0.2195\n",
      "1 ts_micro_f1=1.0000, ts_macro_f1=1.0000\n",
      "3 ts_micro_f1=0.6061, ts_macro_f1=0.2770\n",
      "4 ts_micro_f1=0.5422, ts_macro_f1=0.1742\n",
      "5 ts_micro_f1=0.5568, ts_macro_f1=0.2069\n",
      "6 ts_micro_f1=0.5795, ts_macro_f1=0.2773\n",
      "7 ts_micro_f1=0.6000, ts_macro_f1=0.3168\n",
      "8 ts_micro_f1=0.5278, ts_macro_f1=0.3011\n",
      "**********************************MMR=0.3, DP=2***************************\n",
      "ts_micro_f1=0.5576, ts_macro_f1=0.2791\n",
      "<8 ts_micro_f1=0.5910, ts_macro_f1=0.2643\n",
      "1 ts_micro_f1=1.0000, ts_macro_f1=1.0000\n",
      "3 ts_micro_f1=0.6061, ts_macro_f1=0.2888\n",
      "4 ts_micro_f1=0.5542, ts_macro_f1=0.2013\n",
      "5 ts_micro_f1=0.5909, ts_macro_f1=0.2366\n",
      "6 ts_micro_f1=0.5795, ts_macro_f1=0.2976\n",
      "7 ts_micro_f1=0.6267, ts_macro_f1=0.3655\n",
      "8 ts_micro_f1=0.5380, ts_macro_f1=0.2804\n",
      "**********************************MMR=0.3, DP=3***************************\n",
      "ts_micro_f1=0.5530, ts_macro_f1=0.2752\n",
      "<8 ts_micro_f1=0.5835, ts_macro_f1=0.2571\n",
      "1 ts_micro_f1=1.0000, ts_macro_f1=1.0000\n",
      "3 ts_micro_f1=0.6061, ts_macro_f1=0.2888\n",
      "4 ts_micro_f1=0.5542, ts_macro_f1=0.2022\n",
      "5 ts_micro_f1=0.5795, ts_macro_f1=0.2204\n",
      "6 ts_micro_f1=0.5568, ts_macro_f1=0.2893\n",
      "7 ts_micro_f1=0.6267, ts_macro_f1=0.3829\n",
      "8 ts_micro_f1=0.5351, ts_macro_f1=0.2749\n",
      "**********************************MMR=0.3, DP=4***************************\n",
      "ts_micro_f1=0.5512, ts_macro_f1=0.2763\n",
      "<8 ts_micro_f1=0.5835, ts_macro_f1=0.2592\n",
      "1 ts_micro_f1=1.0000, ts_macro_f1=1.0000\n",
      "3 ts_micro_f1=0.6061, ts_macro_f1=0.2888\n",
      "4 ts_micro_f1=0.5542, ts_macro_f1=0.2022\n",
      "5 ts_micro_f1=0.5795, ts_macro_f1=0.2204\n",
      "6 ts_micro_f1=0.5568, ts_macro_f1=0.2953\n",
      "7 ts_micro_f1=0.6267, ts_macro_f1=0.3926\n",
      "8 ts_micro_f1=0.5322, ts_macro_f1=0.2768\n",
      "**********************************MMR=0.3, DP=5***************************\n",
      "ts_micro_f1=0.5475, ts_macro_f1=0.2767\n",
      "<8 ts_micro_f1=0.5810, ts_macro_f1=0.2586\n",
      "1 ts_micro_f1=1.0000, ts_macro_f1=1.0000\n",
      "3 ts_micro_f1=0.6061, ts_macro_f1=0.2888\n",
      "4 ts_micro_f1=0.5542, ts_macro_f1=0.2022\n",
      "5 ts_micro_f1=0.5795, ts_macro_f1=0.2204\n",
      "6 ts_micro_f1=0.5568, ts_macro_f1=0.2953\n",
      "7 ts_micro_f1=0.6133, ts_macro_f1=0.3854\n",
      "8 ts_micro_f1=0.5278, ts_macro_f1=0.2770\n",
      "**********************************MMR=0.3, DP=6***************************\n",
      "ts_micro_f1=0.5484, ts_macro_f1=0.2769\n",
      "<8 ts_micro_f1=0.5810, ts_macro_f1=0.2586\n",
      "1 ts_micro_f1=1.0000, ts_macro_f1=1.0000\n",
      "3 ts_micro_f1=0.6061, ts_macro_f1=0.2888\n",
      "4 ts_micro_f1=0.5542, ts_macro_f1=0.2022\n",
      "5 ts_micro_f1=0.5795, ts_macro_f1=0.2204\n",
      "6 ts_micro_f1=0.5568, ts_macro_f1=0.2953\n",
      "7 ts_micro_f1=0.6133, ts_macro_f1=0.3854\n",
      "8 ts_micro_f1=0.5292, ts_macro_f1=0.2783\n"
     ]
    }
   ],
   "source": [
    "\n",
    "import itertools\n",
    "\n",
    "def is_sublist(A, B):\n",
    "    it = iter(B)\n",
    "    return all(x in it for x in A)\n",
    "def get_permutation(x):\n",
    "    new = []\n",
    "    x = x.tolist()\n",
    "    if len(x) == 1:\n",
    "        x = x[0]\n",
    "    for k in x:\n",
    "        if k not in new:\n",
    "            new.append(k)\n",
    "    return new\n",
    "\n",
    "\n",
    "micro_all = []\n",
    "macro_all = []\n",
    "# for lbd in [0.5]:\n",
    "for depth in [1, 2, 3, 4, 5, 6]:\n",
    "    print(\"**********************************MMR={}, DP={}***************************\".format(lbd, depth))\n",
    "    src = None\n",
    "    # test_dataset_mmr = GittablesTablewiseIterateMMRDataset(cv=cv,\n",
    "    #                             split=\"test\", src=src,\n",
    "    #                             tokenizer=tokenizer,\n",
    "    #                             max_length=max_length,\n",
    "    #                             gt_only='all' not in task,\n",
    "    #                             device=device,\n",
    "    #                             base_dirpath=os.path.join(args.data_path, \"GitTables/semtab_gittables/2022\"),\n",
    "    #                             small_tag=\"semi1\",\n",
    "    #                             lbd=lbd,\n",
    "    #                             model=model)\n",
    "    # padder = collate_fn(tokenizer.pad_token_id)\n",
    "    # test_dataloader_mmr = DataLoader(test_dataset_mmr,\n",
    "    #                                 batch_size=1,\n",
    "    #                             #   collate_fn=collate_fn)\n",
    "    #                             collate_fn=padder)\n",
    "\n",
    "\n",
    "\n",
    "    model.load_state_dict(best_state_dict, strict=False)\n",
    "    model.eval()\n",
    "    model = model.to(device)\n",
    "    verifier.load_state_dict(veri_state_dict)\n",
    "    verifier.eval()\n",
    "    verifier = verifier.to(device)\n",
    "    labels_test = []\n",
    "    logits_test = []\n",
    "    init_scores = []\n",
    "    init_correctness = []\n",
    "    final_scores = []\n",
    "    final_correctness = []\n",
    "    num_cols = []\n",
    "    total_steps = []\n",
    "    final_steps = []\n",
    "    correct_scores = defaultdict(list)\n",
    "    correct_steps = defaultdict(list)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for batch_idx, batch in enumerate(test_dataloader_mmr):\n",
    "            target_col_mask = batch[\"target_col_mask\"].T\n",
    "            max_score = -float(\"inf\")\n",
    "\n",
    "\n",
    "            init_permutation_i = get_permutation(target_col_mask)\n",
    "            col_idx_set = target_col_mask.unique().tolist()\n",
    "            early_stop = False\n",
    "            i = 0\n",
    "            final_step = 0\n",
    "            total_steps.append(len(init_permutation_i) - max(len(init_permutation_i)//2, len(init_permutation_i)-2) -1)\n",
    "            num_cols.append(len(init_permutation_i))\n",
    "            for r in range(len(init_permutation_i), max(len(init_permutation_i)-depth, 1)-1, -1): # not \n",
    "                for x in itertools.combinations(init_permutation_i, r):\n",
    "                    if 0 not in x:\n",
    "                        continue\n",
    "\n",
    "                    new_batch_data = []\n",
    "                    if len(x) != len(init_permutation_i):\n",
    "                        drop_idx = (set(init_permutation_i)-set(x)).pop()\n",
    "                    else:\n",
    "                        drop_idx = -1\n",
    "                    for col_i in x:\n",
    "                        if col_i == 0:\n",
    "                            if len(new_batch_data) == 0:\n",
    "                                cls_indexes_value = 0\n",
    "                            else:\n",
    "                                cls_indexes_value = sum([len(new_batch_data[i]) for i in range(len(new_batch_data))])\n",
    "                        new_batch_data.append(batch[\"data\"].T[target_col_mask==col_i])\n",
    "                    new_batch_data = torch.cat(new_batch_data, dim=-1).reshape(1, -1)\n",
    "                    cls_indexes = torch.tensor([0, cls_indexes_value]).reshape(1, -1).to(device)\n",
    "                    logits_temp, embs_temp = model(new_batch_data, cls_indexes=cls_indexes, get_enc=True)\n",
    "                    # assert test_embs[batch_idx][i].equal(embs_temp.cpu())\n",
    "                    \n",
    "                    logits_temp = logits_temp.detach().cpu()\n",
    "                    embs_temp = embs_temp.reshape(1,-1)\n",
    "                    scores_temp = F.softmax(verifier(embs_temp))[0,1].item()\n",
    "                    predict_temp = logits_temp.argmax().item()\n",
    "                    i += 1\n",
    "                    if scores_temp > max_score:\n",
    "                        max_score = scores_temp\n",
    "                        logits = logits_temp.clone()\n",
    "                        final_step = len(init_permutation_i) - r\n",
    "                    if logits_temp.argmax().item() == batch[\"label\"].item():\n",
    "                        correct_scores[batch_idx].append(scores_temp)\n",
    "                        correct_steps[batch_idx].append(len(init_permutation_i) - r)\n",
    "                        \n",
    "            final_scores.append(max_score)\n",
    "            final_correctness.append(logits.argmax().item() == batch[\"label\"].item())\n",
    "            final_steps.append(final_step)\n",
    "            labels_test.append(batch[\"label\"].cpu())\n",
    "            logits_test.append(logits.detach().cpu())\n",
    "        labels_test = torch.cat(labels_test, dim=0)\n",
    "        logits_test = torch.stack(logits_test, dim=0)\n",
    "        preds_test = torch.argmax(logits_test, dim=1)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "    from sklearn.metrics import confusion_matrix, f1_score\n",
    "    ts_pred_list = logits_test.argmax(\n",
    "                                1).cpu().detach().numpy().tolist()\n",
    "    ts_micro_f1 = f1_score(labels_test.reshape(-1).numpy().tolist(),\n",
    "                        ts_pred_list,\n",
    "                        average=\"micro\")\n",
    "    ts_macro_f1 = f1_score(labels_test.reshape(-1).numpy().tolist(),\n",
    "                        ts_pred_list,\n",
    "                        average=\"macro\")\n",
    "    ts_f1_full = f1_score(labels_test.reshape(-1).numpy().tolist(),\n",
    "                        ts_pred_list,\n",
    "                        average=None)\n",
    "    # full_f1_init\n",
    "    print(\"ts_micro_f1={:.4f}, ts_macro_f1={:.4f}\".format(ts_micro_f1, ts_macro_f1))\n",
    "\n",
    "    num_cols = torch.tensor(num_cols)\n",
    "    \n",
    "    mask = num_cols < 8\n",
    "    ts_pred_list = logits_test[mask].argmax(\n",
    "                                1).cpu().detach().numpy().tolist()\n",
    "    ts_micro_f1 = f1_score(labels_test.reshape(-1)[mask].numpy().tolist(),\n",
    "                        ts_pred_list,\n",
    "                        average=\"micro\")\n",
    "    ts_macro_f1 = f1_score(labels_test.reshape(-1)[mask].numpy().tolist(),\n",
    "                        ts_pred_list,\n",
    "                        average=\"macro\")\n",
    "        # full_f1_init\n",
    "    print(f\"<8\", \"ts_micro_f1={:.4f}, ts_macro_f1={:.4f}\".format(ts_micro_f1, ts_macro_f1))   \n",
    "\n",
    "    num_cols = torch.tensor(num_cols)\n",
    "    for col_num in num_cols.unique():\n",
    "        mask = num_cols == col_num\n",
    "        ts_pred_list = logits_test[mask].argmax(\n",
    "                                    1).cpu().detach().numpy().tolist()\n",
    "        ts_micro_f1 = f1_score(labels_test.reshape(-1)[mask].numpy().tolist(),\n",
    "                            ts_pred_list,\n",
    "                            average=\"micro\")\n",
    "        ts_macro_f1 = f1_score(labels_test.reshape(-1)[mask].numpy().tolist(),\n",
    "                            ts_pred_list,\n",
    "                            average=\"macro\")\n",
    "        # full_f1_init\n",
    "        print(f\"{col_num}\", \"ts_micro_f1={:.4f}, ts_macro_f1={:.4f}\".format(ts_micro_f1, ts_macro_f1))   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test\n",
      "test 401\n",
      "59.19950124688279 16 223\n",
      "**********************************MMR=0.5, K=5***************************\n"
     ]
    }
   ],
   "source": [
    "lbd = 0.5\n",
    "df_train_groups = []\n",
    "num_cols_train = []\n",
    "for i, (index, group_df) in enumerate(df_train.groupby(\"table_id\")):\n",
    "    group_df.sort_values(by=['col_idx'], inplace=True)\n",
    "    group_df.reset_index(drop=True, inplace=True)\n",
    "    labeled_columns = group_df[group_df['class_id'] > -1]\n",
    "    for index, target_column in labeled_columns.iterrows():\n",
    "        context_columns = group_df.drop(index)\n",
    "        new_df = pd.concat([target_column.to_frame().T, context_columns], ignore_index=True) # put target column in the first row\n",
    "        new_df.reset_index(drop=True, inplace=True)\n",
    "        df_train_groups.append(new_df)\n",
    "        num_cols_train.append(len(new_df))\n",
    "        \n",
    "        \n",
    "K = 5   \n",
    "test_dataset_ret = GittablesTablewiseIterateMMRRetrieveDataset(cv=cv,\n",
    "                            split=\"test\", src=src,\n",
    "                            tokenizer=tokenizer,\n",
    "                            max_length=max_length,\n",
    "                            gt_only='all' not in task,\n",
    "                            device=device,\n",
    "                            base_dirpath=os.path.join(args.data_path, \"GitTables/semtab_gittables/2022\"),\n",
    "                            small_tag=\"semi1\",\n",
    "                            lbd=lbd,\n",
    "                            model=model, \n",
    "                            training_corpus_embs=training_corpus_embs_doduo,\n",
    "                            K=K,\n",
    "                            df_train_groups=df_train_groups)\n",
    "print(np.mean(test_dataset_ret.extra_col_num), np.min(test_dataset_ret.extra_col_num), np.max(test_dataset_ret.extra_col_num))\n",
    "padder = collate_fn(tokenizer.pad_token_id)\n",
    "test_dataloader_ret = DataLoader(test_dataset_ret,\n",
    "                                batch_size=1,\n",
    "                            #   collate_fn=collate_fn)\n",
    "                            collate_fn=padder)\n",
    "print(\"**********************************MMR={}, K={}***************************\".format(lbd, K))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "**********************************Retrieve MMR=0.5, DP=1***************************\n",
      "ts_micro_f1=0.5736, ts_macro_f1=0.2332\n",
      "<8 ts_micro_f1=0.5736, ts_macro_f1=0.2332\n",
      "1 ts_micro_f1=0.0000, ts_macro_f1=0.0000\n",
      "3 ts_micro_f1=0.6212, ts_macro_f1=0.3019\n",
      "4 ts_micro_f1=0.5422, ts_macro_f1=0.2028\n",
      "5 ts_micro_f1=0.6250, ts_macro_f1=0.2479\n",
      "6 ts_micro_f1=0.5000, ts_macro_f1=0.2322\n",
      "7 ts_micro_f1=0.6000, ts_macro_f1=0.3134\n",
      "**********************************Retrieve MMR=0.5, DP=2***************************\n",
      "ts_micro_f1=0.5761, ts_macro_f1=0.2511\n",
      "<8 ts_micro_f1=0.5761, ts_macro_f1=0.2511\n",
      "1 ts_micro_f1=1.0000, ts_macro_f1=1.0000\n",
      "3 ts_micro_f1=0.6061, ts_macro_f1=0.3297\n",
      "4 ts_micro_f1=0.5060, ts_macro_f1=0.2032\n",
      "5 ts_micro_f1=0.6364, ts_macro_f1=0.2726\n",
      "6 ts_micro_f1=0.5227, ts_macro_f1=0.2672\n",
      "7 ts_micro_f1=0.6133, ts_macro_f1=0.3199\n",
      "**********************************Retrieve MMR=0.5, DP=3***************************\n",
      "ts_micro_f1=0.5835, ts_macro_f1=0.2583\n",
      "<8 ts_micro_f1=0.5835, ts_macro_f1=0.2583\n",
      "1 ts_micro_f1=1.0000, ts_macro_f1=1.0000\n",
      "3 ts_micro_f1=0.6061, ts_macro_f1=0.3182\n",
      "4 ts_micro_f1=0.5181, ts_macro_f1=0.1993\n",
      "5 ts_micro_f1=0.6250, ts_macro_f1=0.2650\n",
      "6 ts_micro_f1=0.5455, ts_macro_f1=0.3130\n",
      "7 ts_micro_f1=0.6267, ts_macro_f1=0.3310\n",
      "**********************************Retrieve MMR=0.5, DP=4***************************\n",
      "ts_micro_f1=0.5835, ts_macro_f1=0.2752\n",
      "<8 ts_micro_f1=0.5835, ts_macro_f1=0.2752\n",
      "1 ts_micro_f1=1.0000, ts_macro_f1=1.0000\n",
      "3 ts_micro_f1=0.6061, ts_macro_f1=0.3182\n",
      "4 ts_micro_f1=0.5181, ts_macro_f1=0.1990\n",
      "5 ts_micro_f1=0.6250, ts_macro_f1=0.2637\n",
      "6 ts_micro_f1=0.5341, ts_macro_f1=0.3002\n",
      "7 ts_micro_f1=0.6400, ts_macro_f1=0.3568\n",
      "**********************************Retrieve MMR=0.5, DP=5***************************\n",
      "ts_micro_f1=0.5810, ts_macro_f1=0.2725\n",
      "<8 ts_micro_f1=0.5810, ts_macro_f1=0.2725\n",
      "1 ts_micro_f1=1.0000, ts_macro_f1=1.0000\n",
      "3 ts_micro_f1=0.6061, ts_macro_f1=0.3182\n",
      "4 ts_micro_f1=0.5181, ts_macro_f1=0.1990\n",
      "5 ts_micro_f1=0.6250, ts_macro_f1=0.2647\n",
      "6 ts_micro_f1=0.5341, ts_macro_f1=0.2996\n",
      "7 ts_micro_f1=0.6267, ts_macro_f1=0.3484\n",
      "**********************************Retrieve MMR=0.5, DP=6***************************\n",
      "ts_micro_f1=0.5810, ts_macro_f1=0.2721\n",
      "<8 ts_micro_f1=0.5810, ts_macro_f1=0.2721\n",
      "1 ts_micro_f1=1.0000, ts_macro_f1=1.0000\n",
      "3 ts_micro_f1=0.6061, ts_macro_f1=0.3182\n",
      "4 ts_micro_f1=0.5181, ts_macro_f1=0.2019\n",
      "5 ts_micro_f1=0.6250, ts_macro_f1=0.2647\n",
      "6 ts_micro_f1=0.5341, ts_macro_f1=0.2996\n",
      "7 ts_micro_f1=0.6267, ts_macro_f1=0.3484\n"
     ]
    }
   ],
   "source": [
    "import itertools\n",
    "\n",
    "def is_sublist(A, B):\n",
    "    it = iter(B)\n",
    "    return all(x in it for x in A)\n",
    "def get_permutation(x):\n",
    "    new = []\n",
    "    x = x.tolist()\n",
    "    if len(x) == 1:\n",
    "        x = x[0]\n",
    "    for k in x:\n",
    "        if k not in new:\n",
    "            new.append(k)\n",
    "    return new\n",
    "\n",
    "# verifier iterate test\n",
    "import itertools\n",
    "\n",
    "def is_sublist(A, B):\n",
    "    it = iter(B)\n",
    "    return all(x in it for x in A)\n",
    "def get_permutation(x):\n",
    "    new = []\n",
    "    x = x.tolist()\n",
    "    if len(x) == 1:\n",
    "        x = x[0]\n",
    "    for k in x:\n",
    "        if k not in new:\n",
    "            new.append(k)\n",
    "    return new\n",
    "\n",
    "num_cols = []\n",
    "for batch_idx, batch in enumerate(test_dataloader_iter):\n",
    "    target_col_mask = batch[\"target_col_mask\"].T\n",
    "    max_score = -float(\"inf\")\n",
    "\n",
    "\n",
    "    init_permutation_i = get_permutation(target_col_mask)\n",
    "    num_cols.append(len(init_permutation_i))\n",
    "num_cols = torch.tensor(num_cols)\n",
    "num_cols = num_cols[num_cols<8]\n",
    "\n",
    "\n",
    "micro_all = []\n",
    "macro_all = []\n",
    "# for lbd in [0.5]:\n",
    "for depth in [1, 2, 3, 4, 5, 6]:\n",
    "    print(\"**********************************Retrieve MMR={}, DP={}***************************\".format(lbd, depth))\n",
    "    src = None\n",
    "\n",
    "\n",
    "    model.load_state_dict(best_state_dict, strict=False)\n",
    "    model.eval()\n",
    "    model = model.to(device)\n",
    "    verifier.load_state_dict(veri_state_dict)\n",
    "    verifier.eval()\n",
    "    verifier = verifier.to(device)\n",
    "    labels_test = []\n",
    "    logits_test = []\n",
    "    init_scores = []\n",
    "    init_correctness = []\n",
    "    final_scores = []\n",
    "    final_correctness = []\n",
    "    total_steps = []\n",
    "    final_steps = []\n",
    "    correct_scores = defaultdict(list)\n",
    "    correct_steps = defaultdict(list)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for batch_idx, batch in enumerate(test_dataloader_ret):\n",
    "            target_col_mask = batch[\"target_col_mask\"].T\n",
    "            max_score = -float(\"inf\")\n",
    "\n",
    "\n",
    "            init_permutation_i = get_permutation(target_col_mask)\n",
    "            col_idx_set = target_col_mask.unique().tolist()\n",
    "            early_stop = False\n",
    "            i = 0\n",
    "            final_step = 0\n",
    "            total_steps.append(len(init_permutation_i) - max(len(init_permutation_i)//2, len(init_permutation_i)-2) -1)\n",
    "            for r in range(len(init_permutation_i), max(len(init_permutation_i)-depth, 1)-1, -1): # not \n",
    "                for x in itertools.combinations(init_permutation_i, r):\n",
    "                    if 0 not in x:\n",
    "                        continue\n",
    "\n",
    "                    new_batch_data = []\n",
    "                    if len(x) != len(init_permutation_i):\n",
    "                        drop_idx = (set(init_permutation_i)-set(x)).pop()\n",
    "                    else:\n",
    "                        drop_idx = -1\n",
    "                    for col_i in x:\n",
    "                        if col_i == 0:\n",
    "                            if len(new_batch_data) == 0:\n",
    "                                cls_indexes_value = 0\n",
    "                            else:\n",
    "                                cls_indexes_value = sum([len(new_batch_data[i]) for i in range(len(new_batch_data))])\n",
    "                        new_batch_data.append(batch[\"data\"].T[target_col_mask==col_i])\n",
    "                    new_batch_data = torch.cat(new_batch_data, dim=-1).reshape(1, -1)\n",
    "                    cls_indexes = torch.tensor([0, cls_indexes_value]).reshape(1, -1).to(device)\n",
    "                    logits_temp, embs_temp = model(new_batch_data, cls_indexes=cls_indexes, get_enc=True)\n",
    "                    # assert test_embs[batch_idx][i].equal(embs_temp.cpu())\n",
    "                    \n",
    "                    logits_temp = logits_temp.detach().cpu()\n",
    "                    embs_temp = embs_temp.reshape(1,-1)\n",
    "                    scores_temp = F.softmax(verifier(embs_temp))[0,1].item()\n",
    "                    predict_temp = logits_temp.argmax().item()\n",
    "                    i += 1\n",
    "                    if scores_temp > max_score:\n",
    "                        max_score = scores_temp\n",
    "                        logits = logits_temp.clone()\n",
    "                        final_step = len(init_permutation_i) - r\n",
    "                    if logits_temp.argmax().item() == batch[\"label\"].item():\n",
    "                        correct_scores[batch_idx].append(scores_temp)\n",
    "                        correct_steps[batch_idx].append(len(init_permutation_i) - r)\n",
    "                        \n",
    "            final_scores.append(max_score)\n",
    "            final_correctness.append(logits.argmax().item() == batch[\"label\"].item())\n",
    "            final_steps.append(final_step)\n",
    "            labels_test.append(batch[\"label\"].cpu())\n",
    "            logits_test.append(logits.detach().cpu())\n",
    "        labels_test = torch.cat(labels_test, dim=0)\n",
    "        logits_test = torch.stack(logits_test, dim=0)\n",
    "        preds_test = torch.argmax(logits_test, dim=1)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "    from sklearn.metrics import confusion_matrix, f1_score\n",
    "    ts_pred_list = logits_test.argmax(\n",
    "                                1).cpu().detach().numpy().tolist()\n",
    "    ts_micro_f1 = f1_score(labels_test.reshape(-1).numpy().tolist(),\n",
    "                        ts_pred_list,\n",
    "                        average=\"micro\")\n",
    "    ts_macro_f1 = f1_score(labels_test.reshape(-1).numpy().tolist(),\n",
    "                        ts_pred_list,\n",
    "                        average=\"macro\")\n",
    "    ts_f1_full = f1_score(labels_test.reshape(-1).numpy().tolist(),\n",
    "                        ts_pred_list,\n",
    "                        average=None)\n",
    "    # full_f1_init\n",
    "    print(\"ts_micro_f1={:.4f}, ts_macro_f1={:.4f}\".format(ts_micro_f1, ts_macro_f1))\n",
    "\n",
    "    \n",
    "    mask = num_cols < 8\n",
    "    ts_pred_list = logits_test[mask].argmax(\n",
    "                                1).cpu().detach().numpy().tolist()\n",
    "    ts_micro_f1 = f1_score(labels_test.reshape(-1)[mask].numpy().tolist(),\n",
    "                        ts_pred_list,\n",
    "                        average=\"micro\")\n",
    "    ts_macro_f1 = f1_score(labels_test.reshape(-1)[mask].numpy().tolist(),\n",
    "                        ts_pred_list,\n",
    "                        average=\"macro\")\n",
    "        # full_f1_init\n",
    "    print(f\"<8\", \"ts_micro_f1={:.4f}, ts_macro_f1={:.4f}\".format(ts_micro_f1, ts_macro_f1))   \n",
    "\n",
    "    num_cols = torch.tensor(num_cols)\n",
    "    for col_num in num_cols.unique():\n",
    "        mask = num_cols == col_num\n",
    "        ts_pred_list = logits_test[mask].argmax(\n",
    "                                    1).cpu().detach().numpy().tolist()\n",
    "        ts_micro_f1 = f1_score(labels_test.reshape(-1)[mask].numpy().tolist(),\n",
    "                            ts_pred_list,\n",
    "                            average=\"micro\")\n",
    "        ts_macro_f1 = f1_score(labels_test.reshape(-1)[mask].numpy().tolist(),\n",
    "                            ts_pred_list,\n",
    "                            average=\"macro\")\n",
    "        # full_f1_init\n",
    "        print(f\"{col_num}\", \"ts_micro_f1={:.4f}, ts_macro_f1={:.4f}\".format(ts_micro_f1, ts_macro_f1))   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ts_micro_f1=0.5558, ts_macro_f1=0.2828\n",
      "<8 ts_micro_f1=0.5835, ts_macro_f1=0.2571\n",
      "1 ts_micro_f1=1.0000, ts_macro_f1=1.0000\n",
      "3 ts_micro_f1=0.6061, ts_macro_f1=0.2888\n",
      "4 ts_micro_f1=0.5542, ts_macro_f1=0.2022\n",
      "5 ts_micro_f1=0.5795, ts_macro_f1=0.2204\n",
      "6 ts_micro_f1=0.5568, ts_macro_f1=0.2893\n",
      "7 ts_micro_f1=0.6267, ts_macro_f1=0.3829\n",
      "8 ts_micro_f1=0.5395, ts_macro_f1=0.2859\n"
     ]
    }
   ],
   "source": [
    "        model.load_state_dict(best_state_dict, strict=False)\n",
    "        model.eval()\n",
    "        model = model.to(device)\n",
    "        verifier.load_state_dict(veri_state_dict)\n",
    "        verifier.eval()\n",
    "        verifier = verifier.to(device)\n",
    "        labels_test = []\n",
    "        logits_test = []\n",
    "        init_scores = []\n",
    "        init_correctness = []\n",
    "        final_scores = []\n",
    "        final_correctness = []\n",
    "        num_cols = []\n",
    "        total_steps = []\n",
    "        final_steps = []\n",
    "        correct_scores = defaultdict(list)\n",
    "        correct_steps = defaultdict(list)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            for batch_idx, batch in enumerate(test_dataloader_near):\n",
    "                target_col_mask = batch[\"target_col_mask\"].T\n",
    "                max_score = -float(\"inf\")\n",
    "\n",
    "\n",
    "                init_permutation_i = get_permutation(target_col_mask)\n",
    "                col_idx_set = target_col_mask.unique().tolist()\n",
    "                early_stop = False\n",
    "                i = 0\n",
    "                final_step = 0\n",
    "                total_steps.append(len(init_permutation_i) - max(len(init_permutation_i)//2, len(init_permutation_i)-2) -1)\n",
    "                num_cols.append(len(init_permutation_i))\n",
    "                for r in range(len(init_permutation_i), max(len(init_permutation_i)-depth, 1)-1, -1): # not \n",
    "                    for x in itertools.combinations(init_permutation_i, r):\n",
    "                        if 0 not in x:\n",
    "                            continue\n",
    "\n",
    "                        new_batch_data = []\n",
    "                        if len(x) != len(init_permutation_i):\n",
    "                            drop_idx = (set(init_permutation_i)-set(x)).pop()\n",
    "                        else:\n",
    "                            drop_idx = -1\n",
    "                        for col_i in x:\n",
    "                            if col_i == 0:\n",
    "                                if len(new_batch_data) == 0:\n",
    "                                    cls_indexes_value = 0\n",
    "                                else:\n",
    "                                    cls_indexes_value = sum([len(new_batch_data[i]) for i in range(len(new_batch_data))])\n",
    "                            new_batch_data.append(batch[\"data\"].T[target_col_mask==col_i])\n",
    "                        new_batch_data = torch.cat(new_batch_data, dim=-1).reshape(1, -1)\n",
    "                        cls_indexes = torch.tensor([0, cls_indexes_value]).reshape(1, -1).to(device)\n",
    "                        logits_temp, embs_temp = model(new_batch_data, cls_indexes=cls_indexes, get_enc=True)\n",
    "                        # assert test_embs[batch_idx][i].equal(embs_temp.cpu())\n",
    "                        \n",
    "                        logits_temp = logits_temp.detach().cpu()\n",
    "                        embs_temp = embs_temp.reshape(1,-1)\n",
    "                        scores_temp = F.softmax(verifier(embs_temp))[0,1].item()\n",
    "                        predict_temp = logits_temp.argmax().item()\n",
    "                        i += 1\n",
    "                        if scores_temp > max_score:\n",
    "                            max_score = scores_temp\n",
    "                            logits = logits_temp.clone()\n",
    "                            final_step = len(init_permutation_i) - r\n",
    "                        if logits_temp.argmax().item() == batch[\"label\"].item():\n",
    "                            correct_scores[batch_idx].append(scores_temp)\n",
    "                            correct_steps[batch_idx].append(len(init_permutation_i) - r)\n",
    "                            \n",
    "                final_scores.append(max_score)\n",
    "                final_correctness.append(logits.argmax().item() == batch[\"label\"].item())\n",
    "                final_steps.append(final_step)\n",
    "                labels_test.append(batch[\"label\"].cpu())\n",
    "                logits_test.append(logits.detach().cpu())\n",
    "            labels_test = torch.cat(labels_test, dim=0)\n",
    "            logits_test = torch.stack(logits_test, dim=0)\n",
    "            preds_test = torch.argmax(logits_test, dim=1)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "        from sklearn.metrics import confusion_matrix, f1_score\n",
    "        ts_pred_list = logits_test.argmax(\n",
    "                                    1).cpu().detach().numpy().tolist()\n",
    "        ts_micro_f1 = f1_score(labels_test.reshape(-1).numpy().tolist(),\n",
    "                            ts_pred_list,\n",
    "                            average=\"micro\")\n",
    "        ts_macro_f1 = f1_score(labels_test.reshape(-1).numpy().tolist(),\n",
    "                            ts_pred_list,\n",
    "                            average=\"macro\")\n",
    "        ts_f1_full = f1_score(labels_test.reshape(-1).numpy().tolist(),\n",
    "                            ts_pred_list,\n",
    "                            average=None)\n",
    "        # full_f1_init\n",
    "        print(\"ts_micro_f1={:.4f}, ts_macro_f1={:.4f}\".format(ts_micro_f1, ts_macro_f1))\n",
    "\n",
    "        num_cols = torch.tensor(num_cols)\n",
    "        \n",
    "        mask = num_cols < 8\n",
    "        ts_pred_list = logits_test[mask].argmax(\n",
    "                                    1).cpu().detach().numpy().tolist()\n",
    "        ts_micro_f1 = f1_score(labels_test.reshape(-1)[mask].numpy().tolist(),\n",
    "                            ts_pred_list,\n",
    "                            average=\"micro\")\n",
    "        ts_macro_f1 = f1_score(labels_test.reshape(-1)[mask].numpy().tolist(),\n",
    "                            ts_pred_list,\n",
    "                            average=\"macro\")\n",
    "            # full_f1_init\n",
    "        print(f\"<8\", \"ts_micro_f1={:.4f}, ts_macro_f1={:.4f}\".format(ts_micro_f1, ts_macro_f1))   \n",
    "\n",
    "        num_cols = torch.tensor(num_cols)\n",
    "        for col_num in num_cols.unique():\n",
    "            mask = num_cols == col_num\n",
    "            ts_pred_list = logits_test[mask].argmax(\n",
    "                                        1).cpu().detach().numpy().tolist()\n",
    "            ts_micro_f1 = f1_score(labels_test.reshape(-1)[mask].numpy().tolist(),\n",
    "                                ts_pred_list,\n",
    "                                average=\"micro\")\n",
    "            ts_macro_f1 = f1_score(labels_test.reshape(-1)[mask].numpy().tolist(),\n",
    "                                ts_pred_list,\n",
    "                                average=\"macro\")\n",
    "            # full_f1_init\n",
    "            print(f\"{col_num}\", \"ts_micro_f1={:.4f}, ts_macro_f1={:.4f}\".format(ts_micro_f1, ts_macro_f1))   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "**********************************K=3, DP=3***************************\n",
      "test\n",
      "test 401\n",
      "36.598503740648376 7 209\n",
      "ts_micro_f1=0.5661, ts_macro_f1=0.2337\n",
      "1 ts_micro_f1=1.0000, ts_macro_f1=1.0000\n",
      "3 ts_micro_f1=0.5909, ts_macro_f1=0.2637\n",
      "4 ts_micro_f1=0.5181, ts_macro_f1=0.1947\n",
      "5 ts_micro_f1=0.6136, ts_macro_f1=0.2636\n",
      "6 ts_micro_f1=0.5114, ts_macro_f1=0.2824\n",
      "7 ts_micro_f1=0.6000, ts_macro_f1=0.3029\n",
      "nan nan\n",
      "**********************************K=5, DP=3***************************\n",
      "test\n",
      "test 401\n",
      "59.19950124688279 16 223\n",
      "ts_micro_f1=0.5835, ts_macro_f1=0.2583\n",
      "1 ts_micro_f1=1.0000, ts_macro_f1=1.0000\n",
      "3 ts_micro_f1=0.6061, ts_macro_f1=0.3182\n",
      "4 ts_micro_f1=0.5181, ts_macro_f1=0.1993\n",
      "5 ts_micro_f1=0.6250, ts_macro_f1=0.2650\n",
      "6 ts_micro_f1=0.5455, ts_macro_f1=0.3130\n",
      "7 ts_micro_f1=0.6267, ts_macro_f1=0.3310\n",
      "nan nan\n"
     ]
    }
   ],
   "source": [
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "import itertools\n",
    "\n",
    "def is_sublist(A, B):\n",
    "    it = iter(B)\n",
    "    return all(x in it for x in A)\n",
    "def get_permutation(x):\n",
    "    new = []\n",
    "    x = x.tolist()\n",
    "    if len(x) == 1:\n",
    "        x = x[0]\n",
    "    for k in x:\n",
    "        if k not in new:\n",
    "            new.append(k)\n",
    "    return new\n",
    "\n",
    "\n",
    "micro_all = []\n",
    "macro_all = []\n",
    "for K in [3, 5]:\n",
    "    for depth in [3]:\n",
    "        print(\"**********************************K={}, DP={}***************************\".format(K, depth))\n",
    "        src = None\n",
    "        src = None\n",
    "        test_dataset_ret = GittablesTablewiseIterateMMRRetrieveDataset(cv=cv,\n",
    "                                    split=\"test\", src=src,\n",
    "                                    tokenizer=tokenizer,\n",
    "                                    max_length=max_length,\n",
    "                                    gt_only='all' not in task,\n",
    "                                    device=device,\n",
    "                                    base_dirpath=os.path.join(args.data_path, \"GitTables/semtab_gittables/2022\"),\n",
    "                                    small_tag=\"semi1\",\n",
    "                                    lbd=0.5,\n",
    "                                    model=model, \n",
    "                                    training_corpus_embs=training_corpus_embs_doduo,\n",
    "                                    K=K,\n",
    "                                    df_train_groups=df_train_groups)\n",
    "        print(np.mean(test_dataset_ret.extra_col_num), np.min(test_dataset_ret.extra_col_num), np.max(test_dataset_ret.extra_col_num))\n",
    "        padder = collate_fn(tokenizer.pad_token_id)\n",
    "        test_dataloader_ret = DataLoader(test_dataset_ret,\n",
    "                                        batch_size=1,\n",
    "                                    #   collate_fn=collate_fn)\n",
    "                                    collate_fn=padder)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "        model.load_state_dict(best_state_dict, strict=False)\n",
    "        model.eval()\n",
    "        model = model.to(device)\n",
    "        verifier.load_state_dict(veri_state_dict)\n",
    "        verifier.eval()\n",
    "        verifier = verifier.to(device)\n",
    "        labels_test = []\n",
    "        logits_test = []\n",
    "        init_scores = []\n",
    "        init_correctness = []\n",
    "        final_scores = []\n",
    "        final_correctness = []\n",
    "        total_steps = []\n",
    "        final_steps = []\n",
    "        correct_scores = defaultdict(list)\n",
    "        correct_steps = defaultdict(list)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            for batch_idx, batch in enumerate(test_dataloader_ret):\n",
    "                target_col_mask = batch[\"target_col_mask\"].T\n",
    "                max_score = -float(\"inf\")\n",
    "\n",
    "\n",
    "                init_permutation_i = get_permutation(target_col_mask)\n",
    "                col_idx_set = target_col_mask.unique().tolist()\n",
    "                early_stop = False\n",
    "                i = 0\n",
    "                final_step = 0\n",
    "                total_steps.append(len(init_permutation_i) - max(len(init_permutation_i)//2, len(init_permutation_i)-2) -1)\n",
    "                for r in range(len(init_permutation_i), max(len(init_permutation_i)-depth, 1)-1, -1): # not \n",
    "                    for x in itertools.combinations(init_permutation_i, r):\n",
    "                        if 0 not in x:\n",
    "                            continue\n",
    "\n",
    "                        new_batch_data = []\n",
    "                        if len(x) != len(init_permutation_i):\n",
    "                            drop_idx = (set(init_permutation_i)-set(x)).pop()\n",
    "                        else:\n",
    "                            drop_idx = -1\n",
    "                        for col_i in x:\n",
    "                            if col_i == 0:\n",
    "                                if len(new_batch_data) == 0:\n",
    "                                    cls_indexes_value = 0\n",
    "                                else:\n",
    "                                    cls_indexes_value = sum([len(new_batch_data[i]) for i in range(len(new_batch_data))])\n",
    "                            new_batch_data.append(batch[\"data\"].T[target_col_mask==col_i])\n",
    "                        new_batch_data = torch.cat(new_batch_data, dim=-1).reshape(1, -1)\n",
    "                        cls_indexes = torch.tensor([0, cls_indexes_value]).reshape(1, -1).to(device)\n",
    "                        logits_temp, embs_temp = model(new_batch_data, cls_indexes=cls_indexes, get_enc=True)\n",
    "                        # assert test_embs[batch_idx][i].equal(embs_temp.cpu())\n",
    "                        \n",
    "                        logits_temp = logits_temp.detach().cpu()\n",
    "                        embs_temp = embs_temp.reshape(1,-1)\n",
    "                        scores_temp = F.softmax(verifier(embs_temp))[0,1].item()\n",
    "                        predict_temp = logits_temp.argmax().item()\n",
    "                        i += 1\n",
    "                        if scores_temp > max_score:\n",
    "                            max_score = scores_temp\n",
    "                            logits = logits_temp.clone()\n",
    "                            final_step = len(init_permutation_i) - r\n",
    "                        if logits_temp.argmax().item() == batch[\"label\"].item():\n",
    "                            correct_scores[batch_idx].append(scores_temp)\n",
    "                            correct_steps[batch_idx].append(len(init_permutation_i) - r)\n",
    "                            \n",
    "                final_scores.append(max_score)\n",
    "                final_correctness.append(logits.argmax().item() == batch[\"label\"].item())\n",
    "                final_steps.append(final_step)\n",
    "                labels_test.append(batch[\"label\"].cpu())\n",
    "                logits_test.append(logits.detach().cpu())\n",
    "            labels_test = torch.cat(labels_test, dim=0)\n",
    "            logits_test = torch.stack(logits_test, dim=0)\n",
    "            preds_test = torch.argmax(logits_test, dim=1)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "        from sklearn.metrics import confusion_matrix, f1_score\n",
    "        ts_pred_list = logits_test.argmax(\n",
    "                                    1).cpu().detach().numpy().tolist()\n",
    "        ts_micro_f1 = f1_score(labels_test.reshape(-1).numpy().tolist(),\n",
    "                            ts_pred_list,\n",
    "                            average=\"micro\")\n",
    "        ts_macro_f1 = f1_score(labels_test.reshape(-1).numpy().tolist(),\n",
    "                            ts_pred_list,\n",
    "                            average=\"macro\")\n",
    "        ts_f1_full = f1_score(labels_test.reshape(-1).numpy().tolist(),\n",
    "                            ts_pred_list,\n",
    "                            average=None)\n",
    "        # full_f1_init\n",
    "        print(\"ts_micro_f1={:.4f}, ts_macro_f1={:.4f}\".format(ts_micro_f1, ts_macro_f1))\n",
    "        \n",
    "        num_cols = torch.tensor(num_cols)\n",
    "        for col_num in num_cols.unique():\n",
    "            mask = num_cols == col_num\n",
    "            ts_pred_list = logits_test[mask].argmax(\n",
    "                                        1).cpu().detach().numpy().tolist()\n",
    "            ts_micro_f1 = f1_score(labels_test.reshape(-1)[mask].numpy().tolist(),\n",
    "                                ts_pred_list,\n",
    "                                average=\"micro\")\n",
    "            ts_macro_f1 = f1_score(labels_test.reshape(-1)[mask].numpy().tolist(),\n",
    "                                ts_pred_list,\n",
    "                                average=\"macro\")\n",
    "            # full_f1_init\n",
    "            print(f\"{col_num}\", \"ts_micro_f1={:.4f}, ts_macro_f1={:.4f}\".format(ts_micro_f1, ts_macro_f1))    \n",
    "\n",
    "    print(np.mean(micro_all), np.mean(macro_all))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "**********************************K=8, DP=3***************************\n",
      "test\n",
      "test 401\n",
      "93.85536159600997 37 264\n",
      "ts_micro_f1=0.5686, ts_macro_f1=0.2422\n",
      "1 ts_micro_f1=1.0000, ts_macro_f1=1.0000\n",
      "3 ts_micro_f1=0.6061, ts_macro_f1=0.3271\n",
      "4 ts_micro_f1=0.5181, ts_macro_f1=0.1751\n",
      "5 ts_micro_f1=0.5682, ts_macro_f1=0.2056\n",
      "6 ts_micro_f1=0.5341, ts_macro_f1=0.2988\n",
      "7 ts_micro_f1=0.6267, ts_macro_f1=0.3352\n",
      "nan nan\n",
      "**********************************K=10, DP=3***************************\n",
      "test\n",
      "test 401\n",
      "117.23441396508728 51 291\n",
      "ts_micro_f1=0.5711, ts_macro_f1=0.2446\n",
      "1 ts_micro_f1=0.0000, ts_macro_f1=0.0000\n",
      "3 ts_micro_f1=0.6061, ts_macro_f1=0.3413\n",
      "4 ts_micro_f1=0.5301, ts_macro_f1=0.2068\n",
      "5 ts_micro_f1=0.5795, ts_macro_f1=0.2111\n",
      "6 ts_micro_f1=0.5455, ts_macro_f1=0.3150\n",
      "7 ts_micro_f1=0.6133, ts_macro_f1=0.3179\n",
      "nan nan\n"
     ]
    }
   ],
   "source": [
    "df_train_groups = []\n",
    "num_cols_train = []\n",
    "for i, (index, group_df) in enumerate(df_train.groupby(\"table_id\")):\n",
    "    group_df.sort_values(by=['col_idx'], inplace=True)\n",
    "    group_df.reset_index(drop=True, inplace=True)\n",
    "    labeled_columns = group_df[group_df['class_id'] > -1]\n",
    "    for index, target_column in labeled_columns.iterrows():\n",
    "        context_columns = group_df.drop(index)\n",
    "        new_df = pd.concat([target_column.to_frame().T, context_columns], ignore_index=True) # put target column in the first row\n",
    "        new_df.reset_index(drop=True, inplace=True)\n",
    "        df_train_groups.append(new_df)\n",
    "        num_cols_train.append(len(new_df))\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "import itertools\n",
    "\n",
    "def is_sublist(A, B):\n",
    "    it = iter(B)\n",
    "    return all(x in it for x in A)\n",
    "def get_permutation(x):\n",
    "    new = []\n",
    "    x = x.tolist()\n",
    "    if len(x) == 1:\n",
    "        x = x[0]\n",
    "    for k in x:\n",
    "        if k not in new:\n",
    "            new.append(k)\n",
    "    return new\n",
    "\n",
    "\n",
    "micro_all = []\n",
    "macro_all = []\n",
    "for K in [8, 10]:\n",
    "    for depth in [3]:\n",
    "        print(\"**********************************K={}, DP={}***************************\".format(K, depth))\n",
    "        src = None\n",
    "        src = None\n",
    "        test_dataset_ret = GittablesTablewiseIterateMMRRetrieveDataset(cv=cv,\n",
    "                                    split=\"test\", src=src,\n",
    "                                    tokenizer=tokenizer,\n",
    "                                    max_length=max_length,\n",
    "                                    gt_only='all' not in task,\n",
    "                                    device=device,\n",
    "                                    base_dirpath=os.path.join(args.data_path, \"GitTables/semtab_gittables/2022\"),\n",
    "                                    small_tag=\"semi1\",\n",
    "                                    lbd=0.5,\n",
    "                                    model=model, \n",
    "                                    training_corpus_embs=training_corpus_embs_doduo,\n",
    "                                    K=K,\n",
    "                                    df_train_groups=df_train_groups)\n",
    "        print(np.mean(test_dataset_ret.extra_col_num), np.min(test_dataset_ret.extra_col_num), np.max(test_dataset_ret.extra_col_num))\n",
    "        padder = collate_fn(tokenizer.pad_token_id)\n",
    "        test_dataloader_ret = DataLoader(test_dataset_ret,\n",
    "                                        batch_size=1,\n",
    "                                    #   collate_fn=collate_fn)\n",
    "                                    collate_fn=padder)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "        model.load_state_dict(best_state_dict, strict=False)\n",
    "        model.eval()\n",
    "        model = model.to(device)\n",
    "        verifier.load_state_dict(veri_state_dict)\n",
    "        verifier.eval()\n",
    "        verifier = verifier.to(device)\n",
    "        labels_test = []\n",
    "        logits_test = []\n",
    "        init_scores = []\n",
    "        init_correctness = []\n",
    "        final_scores = []\n",
    "        final_correctness = []\n",
    "        total_steps = []\n",
    "        final_steps = []\n",
    "        correct_scores = defaultdict(list)\n",
    "        correct_steps = defaultdict(list)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            for batch_idx, batch in enumerate(test_dataloader_ret):\n",
    "                target_col_mask = batch[\"target_col_mask\"].T\n",
    "                max_score = -float(\"inf\")\n",
    "\n",
    "\n",
    "                init_permutation_i = get_permutation(target_col_mask)\n",
    "                col_idx_set = target_col_mask.unique().tolist()\n",
    "                early_stop = False\n",
    "                i = 0\n",
    "                final_step = 0\n",
    "                total_steps.append(len(init_permutation_i) - max(len(init_permutation_i)//2, len(init_permutation_i)-2) -1)\n",
    "                for r in range(len(init_permutation_i), max(len(init_permutation_i)-depth, 1)-1, -1): # not \n",
    "                    for x in itertools.combinations(init_permutation_i, r):\n",
    "                        if 0 not in x:\n",
    "                            continue\n",
    "\n",
    "                        new_batch_data = []\n",
    "                        if len(x) != len(init_permutation_i):\n",
    "                            drop_idx = (set(init_permutation_i)-set(x)).pop()\n",
    "                        else:\n",
    "                            drop_idx = -1\n",
    "                        for col_i in x:\n",
    "                            if col_i == 0:\n",
    "                                if len(new_batch_data) == 0:\n",
    "                                    cls_indexes_value = 0\n",
    "                                else:\n",
    "                                    cls_indexes_value = sum([len(new_batch_data[i]) for i in range(len(new_batch_data))])\n",
    "                            new_batch_data.append(batch[\"data\"].T[target_col_mask==col_i])\n",
    "                        new_batch_data = torch.cat(new_batch_data, dim=-1).reshape(1, -1)\n",
    "                        cls_indexes = torch.tensor([0, cls_indexes_value]).reshape(1, -1).to(device)\n",
    "                        logits_temp, embs_temp = model(new_batch_data, cls_indexes=cls_indexes, get_enc=True)\n",
    "                        # assert test_embs[batch_idx][i].equal(embs_temp.cpu())\n",
    "                        \n",
    "                        logits_temp = logits_temp.detach().cpu()\n",
    "                        embs_temp = embs_temp.reshape(1,-1)\n",
    "                        scores_temp = F.softmax(verifier(embs_temp))[0,1].item()\n",
    "                        predict_temp = logits_temp.argmax().item()\n",
    "                        i += 1\n",
    "                        if scores_temp > max_score:\n",
    "                            max_score = scores_temp\n",
    "                            logits = logits_temp.clone()\n",
    "                            final_step = len(init_permutation_i) - r\n",
    "                        if logits_temp.argmax().item() == batch[\"label\"].item():\n",
    "                            correct_scores[batch_idx].append(scores_temp)\n",
    "                            correct_steps[batch_idx].append(len(init_permutation_i) - r)\n",
    "                            \n",
    "                final_scores.append(max_score)\n",
    "                final_correctness.append(logits.argmax().item() == batch[\"label\"].item())\n",
    "                final_steps.append(final_step)\n",
    "                labels_test.append(batch[\"label\"].cpu())\n",
    "                logits_test.append(logits.detach().cpu())\n",
    "            labels_test = torch.cat(labels_test, dim=0)\n",
    "            logits_test = torch.stack(logits_test, dim=0)\n",
    "            preds_test = torch.argmax(logits_test, dim=1)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "        from sklearn.metrics import confusion_matrix, f1_score\n",
    "        ts_pred_list = logits_test.argmax(\n",
    "                                    1).cpu().detach().numpy().tolist()\n",
    "        ts_micro_f1 = f1_score(labels_test.reshape(-1).numpy().tolist(),\n",
    "                            ts_pred_list,\n",
    "                            average=\"micro\")\n",
    "        ts_macro_f1 = f1_score(labels_test.reshape(-1).numpy().tolist(),\n",
    "                            ts_pred_list,\n",
    "                            average=\"macro\")\n",
    "        ts_f1_full = f1_score(labels_test.reshape(-1).numpy().tolist(),\n",
    "                            ts_pred_list,\n",
    "                            average=None)\n",
    "        # full_f1_init\n",
    "        print(\"ts_micro_f1={:.4f}, ts_macro_f1={:.4f}\".format(ts_micro_f1, ts_macro_f1))\n",
    "        \n",
    "        num_cols = torch.tensor(num_cols)\n",
    "        for col_num in num_cols.unique():\n",
    "            mask = num_cols == col_num\n",
    "            ts_pred_list = logits_test[mask].argmax(\n",
    "                                        1).cpu().detach().numpy().tolist()\n",
    "            ts_micro_f1 = f1_score(labels_test.reshape(-1)[mask].numpy().tolist(),\n",
    "                                ts_pred_list,\n",
    "                                average=\"micro\")\n",
    "            ts_macro_f1 = f1_score(labels_test.reshape(-1)[mask].numpy().tolist(),\n",
    "                                ts_pred_list,\n",
    "                                average=\"macro\")\n",
    "            # full_f1_init\n",
    "            print(f\"{col_num}\", \"ts_micro_f1={:.4f}, ts_macro_f1={:.4f}\".format(ts_micro_f1, ts_macro_f1))    \n",
    "\n",
    "    print(np.mean(micro_all), np.mean(macro_all))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "**********************************K=3, DP=3, Tau=0.5***************************\n",
      "test\n",
      "test 401\n",
      "36.598503740648376 7 209\n",
      "ts_micro_f1=0.5736, ts_macro_f1=0.2437\n",
      "1 ts_micro_f1=1.0000, ts_macro_f1=1.0000\n",
      "3 ts_micro_f1=0.6061, ts_macro_f1=0.2945\n",
      "4 ts_micro_f1=0.5301, ts_macro_f1=0.2097\n",
      "5 ts_micro_f1=0.6250, ts_macro_f1=0.2764\n",
      "6 ts_micro_f1=0.5000, ts_macro_f1=0.2640\n",
      "7 ts_micro_f1=0.6133, ts_macro_f1=0.3037\n",
      "**********************************K=5, DP=3, Tau=0.5***************************\n",
      "test\n",
      "test 401\n",
      "59.19950124688279 16 223\n",
      "ts_micro_f1=0.5810, ts_macro_f1=0.2545\n",
      "1 ts_micro_f1=1.0000, ts_macro_f1=1.0000\n",
      "3 ts_micro_f1=0.6061, ts_macro_f1=0.3297\n",
      "4 ts_micro_f1=0.5181, ts_macro_f1=0.2013\n",
      "5 ts_micro_f1=0.6364, ts_macro_f1=0.2653\n",
      "6 ts_micro_f1=0.5114, ts_macro_f1=0.2632\n",
      "7 ts_micro_f1=0.6400, ts_macro_f1=0.3299\n",
      "**********************************K=3, DP=3, Tau=2.0***************************\n",
      "test\n",
      "test 401\n",
      "36.598503740648376 7 209\n",
      "ts_micro_f1=0.5736, ts_macro_f1=0.2450\n",
      "1 ts_micro_f1=1.0000, ts_macro_f1=1.0000\n",
      "3 ts_micro_f1=0.6061, ts_macro_f1=0.2854\n",
      "4 ts_micro_f1=0.5301, ts_macro_f1=0.1967\n",
      "5 ts_micro_f1=0.6136, ts_macro_f1=0.2652\n",
      "6 ts_micro_f1=0.5227, ts_macro_f1=0.2908\n",
      "7 ts_micro_f1=0.6000, ts_macro_f1=0.3239\n",
      "**********************************K=5, DP=3, Tau=2.0***************************\n",
      "test\n",
      "test 401\n",
      "59.19950124688279 16 223\n",
      "ts_micro_f1=0.5810, ts_macro_f1=0.2647\n",
      "1 ts_micro_f1=1.0000, ts_macro_f1=1.0000\n",
      "3 ts_micro_f1=0.6061, ts_macro_f1=0.3182\n",
      "4 ts_micro_f1=0.5181, ts_macro_f1=0.1993\n",
      "5 ts_micro_f1=0.6136, ts_macro_f1=0.2578\n",
      "6 ts_micro_f1=0.5568, ts_macro_f1=0.3218\n",
      "7 ts_micro_f1=0.6133, ts_macro_f1=0.3320\n",
      "**********************************K=3, DP=3, Tau=5.0***************************\n",
      "test\n",
      "test 401\n",
      "36.598503740648376 7 209\n",
      "ts_micro_f1=0.5736, ts_macro_f1=0.2450\n",
      "1 ts_micro_f1=1.0000, ts_macro_f1=1.0000\n",
      "3 ts_micro_f1=0.6061, ts_macro_f1=0.2854\n",
      "4 ts_micro_f1=0.5301, ts_macro_f1=0.1967\n",
      "5 ts_micro_f1=0.6136, ts_macro_f1=0.2652\n",
      "6 ts_micro_f1=0.5227, ts_macro_f1=0.2908\n",
      "7 ts_micro_f1=0.6000, ts_macro_f1=0.3239\n",
      "**********************************K=5, DP=3, Tau=5.0***************************\n",
      "test\n",
      "test 401\n",
      "59.19950124688279 16 223\n",
      "ts_micro_f1=0.5810, ts_macro_f1=0.2620\n",
      "1 ts_micro_f1=1.0000, ts_macro_f1=1.0000\n",
      "3 ts_micro_f1=0.6061, ts_macro_f1=0.3182\n",
      "4 ts_micro_f1=0.5181, ts_macro_f1=0.1936\n",
      "5 ts_micro_f1=0.6136, ts_macro_f1=0.2578\n",
      "6 ts_micro_f1=0.5568, ts_macro_f1=0.3218\n",
      "7 ts_micro_f1=0.6133, ts_macro_f1=0.3320\n",
      "**********************************K=3, DP=3, Tau=10.0***************************\n",
      "test\n",
      "test 401\n",
      "36.598503740648376 7 209\n",
      "ts_micro_f1=0.5736, ts_macro_f1=0.2450\n",
      "1 ts_micro_f1=1.0000, ts_macro_f1=1.0000\n",
      "3 ts_micro_f1=0.6061, ts_macro_f1=0.2854\n",
      "4 ts_micro_f1=0.5301, ts_macro_f1=0.1967\n",
      "5 ts_micro_f1=0.6136, ts_macro_f1=0.2652\n",
      "6 ts_micro_f1=0.5227, ts_macro_f1=0.2908\n",
      "7 ts_micro_f1=0.6000, ts_macro_f1=0.3239\n",
      "**********************************K=5, DP=3, Tau=10.0***************************\n",
      "test\n",
      "test 401\n",
      "59.19950124688279 16 223\n",
      "ts_micro_f1=0.5810, ts_macro_f1=0.2620\n",
      "1 ts_micro_f1=1.0000, ts_macro_f1=1.0000\n",
      "3 ts_micro_f1=0.6061, ts_macro_f1=0.3182\n",
      "4 ts_micro_f1=0.5181, ts_macro_f1=0.1936\n",
      "5 ts_micro_f1=0.6136, ts_macro_f1=0.2578\n",
      "6 ts_micro_f1=0.5568, ts_macro_f1=0.3218\n",
      "7 ts_micro_f1=0.6133, ts_macro_f1=0.3320\n"
     ]
    }
   ],
   "source": [
    "df_train_groups = []\n",
    "num_cols_train = []\n",
    "for i, (index, group_df) in enumerate(df_train.groupby(\"table_id\")):\n",
    "    group_df.sort_values(by=['col_idx'], inplace=True)\n",
    "    group_df.reset_index(drop=True, inplace=True)\n",
    "    labeled_columns = group_df[group_df['class_id'] > -1]\n",
    "    for index, target_column in labeled_columns.iterrows():\n",
    "        context_columns = group_df.drop(index)\n",
    "        new_df = pd.concat([target_column.to_frame().T, context_columns], ignore_index=True) # put target column in the first row\n",
    "        new_df.reset_index(drop=True, inplace=True)\n",
    "        df_train_groups.append(new_df)\n",
    "        num_cols_train.append(len(new_df))\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "import itertools\n",
    "\n",
    "def is_sublist(A, B):\n",
    "    it = iter(B)\n",
    "    return all(x in it for x in A)\n",
    "def get_permutation(x):\n",
    "    new = []\n",
    "    x = x.tolist()\n",
    "    if len(x) == 1:\n",
    "        x = x[0]\n",
    "    for k in x:\n",
    "        if k not in new:\n",
    "            new.append(k)\n",
    "    return new\n",
    "\n",
    "\n",
    "micro_all = []\n",
    "macro_all = []\n",
    "for tau in [0.5, 2.0, 5.0, 10.0]:\n",
    "    for K in [3, 5]:\n",
    "        for depth in [3]:\n",
    "            print(\"**********************************K={}, DP={}, Tau={}***************************\".format(K, depth, tau))\n",
    "            src = None\n",
    "            src = None\n",
    "            test_dataset_ret = GittablesTablewiseIterateMMRRetrieveDataset(cv=cv,\n",
    "                                        split=\"test\", src=src,\n",
    "                                        tokenizer=tokenizer,\n",
    "                                        max_length=max_length,\n",
    "                                        gt_only='all' not in task,\n",
    "                                        device=device,\n",
    "                                        base_dirpath=os.path.join(args.data_path, \"GitTables/semtab_gittables/2022\"),\n",
    "                                        small_tag=\"semi1\",\n",
    "                                        lbd=0.5,\n",
    "                                        model=model, \n",
    "                                        training_corpus_embs=training_corpus_embs_doduo,\n",
    "                                        K=K,\n",
    "                                        df_train_groups=df_train_groups)\n",
    "            print(np.mean(test_dataset_ret.extra_col_num), np.min(test_dataset_ret.extra_col_num), np.max(test_dataset_ret.extra_col_num))\n",
    "            padder = collate_fn(tokenizer.pad_token_id)\n",
    "            test_dataloader_ret = DataLoader(test_dataset_ret,\n",
    "                                            batch_size=1,\n",
    "                                        #   collate_fn=collate_fn)\n",
    "                                        collate_fn=padder)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "            model.load_state_dict(best_state_dict, strict=False)\n",
    "            model.eval()\n",
    "            model = model.to(device)\n",
    "            verifier.load_state_dict(veri_state_dict)\n",
    "            verifier.eval()\n",
    "            verifier = verifier.to(device)\n",
    "            labels_test = []\n",
    "            logits_test = []\n",
    "            init_scores = []\n",
    "            init_correctness = []\n",
    "            final_scores = []\n",
    "            final_correctness = []\n",
    "            total_steps = []\n",
    "            final_steps = []\n",
    "            correct_scores = defaultdict(list)\n",
    "            correct_steps = defaultdict(list)\n",
    "\n",
    "            with torch.no_grad():\n",
    "                for batch_idx, batch in enumerate(test_dataloader_ret):\n",
    "                    target_col_mask = batch[\"target_col_mask\"].T\n",
    "                    max_score = -float(\"inf\")\n",
    "\n",
    "\n",
    "                    init_permutation_i = get_permutation(target_col_mask)\n",
    "                    col_idx_set = target_col_mask.unique().tolist()\n",
    "                    early_stop = False\n",
    "                    i = 0\n",
    "                    final_step = 0\n",
    "                    total_steps.append(len(init_permutation_i) - max(len(init_permutation_i)//2, len(init_permutation_i)-2) -1)\n",
    "                    for r in range(len(init_permutation_i), max(len(init_permutation_i)-depth, 1)-1, -1): # not \n",
    "                        for x in itertools.combinations(init_permutation_i, r):\n",
    "                            if 0 not in x:\n",
    "                                continue\n",
    "\n",
    "                            new_batch_data = []\n",
    "                            if len(x) != len(init_permutation_i):\n",
    "                                drop_idx = (set(init_permutation_i)-set(x)).pop()\n",
    "                            else:\n",
    "                                drop_idx = -1\n",
    "                            for col_i in x:\n",
    "                                if col_i == 0:\n",
    "                                    if len(new_batch_data) == 0:\n",
    "                                        cls_indexes_value = 0\n",
    "                                    else:\n",
    "                                        cls_indexes_value = sum([len(new_batch_data[i]) for i in range(len(new_batch_data))])\n",
    "                                new_batch_data.append(batch[\"data\"].T[target_col_mask==col_i])\n",
    "                            new_batch_data = torch.cat(new_batch_data, dim=-1).reshape(1, -1)\n",
    "                            cls_indexes = torch.tensor([0, cls_indexes_value]).reshape(1, -1).to(device)\n",
    "                            logits_temp, embs_temp = model(new_batch_data, cls_indexes=cls_indexes, get_enc=True)\n",
    "                            # assert test_embs[batch_idx][i].equal(embs_temp.cpu())\n",
    "                            \n",
    "                            logits_temp = logits_temp.detach().cpu()\n",
    "                            embs_temp = embs_temp.reshape(1,-1)\n",
    "                            scores_temp = F.softmax(verifier(embs_temp)/tau)[0,1].item()\n",
    "                            predict_temp = logits_temp.argmax().item()\n",
    "                            i += 1\n",
    "                            if scores_temp > max_score:\n",
    "                                max_score = scores_temp\n",
    "                                logits = logits_temp.clone()\n",
    "                                final_step = len(init_permutation_i) - r\n",
    "                            if logits_temp.argmax().item() == batch[\"label\"].item():\n",
    "                                correct_scores[batch_idx].append(scores_temp)\n",
    "                                correct_steps[batch_idx].append(len(init_permutation_i) - r)\n",
    "                                \n",
    "                    final_scores.append(max_score)\n",
    "                    final_correctness.append(logits.argmax().item() == batch[\"label\"].item())\n",
    "                    final_steps.append(final_step)\n",
    "                    labels_test.append(batch[\"label\"].cpu())\n",
    "                    logits_test.append(logits.detach().cpu())\n",
    "                labels_test = torch.cat(labels_test, dim=0)\n",
    "                logits_test = torch.stack(logits_test, dim=0)\n",
    "                preds_test = torch.argmax(logits_test, dim=1)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "            from sklearn.metrics import confusion_matrix, f1_score\n",
    "            ts_pred_list = logits_test.argmax(\n",
    "                                        1).cpu().detach().numpy().tolist()\n",
    "            ts_micro_f1 = f1_score(labels_test.reshape(-1).numpy().tolist(),\n",
    "                                ts_pred_list,\n",
    "                                average=\"micro\")\n",
    "            ts_macro_f1 = f1_score(labels_test.reshape(-1).numpy().tolist(),\n",
    "                                ts_pred_list,\n",
    "                                average=\"macro\")\n",
    "            ts_f1_full = f1_score(labels_test.reshape(-1).numpy().tolist(),\n",
    "                                ts_pred_list,\n",
    "                                average=None)\n",
    "            # full_f1_init\n",
    "            print(\"ts_micro_f1={:.4f}, ts_macro_f1={:.4f}\".format(ts_micro_f1, ts_macro_f1))\n",
    "            \n",
    "            num_cols = torch.tensor(num_cols)\n",
    "            for col_num in num_cols.unique():\n",
    "                mask = num_cols == col_num\n",
    "                ts_pred_list = logits_test[mask].argmax(\n",
    "                                            1).cpu().detach().numpy().tolist()\n",
    "                ts_micro_f1 = f1_score(labels_test.reshape(-1)[mask].numpy().tolist(),\n",
    "                                    ts_pred_list,\n",
    "                                    average=\"micro\")\n",
    "                ts_macro_f1 = f1_score(labels_test.reshape(-1)[mask].numpy().tolist(),\n",
    "                                    ts_pred_list,\n",
    "                                    average=\"macro\")\n",
    "                # full_f1_init\n",
    "                print(f\"{col_num}\", \"ts_micro_f1={:.4f}, ts_macro_f1={:.4f}\".format(ts_micro_f1, ts_macro_f1))    \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "401.9887111186981\n",
      "ts_micro_f1=0.5641, ts_macro_f1=0.2924\n",
      "<8 ts_micro_f1=0.5810, ts_macro_f1=0.2549\n",
      "1 ts_micro_f1=1.0000, ts_macro_f1=1.0000\n",
      "3 ts_micro_f1=0.6061, ts_macro_f1=0.2888\n",
      "4 ts_micro_f1=0.5422, ts_macro_f1=0.1955\n",
      "5 ts_micro_f1=0.5795, ts_macro_f1=0.2204\n",
      "6 ts_micro_f1=0.5568, ts_macro_f1=0.2798\n",
      "7 ts_micro_f1=0.6267, ts_macro_f1=0.3829\n",
      "8 ts_micro_f1=0.5541, ts_macro_f1=0.3008\n"
     ]
    }
   ],
   "source": [
    "dp = 6\n",
    "model.load_state_dict(best_state_dict, strict=False)\n",
    "model.eval()\n",
    "model = model.to(device)\n",
    "verifier.load_state_dict(veri_state_dict)\n",
    "verifier.eval()\n",
    "verifier = verifier.to(device)\n",
    "labels_test = []\n",
    "logits_test = []\n",
    "init_scores = []\n",
    "init_correctness = []\n",
    "final_scores = []\n",
    "final_correctness = []\n",
    "num_cols = []\n",
    "total_steps = []\n",
    "final_steps = []\n",
    "correct_scores = defaultdict(list)\n",
    "correct_steps = defaultdict(list)\n",
    "ts = time.time()\n",
    "with torch.no_grad():\n",
    "    for batch_idx, batch in enumerate(test_dataloader_mmr):\n",
    "        target_col_mask = batch[\"target_col_mask\"].T\n",
    "        max_score = -float(\"inf\")\n",
    "\n",
    "\n",
    "        init_permutation_i = get_permutation(target_col_mask)\n",
    "        col_idx_set = target_col_mask.unique().tolist()\n",
    "        early_stop = False\n",
    "        i = 0\n",
    "        final_step = 0\n",
    "        total_steps.append(len(init_permutation_i) - max(len(init_permutation_i)//2, len(init_permutation_i)-2) -1)\n",
    "        num_cols.append(len(init_permutation_i))\n",
    "        for r in range(len(init_permutation_i), max(len(init_permutation_i)-depth, 1)-1, -1): # not \n",
    "            for x in itertools.combinations(init_permutation_i, r):\n",
    "                if 0 not in x:\n",
    "                    continue\n",
    "\n",
    "                new_batch_data = []\n",
    "                if len(x) != len(init_permutation_i):\n",
    "                    drop_idx = (set(init_permutation_i)-set(x)).pop()\n",
    "                else:\n",
    "                    drop_idx = -1\n",
    "                for col_i in x:\n",
    "                    if col_i == 0:\n",
    "                        if len(new_batch_data) == 0:\n",
    "                            cls_indexes_value = 0\n",
    "                        else:\n",
    "                            cls_indexes_value = sum([len(new_batch_data[i]) for i in range(len(new_batch_data))])\n",
    "                    new_batch_data.append(batch[\"data\"].T[target_col_mask==col_i])\n",
    "                new_batch_data = torch.cat(new_batch_data, dim=-1).reshape(1, -1)\n",
    "                cls_indexes = torch.tensor([0, cls_indexes_value]).reshape(1, -1).to(device)\n",
    "                logits_temp, embs_temp = model(new_batch_data, cls_indexes=cls_indexes, get_enc=True)\n",
    "                # assert test_embs[batch_idx][i].equal(embs_temp.cpu())\n",
    "                \n",
    "                logits_temp = logits_temp.detach().cpu()\n",
    "                embs_temp = embs_temp.reshape(1,-1)\n",
    "                scores_temp = F.softmax(verifier(embs_temp)/tau)[0,1].item()\n",
    "                predict_temp = logits_temp.argmax().item()\n",
    "                i += 1\n",
    "                if scores_temp > max_score:\n",
    "                    max_score = scores_temp\n",
    "                    logits = logits_temp.clone()\n",
    "                    final_step = len(init_permutation_i) - r\n",
    "                if logits_temp.argmax().item() == batch[\"label\"].item():\n",
    "                    correct_scores[batch_idx].append(scores_temp)\n",
    "                    correct_steps[batch_idx].append(len(init_permutation_i) - r)\n",
    "                    \n",
    "        final_scores.append(max_score)\n",
    "        final_correctness.append(logits.argmax().item() == batch[\"label\"].item())\n",
    "        final_steps.append(final_step)\n",
    "        labels_test.append(batch[\"label\"].cpu())\n",
    "        logits_test.append(logits.detach().cpu())\n",
    "    labels_test = torch.cat(labels_test, dim=0)\n",
    "    logits_test = torch.stack(logits_test, dim=0)\n",
    "    preds_test = torch.argmax(logits_test, dim=1)\n",
    "\n",
    "print(time.time()-ts)\n",
    "\n",
    "\n",
    "from sklearn.metrics import confusion_matrix, f1_score\n",
    "ts_pred_list = logits_test.argmax(\n",
    "                            1).cpu().detach().numpy().tolist()\n",
    "ts_micro_f1 = f1_score(labels_test.reshape(-1).numpy().tolist(),\n",
    "                    ts_pred_list,\n",
    "                    average=\"micro\")\n",
    "ts_macro_f1 = f1_score(labels_test.reshape(-1).numpy().tolist(),\n",
    "                    ts_pred_list,\n",
    "                    average=\"macro\")\n",
    "ts_f1_full = f1_score(labels_test.reshape(-1).numpy().tolist(),\n",
    "                    ts_pred_list,\n",
    "                    average=None)\n",
    "# full_f1_init\n",
    "print(\"ts_micro_f1={:.4f}, ts_macro_f1={:.4f}\".format(ts_micro_f1, ts_macro_f1))\n",
    "\n",
    "num_cols = torch.tensor(num_cols)\n",
    "\n",
    "mask = num_cols < 8\n",
    "ts_pred_list = logits_test[mask].argmax(\n",
    "                            1).cpu().detach().numpy().tolist()\n",
    "ts_micro_f1 = f1_score(labels_test.reshape(-1)[mask].numpy().tolist(),\n",
    "                    ts_pred_list,\n",
    "                    average=\"micro\")\n",
    "ts_macro_f1 = f1_score(labels_test.reshape(-1)[mask].numpy().tolist(),\n",
    "                    ts_pred_list,\n",
    "                    average=\"macro\")\n",
    "    # full_f1_init\n",
    "print(f\"<8\", \"ts_micro_f1={:.4f}, ts_macro_f1={:.4f}\".format(ts_micro_f1, ts_macro_f1))   \n",
    "\n",
    "num_cols = torch.tensor(num_cols)\n",
    "for col_num in num_cols.unique():\n",
    "    mask = num_cols == col_num\n",
    "    ts_pred_list = logits_test[mask].argmax(\n",
    "                                1).cpu().detach().numpy().tolist()\n",
    "    ts_micro_f1 = f1_score(labels_test.reshape(-1)[mask].numpy().tolist(),\n",
    "                        ts_pred_list,\n",
    "                        average=\"micro\")\n",
    "    ts_macro_f1 = f1_score(labels_test.reshape(-1)[mask].numpy().tolist(),\n",
    "                        ts_pred_list,\n",
    "                        average=\"macro\")\n",
    "    # full_f1_init\n",
    "    print(f\"{col_num}\", \"ts_micro_f1={:.4f}, ts_macro_f1={:.4f}\".format(ts_micro_f1, ts_macro_f1))   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "389.56192684173584\n",
      "ts_micro_f1=0.5567, ts_macro_f1=0.2902\n",
      "1 ts_micro_f1=1.0000, ts_macro_f1=1.0000\n",
      "3 ts_micro_f1=0.6061, ts_macro_f1=0.2888\n",
      "4 ts_micro_f1=0.5422, ts_macro_f1=0.1955\n",
      "5 ts_micro_f1=0.5795, ts_macro_f1=0.2204\n",
      "6 ts_micro_f1=0.5568, ts_macro_f1=0.2856\n",
      "7 ts_micro_f1=0.6133, ts_macro_f1=0.3854\n",
      "8 ts_micro_f1=0.5439, ts_macro_f1=0.2966\n"
     ]
    }
   ],
   "source": [
    "depth = 6\n",
    "model.load_state_dict(best_state_dict, strict=False)\n",
    "model.eval()\n",
    "model = model.to(device)\n",
    "verifier.load_state_dict(veri_state_dict)\n",
    "verifier.eval()\n",
    "verifier = verifier.to(device)\n",
    "labels_test = []\n",
    "logits_test = []\n",
    "init_scores = []\n",
    "init_correctness = []\n",
    "final_scores = []\n",
    "final_correctness = []\n",
    "total_steps = []\n",
    "final_steps = []\n",
    "correct_scores = defaultdict(list)\n",
    "correct_steps = defaultdict(list)\n",
    "ts = time.time()\n",
    "with torch.no_grad():\n",
    "    for batch_idx, batch in enumerate(test_dataloader_mmr):\n",
    "        target_col_mask = batch[\"target_col_mask\"].T\n",
    "        max_score = -float(\"inf\")\n",
    "\n",
    "\n",
    "        init_permutation_i = get_permutation(target_col_mask)\n",
    "        col_idx_set = target_col_mask.unique().tolist()\n",
    "        early_stop = False\n",
    "        i = 0\n",
    "        final_step = 0\n",
    "        total_steps.append(len(init_permutation_i) - max(len(init_permutation_i)//2, len(init_permutation_i)-2) -1)\n",
    "        batch_data_all = []\n",
    "        cls_indexes_all = []\n",
    "        count = 0\n",
    "        for r in range(len(init_permutation_i), max(len(init_permutation_i)-depth, 1)-1, -1): # not \n",
    "            for x in itertools.combinations(init_permutation_i, r):\n",
    "                if 0 not in x:\n",
    "                    continue\n",
    "\n",
    "                new_batch_data = []\n",
    "                if len(x) != len(init_permutation_i):\n",
    "                    drop_idx = (set(init_permutation_i)-set(x)).pop()\n",
    "                else:\n",
    "                    drop_idx = -1\n",
    "                for col_i in x:\n",
    "                    if col_i == 0:\n",
    "                        if len(new_batch_data) == 0:\n",
    "                            cls_indexes_value = 0\n",
    "                        else:\n",
    "                            cls_indexes_value = sum([len(new_batch_data[i]) for i in range(len(new_batch_data))])\n",
    "                    new_batch_data.append(batch[\"data\"].T[target_col_mask==col_i])\n",
    "                new_batch_data = torch.cat(new_batch_data, dim=-1).reshape(-1)\n",
    "                batch_data_all.append(new_batch_data)\n",
    "                cls_indexes = torch.tensor([count, cls_indexes_value]).reshape(1, -1).to(device)\n",
    "                cls_indexes_all.append(cls_indexes)\n",
    "                count += 1\n",
    "        batch_data_all = torch.nn.utils.rnn.pad_sequence(\n",
    "            batch_data_all, padding_value=tokenizer.pad_token_id, batch_first=True).reshape(count, -1)\n",
    "        cls_indexes_all = torch.cat(cls_indexes_all, dim=0).reshape(count, -1)\n",
    "        logits_temp, embs_temp = model(batch_data_all, cls_indexes=cls_indexes_all, get_enc=True)\n",
    "        # assert test_embs[batch_idx][i].equal(embs_temp.cpu())\n",
    "        embs_temp = embs_temp.reshape(count,-1)\n",
    "        logits_temp = logits_temp.detach().cpu().reshape(count, -1)\n",
    "        scores_temp = F.softmax(verifier(embs_temp)/tau)[:,1].detach().cpu()\n",
    "        max_index = scores_temp.argmax().item()\n",
    "        max_score = scores_temp[max_index].item()\n",
    "        logits = logits_temp[max_index]\n",
    "        predict_temp = logits_temp.argmax().item()\n",
    "        i += 1\n",
    "\n",
    "        labels_test.append(batch[\"label\"].cpu())\n",
    "        logits_test.append(logits.detach().cpu())\n",
    "    labels_test = torch.cat(labels_test, dim=0)\n",
    "    logits_test = torch.stack(logits_test, dim=0)\n",
    "    preds_test = torch.argmax(logits_test, dim=1)\n",
    "\n",
    "\n",
    "print(time.time()-ts)\n",
    "from sklearn.metrics import confusion_matrix, f1_score\n",
    "ts_pred_list = logits_test.argmax(\n",
    "                            1).cpu().detach().numpy().tolist()\n",
    "ts_micro_f1 = f1_score(labels_test.reshape(-1).numpy().tolist(),\n",
    "                    ts_pred_list,\n",
    "                    average=\"micro\")\n",
    "ts_macro_f1 = f1_score(labels_test.reshape(-1).numpy().tolist(),\n",
    "                    ts_pred_list,\n",
    "                    average=\"macro\")\n",
    "ts_f1_full = f1_score(labels_test.reshape(-1).numpy().tolist(),\n",
    "                    ts_pred_list,\n",
    "                    average=None)\n",
    "# full_f1_init\n",
    "print(\"ts_micro_f1={:.4f}, ts_macro_f1={:.4f}\".format(ts_micro_f1, ts_macro_f1))\n",
    "\n",
    "num_cols = torch.tensor(num_cols)\n",
    "for col_num in num_cols.unique():\n",
    "    mask = num_cols == col_num\n",
    "    ts_pred_list = logits_test[mask].argmax(\n",
    "                                1).cpu().detach().numpy().tolist()\n",
    "    ts_micro_f1 = f1_score(labels_test.reshape(-1)[mask].numpy().tolist(),\n",
    "                        ts_pred_list,\n",
    "                        average=\"micro\")\n",
    "    ts_macro_f1 = f1_score(labels_test.reshape(-1)[mask].numpy().tolist(),\n",
    "                        ts_pred_list,\n",
    "                        average=\"macro\")\n",
    "    # full_f1_init\n",
    "    print(f\"{col_num}\", \"ts_micro_f1={:.4f}, ts_macro_f1={:.4f}\".format(ts_micro_f1, ts_macro_f1))    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([768])"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embs_temp.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "count"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "**********************************K=10, DP=3, Tau=1.0***************************\n",
      "test\n",
      "test 1085\n",
      "43.32811059907834 0 291\n",
      "290.12770652770996\n",
      "ts_micro_f1=0.5613, ts_macro_f1=0.2896\n",
      "8 ts_micro_f1=0.5613, ts_macro_f1=0.2896\n"
     ]
    }
   ],
   "source": [
    "df_train_groups = []\n",
    "num_cols_train = []\n",
    "for i, (index, group_df) in enumerate(df_train.groupby(\"table_id\")):\n",
    "    group_df.sort_values(by=['col_idx'], inplace=True)\n",
    "    group_df.reset_index(drop=True, inplace=True)\n",
    "    labeled_columns = group_df[group_df['class_id'] > -1]\n",
    "    for index, target_column in labeled_columns.iterrows():\n",
    "        context_columns = group_df.drop(index)\n",
    "        new_df = pd.concat([target_column.to_frame().T, context_columns], ignore_index=True) # put target column in the first row\n",
    "        new_df.reset_index(drop=True, inplace=True)\n",
    "        df_train_groups.append(new_df)\n",
    "        num_cols_train.append(len(new_df))\n",
    "        \n",
    "        \n",
    "import itertools\n",
    "\n",
    "def is_sublist(A, B):\n",
    "    it = iter(B)\n",
    "    return all(x in it for x in A)\n",
    "def get_permutation(x):\n",
    "    new = []\n",
    "    x = x.tolist()\n",
    "    if len(x) == 1:\n",
    "        x = x[0]\n",
    "    for k in x:\n",
    "        if k not in new:\n",
    "            new.append(k)\n",
    "    return new\n",
    "\n",
    "\n",
    "import time\n",
    "for tau in [1.0]:\n",
    "    for K in [10]:\n",
    "        for depth in [3]:\n",
    "            print(\"**********************************K={}, DP={}, Tau={}***************************\".format(K, depth, tau))\n",
    "            src = None\n",
    "            src = None\n",
    "            test_dataset_ret = GittablesTablewiseIterateMMRRetrieveDataset(cv=cv,\n",
    "                                        split=\"test\", src=src,\n",
    "                                        tokenizer=tokenizer,\n",
    "                                        max_length=max_length,\n",
    "                                        gt_only='all' not in task,\n",
    "                                        device=device,\n",
    "                                        base_dirpath=os.path.join(args.data_path, \"GitTables/semtab_gittables/2022\"),\n",
    "                                        small_tag=\"semi1\",\n",
    "                                        lbd=0.5,\n",
    "                                        model=model, \n",
    "                                        training_corpus_embs=training_corpus_embs_doduo,\n",
    "                                        K=K,\n",
    "                                        df_train_groups=df_train_groups)\n",
    "            print(np.mean(test_dataset_ret.extra_col_num), np.min(test_dataset_ret.extra_col_num), np.max(test_dataset_ret.extra_col_num))\n",
    "            padder = collate_fn(tokenizer.pad_token_id)\n",
    "            test_dataloader_ret = DataLoader(test_dataset_ret,\n",
    "                                            batch_size=1,\n",
    "                                        #   collate_fn=collate_fn)\n",
    "                                        collate_fn=padder)\n",
    "            \n",
    "            \n",
    "            model.load_state_dict(best_state_dict, strict=False)\n",
    "            model.eval()\n",
    "            model = model.to(device)\n",
    "            verifier.load_state_dict(veri_state_dict)\n",
    "            verifier.eval()\n",
    "            verifier = verifier.to(device)\n",
    "            labels_test = []\n",
    "            logits_test = []\n",
    "            init_scores = []\n",
    "            init_correctness = []\n",
    "            final_scores = []\n",
    "            final_correctness = []\n",
    "            total_steps = []\n",
    "            final_steps = []\n",
    "            correct_scores = defaultdict(list)\n",
    "            correct_steps = defaultdict(list)\n",
    "            num_cols = []\n",
    "            ts = time.time()\n",
    "            with torch.no_grad():\n",
    "                for batch_idx, batch in enumerate(test_dataloader_ret):\n",
    "                    target_col_mask = batch[\"target_col_mask\"].T\n",
    "                    max_score = -float(\"inf\")\n",
    "\n",
    "\n",
    "                    init_permutation_i = get_permutation(target_col_mask)\n",
    "                    num_cols.append(len(init_permutation_i))\n",
    "                    col_idx_set = target_col_mask.unique().tolist()\n",
    "                    early_stop = False\n",
    "                    i = 0\n",
    "                    final_step = 0\n",
    "                    total_steps.append(len(init_permutation_i) - max(len(init_permutation_i)//2, len(init_permutation_i)-2) -1)\n",
    "                    batch_data_all = []\n",
    "                    cls_indexes_all = []\n",
    "                    count = 0\n",
    "                    for r in range(len(init_permutation_i), max(len(init_permutation_i)-depth, 1)-1, -1): # not \n",
    "                        for x in itertools.combinations(init_permutation_i, r):\n",
    "                            if 0 not in x:\n",
    "                                continue\n",
    "\n",
    "                            new_batch_data = []\n",
    "                            if len(x) != len(init_permutation_i):\n",
    "                                drop_idx = (set(init_permutation_i)-set(x)).pop()\n",
    "                            else:\n",
    "                                drop_idx = -1\n",
    "                            for col_i in x:\n",
    "                                if col_i == 0:\n",
    "                                    if len(new_batch_data) == 0:\n",
    "                                        cls_indexes_value = 0\n",
    "                                    else:\n",
    "                                        cls_indexes_value = sum([len(new_batch_data[i]) for i in range(len(new_batch_data))])\n",
    "                                new_batch_data.append(batch[\"data\"].T[target_col_mask==col_i])\n",
    "                            new_batch_data = torch.cat(new_batch_data, dim=-1).reshape(-1)\n",
    "                            batch_data_all.append(new_batch_data)\n",
    "                            cls_indexes = torch.tensor([count, cls_indexes_value]).reshape(1, -1).to(device)\n",
    "                            cls_indexes_all.append(cls_indexes)\n",
    "                            count += 1\n",
    "                    batch_data_all = torch.nn.utils.rnn.pad_sequence(\n",
    "                        batch_data_all, padding_value=tokenizer.pad_token_id, batch_first=True)\n",
    "                    cls_indexes_all = torch.cat(cls_indexes_all, dim=0)\n",
    "                    logits_temp, embs_temp = model(batch_data_all, cls_indexes=cls_indexes_all, get_enc=True)\n",
    "                    # assert test_embs[batch_idx][i].equal(embs_temp.cpu())\n",
    "                    \n",
    "                    logits_temp = logits_temp.detach().cpu()\n",
    "                    scores_temp = F.softmax(verifier(embs_temp)/tau)[:,1].detach().cpu()\n",
    "                    max_index = scores_temp.argmax().item()\n",
    "                    max_score = scores_temp[max_index].item()\n",
    "                    logits = logits_temp[max_index]\n",
    "                    predict_temp = logits_temp.argmax().item()\n",
    "                    i += 1\n",
    "\n",
    "                    labels_test.append(batch[\"label\"].cpu())\n",
    "                    logits_test.append(logits.detach().cpu())\n",
    "                labels_test = torch.cat(labels_test, dim=0)\n",
    "                logits_test = torch.stack(logits_test, dim=0)\n",
    "                preds_test = torch.argmax(logits_test, dim=1)\n",
    "\n",
    "            num_cols = torch.tensor(num_cols)\n",
    "            print(time.time()-ts)\n",
    "            from sklearn.metrics import confusion_matrix, f1_score\n",
    "            ts_pred_list = logits_test.argmax(\n",
    "                                        1).cpu().detach().numpy().tolist()\n",
    "            ts_micro_f1 = f1_score(labels_test.reshape(-1).numpy().tolist(),\n",
    "                                ts_pred_list,\n",
    "                                average=\"micro\")\n",
    "            ts_macro_f1 = f1_score(labels_test.reshape(-1).numpy().tolist(),\n",
    "                                ts_pred_list,\n",
    "                                average=\"macro\")\n",
    "            ts_f1_full = f1_score(labels_test.reshape(-1).numpy().tolist(),\n",
    "                                ts_pred_list,\n",
    "                                average=None)\n",
    "            # full_f1_init\n",
    "            print(\"ts_micro_f1={:.4f}, ts_macro_f1={:.4f}\".format(ts_micro_f1, ts_macro_f1))\n",
    "\n",
    "\n",
    "            for col_num in num_cols.unique():\n",
    "                mask = num_cols == col_num\n",
    "                ts_pred_list = logits_test[mask].argmax(\n",
    "                                            1).cpu().detach().numpy().tolist()\n",
    "                ts_micro_f1 = f1_score(labels_test.reshape(-1)[mask].numpy().tolist(),\n",
    "                                    ts_pred_list,\n",
    "                                    average=\"micro\")\n",
    "                ts_macro_f1 = f1_score(labels_test.reshape(-1)[mask].numpy().tolist(),\n",
    "                                    ts_pred_list,\n",
    "                                    average=\"macro\")\n",
    "                # full_f1_init\n",
    "                print(f\"{col_num}\", \"ts_micro_f1={:.4f}, ts_macro_f1={:.4f}\".format(ts_micro_f1, ts_macro_f1))    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "veri_results = torch.load(f\"/data/zhihao/TU/Watchog/verification/{args.task}/random1_veri_data_rand_seed2_mmr0.5_dp3.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "veri_embs = veri_results[\"embs\"]\n",
    "veri_label = veri_results[\"label\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MahalanobisDistance(object):\n",
    "    def __init__(self, ftrain, labels_train):\n",
    "        self.xc = [ftrain[labels_train == i] for i in np.unique(labels_train)]\n",
    "        self.cov_inv = torch.linalg.pinv(torch.Tensor(np.cov(ftrain.T, bias=True))).numpy()\n",
    "    def __call__(self, ftest):\n",
    "        dtest = [\n",
    "            np.sum(\n",
    "                (ftest - np.mean(x, axis=0, keepdims=True))\n",
    "                * (\n",
    "                    self.cov_inv.dot(\n",
    "                        (ftest - np.mean(x, axis=0, keepdims=True)).T\n",
    "                    )\n",
    "                ).T,\n",
    "                axis=-1,\n",
    "            )\n",
    "            for x in self.xc\n",
    "        ]\n",
    "\n",
    "        dtest = np.min(dtest, axis=0)\n",
    "\n",
    "        return dtest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "mah_dist = MahalanobisDistance(veri_embs, veri_label)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'gt-semtab22-dbpedia-all0'"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "args.task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dataset = torch.load(f\"/data/zhihao/TU/Watchog/verification/{args.task}/random1_test_data_mmr0.5_dp3.pth\")\n",
    "test_embs = test_dataset['embs']\n",
    "test_logits = test_dataset['logits']\n",
    "test_class = test_dataset['class']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "                labels_test = []\n",
    "                logits_test = []\n",
    "\n",
    "                with torch.no_grad():\n",
    "                    for batch_idx, batch in enumerate(test_embs):\n",
    "                        embs = test_embs[batch_idx][0].reshape(-1).to(device)\n",
    "                        logits = test_logits[batch_idx][0].reshape(-1).to(device)\n",
    "                        # scores_init = verifier(embs)\n",
    "                        max_score = -float(\"inf\")\n",
    "                        embs_temp_all = []\n",
    "                        logits_temp_all = []\n",
    "                        count = 0\n",
    "                        for i in range(len(test_embs[batch_idx])):\n",
    "                            logits_temp = test_logits[batch_idx][i].reshape(1, -1).to(device)\n",
    "                            embs_temp = test_embs[batch_idx][i].reshape(1, -1).to(device)\n",
    "                            # if args.context == \"init\":\n",
    "                            #     embs_temp = torch.cat([embs, embs_temp], dim=-1)\n",
    "                            # elif args.context == \"target\":\n",
    "                            #     embs_temp = torch.cat([embs_target, embs_temp], dim=-1)\n",
    "                            embs_temp_all.append(embs_temp)\n",
    "                            logits_temp_all.append(logits_temp)\n",
    "                            count += 1\n",
    "                        embs_temp_all = torch.cat(embs_temp_all, dim=0)\n",
    "                        logits_temp_all = torch.cat(logits_temp_all, dim=0)\n",
    "                        scores_temp = F.softmax(verifier(embs_temp))[:,1].detach().cpu()\n",
    "                        max_index = scores_temp.argmax().item()\n",
    "                        max_score = scores_temp[max_index].item()\n",
    "                        logits = logits_temp[max_index]\n",
    "                            # if len(x) == 1 and 0 in x:\n",
    "                            #     predict_target = predict_temp\n",
    "                            #     msp_target = msp_temp\n",
    "                            # # print(x, msp_temp, predict_temp)\n",
    "                            # if 0 not in x and msp_temp > debias_threshold and (predict_temp != predict_target):\n",
    "                            #     debias_classes.append(predict_temp)\n",
    "                            #     continue\n",
    "\n",
    "                        labels_test.append(torch.tensor([test_class[batch_idx][0]]))\n",
    "                        logits_test.append(logits.detach().cpu())\n",
    "                    labels_test = torch.cat(labels_test, dim=0)\n",
    "                    logits_test = torch.stack(logits_test, dim=0)\n",
    "                    preds_test = torch.argmax(logits_test, dim=1)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "                from sklearn.metrics import confusion_matrix, f1_score\n",
    "                ts_pred_list = logits_test.argmax(\n",
    "                                            1).cpu().detach().numpy().tolist()\n",
    "                ts_micro_f1 = f1_score(labels_test.reshape(-1).numpy().tolist(),\n",
    "                                    ts_pred_list,\n",
    "                                    average=\"micro\")\n",
    "                ts_macro_f1 = f1_score(labels_test.reshape(-1).numpy().tolist(),\n",
    "                                    ts_pred_list,\n",
    "                                    average=\"macro\")\n",
    "                print(\"ts_micro_f1={:.4f}, ts_macro_f1={:.4f}\".format(ts_micro_f1, ts_macro_f1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ts_micro_f1=0.4885, ts_macro_f1=0.2163\n"
     ]
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "jupyter",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}