Building a Powerful Question Answering API with Flask, Docker, BERT, and Haystack

Introduction

As natural language processing (NLP) technologies continue to advance, there is growing demand for tools and applications that can intelligently process text data and provide human-like responses. One area that has seen rapid progress in recent years is question answering (QA) – the ability for a computer to understand a question posed in natural language and locate or generate a relevant answer from a corpus of text.

QA has many practical applications, from powering chatbots and virtual assistants to helping users quickly find information within large documents or databases. While building a QA system from scratch requires significant expertise in NLP and machine learning, open source frameworks like Haystack now make it easier than ever for developers to spin up powerful, production-ready QA APIs.

In this tutorial, we‘ll walk through the process of building a QA API using Flask, Docker, a BERT-based language model, and the Haystack framework. By the end, you‘ll have a fully containerized API that can ingest a corpus of documents, fine-tune a QA model on your own custom dataset, and provide accurate, extractive answers to natural language questions. Let‘s dive in!

Overview of the Technologies

Before we get started with the hands-on portion of this tutorial, let‘s take a quick look at the key technologies we‘ll be using:

• Flask: Flask is a lightweight, Python-based web application framework that makes it easy to build APIs and web apps. We‘ll use Flask to define the endpoints of our QA API and handle requests.

• Docker: Docker is a platform for building, shipping, and running applications in containers. Containerizing our QA API with Docker will make it easy to deploy and scale in any environment.

• BERT: BERT (Bidirectional Encoder Representations from Transformers) is a state-of-the-art language model developed by Google. Pre-trained on a massive corpus of text data, BERT can be fine-tuned for a variety of NLP tasks, including question answering. We‘ll be using a pre-trained BERT model as the foundation of our QA system.

• Haystack: Haystack is an open source framework for building end-to-end question answering systems. It provides a set of flexible, extensible tools for processing documents, retrieving relevant passages, and extracting granular answers to questions. Haystack will serve as the backbone of our QA pipeline.

With these powerful tools in our arsenal, let‘s start building our QA API!

Setting Up the Development Environment

To keep things simple and reproducible, we‘ll be using Docker to containerize our API and all of its dependencies. This will allow us to develop and test locally in a consistent environment that matches what we‘ll eventually deploy to production.

First, make sure you have Docker installed on your machine. If you don‘t already have it, you can download Docker Desktop from the official website: https://www.docker.com/products/docker-desktop

Next, create a new directory for your project and navigate to it in your terminal:

mkdir qa-api
cd qa-api

In this directory, create a file called Dockerfile and paste in the following contents:

FROM python:3.8-slim-buster

WORKDIR /app

COPY requirements.txt .
RUN pip install -r requirements.txt

COPY . .

EXPOSE 5000

CMD ["gunicorn", "-b", "0.0.0.0:5000", "main:app"]

This Dockerfile specifies the base image for our container (Python 3.8), copies our application code into the /app directory, installs the required Python packages from a requirements.txt file, exposes port 5000, and defines the command to run when the container starts up.

Create a requirements.txt file and add the following dependencies:

flask
gunicorn
farm-haystack

This will install Flask for our API, Gunicorn as a production-grade web server, and the farm-haystack package which contains the core Haystack framework.

We‘re now ready to write the code for our API!

Building the API with Flask and Haystack

Create a new file called main.py which will contain the Flask app and API endpoints:

from flask import Flask, request, jsonify
from haystack.document_store.elasticsearch import ElasticsearchDocumentStore
from haystack.retriever.dense import DensePassageRetriever
from haystack.reader.farm import FARMReader

# Connect to Elasticsearch 
document_store = ElasticsearchDocumentStore(
    host="localhost",
    username="",
    password="",
    index="document",
    embedding_field="emb",
    embedding_dim=768,
    excluded_meta_data=["emb"]
)

# Initialize retriever and reader 
retriever = DensePassageRetriever(document_store)
reader = FARMReader(model_name_or_path="distilbert-base-uncased-distilled-squad", use_gpu=True)

app = Flask(__name__)

@app.route("/upload", methods=["POST"])
def upload_document():
    # Get document from request
    document = request.files["document"]

    # Add document to DocumentStore
    document_store.add_document(document)

    return jsonify({"status": "success"})

@app.route("/query", methods=["GET"])
def query():
    # Get question from request
    question = request.args.get("question", None)
    top_k = request.args.get("top_k", 3)

    # Retrieve relevant documents from store
    retrieved_docs = retriever.retrieve(query=question, top_k=top_k)

    # Get predicted answers from reader
    predictions = reader.predict(question=question, documents=retrieved_docs)

    return jsonify(predictions)

if __name__ == "__main__":  
    app.run(host="0.0.0.0")

Let‘s break this down:

• First, we import the necessary modules from Flask and Haystack. The ElasticsearchDocumentStore will serve as our knowledge base for storing and indexing documents, while the DensePassageRetriever and FARMReader will handle retrieving relevant documents and extracting granular answers from them.

