Table Question Answering (Table QA) refers to providing precise answers from tables to answer a user's question. With recent works on Table QA, is it now possible to answer natural language queries from tabular data. This notebook demonstrates how you can build a Table QA system that can answer your natural language queries using the Pinecone vector database.

We need three main components to build the Table QA system:

• A vector index to store table embeddings
• A retriever model for embedding queries and tables
• A reader model to read the tables and extract answers

Install Dependencies

# torch-scatter may take few minutes to install
!pip install datasets pinecone-client sentence_transformers torch-scatter

Load the Dataset

We will work with a subset of the Open Table-and-Text Question Answering (OTT-QA) dataset, consisting of texts and tables from Wikipedia. The subset contains 20,000 tables, and it can be loaded from the Huggigface Datasets hub as follows:

from datasets import load_dataset

# load the dataset from huggingface datasets hub
data = load_dataset("ashraq/ott-qa-20k", split="train")
data

Dataset({
    features: ['url', 'title', 'header', 'data', 'section_title', 'section_text', 'uid', 'intro'],
    num_rows: 20000
})

data[2]

{'url': 'https://en.wikipedia.org/wiki/1976_New_York_Mets_season',
 'title': '1976 New York Mets season',
 'header': ['Level', 'Team', 'League', 'Manager'],
 'data': [['AAA', 'Tidewater Tides', 'International League', 'Tom Burgess'],
  ['AA', 'Jackson Mets', 'Texas League', 'John Antonelli'],
  ['A', 'Lynchburg Mets', 'Carolina League', 'Jack Aker'],
  ['A', 'Wausau Mets', 'Midwest League', 'Bill Monbouquette'],
  ['Rookie', 'Marion Mets', 'Appalachian League', 'Al Jackson']],
 'section_title': 'Farm system',
 'section_text': 'See also : Minor League Baseball',
 'uid': '1976_New_York_Mets_season_7',
 'intro': 'The New York Mets season was the 15th regular season for the Mets, who played home games at Shea Stadium. Led by manager Joe Frazier, the team had an 86-76 record and finished in third place in the National League East.'}

As we can see, the dataset includes both textual and tabular data that are related to one another. Let's extract and transform the dataset's tables into pandas dataframes as we will only be using the tables in this example.

import pandas as pd

# store all tables in the tables list
tables = []
# loop through the dataset and convert tabular data to pandas dataframes
for doc in data:
    table = pd.DataFrame(doc["data"], columns=doc["header"])
    tables.append(table)

tables[2]

Initialize Retriever

The retriever transforms natural language queries and tabular data into embeddings/vectors. It will generate embeddings in a way that the natural language questions and tables containing answers to our questions are nearby in the vector space.

We will use a SentenceTransformer model trained specifically for embedding tabular data for retrieval tasks. The model can be loaded from the Huggingface Models hub as follows:

import torch
from sentence_transformers import SentenceTransformer

# set device to GPU if available
device = 'cuda' if torch.cuda.is_available() else 'cpu'
# load the table embedding model from huggingface models hub
retriever = SentenceTransformer("deepset/all-mpnet-base-v2-table", device=device)
retriever

SentenceTransformer(
  (0): Transformer({'max_seq_length': 384, 'do_lower_case': False}) with Transformer model: MPNetModel 
  (1): Pooling({'word_embedding_dimension': 768, 'pooling_mode_cls_token': False, 'pooling_mode_mean_tokens': True, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False})
  (2): Normalize()
)

The retriever expects tables to be in a particular format. Let's write a function to convert the tables to this format.

def _preprocess_tables(tables: list):
    processed = []
    # loop through all tables
    for table in tables:
        # convert the table to csv and 
        processed_table = "\n".join([table.to_csv(index=False)])
        # add the processed table to processed list
        processed.append(processed_table)
    return processed

Notice that we are only using tables here. However, if you want the retriever to take the metadata into account while retrieving the tables, you can join any metadata strings, such as title, section_title, etc., separated by new line characters at the beginning of the processed table.

Let's take a look at the formatted tables.

# format all the dataframes in the tables list
processed_tables = _preprocess_tables(tables)
# display the formatted table
processed_tables[2]

'Level,Team,League,Manager\nAAA,Tidewater Tides,International League,Tom Burgess\nAA,Jackson Mets,Texas League,John Antonelli\nA,Lynchburg Mets,Carolina League,Jack Aker\nA,Wausau Mets,Midwest League,Bill Monbouquette\nRookie,Marion Mets,Appalachian League,Al Jackson\n'

The formatted table may not make sense to us, but the embedding model is trained to understand it and generate accurate embeddings.

Initialize Pinecone Index

We will use the Pinecone vector database as our vector index. The Pinecone index stores vector representations of our tables which we can retrieve using a natural language query (query vector). Pinecone does this by computing the similarity between the query vector and the embedded tables stored in the vector index.

To use Pinecone, we first need to initialize a connection to Pinecone. For this, we need a free API key. You can find your environment in the Pinecone console under API Keys. We initialize the connection like so:

import pinecone

# connect to pinecone environment
pinecone.init(
    api_key="YOUR API KEY",
    environment="YOUR_ENVIRONMENT"
)

Now we create a new index. We specify the metric type as "cosine" and dimension as 768 because the retriever we use to generate context embeddings outputs 768-dimension vectors. Pinecone will use cosine similarity to compute the similarity between the query and table embeddings.

# you can choose any name for the index
index_name = "table-qa"

# check if the table-qa index exists
if index_name not in pinecone.list_indexes():
    # create the index if it does not exist
    pinecone.create_index(
        index_name,
        dimension=768,
        metric="cosine"
    )

# connect to table-qa index we created
index = pinecone.Index(index_name)

