public_PLDI22AE_Poe_TaPas_on_VisQA.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "public-PLDI22AE-Poe-TaPas-on-VisQA",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/chyanju/Poe/blob/main/public_PLDI22AE_Poe_TaPas_on_VisQA.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AZGoIr5tvdIx"
      },
      "source": [
        "# Poe: Collecting TaPas results on VisQA dataset.\n",
        "This colab notebook is modified by ***Poe*** from the original TaPas tool. Follow the instructions to get the prediction results."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-07bRHwv0C7L"
      },
      "source": [
        "##### Copyright 2020 The Google AI Language Team Authors\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SSpOxRRH0BCU"
      },
      "source": [
        "# Copyright 2019 The Google AI Language Team Authors.\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     http://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y-m_JoVCFCV0"
      },
      "source": [
        "# 1. Clone and install the repository\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lF84Z-KayR3Z"
      },
      "source": [
        "First, let's install the code."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uI6zyIM20Kw4"
      },
      "source": [
        "! pip install tapas-table-parsing==0.0.1.dev0"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1rHkOZhQuVii"
      },
      "source": [
        "**Poe**: You need to restart the runtime and resume from here."
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Poe**: Replace the `run_task_main.py` with a customized one for Poe. This block only needs to be run **once** during the runtime."
      ],
      "metadata": {
        "id": "icCzFsZ9sgIr"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "! git clone https://github.com/chyanju/Poe.git\n",
        "! cp /content/Poe/benchmarks/VisQA/shared/run_task_main.py /usr/local/lib/python3.7/dist-packages/tapas/\n",
        "! cp /content/Poe/benchmarks/VisQA/shared/tapas_on_visqa_inputs.pkl /content/"
      ],
      "metadata": {
        "id": "tYEnF1i-v4jE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JfGBVnJroVgF"
      },
      "source": [
        "import tapas\n",
        "print(tapas.__file__)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7We9ofHuFMuk"
      },
      "source": [
        "# 2. Fetch models fom Google Storage"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Poe**: This block only needs to be run once during the runtime."
      ],
      "metadata": {
        "id": "K6qJpqAFyEH1"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "B10C0Yz6gQyD"
      },
      "source": [
        "! gsutil cp \"gs://tapas_models/2020_08_05/tapas_wtq_wikisql_sqa_masklm_medium_reset.zip\" \"tapas_model.zip\" && unzip tapas_model.zip\n",
        "! mv tapas_wtq_wikisql_sqa_masklm_medium_reset tapas_model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E3107bGlGm7d"
      },
      "source": [
        "# 3. Imports"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pnUjDlLqDd3m"
      },
      "source": [
        "import tensorflow.compat.v1 as tf\n",
        "import os \n",
        "import shutil\n",
        "import csv\n",
        "import pandas as pd\n",
        "import IPython\n",
        "\n",
        "tf.get_logger().setLevel('ERROR')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aml6oLFl1dSt"
      },
      "source": [
        "from tapas.utils import tf_example_utils\n",
        "from tapas.protos import interaction_pb2\n",
        "from tapas.utils import number_annotation_utils\n",
        "from tapas.scripts import prediction_utils"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AbMUYT1bKMp9"
      },
      "source": [
        "# 4. Load checkpoint for prediction"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IO0d_wFMy82O"
      },
      "source": [
        "Here's the prediction code, which will create and `interaction_pb2.Interaction` protobuf object, which is the datastructure we use to store examples, and then call the prediction script."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UKfxspnVFPsc"
      },
      "source": [
        "os.makedirs('results/wtq/tf_examples', exist_ok=True)\n",
        "os.makedirs('results/wtq/model', exist_ok=True)\n",
        "with open('results/wtq/model/checkpoint', 'w') as f:\n",
        "  f.write('model_checkpoint_path: \"model.ckpt-0\"')\n",
        "for suffix in ['.data-00000-of-00001', '.index', '.meta']:\n",
        "  shutil.copyfile(f'tapas_model/model.ckpt{suffix}', f'results/wtq/model/model.ckpt-0{suffix}')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9RlvgDAmCNtP"
      },
      "source": [
        "max_seq_length = 512\n",
        "vocab_file = \"tapas_model/vocab.txt\"\n",
        "config = tf_example_utils.ClassifierConversionConfig(\n",
        "    vocab_file=vocab_file,\n",
        "    max_seq_length=max_seq_length,\n",
        "    max_column_id=max_seq_length,\n",
        "    max_row_id=max_seq_length,\n",
        "    strip_column_names=False,\n",
        "    add_aggregation_candidates=False,\n",
        ")\n",
        "converter = tf_example_utils.ToClassifierTensorflowExample(config)\n",
        "\n",
        "def convert_interactions_to_examples(tables_and_queries):\n",
        "  \"\"\"Calls Tapas converter to convert interaction to example.\"\"\"\n",
        "  for idx, (table, queries) in enumerate(tables_and_queries):\n",
        "    interaction = interaction_pb2.Interaction()\n",
        "    for position, query in enumerate(queries):\n",
        "      question = interaction.questions.add()\n",
        "      question.original_text = query\n",
        "      question.id = f\"{idx}-0_{position}\"\n",
        "    for header in table[0]:\n",
        "      interaction.table.columns.add().text = header\n",
        "    for line in table[1:]:\n",
        "      row = interaction.table.rows.add()\n",
        "      for cell in line:\n",
        "        row.cells.add().text = cell\n",
        "    number_annotation_utils.add_numeric_values(interaction)\n",
        "    for i in range(len(interaction.questions)):\n",
        "      try:\n",
        "        yield converter.convert(interaction, i)\n",
        "      except ValueError as e:\n",
        "        print(f\"Can't convert interaction: {interaction.id} error: {e}\")\n",
        "        \n",
        "def write_tf_example(filename, examples):\n",
        "  with tf.io.TFRecordWriter(filename) as writer:\n",
        "    for example in examples:\n",
        "      writer.write(example.SerializeToString())\n",
        "\n",
        "def aggregation_to_string(index):\n",
        "  if index == 0:\n",
        "    return \"NONE\"\n",
        "  if index == 1:\n",
        "    return \"SUM\"\n",
        "  if index == 2:\n",
        "    return \"AVERAGE\"\n",
        "  if index == 3:\n",
        "    return \"COUNT\"\n",
        "  raise ValueError(f\"Unknown index: {index}\")\n",
        "\n",
        "def predict(table_data, queries):\n",
        "  table = [list(map(lambda s: s.strip(), row.split(\"|\"))) \n",
        "           for row in table_data.split(\"\\n\") if row.strip()]\n",
        "  examples = convert_interactions_to_examples([(table, queries)])\n",
        "  write_tf_example(\"results/wtq/tf_examples/test.tfrecord\", examples)\n",
        "  write_tf_example(\"results/wtq/tf_examples/random-split-1-dev.tfrecord\", [])\n",
        "  \n",
        "  ! python -m tapas.run_task_main \\\n",
        "    --task=\"WTQ\" \\\n",
        "    --output_dir=\"results\" \\\n",
        "    --noloop_predict \\\n",
        "    --test_batch_size={len(queries)} \\\n",
        "    --tapas_verbosity=\"ERROR\" \\\n",
        "    --compression_type= \\\n",
        "    --reset_position_index_per_cell \\\n",
        "    --init_checkpoint=\"tapas_model/model.ckpt\" \\\n",
        "    --bert_config_file=\"tapas_model/bert_config.json\" \\\n",
        "    --mode=\"predict\" 2> error\n",
        "\n",
        "\n",
        "  results_path = \"results/wtq/model/test.tsv\"\n",
        "  all_coordinates = []\n",
        "  df = pd.DataFrame(table[1:], columns=table[0])\n",
        "  # display(IPython.display.HTML(df.to_html(index=False)))\n",
        "  # print()\n",
        "  with open(results_path) as csvfile:\n",
        "    reader = csv.DictReader(csvfile, delimiter='\\t')\n",
        "    for row in reader:\n",
        "      coordinates = sorted(prediction_utils.parse_coordinates(row[\"answer_coordinates\"]))\n",
        "      all_coordinates.append(coordinates)\n",
        "      answers = ', '.join([table[row + 1][col] for row, col in coordinates])\n",
        "      position = int(row['position'])\n",
        "      aggregation = aggregation_to_string(int(row[\"pred_aggr\"]))\n",
        "      print(\">\", queries[position])\n",
        "      answer_text = str(answers)\n",
        "      if aggregation != \"NONE\":\n",
        "        answer_text = f\"{aggregation} of {answer_text}\"\n",
        "      print(answer_text)\n",
        "  return all_coordinates"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gqu-I-M9QaoA"
      },
      "source": [
        "# 5. Predict"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MSCSi6D_4ODm"
      },
      "source": [
        "**Poe**: Make an initial (empty) output file first."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KGLQBNvq4VEm"
      },
      "source": [
        "import pickle\n",
        "with open(\"/content/tapas_on_visqa_outputs.pkl\", \"wb\") as f:\n",
        "  pickle.dump([],f)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yI6nJ4nc4eMd"
      },
      "source": [
        "**Poe**: Then we read the inputs, predict and collect all the raw results."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "o4VujmAw4nde"
      },
      "source": [
        "import pickle\n",
        "with open(\"/content/tapas_on_visqa_inputs.pkl\", \"rb\") as f:\n",
        "  dt = pickle.load(f)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Poe**: Start collection. This may run for several hours."
      ],
      "metadata": {
        "id": "EuKMRAijyVKo"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "coJrvTk75CnG"
      },
      "source": [
        "for i in range(len(dt)):\n",
        "  print(\"# processing benchmark id={}/{}, query={}\".format(i, dt[i][0], dt[i][1]))\n",
        "  predict(dt[i][2], [dt[i][1]])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SIE7bTJMVuSh"
      },
      "source": [
        "print(\"# done\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Fk4n0qCpP1sB"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}