tapas-on-visqa-helper.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "55a60603-c4ed-4362-ad47-a1360d4c2eba",
   "metadata": {},
   "source": [
    "#### This prepares the `tapas_on_visqa_inputs.pkl` as the input to the TaPas model\n",
    "- Put it to the repo root before using, since this notebook requires some main tool components.\n",
    "- This requires the `visqa_dataset.pkl` ready before running."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d079203f-b132-4078-9154-3264d36166fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2521e5b6-5079-450b-ad4b-eb7d1013502e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# path to the original VisQA dataset\n",
    "data_path = \"./benchmarks/VisQA/shared/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26245ec1-2729-4be5-bc01-562eac8fc6db",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"{}/visqa_dataset.pkl\".format(data_path), \"rb\") as f:\n",
    "    dt = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d242a6af-175b-4b67-896e-8374a9f9e835",
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_table(arg_df):\n",
    "    return arg_df.to_markdown(index=False,tablefmt=\"jira\",numalign=\"left\").replace(\"||\",\"|\").replace(\"|\\n|\",\"\\n\").strip(\"|\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ffdf152-2df5-4e2d-bfa2-bc0761277210",
   "metadata": {},
   "outputs": [],
   "source": [
    "# construct inputs for every benchmark\n",
    "tapas_inputs = []\n",
    "for i in range(len(dt)):\n",
    "    p = dt[i]\n",
    "    str_table = format_table(p[\"rendered_table\"])\n",
    "    tapas_inputs.append((p[\"short_id\"], p[\"query\"], str_table))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd749412-c535-4394-a087-90797c73fb42",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"{}/tapas_on_visqa_inputs.pkl\".format(data_path), \"wb\") as f:\n",
    "    pickle.dump(tapas_inputs, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a507463d-55ec-4010-9d62-56e062eba816",
   "metadata": {},
   "outputs": [],
   "source": [
    "# alternative old inputs construction (for debugging only)\n",
    "# tapas_inputs = []\n",
    "# for i in range(len(dt)):\n",
    "#     p = dt[i]\n",
    "#     str_table = format_table(p[\"table\"])\n",
    "#     tapas_inputs.append((p[\"short_id\"], p[\"query\"], str_table))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18bbfbdf-7cab-4088-ae5f-27a965d56bb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# with open(\"{}/tapas_on_visqa_inputs_old.pkl\".format(data_path), \"wb\") as f:\n",
    "#     pickle.dump(tapas_inputs, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8fc8679-ce19-4716-a6e4-2927502a73ca",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "9d3a3e70-10b7-425b-aff9-9cc61e5d91b6",
   "metadata": {},
   "source": [
    "#### This processes `tapas_on_visqa_outputs.pkl` to generate `tapas_on_visqa_dataset.pkl`\n",
    "- Performs top-k strategy and generate candidate outputs\n",
    "- Based on `visqa_dataset.pkl`, and merge the tapas results with it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69623d31-763d-4fd2-82df-8b0cad42d798",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from trinity.utils.visqa import normalize_table, parse_value\n",
    "from trinity.utils.visqa_strategy import strategy_TaPas_A, strategy_TaPas_B, strategy_TaPas_C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca4d09ac-06a4-4668-9f6a-df97ade523a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# path to the original VisQA dataset\n",
    "data_path = \"./benchmarks/VisQA/shared/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9afae8f-8707-4500-bcf3-2145262545a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def interpret_answer(arg_line):\n",
    "    if arg_line.startswith(\"COUNT of \"):\n",
    "        tmp_operands = [parse_value(p) for p in arg_line[len(\"COUNT of \"):].split(\", \")]\n",
    "        return [len(tmp_operands)]\n",
    "    elif arg_line.startswith(\"SUM of \"):\n",
    "        tmp_operands = [parse_value(p) for p in arg_line[len(\"SUM of \"):].split(\", \")]\n",
    "        return [sum(tmp_operands)]\n",
    "    elif arg_line.startswith(\"AVERAGE of \"):\n",
    "        tmp_operands = [parse_value(p) for p in arg_line[len(\"AVERAGE of \"):].split(\", \")]\n",
    "        return [sum(tmp_operands)/len(tmp_operands)]\n",
    "    else:\n",
    "        # no ops\n",
    "        tmp_operands = [parse_value(p) for p in arg_line.split(\", \")]\n",
    "        if len(tmp_operands)==0:\n",
    "            return [\"<no answer>\"]\n",
    "        elif len(tmp_operands)==1:\n",
    "            if isinstance(tmp_operands[0], str) and tmp_operands[0].strip()==\"\":\n",
    "                return [\"<no answer>\"]\n",
    "            else:\n",
    "                return [tmp_operands[0]]\n",
    "        else:\n",
    "            # len>1\n",
    "            return sorted([p for p in tmp_operands], key=lambda x:str(x))\n",
    "        \n",
    "def extract_answers_from_logs(arg_logs):\n",
    "    tmp_answers = []\n",
    "    for i in range(len(arg_logs)):\n",
    "        if arg_logs[i].startswith(\"Evaluation finished\"):\n",
    "            if arg_logs[i+1].startswith(\">\"):\n",
    "                try:\n",
    "                    tmp_answers.append(interpret_answer(arg_logs[i+2]))\n",
    "                except TypeError:\n",
    "                    tmp_answers.append([\"<type error>\"])\n",
    "            else:\n",
    "                # TaPas exception/error\n",
    "                tmp_answers.append([\"<tapas exception>\"])\n",
    "    return tmp_answers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb18817d-5190-4d00-a255-b7903b240a84",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"{}/visqa_dataset.pkl\".format(data_path), \"rb\") as f:\n",
    "    dt = pickle.load(f)\n",
    "\n",
    "with open(\"{}/tapas_on_visqa_outputs.log\".format(data_path), \"r\") as f:\n",
    "    tapas_logs = f.readlines()\n",
    "tapas_logs = extract_answers_from_logs(tapas_logs)\n",
    "with open(\"{}/tapas_on_visqa_outputs.pkl\".format(data_path), \"rb\") as f:\n",
    "    tapas_outputs = pickle.load(f)\n",
    "tapas_outputs = [tapas_outputs[i] for i in range(len(tapas_outputs)) if i%2!=0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d6a4c33-88aa-47fd-8844-9bf256a337f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(tapas_outputs)==len(dt)\n",
    "assert len(tapas_logs)==len(dt)\n",
    "len(tapas_outputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d498b33e-96a6-44d1-924a-a3fd84892b14",
   "metadata": {},
   "outputs": [],
   "source": [
    "# first extract all the cell pointers with probs\n",
    "tapas_parsed_outputs = []\n",
    "for i in range(len(tapas_outputs)):\n",
    "    # print(\"# i={}\".format(i))\n",
    "    if len(tapas_outputs[i])>0:\n",
    "        p = tapas_outputs[i][0] # always at 0 since we pass 1 benchmark to TaPas at a time\n",
    "        dop = p[\"pred_aggr\"] # predicted operator\n",
    "        qlist = p[\"probabilities\"]>0 # find all cells with prob>0\n",
    "        cpps = []\n",
    "        for j in range(len(qlist)):\n",
    "            if qlist[j]:\n",
    "                drow = p[\"row_ids\"][j]-1\n",
    "                dcol = p[\"column_ids\"][j]-1\n",
    "                dprob= p[\"probabilities\"][j]\n",
    "                cpps.append((drow,dcol,dprob))\n",
    "        cpps = sorted(cpps, key=lambda x:x[2], reverse=True)\n",
    "        tapas_parsed_outputs.append((dop,cpps)) # (aggr, cpps)\n",
    "    else:\n",
    "        # no outputs, could be something wrong?\n",
    "        print(\"# warning: no output for i={}\".format(i))\n",
    "        tapas_parsed_outputs.append((0,[])) # (aggr, cpps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c25f4ad-e1c9-4c1b-986b-12223f36e7cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# then build \"expected_output\" table and \"candidate_outputs\" table\n",
    "for i in range(len(dt)):\n",
    "    print(\"\\r# processing {}/{}\".format(i, len(dt)), end=\"\")\n",
    "    p = dt[i]\n",
    "\n",
    "    if isinstance(p[\"repr_answer\"], list):\n",
    "        tmp_expected_output = normalize_table(pd.DataFrame.from_records(\n",
    "            np.asarray([p[\"repr_answer\"]]).T, columns=[\"ANSWER\"],\n",
    "        ))\n",
    "    elif isinstance(p[\"repr_answer\"], (int, float, str)):\n",
    "        tmp_expected_output = normalize_table(pd.DataFrame.from_records(\n",
    "            [[p[\"repr_answer\"]]], columns=[\"ANSWER\"],\n",
    "        ))\n",
    "    else:\n",
    "        raise NotImplementedError(\"Unsupported type of answer, got: {}.\".format(type(p[\"repr_answer\"])))\n",
    "\n",
    "    dt[i][\"expected_output\"] = tmp_expected_output\n",
    "\n",
    "    tmp_outputs_original = tapas_logs[i]\n",
    "    tmp_outputs_TaPas_A = strategy_TaPas_A(tapas_parsed_outputs[i], p[\"rendered_table\"])\n",
    "    tmp_outputs_TaPas_B = strategy_TaPas_B(tapas_parsed_outputs[i], p[\"rendered_table\"])\n",
    "    tmp_probs_TaPas_C, tmp_outputs_TaPas_C = strategy_TaPas_C(tapas_parsed_outputs[i], p[\"rendered_table\"])\n",
    "    dt[i][\"candidate_outputs\"] = {\n",
    "        \"TaPas_original\": tmp_outputs_original,\n",
    "        \"TaPas_A\": tmp_outputs_TaPas_A,\n",
    "        \"TaPas_B\": tmp_outputs_TaPas_B,\n",
    "        \"TaPas_C\": tmp_outputs_TaPas_C,\n",
    "        \"TaPas_probs_C\": tmp_probs_TaPas_C,\n",
    "    }\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14face60-5a9d-4166-bb10-0f7ad38324b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"{}/tapas_on_visqa_dataset.pkl\".format(data_path), \"wb\") as f:\n",
    "    pickle.dump(dt, f)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "base"
  },
  "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}