Generate Embeddings and Upsert

Next we need to generate the table embeddings and upload it to the Pinecone index. We can easily do that as follows:

from tqdm.auto import tqdm

# we will use batches of 64
batch_size = 64

for i in tqdm(range(0, len(processed_tables), batch_size)):
    # find end of batch
    i_end = min(i+batch_size, len(processed_tables))
    # extract batch
    batch = processed_tables[i:i_end]
    # generate embeddings for batch
    emb = retriever.encode(batch).tolist()
    # create unique IDs ranging from zero to the total number of tables in the dataset
    ids = [f"{idx}" for idx in range(i, i_end)]
    # add all to upsert list
    to_upsert = list(zip(ids, emb))
    # upsert/insert these records to pinecone
    _ = index.upsert(vectors=to_upsert)

# check that we have all vectors in index
index.describe_index_stats()

  100%|██████████| 313/313 [09:12<00:00, 1.49s/it]

{'dimension': 768,
 'index_fullness': 0.0,
 'namespaces': {'': {'vector_count': 20000}},
 'total_vector_count': 20000}

Now the Pinecone index is ready for querying. Let's test to see if it returns tables relevant to our queries.

query = "which country has the highest GDP in 2020?"
# generate embedding for the query
xq = retriever.encode([query]).tolist()
# query pinecone index to find the table containing answer to the query
result = index.query(xq, top_k=1)
result

{'matches': [{'id': '19931', 'score': 0.822087, 'values': []}], 'namespace': ''}

The Pinecone index has returned the id of a table that would contain the answer to our query with 82.2% confidence. Let's see if this table actually contains the answer. We can use the returned id as an index to get the relevant pandas dataframe from the tables list.

id = int(result["matches"][0]["id"])
tables[id].head()

The table returned by the Pinecone index indeed has the answer to our query. Now we need a model that can read this table and extract the precise answer.

Initialize Table Reader

As the reader, we will use a TAPAS model fine-tuned for the Table QA task. TAPAS is a BERT-like Transformer model pretrained in a self-supervised manner on a large corpus of English language data from Wikipedia. We load the model and tokenizer from the Huggingface model hub into a question-answering pipeline.

from transformers import pipeline, TapasTokenizer, TapasForQuestionAnswering

model_name = "google/tapas-base-finetuned-wtq"
# load the tokenizer and the model from huggingface model hub
tokenizer = TapasTokenizer.from_pretrained(model_name)
model = TapasForQuestionAnswering.from_pretrained(model_name, local_files_only=False)
# load the model and tokenizer into a question-answering pipeline
pipe = pipeline("table-question-answering",  model=model, tokenizer=tokenizer, device=device)

Let's run the table returned by the Pinecone index and the query we used before into the question-answering pipeline to extract the answer.

pipe(table=tables[id], query=query)

{'answer': 'China',
 'coordinates': [(0, 1)],
 'cells': ['China'],
 'aggregator': 'NONE'}

The model has precisely answered our query. Let's run some more queries.

Querying

First, we will define two function to handle our queries and extract answers from tables.

def query_pinecone(query):
    # generate embedding for the query
    xq = retriever.encode([query]).tolist()
    # query pinecone index to find the table containing answer to the query
    result = index.query(xq, top_k=1)
    # return the relevant table from the tables list
    return tables[int(result["matches"][0]["id"])]

def get_answer_from_table(table, query):
    # run the table and query through the question-answering pipeline
    answers = pipe(table=table, query=query)
    return answers

query = "which car manufacturers produce cars with a top speed of above 180 kph?"
table = query_pinecone(query)
table

get_answer_from_table(table, query)

{'answer': 'Fiat, Bugatti, Benz, Miller',
 'coordinates': [(0, 0), (3, 0), (7, 0), (8, 0)],
 'cells': ['Fiat', 'Bugatti', 'Benz', 'Miller'],
 'aggregator': 'NONE'}

query = "which scientist is known for improving the steam engine?"
table = query_pinecone(query)
table.head()

get_answer_from_table(table, query)

{'answer': 'George Henry Corliss',
 'coordinates': [(4, 1)],
 'cells': ['George Henry Corliss'],
 'aggregator': 'NONE'}

query = "What is the Maldivian island name for Oblu Select at Sangeli   resort?"
table = query_pinecone(query)
table.head()

get_answer_from_table(table, query)

{'answer': 'Akirifushi',
 'coordinates': [(1, 0)],
 'cells': ['Akirifushi'],
 'aggregator': 'NONE'}

As we can see, our Table QA system can retrieve the correct table from the Pinecone index and extract precise answers from the table. The TAPAS model we use supports more advanced queries. It has an aggregation head which indicates whether we need to count, sum, or average cells to answer the questions. Let's run some advanced queries that require aggregation to answer.

query = "what was the total GDP of China and Indonesia in 2020?"
table = query_pinecone(query)
table.head()

get_answer_from_table(table, query)

{'answer': 'SUM > 27,804,953, 3,778,134',
 'coordinates': [(0, 2), (3, 2)],
 'cells': ['27,804,953', '3,778,134'],
 'aggregator': 'SUM'}

Here the QA system suggests the correct cells to add in order to get the total GDP of China and Indonesia in 2020.

query = "what is the average carbon emission of power stations in australia, canada and germany?"
table = query_pinecone(query)
table.head()

get_answer_from_table(table, query)

{'answer': 'AVERAGE > 1.58, 1.27, 1.25',
 'coordinates': [(0, 0), (2, 0), (3, 0)],
 'cells': ['1.58', '1.27', '1.25'],
 'aggregator': 'AVERAGE'}

As we can see, the QA system correctly identified which cells to average to answer our question.