Abstractive Question Answering

Open In Colab Open nbviewer Open Github

Abstractive question-answering focuses on the generation of multi-sentence answers to open-ended questions. It usually works by searching massive document stores for relevant information and then using this information to synthetically generate answers. This notebook demonstrates how Pinecone helps you build an abstractive question-answering system. We need three main components:

  • A vector index to store and run semantic search
  • A retriever model for embedding context passages
  • A generator model to generate answers

Install Dependencies

!pip install -qU datasets pinecone-client sentence-transformers torch

Load and Prepare Dataset

Our source data will be taken from the Wiki Snippets dataset, which contains over 17 million passages from Wikipedia. But, since indexing the entire dataset may take some time, we will only utilize 50,000 passages in this demo that include "History" in the "section title" column. If you want, you may utilize the complete dataset. Pinecone vector database can effortlessly manage millions of documents for you.

from datasets import load_dataset

# load the dataset from huggingface in streaming mode and shuffle it
wiki_data = load_dataset(
    'vblagoje/wikipedia_snippets_streamed',
    split='train',
    streaming=True
).shuffle(seed=960)

We are loading the dataset in the streaming mode so that we don't have to wait for the whole dataset to download (which is over 9GB). Instead, we iteratively download records one at a time.

# show the contents of a single document in the dataset
next(iter(wiki_data))
{'wiki_id': 'Q7649565',
 'start_paragraph': 20,
 'start_character': 272,
 'end_paragraph': 24,
 'end_character': 380,
 'article_title': 'Sustainable Agriculture Research and Education',
 'section_title': "2000s & Evaluation of the program's effectiveness",
 'passage_text': "preserving the surrounding prairies. It ran until March 31, 2001.\nIn 2008, SARE celebrated its 20th anniversary. To that date, the program had funded 3,700 projects and was operating with an annual budget of approximately $19 million. Evaluation of the program's effectiveness As of 2008, 64% of farmers who had received SARE grants stated that they had been able to earn increased profits as a result of the funding they received and utilization of sustainable agriculture methods. Additionally, 79% of grantees said that they had experienced a significant improvement in soil quality though the environmentally friendly, sustainable methods that they were"}
# filter only documents with History as section_title
history = wiki_data.filter(
    lambda d: d['section_title'].startswith('History')
)

Let's iterate through the dataset and apply our filter to select the 50,000 historical passages. We will extract article_title, section_title and passage_text from each document.

from tqdm.auto import tqdm  # progress bar

total_doc_count = 50000

counter = 0
docs = []
# iterate through the dataset and apply our filter
for d in tqdm(history, total=total_doc_count):
    # extract the fields we need
    doc = {
        "article_title": d["article_title"],
        "section_title": d["section_title"],
        "passage_text": d["passage_text"]
    }
    # add the dict containing fields we need to docs list
    docs.append(doc)

    # stop iteration once we reach 50k
    if counter == total_doc_count:
        break

    # increase the counter on every iteration
    counter += 1
  100%|██████████| 50000/50000 [05:18<00:00, 145.03it/s]
import pandas as pd

# create a pandas dataframe with the documents we extracted
df = pd.DataFrame(docs)
df.head()
article_title section_title passage_text
0 Taupo District History was not until the 1950s that the region starte...
1 Sutarfeni History & Western asian analogues Sutarfeni History strand-like pheni were Phena...
2 The Bishop Wand Church of England School History The Bishop Wand Church of England School Histo...
3 Teufelsmoor History & Situation today made to preserve the original landscape, altho...
4 Surface Hill Uniting Church History in perpetual reminder that work and worship go...

Initialize Pinecone Index

The Pinecone index stores vector representations of our historical passages which we can retrieve later using another vector (query vector). To build our vector index, we must first establish a connection with Pinecone. For this, we need an API from Pinecone. You can get one for free from here. You also need to know the environment for your index; for new accounts, the default environment is us-east1-gcp.

We initialize the connection as follows:

import pinecone

# connect to pinecone environment
pinecone.init(
    api_key="<<YOUR_API_KEY>>",
    environment="<<YOUR_ENVIRONMENT>>"
)

Now we create a new index. We will name it "abstractive-question-answering" — you can name it anything we want. We specify the metric type as "cosine" and dimension as 768 because the retriever we use to generate context embeddings is optimized for cosine similarity and outputs 768-dimension vectors.

