build-visqa-dataset.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "84876017-deff-421b-afe9-ef18a5cdaa31",
   "metadata": {},
   "source": [
    "#### This prepares the `visqa_dataset.pkl` for VisQA dataset by packing up necessary information\n",
    "- Put it to the repo root before using, since this notebook requires some main tool components."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "baa2946a-571c-4e3a-ac9b-407d9b48493b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# path to the original VisQA dataset\n",
    "data_path = \"./benchmarks/VisQA/shared/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a511dd3a-0cc3-402f-a9b2-086e09e17c4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import pickle\n",
    "import pandas as pd\n",
    "pd.set_option('display.max_rows', 10)\n",
    "from io import StringIO\n",
    "from trinity.utils.visqa import vl_to_svg, svg_to_table, normalize_table, grouped_svg_to_table, vl_to_grouped_svg, parse_value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92318cd5-6392-4b52-b5a0-4d8f986b2cf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"{}/qadata.json\".format(data_path), \"r\") as f:\n",
    "    dt = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a318450-ada7-4cae-98a5-9f339399b773",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_spec_name(arg_dataset, arg_chart_name):\n",
    "    if arg_dataset==\"kong\":\n",
    "        return arg_chart_name.split(\"_\")[1]\n",
    "    elif arg_dataset==\"d3\":\n",
    "        return arg_chart_name.split(\"_\")[1]\n",
    "        raise None\n",
    "    elif arg_dataset==\"vega-lite-example-gallery\":\n",
    "        return arg_chart_name\n",
    "    elif \"wikitables\" in arg_dataset:\n",
    "        return arg_chart_name\n",
    "    else:\n",
    "        raise NotImplementedError(\"Unsupported data series, got: {}, {}.\".format(arg_dataset, arg_chart_name))\n",
    "\n",
    "def get_repr_answer(arg_refined_answer):\n",
    "    if \"|\" in arg_refined_answer:\n",
    "        return sorted([parse_value(p) for p in arg_refined_answer.split(\"|\")], key=lambda x:str(x))\n",
    "    else:\n",
    "        # parse as a single value\n",
    "        return [parse_value(arg_refined_answer)]\n",
    "\n",
    "# start collecting\n",
    "cached_utils = {}\n",
    "benchmark_collections = []\n",
    "tmp_count = 0\n",
    "for dkey in dt.keys():\n",
    "    tmp_count += 1\n",
    "    \n",
    "#     if tmp_count < 195:\n",
    "#         continue\n",
    "    \n",
    "    print(\"\\r# processing {}/{}, {}\".format(tmp_count, len(dt), dkey), end=\"\")\n",
    "    \n",
    "    bitem = {}\n",
    "    bitem[\"id\"] = dkey\n",
    "    bitem[\"short_id\"] = dkey[:6]\n",
    "    bitem[\"query\"] = dt[dkey][\"question\"]\n",
    "    bitem[\"answer\"] = dt[dkey][\"answer\"] # original answer from the dataset\n",
    "    \n",
    "    # refined_answer: some original answers are wrong, which will be corrected as refined answers\n",
    "    #                 see block below for detailed refinements\n",
    "    bitem[\"refined_answer\"] = dt[dkey][\"answer\"]\n",
    "    # repr_answer is for quick comparison for accuracy; its data structure should support equality comparison\n",
    "    bitem[\"repr_answer\"] = get_repr_answer(bitem[\"refined_answer\"])\n",
    "    \n",
    "    bitem[\"data_series\"] = dt[dkey][\"dataset\"]\n",
    "    \n",
    "    bitem[\"spec_name\"] = get_spec_name(dt[dkey][\"dataset\"], dt[dkey][\"chartName\"])\n",
    "    # load full spec\n",
    "    with open(\"{}/dataset/{}/specs/{}.json\".format(data_path, bitem[\"data_series\"], bitem[\"spec_name\"]), \"r\") as f:\n",
    "        tmp_spec = json.load(f)\n",
    "    bitem[\"spec\"] = tmp_spec\n",
    "    \n",
    "    bitem[\"csv_name\"] = bitem[\"spec\"][\"data\"][\"url\"].split(\"/\")[-1]\n",
    "    if bitem[\"csv_name\"].endswith(\".csv\"):\n",
    "        bitem[\"csv_name\"] = bitem[\"csv_name\"].replace(\".csv\", \"\")\n",
    "        # load csv and remove trailing delimiters\n",
    "        with open(\"{}/dataset/{}/data/{}.csv\".format(data_path, bitem[\"data_series\"], bitem[\"csv_name\"]), \"r\") as f:\n",
    "            tmp_csv = f.readlines()\n",
    "        for i in range(len(tmp_csv)):\n",
    "            tmp_csv[i] = tmp_csv[i].strip()\n",
    "            while tmp_csv[i].endswith(\",\"):\n",
    "                tmp_csv[i] = tmp_csv[i][:-1]\n",
    "        bitem[\"csv\"] = \"\\n\".join(tmp_csv)\n",
    "        bitem[\"table\"] = normalize_table(pd.read_csv(StringIO(bitem[\"csv\"])))\n",
    "    elif bitem[\"csv_name\"].endswith(\".json\"):\n",
    "        # some of the data comes in json format\n",
    "        # convert to csv\n",
    "        bitem[\"csv_name\"] = bitem[\"csv_name\"].replace(\".json\", \"\")\n",
    "        with open(\"{}/dataset/{}/data/{}.json\".format(data_path, bitem[\"data_series\"], bitem[\"csv_name\"]), \"r\") as f:\n",
    "            tmp_json = json.load(f)\n",
    "        bitem[\"table\"] = normalize_table(pd.DataFrame.from_records(tmp_json))\n",
    "        bitem[\"csv\"] = bitem[\"table\"].to_csv(index=False)\n",
    "    else:\n",
    "        raise NotImplementedError(\"Unsupported data file type, got: {}.\".format(bitem[\"csv_name\"]))\n",
    "    \n",
    "    # update spec\n",
    "    del bitem[\"spec\"][\"data\"][\"url\"]\n",
    "    bitem[\"spec\"][\"data\"][\"values\"] = bitem[\"table\"].to_dict(\"records\")\n",
    "    \n",
    "    # render table\n",
    "    # check cache first\n",
    "    tmp_key = (bitem[\"data_series\"], bitem[\"spec_name\"])\n",
    "    if tmp_key in cached_utils.keys():\n",
    "        bitem[\"grouped_svg\"] = cached_utils[tmp_key][\"grouped_svg\"]\n",
    "        bitem[\"rendered_table\"] = cached_utils[tmp_key][\"rendered_table\"]\n",
    "    else:\n",
    "        bitem[\"grouped_svg\"] = vl_to_grouped_svg(bitem[\"spec\"])\n",
    "        bitem[\"rendered_table\"] = grouped_svg_to_table(bitem[\"grouped_svg\"], bitem[\"spec\"])\n",
    "        cached_utils[tmp_key] = {}\n",
    "        cached_utils[tmp_key][\"grouped_svg\"] = bitem[\"grouped_svg\"]\n",
    "        cached_utils[tmp_key][\"rendered_table\"] = bitem[\"rendered_table\"]\n",
    "        # for debugging only\n",
    "        cached_utils[tmp_key][\"table\"] = bitem[\"table\"]\n",
    "        # merge table to prevent data loss due to spec filter, e.g., kong:15\n",
    "        # bitem[\"merged_table\"] = pd.merge(bitem[\"table\"], bitem[\"rendered_table\"])\n",
    "    \n",
    "    # re-order the column order\n",
    "    tmp_cs = [p for p in bitem[\"table\"].columns if p in bitem[\"rendered_table\"].columns]\n",
    "    for q in bitem[\"rendered_table\"].columns:\n",
    "        if q not in tmp_cs:\n",
    "            tmp_cs.append(q)\n",
    "    bitem[\"rendered_table\"] = bitem[\"rendered_table\"][tmp_cs]\n",
    "    \n",
    "    # put to collection\n",
    "    benchmark_collections.append(bitem)\n",
    "#     break\n",
    "    \n",
    "print(\"\\n# done\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bca5ad54-e215-4685-be44-e648adf9f60f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "134363b6-2b26-4652-a736-78bcbce01665",
   "metadata": {},
   "outputs": [],
   "source": [
    "# additional patches to apply to data, see README or Google Sheets for details\n",
    "benchmark_collections[74][\"refined_answer\"] = \"Hindus\"\n",
    "benchmark_collections[74][\"repr_answer\"] = get_repr_answer(\"Hindus\")\n",
    "\n",
    "benchmark_collections[83][\"refined_answer\"] = \"Hindus\"\n",
    "benchmark_collections[83][\"repr_answer\"] = get_repr_answer(\"Hindus\")\n",
    "\n",
    "# yes, refined_answers are all strings\n",
    "benchmark_collections[244][\"refined_answer\"] = \"40221893\"\n",
    "benchmark_collections[244][\"repr_answer\"] = get_repr_answer(\"40221893\")\n",
    "\n",
    "benchmark_collections[288][\"refined_answer\"] = \"Aug 1 2004\"\n",
    "benchmark_collections[288][\"repr_answer\"] = get_repr_answer(\"Aug 1 2004\")\n",
    "\n",
    "benchmark_collections[348][\"refined_answer\"] = \"Aug\"\n",
    "benchmark_collections[348][\"repr_answer\"] = get_repr_answer(\"Aug\")\n",
    "\n",
    "benchmark_collections[353][\"refined_answer\"] = \"Jul\"\n",
    "benchmark_collections[353][\"repr_answer\"] = get_repr_answer(\"Jul\")\n",
    "\n",
    "benchmark_collections[354][\"refined_answer\"] = \"Feb\"\n",
    "benchmark_collections[354][\"repr_answer\"] = get_repr_answer(\"Feb\")\n",
    "\n",
    "benchmark_collections[355][\"refined_answer\"] = \"Aug\"\n",
    "benchmark_collections[355][\"repr_answer\"] = get_repr_answer(\"Aug\")\n",
    "\n",
    "# yes, refined_answers are all strings\n",
    "benchmark_collections[438][\"refined_answer\"] = \"2915000\"\n",
    "benchmark_collections[438][\"repr_answer\"] = get_repr_answer(\"2915000\")\n",
    "\n",
    "benchmark_collections[520][\"refined_answer\"] = \"gojō|yoshino\"\n",
    "benchmark_collections[520][\"repr_answer\"] = get_repr_answer(\"gojō|yoshino\")\n",
    "\n",
    "benchmark_collections[537][\"refined_answer\"] = \"north-east skåne\"\n",
    "benchmark_collections[537][\"repr_answer\"] = get_repr_answer(\"north-east skåne\")\n",
    "\n",
    "benchmark_collections[541][\"refined_answer\"] = \"north-east skåne\"\n",
    "benchmark_collections[541][\"repr_answer\"] = get_repr_answer(\"north-east skåne\")\n",
    "\n",
    "benchmark_collections[616][\"refined_answer\"] = \"x-Houston Rockets|x-San Antonio Spurs\"\n",
    "benchmark_collections[616][\"repr_answer\"] = get_repr_answer(\"x-Houston Rockets|x-San Antonio Spurs\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4ba6165-2bba-41dc-8d5d-b1a15e91bff1",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(benchmark_collections)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18fafd86-9fee-40f0-9c85-e0d01dbac798",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"{}/visqa_dataset.pkl\".format(data_path), \"wb\") as f:\n",
    "    pickle.dump(benchmark_collections, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a5aaf54-4c51-4918-b37e-e180a80a4bf7",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd487b98-609b-4b1a-8944-9c9d77594ee7",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "d07d096f-069b-48f4-9017-cd3954d48105",
   "metadata": {},
   "source": [
    "### inspect the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcd324c6-5980-438b-8a93-a88ce9f556a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "pd.set_option('display.max_rows', 10)\n",
    "\n",
    "from trinity.utils.visqa import normalize_table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5c6839f-aa41-40c7-bcf5-64eb2f69413a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# path to the original VisQA dataset\n",
    "data_path = \"./benchmarks/VisQA/shared/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b2ce7f8-748a-4bdb-94ac-137af6cb0779",
   "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": "3be2fd31-9f1c-456b-b5b6-7b682d28d1ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "for dkey in cached_utils.keys():\n",
    "    print(\"==============================\")\n",
    "    print(dkey)\n",
    "    print(cached_utils[dkey][\"table\"])\n",
    "    print(cached_utils[dkey][\"rendered_table\"].dtypes)\n",
    "    print(cached_utils[dkey][\"rendered_table\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6f97861-5ae9-4c1e-86ce-150e32573dbd",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d74a961f-ee32-4302-b4c6-9ff52f780e4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "cached_utils[('kong', '18')][\"table\"].columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80d89536-d757-445a-adaf-e8c590f86136",
   "metadata": {},
   "outputs": [],
   "source": [
    "cached_utils[('kong', '18')][\"rendered_table\"].columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78f69b93-7237-4aa4-bd95-bf8d159cf381",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "716f8d92-25ce-4ba4-b46a-e46a7940b7b1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for i in range(len(dt)):\n",
    "    # display(dt[i][\"rendered_table\"])\n",
    "    print(\"i={}\".format(i))\n",
    "    print(\"{}\".format(dt[i][\"rendered_table\"].dtypes))\n",
    "    print(\"{}\".format(dt[i][\"rendered_table\"]))\n",
    "    # input(\"i={}\".format(i))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e73dd614-2e11-4c32-8ef1-a461a79aaa69",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}