From 240fb61a8e78094836caecb315fde294f93efeb1 Mon Sep 17 00:00:00 2001 From: bharatji30 <43114202+bharatji30@users.noreply.github.com> Date: Thu, 21 Jul 2022 22:07:11 +0530 Subject: [PATCH] Created using Colaboratory --- notebooks/retrieval_predictions.ipynb | 513 +++++++++++++++++++++++++- 1 file changed, 512 insertions(+), 1 deletion(-) diff --git a/notebooks/retrieval_predictions.ipynb b/notebooks/retrieval_predictions.ipynb index 577cc7c..6e151b4 100644 --- a/notebooks/retrieval_predictions.ipynb +++ b/notebooks/retrieval_predictions.ipynb @@ -1 +1,512 @@ -{"nbformat":4,"nbformat_minor":0,"metadata":{"colab":{"name":"retrieval_predictions.ipynb","provenance":[],"collapsed_sections":[]},"kernelspec":{"display_name":"Python 3","name":"python3"},"language_info":{"name":"python"}},"cells":[{"cell_type":"markdown","metadata":{"id":"Bf-86KknHHps"},"source":["\"Open"]},{"cell_type":"markdown","metadata":{"id":"M8XF59AeJask"},"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":"rnPwClY0taWp"},"source":["# Copyright 2021 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":"3jiJ5VvrJkih"},"source":["# Run TAPAS retrieval models description\n","This notebook shows how to use retrieval models, which was introduced in the paper: [Open Domain Question Answering over Tables via Dense Retrieval](https://arxiv.org/pdf/2103.12011.pdf).\n","1. Load pre-trained and fine-tuned models.\n"," > * the dual encoder. (called tapas_retriever)\n"," > * the reader models. (called tapas_reader)\n","2. Add handcrafted query and extract the interactions and tf-examples.\n","3. Get nearest neighbors for each query, and extract the interactions to pass to the reader.\n","4. Call the reader on the new interactions and print \n"," > * the query.\n"," > * the probability of this table containing the answer.\n"," > * the table with a highlighted answer found by the reader.\n","\n","\n"]},{"cell_type":"code","metadata":{"id":"G9Z8tP11vCOl"},"source":["# Clone and install the repository\n","! sudo apt-get install protobuf-compiler\n","! git clone https://github.com/google-research/tapas\n","! pip install -e ./tapas\n","\n","\n","# Run the imports needed for all the colab\n","import tensorflow.compat.v1 as tf\n","import os \n","import shutil\n","import csv\n","import pandas as pd\n","import IPython\n","import ast\n","\n","tf.get_logger().setLevel('ERROR')\n","\n","from tapas.utils import tf_example_utils\n","from tapas.utils import beam_runner\n","from tapas.utils import create_data\n","from tapas.protos import interaction_pb2\n","from tapas.utils import number_annotation_utils\n","from tapas.scripts import prediction_utils\n","from tapas.scripts import eval_table_retriever_utils\n","from tapas.retrieval import tf_example_utils as retrieval_utils\n","from tapas.experiments import table_retriever_experiment\n"],"execution_count":null,"outputs":[]},{"cell_type":"code","metadata":{"id":"6fiPv3GV7yXZ"},"source":["# 1. Load pre-trained and fine-tuned models fom Google Storage.\n","# The dual encoder model\n","\n","# The dual encoder model\n","! gsutil cp \"gs://tapas_models/2021_04_27/tapas_nq_hn_retriever_medium.zip\" \"tapas_retriever.zip\" && unzip tapas_retriever.zip\n","! mv tapas_nq_hn_retriever_medium/ tapas_retriever\n","\n","\n","# The reader model\n","! gsutil cp \"gs://tapas_models/2021_04_27/tapas_nq_hn_reader_large.zip\" \"tapas_reader.zip\" && unzip tapas_reader.zip\n","! mv tapas_nq_hn_reader_large tapas_reader\n","\n","# Load the released nq_tables data.\n","os.makedirs('tapas_models_nq_tables', exist_ok=True)\n","! gsutil -m cp -R gs://tapas_models/2021_07_22/nq_tables/* tapas_models_nq_tables/\n"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"9vGb8QmE6BsH"},"source":["# Code"]},{"cell_type":"code","metadata":{"id":"9mN0WRCQv7hr"},"source":["# 2. Add handcrafted tables, and queries\n","# 2.1. Create the needed directories.\n","def create_directories():\n"," \"\"\"Create directories.\"\"\"\n"," # To be used for the dual encoder.\n"," os.makedirs('results/nq_retrieval/model', exist_ok=True)\n"," with open('results/nq_retrieval/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_retriever/model.ckpt{suffix}', f'results/nq_retrieval/model/model.ckpt-0{suffix}')\n"," shutil.copyfile(f'tapas_retriever/tables.tsv', f'results/nq_retrieval/model/tables.tsv')\n"," shutil.copyfile(f'tapas_retriever/bert_config.json', f'results/nq_retrieval/model/bert_config.json')\n"," # To be used for nq_reder.\n"," os.makedirs('results/nq_retrieval/tf_examples', exist_ok=True)\n"," os.makedirs('results/nq_retrieval/queries', exist_ok=True)\n"," # os.makedirs('results/nq_retrieval/tables', exist_ok=True)\n","\n"," os.makedirs('results/nq_reader/model', exist_ok=True)\n"," os.makedirs('results/nq_reader/queries', exist_ok=True)\n"," os.makedirs('results/nq_reader/nq_retrieval/tf_examples', exist_ok=True)\n"," os.makedirs('results/nq_reader/nq_retrieval/model', exist_ok=True)\n","\n"," with open('results/nq_reader/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_reader/model.ckpt{suffix}', f'results/nq_reader/model/model.ckpt-0{suffix}')\n","\n","# 2.2. Code to extract the data: the interactions than the tf_examples.\n","# interaction_pb2.Interaction` protobuf object is the data structure we use to\n","# store examples, and then to call the prediction script.\n","def get_table(document_title, table_data):\n"," \"\"\"Extracts the interaction for an str table.\n"," \n"," Args:\n"," table_data: str table where the columns are separated by '|'.\n"," document_title: str title of the page containing the table or a table title\n"," it also could be empty str.\"\"\"\n"," table = [list(map(lambda s: s.strip(), row.split(\"|\"))) \n"," for row in table_data.split(\"\\n\") if row.strip()]\n"," table_interaction = interaction_pb2.Table()\n"," table_interaction.document_title = document_title\n"," table_interaction.table_id = document_title\n"," if not table:\n"," return table_interaction\n"," for header in table[0]:\n"," table_interaction.columns.add().text = header\n"," for line in table[1:]:\n"," row = table_interaction.rows.add()\n"," for cell in line:\n"," row.cells.add().text = cell\n"," return table_interaction\n","\n","def extract_queries(queries):\n"," \"\"\"Extracts the interaction for a list of queries.\n"," \n"," This is used to create the interaction queries file.\n"," Args:\n"," queries: list of str queries.\"\"\"\n"," for idx, query in enumerate(queries):\n"," interaction = interaction_pb2.Interaction()\n"," interaction.id = f\"queries_{idx}\"\n"," question = interaction.questions.add()\n"," question.original_text = query\n"," question.id = f\"{interaction.id}-0_0\"\n"," interaction.table.CopyFrom(get_table(\"FAKE\", \" | \\n | \\n\"))\n"," number_annotation_utils.add_numeric_values(interaction) \n"," yield interaction\n","\n","def write_tfrecord(filename, examples):\n"," \"\"\"From interactions examples to tfrecord.\"\"\"\n"," with tf.io.TFRecordWriter(filename) as writer:\n"," for example in examples:\n"," writer.write(example.SerializeToString())\n","\n","def get_config():\n"," max_seq_length = 512\n"," vocab_file = \"tapas_retriever/vocab.txt\"\n"," config=tf_example_utils.RetrievalConversionConfig(\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"," cell_trim_length=-1,\n"," use_document_title=True,\n"," )\n"," return config\n","\n","def extract_queries_data(queries):\n"," \"\"\"Extracts the interactions then the tf_examples.\n"," \n"," Args:\n"," queries: list of str queries.\n"," \"\"\"\n"," examples = extract_queries(queries)\n"," input_queries = \"results/nq_retrieval/queries/queries.tfrecord\"\n"," write_tfrecord(input_queries, examples)\n"," config = get_config()\n"," beam_runner.run_type(\n"," create_data.build_retrieval_pipeline(\n"," input_files=[input_queries],\n"," output_files=[os.path.join(\"results/nq_retrieval/tf_examples\",\n"," \"queries.tfrecord\")],\n"," input_format=create_data.InputFormat.INTERACTION,\n"," config=config,\n"," ), beam_runner.RunnerType.DIRECT).wait_until_finish()\n","\n","# 3. Retrieval: Extract the queries interactions to pass to the reader\n","# 3.1. Extracts the queries embeddings.\n","def get_queries_embeddings():\n"," ! python -m tapas.experiments.table_retriever_experiment \\\n"," --do_predict \\\n"," --eval_name=\"dual_encoder_queries\" \\\n"," --minutes_to_sleep_before_predictions=0 \\\n"," --num_eval_steps=0 \\\n"," --model_dir=\"results/nq_retrieval/model\" \\\n"," --prediction_output_dir=\"results/nq_retrieval/model/queries\" \\\n"," --evaluated_checkpoint_step=0 \\\n"," --input_file_predict=\"results/nq_retrieval/tf_examples/queries.tfrecord\" \\\n"," --bert_config_file=\"tapas_retriever/bert_config.json\" \\\n"," --init_from_single_encoder=false \\\n"," --tf_random_seed=\"1\" \\\n"," --compression_type= \\\n"," --down_projection_dim=256 \\\n"," --eval_batch_size=1 \\\n"," --max_seq_length=512 2> error\n","\n"," with open(\"results/nq_retrieval/model/queries/predict_results_0.tsv\") as csvfile_reader:\n"," reader = csv.DictReader(csvfile_reader, delimiter='\\t')\n"," for i, row in enumerate(reader):\n"," print(\"Adding query_id: \", row[\"query_id\"], \":\", queries[i])\n","\n","# 3.2. Get nearest tables neighbors for each query.\n","def get_nearest_neighbors(num_neighbors):\n"," queries_pred = eval_table_retriever_utils.read_queries(\"results/nq_retrieval/model/queries/predict_results_0.tsv\")\n"," tables = eval_table_retriever_utils.read_tables(\"results/nq_retrieval/model/tables.tsv\", make_tables_unique=False)\n"," index = eval_table_retriever_utils.build_table_index(tables)\n"," similarities, neighbors = eval_table_retriever_utils._retrieve(queries_pred, index)\n"," selected_tables = {}\n"," for i, s in enumerate(similarities):\n"," print(\"Query index\", i, \":\", queries[i])\n"," selected_tables[i]={}\n"," for pos in range(num_neighbors):\n"," table_id = tables[neighbors[i][pos]].table_id\n"," selected_tables[i][table_id] = (s[pos], neighbors[i][pos])\n"," print(\" Related table id:\", table_id)\n"," print(\" Table's score \", s[pos])\n"," print(\" ----------------------------------------------\")\n"," return selected_tables\n","# 3.3. Extract the interactions to pass to the reader.\n","def iterate_tables(input_file):\n"," \"\"\"Reads interaction_pb2.Table().\"\"\"\n"," for value in tf.python_io.tf_record_iterator(input_file):\n"," table = interaction_pb2.Table()\n"," table.ParseFromString(value)\n"," yield table\n","\n","def create_queries_tables_interactions(selected_tables):\n"," \"\"\"Creates the interaction by linking the query to the selected table.\"\"\"\n"," queries_interactions = prediction_utils.iterate_interactions(\n"," \"results/nq_retrieval/queries/queries.tfrecord\")\n"," all_tables = iterate_tables(\"tapas_models_nq_tables/tables/tables.tfrecord\")\n"," tables = {table.table_id: table for table in all_tables}\n"," for i, q in enumerate(queries_interactions):\n"," print(\"\\n Query index:\", i, \":\", queries[i])\n"," t = selected_tables[i]\n"," for table_id in t.keys():\n"," if table_id in tables.keys():\n"," table = tables[table_id]\n"," print(\" > Converted query:\", q.questions[0].original_text)\n"," new_interaction = interaction_pb2.Interaction()\n"," new_interaction.CopyFrom(q)\n"," new_interaction.id = f\"{new_interaction.id}_{table.table_id}\"\n"," new_interaction.questions[0].id = f\"{new_interaction.id}_0\"\n"," new_interaction.table.CopyFrom(table)\n"," yield new_interaction\n"," else:\n"," print(\" > Not found in table file.\")\n"," print(\" Related table id: \", table_id)\n"," print(\" Table's score: \", t[table_id][0])\n"," print(\" ----------------------------------------------\")\n","\n","def create_interactions_for_reader(selected_tables):\n"," examples = create_queries_tables_interactions(selected_tables)\n"," write_tfrecord(\"results/nq_reader/queries/reader_queries.tfrecord\", examples)\n","\n","# 4. Reader: Get the answer given the question and the table\n","def get_converter(max_seq_length):\n"," \"\"\"Get a clssifier conferter.\"\"\"\n"," config = tf_example_utils.ClassifierConversionConfig(\n"," vocab_file=\"tapas_reader/vocab.txt\",\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"," return tf_example_utils.ToClassifierTensorflowExample(config)\n","\n","def convert_interactions_to_examples(converter):\n"," \"\"\"Calls Tapas converter to convert interaction to example.\"\"\"\n"," interactions = prediction_utils.iterate_interactions(\n"," \"results/nq_reader/queries/reader_queries.tfrecord\")\n"," for interaction in interactions:\n"," try:\n"," yield converter.convert(interaction, 0)\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","class Colors:\n"," \"\"\"Used to highlight the answers.\"\"\"\n"," ANSWER = '\\033[94m'\n"," BASE = '\\033[95m'\n"," BLACK = '\\033[0m'\n","\n"," \n","def set_answer_color(input, begin, end):\n"," \"\"\"Highlights the answers.\"\"\"\n"," list_output = [i.original_text for i in input]\n"," list_output[begin] = Colors.ANSWER + list_output[begin]\n"," list_output[end - 1] = list_output[end - 1] + Colors.BASE\n"," return \" \".join(list_output)\n","\n","\n","def get_table_df(table):\n"," \"\"\"Extracts a dataframe table for a better visualisation.\"\"\"\n"," printabe_table = [[Colors.BASE + c.text + Colors.BASE for c in table.columns]] \n"," for r in table.rows:\n"," printabe_table.append([Colors.BASE + c.text + Colors.BASE for c in r.cells])\n"," return pd.DataFrame(printabe_table)\n","\n","def predict():\n"," \"\"\"Predict the answer given the query and the table.\"\"\"\n"," max_seq_length = 512\n"," # Extracts the tf examples given the interactions.\n"," converter = get_converter(max_seq_length)\n"," examples = convert_interactions_to_examples(converter)\n"," write_tf_example(\"results/nq_reader/nq_retrieval/tf_examples/test.tfrecord\", examples)\n"," write_tf_example(\"results/nq_reader/nq_retrieval/tf_examples/dev.tfrecord\", [])\n"," # Run prediction\n"," ! python -m tapas.run_task_main \\\n"," --task=\"NQ_RETRIEVAL\" \\\n"," --output_dir=\"results/nq_reader\" \\\n"," --model_dir=\"results/nq_reader/model\" \\\n"," --noloop_predict \\\n"," --tapas_verbosity=\"ERROR\" \\\n"," --test_batch_size={len(queries)} \\\n"," --reset_position_index_per_cell \\\n"," --init_checkpoint=\"tapas_reader/model.ckpt\" \\\n"," --bert_config_file=\"tapas_reader/bert_config.json\" \\\n"," --bert_vocab_file=\"tapas_reader/vocab.txt\" \\\n"," --compression_type= \\\n"," --mode=\"predict\" 2> error\n"," # Display results\n"," results_path = \"results/nq_reader/model/test.tsv\"\n"," \n"," interactions = prediction_utils.iterate_interactions(\n"," \"results/nq_reader/queries/reader_queries.tfrecord\")\n"," tables = {\n"," interaction.questions[0].id : (get_table_df(interaction.table),\n"," interaction.table.table_id,\n"," converter._tokenize_table(interaction.table),\n"," interaction.questions[0].original_text)\n"," for interaction in interactions}\n","\n"," with open(results_path) as csvfile:\n"," reader = csv.DictReader(csvfile, delimiter='\\t')\n"," \n"," for row in reader:\n"," # question_id\n"," df, table_id, table_tokens, query_text = tables[row[\"question_id\"]]\n"," print(Colors.BLACK)\n"," print(\"query >\", query_text)\n"," print(\" > table id: \", table_id)\n"," print(\" > table prediction score: \", row[\"logits_cls\"])\n"," answers = ast.literal_eval(row[\"answers\"])\n"," for a in answers:\n"," index_r = a[\"row_index\"]\n"," index_c = a[\"column_index\"]\n"," colored_answer = set_answer_color(\n"," table_tokens.rows[index_r+1][index_c],\n"," a[\"begin_token_index\"], a[\"end_token_index\"])\n"," df.iat[index_r+1, index_c] = colored_answer\n"," print(\" > Answer cell:\", Colors.BASE + colored_answer)\n"," print(Colors.BLACK + \" > Answer score:\", a[\"score\"], \"\\n\")\n"," with pd.option_context(\n"," 'display.max_rows', None, 'display.max_columns', None,\n"," 'expand_frame_repr', False, 'display.unicode.ambiguous_as_wide', False,\n"," 'display.max_colwidth', None):\n"," print(\"Table:\\n\")\n"," print(Colors.BASE + df.to_string(index=False, header=False))\n"," print(Colors.BLACK +\"-------------------------------------------------------------------------------------\\n\")\n"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"uOXJ_JZJFmZ6"},"source":["## Run predict"]},{"cell_type":"code","metadata":{"id":"aAyvPvoeyxAH"},"source":["create_directories()\n","\n","new_query = \"Write your question\" #@param {type:\"string\"}\n","# You can add multiple queries.\n","queries = [new_query]\n","\n","extract_queries_data(queries)\n","get_queries_embeddings()\n","\n","num_neighbors = 4 #@param {type:\"integer\"}\n","selected_tables = get_nearest_neighbors(num_neighbors=num_neighbors)\n","create_interactions_for_reader(selected_tables=selected_tables)\n","\n","predict()\n"],"execution_count":null,"outputs":[]}]} +{ + "nbformat": 4, + "nbformat_minor": 0, + "metadata": { + "colab": { + "name": "retrieval_predictions.ipynb", + "provenance": [], + "collapsed_sections": [] + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "Bf-86KknHHps" + }, + "source": [ + "\"Open" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "M8XF59AeJask" + }, + "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": "rnPwClY0taWp" + }, + "source": [ + "# Copyright 2021 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": "3jiJ5VvrJkih" + }, + "source": [ + "# Run TAPAS retrieval models description\n", + "This notebook shows how to use retrieval models, which was introduced in the paper: [Open Domain Question Answering over Tables via Dense Retrieval](https://arxiv.org/pdf/2103.12011.pdf).\n", + "1. Load pre-trained and fine-tuned models.\n", + " > * the dual encoder. (called tapas_retriever)\n", + " > * the reader models. (called tapas_reader)\n", + "2. Add handcrafted query and extract the interactions and tf-examples.\n", + "3. Get nearest neighbors for each query, and extract the interactions to pass to the reader.\n", + "4. Call the reader on the new interactions and print \n", + " > * the query.\n", + " > * the probability of this table containing the answer.\n", + " > * the table with a highlighted answer found by the reader.\n", + "\n", + "\n" + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "G9Z8tP11vCOl" + }, + "source": [ + "# Clone and install the repository\n", + "! sudo apt-get install protobuf-compiler\n", + "! git clone https://github.com/google-research/tapas\n", + "! pip install -e ./tapas" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "### NOTE: \n", + "Please ensure you have restarted the runtime before continuing executing the cells below. " + ], + "metadata": { + "id": "l3fyrVm2nUd6" + } + }, + { + "cell_type": "code", + "source": [ + "# Run the imports needed for all the colab\n", + "import tensorflow.compat.v1 as tf\n", + "import os \n", + "import shutil\n", + "import csv\n", + "import pandas as pd\n", + "import IPython\n", + "import ast\n", + "\n", + "tf.get_logger().setLevel('ERROR')\n", + "\n", + "from tapas.utils import tf_example_utils\n", + "from tapas.utils import beam_runner\n", + "from tapas.utils import create_data\n", + "from tapas.protos import interaction_pb2\n", + "from tapas.utils import number_annotation_utils\n", + "from tapas.scripts import prediction_utils\n", + "from tapas.scripts import eval_table_retriever_utils\n", + "from tapas.retrieval import tf_example_utils as retrieval_utils\n", + "from tapas.experiments import table_retriever_experiment" + ], + "metadata": { + "id": "QFk6REz6nRT3" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "id": "6fiPv3GV7yXZ" + }, + "source": [ + "# 1. Load pre-trained and fine-tuned models fom Google Storage.\n", + "# The dual encoder model\n", + "\n", + "# The dual encoder model\n", + "! gsutil cp \"gs://tapas_models/2021_04_27/tapas_nq_hn_retriever_medium.zip\" \"tapas_retriever.zip\" && unzip tapas_retriever.zip\n", + "! mv tapas_nq_hn_retriever_medium/ tapas_retriever\n", + "\n", + "\n", + "# The reader model\n", + "! gsutil cp \"gs://tapas_models/2021_04_27/tapas_nq_hn_reader_large.zip\" \"tapas_reader.zip\" && unzip tapas_reader.zip\n", + "! mv tapas_nq_hn_reader_large tapas_reader\n", + "\n", + "# Load the released nq_tables data.\n", + "os.makedirs('tapas_models_nq_tables', exist_ok=True)\n", + "! gsutil -m cp -R gs://tapas_models/2021_07_22/nq_tables/* tapas_models_nq_tables/\n" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "9vGb8QmE6BsH" + }, + "source": [ + "# Code" + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "9mN0WRCQv7hr" + }, + "source": [ + "# 2. Add handcrafted tables, and queries\n", + "# 2.1. Create the needed directories.\n", + "def create_directories():\n", + " \"\"\"Create directories.\"\"\"\n", + " # To be used for the dual encoder.\n", + " os.makedirs('results/nq_retrieval/model', exist_ok=True)\n", + " with open('results/nq_retrieval/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_retriever/model.ckpt{suffix}', f'results/nq_retrieval/model/model.ckpt-0{suffix}')\n", + " shutil.copyfile(f'tapas_retriever/tables.tsv', f'results/nq_retrieval/model/tables.tsv')\n", + " shutil.copyfile(f'tapas_retriever/bert_config.json', f'results/nq_retrieval/model/bert_config.json')\n", + " # To be used for nq_reder.\n", + " os.makedirs('results/nq_retrieval/tf_examples', exist_ok=True)\n", + " os.makedirs('results/nq_retrieval/queries', exist_ok=True)\n", + " # os.makedirs('results/nq_retrieval/tables', exist_ok=True)\n", + "\n", + " os.makedirs('results/nq_reader/model', exist_ok=True)\n", + " os.makedirs('results/nq_reader/queries', exist_ok=True)\n", + " os.makedirs('results/nq_reader/nq_retrieval/tf_examples', exist_ok=True)\n", + " os.makedirs('results/nq_reader/nq_retrieval/model', exist_ok=True)\n", + "\n", + " with open('results/nq_reader/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_reader/model.ckpt{suffix}', f'results/nq_reader/model/model.ckpt-0{suffix}')\n", + "\n", + "# 2.2. Code to extract the data: the interactions than the tf_examples.\n", + "# interaction_pb2.Interaction` protobuf object is the data structure we use to\n", + "# store examples, and then to call the prediction script.\n", + "def get_table(document_title, table_data):\n", + " \"\"\"Extracts the interaction for an str table.\n", + " \n", + " Args:\n", + " table_data: str table where the columns are separated by '|'.\n", + " document_title: str title of the page containing the table or a table title\n", + " it also could be empty str.\"\"\"\n", + " table = [list(map(lambda s: s.strip(), row.split(\"|\"))) \n", + " for row in table_data.split(\"\\n\") if row.strip()]\n", + " table_interaction = interaction_pb2.Table()\n", + " table_interaction.document_title = document_title\n", + " table_interaction.table_id = document_title\n", + " if not table:\n", + " return table_interaction\n", + " for header in table[0]:\n", + " table_interaction.columns.add().text = header\n", + " for line in table[1:]:\n", + " row = table_interaction.rows.add()\n", + " for cell in line:\n", + " row.cells.add().text = cell\n", + " return table_interaction\n", + "\n", + "def extract_queries(queries):\n", + " \"\"\"Extracts the interaction for a list of queries.\n", + " \n", + " This is used to create the interaction queries file.\n", + " Args:\n", + " queries: list of str queries.\"\"\"\n", + " for idx, query in enumerate(queries):\n", + " interaction = interaction_pb2.Interaction()\n", + " interaction.id = f\"queries_{idx}\"\n", + " question = interaction.questions.add()\n", + " question.original_text = query\n", + " question.id = f\"{interaction.id}-0_0\"\n", + " interaction.table.CopyFrom(get_table(\"FAKE\", \" | \\n | \\n\"))\n", + " number_annotation_utils.add_numeric_values(interaction) \n", + " yield interaction\n", + "\n", + "def write_tfrecord(filename, examples):\n", + " \"\"\"From interactions examples to tfrecord.\"\"\"\n", + " with tf.io.TFRecordWriter(filename) as writer:\n", + " for example in examples:\n", + " writer.write(example.SerializeToString())\n", + "\n", + "def get_config():\n", + " max_seq_length = 512\n", + " vocab_file = \"tapas_retriever/vocab.txt\"\n", + " config=tf_example_utils.RetrievalConversionConfig(\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", + " cell_trim_length=-1,\n", + " use_document_title=True,\n", + " )\n", + " return config\n", + "\n", + "def extract_queries_data(queries):\n", + " \"\"\"Extracts the interactions then the tf_examples.\n", + " \n", + " Args:\n", + " queries: list of str queries.\n", + " \"\"\"\n", + " examples = extract_queries(queries)\n", + " input_queries = \"results/nq_retrieval/queries/queries.tfrecord\"\n", + " write_tfrecord(input_queries, examples)\n", + " config = get_config()\n", + " beam_runner.run_type(\n", + " create_data.build_retrieval_pipeline(\n", + " input_files=[input_queries],\n", + " output_files=[os.path.join(\"results/nq_retrieval/tf_examples\",\n", + " \"queries.tfrecord\")],\n", + " input_format=create_data.InputFormat.INTERACTION,\n", + " config=config,\n", + " ), beam_runner.RunnerType.DIRECT).wait_until_finish()\n", + "\n", + "# 3. Retrieval: Extract the queries interactions to pass to the reader\n", + "# 3.1. Extracts the queries embeddings.\n", + "def get_queries_embeddings():\n", + " ! python -m tapas.experiments.table_retriever_experiment \\\n", + " --do_predict \\\n", + " --eval_name=\"dual_encoder_queries\" \\\n", + " --minutes_to_sleep_before_predictions=0 \\\n", + " --num_eval_steps=0 \\\n", + " --model_dir=\"results/nq_retrieval/model\" \\\n", + " --prediction_output_dir=\"results/nq_retrieval/model/queries\" \\\n", + " --evaluated_checkpoint_step=0 \\\n", + " --input_file_predict=\"results/nq_retrieval/tf_examples/queries.tfrecord\" \\\n", + " --bert_config_file=\"tapas_retriever/bert_config.json\" \\\n", + " --init_from_single_encoder=false \\\n", + " --tf_random_seed=\"1\" \\\n", + " --compression_type= \\\n", + " --down_projection_dim=256 \\\n", + " --eval_batch_size=1 \\\n", + " --max_seq_length=512 2> error\n", + "\n", + " with open(\"results/nq_retrieval/model/queries/predict_results_0.tsv\") as csvfile_reader:\n", + " reader = csv.DictReader(csvfile_reader, delimiter='\\t')\n", + " for i, row in enumerate(reader):\n", + " print(\"Adding query_id: \", row[\"query_id\"], \":\", queries[i])\n", + "\n", + "# 3.2. Get nearest tables neighbors for each query.\n", + "def get_nearest_neighbors(num_neighbors):\n", + " queries_pred = eval_table_retriever_utils.read_queries(\"results/nq_retrieval/model/queries/predict_results_0.tsv\")\n", + " tables = eval_table_retriever_utils.read_tables(\"results/nq_retrieval/model/tables.tsv\", make_tables_unique=False)\n", + " index = eval_table_retriever_utils.build_table_index(tables)\n", + " similarities, neighbors = eval_table_retriever_utils._retrieve(queries_pred, index)\n", + " selected_tables = {}\n", + " for i, s in enumerate(similarities):\n", + " print(\"Query index\", i, \":\", queries[i])\n", + " selected_tables[i]={}\n", + " for pos in range(num_neighbors):\n", + " table_id = tables[neighbors[i][pos]].table_id\n", + " selected_tables[i][table_id] = (s[pos], neighbors[i][pos])\n", + " print(\" Related table id:\", table_id)\n", + " print(\" Table's score \", s[pos])\n", + " print(\" ----------------------------------------------\")\n", + " return selected_tables\n", + "# 3.3. Extract the interactions to pass to the reader.\n", + "def iterate_tables(input_file):\n", + " \"\"\"Reads interaction_pb2.Table().\"\"\"\n", + " for value in tf.python_io.tf_record_iterator(input_file):\n", + " table = interaction_pb2.Table()\n", + " table.ParseFromString(value)\n", + " yield table\n", + "\n", + "def create_queries_tables_interactions(selected_tables):\n", + " \"\"\"Creates the interaction by linking the query to the selected table.\"\"\"\n", + " queries_interactions = prediction_utils.iterate_interactions(\n", + " \"results/nq_retrieval/queries/queries.tfrecord\")\n", + " all_tables = iterate_tables(\"tapas_models_nq_tables/tables/tables.tfrecord\")\n", + " tables = {table.table_id: table for table in all_tables}\n", + " for i, q in enumerate(queries_interactions):\n", + " print(\"\\n Query index:\", i, \":\", queries[i])\n", + " t = selected_tables[i]\n", + " for table_id in t.keys():\n", + " if table_id in tables.keys():\n", + " table = tables[table_id]\n", + " print(\" > Converted query:\", q.questions[0].original_text)\n", + " new_interaction = interaction_pb2.Interaction()\n", + " new_interaction.CopyFrom(q)\n", + " new_interaction.id = f\"{new_interaction.id}_{table.table_id}\"\n", + " new_interaction.questions[0].id = f\"{new_interaction.id}_0\"\n", + " new_interaction.table.CopyFrom(table)\n", + " yield new_interaction\n", + " else:\n", + " print(\" > Not found in table file.\")\n", + " print(\" Related table id: \", table_id)\n", + " print(\" Table's score: \", t[table_id][0])\n", + " print(\" ----------------------------------------------\")\n", + "\n", + "def create_interactions_for_reader(selected_tables):\n", + " examples = create_queries_tables_interactions(selected_tables)\n", + " write_tfrecord(\"results/nq_reader/queries/reader_queries.tfrecord\", examples)\n", + "\n", + "# 4. Reader: Get the answer given the question and the table\n", + "def get_converter(max_seq_length):\n", + " \"\"\"Get a clssifier conferter.\"\"\"\n", + " config = tf_example_utils.ClassifierConversionConfig(\n", + " vocab_file=\"tapas_reader/vocab.txt\",\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", + " return tf_example_utils.ToClassifierTensorflowExample(config)\n", + "\n", + "def convert_interactions_to_examples(converter):\n", + " \"\"\"Calls Tapas converter to convert interaction to example.\"\"\"\n", + " interactions = prediction_utils.iterate_interactions(\n", + " \"results/nq_reader/queries/reader_queries.tfrecord\")\n", + " for interaction in interactions:\n", + " try:\n", + " yield converter.convert(interaction, 0)\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", + "class Colors:\n", + " \"\"\"Used to highlight the answers.\"\"\"\n", + " ANSWER = '\\033[94m'\n", + " BASE = '\\033[95m'\n", + " BLACK = '\\033[0m'\n", + "\n", + " \n", + "def set_answer_color(input, begin, end):\n", + " \"\"\"Highlights the answers.\"\"\"\n", + " list_output = [i.original_text for i in input]\n", + " list_output[begin] = Colors.ANSWER + list_output[begin]\n", + " list_output[end - 1] = list_output[end - 1] + Colors.BASE\n", + " return \" \".join(list_output)\n", + "\n", + "\n", + "def get_table_df(table):\n", + " \"\"\"Extracts a dataframe table for a better visualisation.\"\"\"\n", + " printabe_table = [[Colors.BASE + c.text + Colors.BASE for c in table.columns]] \n", + " for r in table.rows:\n", + " printabe_table.append([Colors.BASE + c.text + Colors.BASE for c in r.cells])\n", + " return pd.DataFrame(printabe_table)\n", + "\n", + "def predict():\n", + " \"\"\"Predict the answer given the query and the table.\"\"\"\n", + " max_seq_length = 512\n", + " # Extracts the tf examples given the interactions.\n", + " converter = get_converter(max_seq_length)\n", + " examples = convert_interactions_to_examples(converter)\n", + " write_tf_example(\"results/nq_reader/nq_retrieval/tf_examples/test.tfrecord\", examples)\n", + " write_tf_example(\"results/nq_reader/nq_retrieval/tf_examples/dev.tfrecord\", [])\n", + " # Run prediction\n", + " ! python -m tapas.run_task_main \\\n", + " --task=\"NQ_RETRIEVAL\" \\\n", + " --output_dir=\"results/nq_reader\" \\\n", + " --model_dir=\"results/nq_reader/model\" \\\n", + " --noloop_predict \\\n", + " --tapas_verbosity=\"ERROR\" \\\n", + " --test_batch_size={len(queries)} \\\n", + " --reset_position_index_per_cell \\\n", + " --init_checkpoint=\"tapas_reader/model.ckpt\" \\\n", + " --bert_config_file=\"tapas_reader/bert_config.json\" \\\n", + " --bert_vocab_file=\"tapas_reader/vocab.txt\" \\\n", + " --compression_type= \\\n", + " --mode=\"predict\" 2> error\n", + " # Display results\n", + " results_path = \"results/nq_reader/model/test.tsv\"\n", + " \n", + " interactions = prediction_utils.iterate_interactions(\n", + " \"results/nq_reader/queries/reader_queries.tfrecord\")\n", + " tables = {\n", + " interaction.questions[0].id : (get_table_df(interaction.table),\n", + " interaction.table.table_id,\n", + " converter._tokenize_table(interaction.table),\n", + " interaction.questions[0].original_text)\n", + " for interaction in interactions}\n", + "\n", + " with open(results_path) as csvfile:\n", + " reader = csv.DictReader(csvfile, delimiter='\\t')\n", + " \n", + " for row in reader:\n", + " # question_id\n", + " df, table_id, table_tokens, query_text = tables[row[\"question_id\"]]\n", + " print(Colors.BLACK)\n", + " print(\"query >\", query_text)\n", + " print(\" > table id: \", table_id)\n", + " print(\" > table prediction score: \", row[\"logits_cls\"])\n", + " answers = ast.literal_eval(row[\"answers\"])\n", + " for a in answers:\n", + " index_r = a[\"row_index\"]\n", + " index_c = a[\"column_index\"]\n", + " colored_answer = set_answer_color(\n", + " table_tokens.rows[index_r+1][index_c],\n", + " a[\"begin_token_index\"], a[\"end_token_index\"])\n", + " df.iat[index_r+1, index_c] = colored_answer\n", + " print(\" > Answer cell:\", Colors.BASE + colored_answer)\n", + " print(Colors.BLACK + \" > Answer score:\", a[\"score\"], \"\\n\")\n", + " with pd.option_context(\n", + " 'display.max_rows', None, 'display.max_columns', None,\n", + " 'expand_frame_repr', False, 'display.unicode.ambiguous_as_wide', False,\n", + " 'display.max_colwidth', None):\n", + " print(\"Table:\\n\")\n", + " print(Colors.BASE + df.to_string(index=False, header=False))\n", + " print(Colors.BLACK +\"-------------------------------------------------------------------------------------\\n\")\n" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "uOXJ_JZJFmZ6" + }, + "source": [ + "## Run predict" + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "aAyvPvoeyxAH" + }, + "source": [ + "create_directories()\n", + "\n", + "new_query = \"Write your question\" #@param {type:\"string\"}\n", + "# You can add multiple queries.\n", + "queries = [new_query]\n", + "\n", + "extract_queries_data(queries)\n", + "get_queries_embeddings()\n", + "\n", + "num_neighbors = 4 #@param {type:\"integer\"}\n", + "selected_tables = get_nearest_neighbors(num_neighbors=num_neighbors)\n", + "create_interactions_for_reader(selected_tables=selected_tables)\n", + "\n", + "predict()\n" + ], + "execution_count": null, + "outputs": [] + } + ] +} \ No newline at end of file