index_name = "abstractive-question-answering"

# check if the abstractive-question-answering 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 abstractive-question-answering index we created
index = pinecone.Index(index_name)

Initialize Retriever

Next, we need to initialize our retriever. The retriever will mainly do two things:

  • Generate embeddings for all historical passages (context vectors/embeddings)
  • Generate embeddings for our questions (query vector/embedding)

The retriever will create embeddings such that the questions and passages that hold the answers to our queries are close to one another in the vector space. We will use a SentenceTransformer model based on Microsoft's MPNet as our retriever. This model performs quite well for comparing the similarity between queries and documents. We can use Cosine Similarity to compute the similarity between query and context vectors generated by this model (Pinecone automatically does this for us).

import torch
from sentence_transformers import SentenceTransformer

# set device to GPU if available
device = 'cuda' if torch.cuda.is_available() else 'cpu'
# load the retriever model from huggingface model hub
retriever = SentenceTransformer("flax-sentence-embeddings/all_datasets_v3_mpnet-base", device=device)
retriever
SentenceTransformer(
  (0): Transformer({'max_seq_length': 128, '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()
)

Generate Embeddings and Upsert

Next, we need to generate embeddings for the context passages. We will do this in batches to help us more quickly generate embeddings and upload them to the Pinecone index. When passing the documents to Pinecone, we need an id (a unique value), context embedding, and metadata for each document representing context passages in the dataset. The metadata is a dictionary containing data relevant to our embeddings, such as the article title, section title, passage text, etc.

# we will use batches of 64
batch_size = 64

for i in tqdm(range(0, len(df), batch_size)):
    # find end of batch
    i_end = min(i+batch_size, len(df))
    # extract batch
    batch = df.iloc[i:i_end]
    # generate embeddings for batch
    emb = retriever.encode(batch["passage_text"].tolist()).tolist()
    # get metadata
    meta = batch.to_dict(orient="records")
    # create unique IDs
    ids = [f"{idx}" for idx in range(i, i_end)]
    # add all to upsert list
    to_upsert = list(zip(ids, emb, meta))
    # upsert/insert these records to pinecone
    _ = index.upsert(vectors=to_upsert)

# check that we have all vectors in index
index.describe_index_stats()
  100%|██████████| 782/782 [13:56<00:00, 1.13it/s]

{'dimension': 768,
 'index_fullness': 0.1,
 'namespaces': {'': {'vector_count': 50001}},
 'total_vector_count': 50001}

Initialize Generator

We will use ELI5 BART for the generator which is a Sequence-To-Sequence model trained using the ‘Explain Like I’m 5’ (ELI5) dataset. Sequence-To-Sequence models can take a text sequence as input and produce a different text sequence as output.

The input to the ELI5 BART model is a single string which is a concatenation of the query and the relevant documents providing the context for the answer. The documents are separated by a special token <P>, so the input string will look as follows:

question: What is a sonic boom? context: <P> A sonic boom is a sound associated with shock waves created when an object travels through the air faster than the speed of sound. <P> Sonic booms generate enormous amounts of sound energy, sounding similar to an explosion or a thunderclap to the human ear. <P> Sonic booms due to large supersonic aircraft can be particularly loud and startling, tend to awaken people, and may cause minor damage to some structures. This led to prohibition of routine supersonic flight overland.

More detail on how the ELI5 dataset was built is available here and how ELI5 BART model was trained is available here.

Let's initialize the BART model using transformers.

from transformers import BartTokenizer, BartForConditionalGeneration

# load bart tokenizer and model from huggingface
tokenizer = BartTokenizer.from_pretrained('vblagoje/bart_lfqa')
generator = BartForConditionalGeneration.from_pretrained('vblagoje/bart_lfqa').to(device)

All the components of our abstract QA system are complete and ready to be queried. But first, let's write some helper functions to retrieve context passages from Pinecone index and to format the query in the way the generator expects the input.

def query_pinecone(query, top_k):
    # generate embeddings for the query
    xq = retriever.encode([query]).tolist()
    # search pinecone index for context passage with the answer
    xc = index.query(xq, top_k=top_k, include_metadata=True)
    return xc
def format_query(query, context):
    # extract passage_text from Pinecone search result and add the <P> tag
    context = [f"<P> {m['metadata']['passage_text']}" for m in context]
    # concatinate all context passages
    context = " ".join(context)
    # contcatinate the query and context passages
    query = f"question: {query} context: {context}"
    return query

Let's test the helper functions. We will query the Pinecone index function we created earlier with the query_pinecone to get context passages and pass them to the format_query function.

query = "when was the first electric power system built?"
result = query_pinecone(query, top_k=1)
result
{'matches': [{'id': '3593',
              'metadata': {'article_title': 'Electric power system',
                           'passage_text': 'Electric power system History In '
                                           '1881, two electricians built the '
                                           "world's first power system at "
                                           'Godalming in England. It was '
                                           'powered by two waterwheels and '
                                           'produced an alternating current '
                                           'that in turn supplied seven '
                                           'Siemens arc lamps at 250 volts and '
                                           '34 incandescent lamps at 40 volts. '
                                           'However, supply to the lamps was '
                                           'intermittent and in 1882 Thomas '
                                           'Edison and his company, The Edison '
                                           'Electric Light Company, developed '
                                           'the first steam-powered electric '
                                           'power station on Pearl Street in '
                                           'New York City. The Pearl Street '
                                           'Station initially powered around '
                                           '3,000 lamps for 59 customers. The '
                                           'power station generated direct '
                                           'current and',
                           'section_title': 'History'},
              'score': 0.69118017,
              'values': []}],
 'namespace': ''}
from pprint import pprint
# format the query in the form generator expects the input
query = format_query(query, result["matches"])
pprint(query)
('question: when was the first electric power system built? context: <P> '
 "Electric power system History In 1881, two electricians built the world's "
 'first power system at Godalming in England. It was powered by two '
 'waterwheels and produced an alternating current that in turn supplied seven '
 'Siemens arc lamps at 250 volts and 34 incandescent lamps at 40 volts. '
 'However, supply to the lamps was intermittent and in 1882 Thomas Edison and '
 'his company, The Edison Electric Light Company, developed the first '
 'steam-powered electric power station on Pearl Street in New York City. The '
 'Pearl Street Station initially powered around 3,000 lamps for 59 customers. '
 'The power station generated direct current and')

The output looks great. Now let's write a function to generate answers.

def generate_answer(query):
    # tokenize the query to get input_ids
    inputs = tokenizer([query], max_length=1024, return_tensors="pt")
    # use generator to predict output ids
    ids = generator.generate(inputs["input_ids"], num_beams=2, min_length=20, max_length=40)
    # use tokenizer to decode the output ids
    answer = tokenizer.batch_decode(ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
    return pprint(answer)
generate_answer(query)
('The first electric power system was built in 1881 at Godalming in England. '
 'It was powered by two waterwheels and produced alternating current that in '
 'turn supplied seven Siemens arc lamps')

As we can see, the generator used the provided context to answer our question. Let's run some more queries.

query = "How was the first wireless message sent?"
context = query_pinecone(query, top_k=5)
query = format_query(query, context["matches"])
generate_answer(query)
('The first wireless message was sent in 1866 by Mahlon Loomis, who had a kite '
 'on a mountaintop 14 miles apart. The kite was connected to a cable')

To confirm that this answer is correct, we can check the contexts used to generate the answer.

for doc in context["matches"]:
    print(doc["metadata"]["passage_text"], end='\n---\n')
by electrostatic induction or electromagnetic induction, which had too short a range to be practical. In 1866 Mahlon Loomis claimed to have transmitted an electrical signal through the atmosphere between two 600 foot wires held aloft by kites on mountaintops 14 miles apart. Thomas Edison had come close to discovering radio in 1875; he had generated and detected radio waves which he called "etheric currents" experimenting with high-voltage spark circuits, but due to lack of time did not pursue the matter. David Edward Hughes in 1879 had also stumbled on radio wave transmission which he received with his carbon microphone
---
the east coast of India, then on to Penang, Malacca, Singapore, Batavia (current Jakarta), to finally reach Darwin, Australia. It was the first direct link between Australia and Great Britain. The company that laid the first part of the cable took the name of Falmouth, Gibraltar and Malta Telegraph Company and had been founded in 1869. This company later operated as the Eastern Telegraph Company from Mount Pleasant in Gibraltar and eventually became Cable & Wireless.
The first telephones were introduced to Gibraltar in 1886 by a private company which was later taken over by the colonial authorities. The first wireless
---
audio distance records, and were heard as far west as Hawaii. They were also received in Paris, France, which marked the first transmission of speech across the Atlantic.
With the entrance of the United States into World War I in April 1917 the federal government took over full control of the radio industry, and it became illegal for civilians to possess an operational radio receiver. However NAA continued to operate during the conflict. In addition to time signals and weather reports, it also broadcast news summaries received by troops on land and aboard ships in the Atlantic. Effective April 15, 1919
---
Message from space (science fiction) For other uses, see Message from Space (disambiguation).
"Message from space" is a type of "first contact" theme in science fiction . Stories of this type involve receiving an interstellar message which reveals the existence of other intelligent life in the universe. History An early short story, A Message from Space (Joseph Schlossel, Weird Tales, March 1926) tells of an amateur who builds a ham TV set and suddenly sees an alien, The latter one realises it is being watched and tells its soap opera story. The verdict of Everett Franklin Bleiler: "original ideas, but clumsy
---
radio operators, interested in a practical benefit from their hobby, and jewelers, who previously had been reliant on time services transmitted over telegraph wires, which had a reputation for being both expensive and of questionable reliability, especially compared to the free and very accurate NAA transmissions.
NAA's original transmitters were only capable of producing the dots-and-dashes of Morse code. The later development of vacuum tube transmitters made audio transmissions practical, and in 1915 the American Telephone and Telegraph Company (AT&T) received permission from the Navy to conduct a series of tests at the NAA facility. These experimental transmissions set impressive new
---

In this case, the answer looks correct. If we ask a question and no relevant contexts are retrieved, the generator will typically return nonsensical or false answers, like with this question about COVID-19:

query = "where did COVID-19 originate?"
context = query_pinecone(query, top_k=3)
query = format_query(query, context["matches"])
generate_answer(query)
('COVID-19 is a zoonotic disease, which means that it is a virus that is '
 'transmitted from one animal to another. It is not a virus that can be '
 'transmitted from person')
for doc in context["matches"]:
    print(doc["metadata"]["passage_text"], end='\n---\n')
to establish with certainty which diseases jumped from other animals to humans, but there is increasing evidence from DNA and RNA sequencing, that measles, smallpox, influenza, HIV, and diphtheria came to humans this way. Various forms of the common cold and tuberculosis also are adaptations of strains originating in other species.
Zoonoses are of interest because they are often previously unrecognized diseases or have increased virulence in populations lacking immunity. The West Nile virus appeared in the United States in 1999 in the New York City area, and moved through the country in the summer of 2002, causing much distress. Bubonic
---
plague is a zoonotic disease, as are salmonellosis, Rocky Mountain spotted fever, and Lyme disease.
A major factor contributing to the appearance of new zoonotic pathogens in human populations is increased contact between humans and wildlife. This can be caused either by encroachment of human activity into wilderness areas or by movement of wild animals into areas of human activity. An example of this is the outbreak of Nipah virus in peninsular Malaysia in 1999, when intensive pig farming began on the habitat of infected fruit bats. Unidentified infection of the pigs amplified the force of infection, eventually transmitting the virus
---
man killed and twenty-nine died of disease.
---

Let’s finish with a final few questions.

query = "what was the war of currents?"
context = query_pinecone(query, top_k=5)
query = format_query(query, context["matches"])
generate_answer(query)
('The War of Currents was a series of events in the early 1900s between Edison '
 'and Westinghouse. The two companies were competing for the market share of '
 'electric power in the United States')
query = "who was the first person on the moon?"
context = query_pinecone(query, top_k=10)
query = format_query(query, context["matches"])
generate_answer(query)
('The first person to walk on the moon was Neil Armstrong in 1969. He walked '
 'on the moon in 1969. He was the first person to walk on the moon.')
query = "what was NASAs most expensive project?"
context = query_pinecone(query, top_k=3)
query = format_query(query, context["matches"])
generate_answer(query)
('The Space Shuttle was the most expensive project in the history of NASA. It '
 'cost about $10 billion to build.')

As we can see, the model can generate some great answers.

Example Application

To try out an example application of abstractive QA, see this demo app.