• We establish a connection to Elasticsearch and initialize our document_store. You‘ll need to have Elasticsearch running locally on the default port (9200).

• We create instances of our retriever and reader, specifying that we want to use a DistilBERT model fine-tuned for QA on the SQuAD dataset. This is a great general purpose model that balances performance and efficiency.

• We define our Flask app and two API endpoints:

  • /upload expects a POST request with a document file. It simply passes this document along to the document_store to be indexed for later retrieval.
  • /query expects a GET request with a question parameter containing a natural language query. It first uses the retriever to fetch the most relevant documents for the query, then passes those documents and the original question to the reader to extract a final answer. The predicted answers are returned in the HTTP response as JSON.

With just this simple skeleton, we already have a functional QA API that can be called via HTTP to upload documents and answer questions! Of course, there are many ways we can extend and optimize this further.

Containerizing and Running the API

To run our newly created API, we simply need to build and run the Docker container. From the qa-api directory, run:

docker build -t qa-api .
docker run -p 5000:5000 qa-api

This will build a Docker image tagged qa-api based on our Dockerfile, then start a container running our Flask app. The -p flag maps port 5000 in the container to port 5000 on our local machine, so we can access the API at http://localhost:5000.

To test it out, first upload a document:

curl -X POST -F ‘document=@/path/to/document.txt‘ http://localhost:5000/upload

Then ask a question:

curl http://localhost:5000/query?question=What+is+the+capital+of+France?

If everything is working, you should get back a JSON response with the predicted answer(s)!

Fine-Tuning the QA Model

While the pre-trained DistilBERT model does a decent job out-of-the-box, we can significantly improve the performance of our QA system by fine-tuning the model on our own domain-specific dataset. This will teach the model to better understand the context and language of the types of questions we expect to receive.

To generate a dataset for fine-tuning, we can use the handy annotation tool provided by deepset, the creators of Haystack. This is a simple web-based UI that allows you to manually label question/answer pairs on your documents.

First, export your documents from the Elasticsearch document_store in JSON format:

import json

docs = document_store.get_all_documents()
with open("documents.json", "w") as f:
    json.dump(docs, f)

Then, head over to https://annotate.deepset.ai/index.html, create a new project, and upload your documents.json file. You can now start highlighting questions and answers in the documents to create training examples. Aim for a few dozen high-quality annotations to start.

When you‘re done, export your annotations in SQuAD JSON format. This is a standard format used for question answering datasets.

To fine-tune your model, we‘ll use the Haystack ModelTrainer API:

from haystack.trainer import ModelTrainer

trainer = ModelTrainer(
    model=reader,
    data_dir="data/squad_format",
    train_filename="annotations.json",
    use_gpu=True,
    n_epochs=1,
    learning_rate=1e-5, 
    batch_size=16
)

trainer.train()

# Save the fine-tuned model
reader.model.model.save_pretrained("my_model")

This will load our annotated data, train the model for a single epoch with a batch size of 16, and save the resulting model to disk. We can now update our Flask app to load the fine-tuned model:

reader = FARMReader(model_name_or_path="my_model", use_gpu=True)

With just a small amount of annotated data, you should notice a significant improvement in the relevance and accuracy of the answers returned by the model. Of course, the more high-quality examples you provide, the better the model will perform on your specific domain.

Deploying to Production

Once you‘re happy with your fine-tuned model, it‘s time to deploy your API to production. Since we‘ve packaged everything neatly in a Docker container, this is relatively straightforward.

One option is to deploy to a managed container platform like Google Cloud Run or AWS Fargate. These platforms will handle the scaling and infrastructure management for you, so you can focus on your application code.

Alternatively, you can deploy to a Kubernetes cluster for maximum flexibility and control. Create a Kubernetes Deployment manifest for your qa-api container, set up a load balancer Service to route traffic to it, and use a tool like Helm or Kustomize to manage the release.

Whichever method you choose, be sure to configure logging and monitoring for your API so you can keep tabs on performance and quickly identify and fix any issues. Tools like Prometheus and Grafana are great for this.

Conclusion

In this tutorial, we‘ve seen how to build a production-ready question answering API with Flask, Docker, BERT, and Haystack. We walked through the key components of Haystack, including the DocumentStore for storing and indexing documents, the Retriever for finding relevant passages, and the Reader for extracting granular answers.

We then showed how to train a custom QA model using annotated data and package everything up in a Docker container for easy deployment. The end result is a powerful, flexible API that can provide intelligent answers to user questions across a wide range of domains.

Of course, this is just the beginning. There are many ways you can extend and improve this system, such as:

• Experimenting with different retriever and reader models to optimize performance
• Implementing a user feedback loop to collect new training examples over time
• Adding support for multiple languages
• Integrating with other data sources and APIs to enrich your knowledge base
• Building a full-fledged chatbot interface on top of the QA API

The possibilities are endless! I encourage you to use this tutorial as a starting point and continue exploring the exciting world of question answering and NLP. With the right tools and a bit of creativity, you can build some truly amazing applications. Happy coding!