Search COVID papers with Deep Learning
By: Francesco Zuppichini
Tags: Deep Learning, Programming, Python, Covid 19, Data Science

As you can see, we have a lot of information. We are obviously interested in the text columns. Working with pandas is not ideal, so let’s create a Dataset . This will allow us to later create a DataLoader to perform batch-wise encoding. If you are not familiar with the Pytorch data loading ecosystem you can read more about here

In order, I have subclassed torch.utils.data.Dataset to create a custom dataset. The dataset is expecting a dataframe as input from which we kept only the interesting columns. Then, we removed some of the rows where the abstract and title columns matched one of the "junk" words in FILTER_TITLE and FILTER_ABSTRACT respectively. This is done because articles were scrapped in an automatic fashion, and many have irrelevant entries instead of title/abstract information.

The dataset returns a dictionary since pd.DataFrame is not a supported type in PyTorch. To give our search engine more context, we merge the title and the abstract together, the result is stored in the title_abstract key.

We can now call the dataset and see if everything is correct

ds = CovidPapersDataset.from_path('./dataset/metadata.csv') ds[0]['title']

Output:

'Sequence requirements for RNA strand transfer during nidovirus discontinuous subgenomic RNA synthesis'

Embed

We now need a way to create a vector ( embedding) from each data point. We define a class Embedder that loads automatically a model from HuggingFace's transformers using the sentence_transformers library.

The model of choice is gsarti/biobert-nli a BioBERT model fine-tuned on the SNLI and the MultiNLI to produce universal sentence embeddings. Gabriele Sarti made the finetuning, the code to reproduce it is available here.

BioBERT is especially fit for our dataset since it was originally trained on biomedical scientific publications. So, it should create better context-aware embeddings given the similarity with our data.

Under the hood, the model first tokenizes the input string in tokens, then it creates one vector for each one of them. So, if we have N tokens in one paper we will get a [N, 768] vector (note that a token often corresponds to a word piece, read more about tokenization strategies here. Thus, if two papers have a different word size, we will have two vectors with two different first dimensions. This is a problem since we need to compare them to search.

To get a fixed embed for each paper, we apply average pooling. This methodology computes the average of each word and outputs a fixed-size vector of dims [1, 768]

So, let’s code an Embedder class

We can try our embedder on a data point

embedder = Embedder() emb = embedder(ds[0]['title_abstract']) emb[0].shape // (768,)

Et voilà! We encoded one paper.

Search

Okay, we know how to embed each paper, but how can we search in the data using a query? Assuming we have embedded all the papers we could also embed the query and compute the cosine similarity between the query and all the embeddings. Then, we can show the results sorted by the distance (score). Intuitively, the closer they are in the embedding space to the query the more context similarity they share.

But, how? First, we need a proper way to manage the data and to run the cosine similarity fast enough. Fortunately, Elastic Search comes to the rescue!

Elastic Search

Elastic Search is a database with one goal, yes you guessed right: search. We will first store all the embedding in elastic and then use its API to perform the searching. If you are lazy like me you can install elastic search with docker

docker pull docker.elastic.co/elasticsearch/elasticsearch:7.6.2

docker run -p 9200:9200 -p 9300:9300 -e "discovery.type=single-node" docker.elastic.co/elasticsearch/elasticsearch:7.6.2

Perfect. The next step is to store the embeddings and the papers’ information on elastic search. It is a very straightforward process. We have to need to create an index (a new database) and then build one entry for each paper.

To create an index we need to describe for elastic what we wish to store. In our case:

You can read more about the index creation on the elastic search doc. The last entry defines the embed field as a dense vector with 768 . This is indeed our embedding. For convenience, I have stored the configuration in a .json file and created a class named ElasticSearchProvider to handle the storing process.

Most of the work is done in create_and_bulk_documents where we just deconstruct one entry at the time and add two elastic search parameters.

Unfortunately, Elastic Search won’t be able to serialize the numpy arrays. So we need to create an adapter for our data. This class takes as input the paper data and the embedding and "adapt" them to work in our ElasticSearchProvider .

Okay, we have everything in place. A way to represent the data, one to encode it in a vector and a method to store the result. Let’s wrap everything up and encode all the papers.

There are two tricks here, first, we use torch.utils.data.DataLoader to create a batch-wise iterator. In general, feeding data to the model in batch rather than as a single point boost the performance (in my case x100). Second, we replace the collate_fn parameter in the DataLoader constructor. This is because, by default, Pytorch will try to cast all our data into a torch.Tensor but it will fail to convert strings. By doing so, we just return an array of dictionaries, the output from CovidPapersDataset . So, a batch is a list of dictionaries with length batch_size . After we finished (~7m on a 1080ti), we can take a look at http://localhost:9200/covid/_search?pretty=true&q=*:* .

If everything works correctly, you should see our data displayed by elastic search

Make a query

We are almost done. The last piece of the puzzle is a way to search in the database. Elastic search can perform cosine similarity between one input vector and a target vector field in all the documents. The syntax is very straightforward:

{

"query": {

"match_all": {}

},

"script": {

"source":

"cosineSimilarity(params.query_vector, doc['embed']) + 1.0",

"params": {

"query_vector": vector

}

}

}

Where vector is our input. So, we created a class that takes a vector as an input and show all the results from the query

Let’s see the first result (I have copied and pasted the first matching paper’s abstract)

es_search = ElasticSearcher()

es_search(embedder(['Effect of the virus on pregnant women'])[0].tolist())

As public health professionals respond to emerging infections, particular attention needs to be paid to pregnant women and their offspring. Pregnant women might be more susceptible to, or more severely affected by, emerging infections. The effects of a new maternal infection on the embryo or fetus are difficult to predict. Some medications recommended for prophylaxis or treatment could harm the embryo or fetus. We discuss the challenges of responding to emerging infections among pregnant women, and we propose strategies for overcoming these challenges.

It worked! I have also created a command-line where the user can type a query. The final result is:

More queries

Finally, we can ask a question and find interesting papers. Empirically, queries with more details work better since they provide more context.

For example, we may want to know the What is the effectiveness of chloroquine for COVID-19. The results are

Or How does COVID-19 bind to the ACE2 receptor?

The search engine seems to work good but it is not perfect, in the next part of this tutorial we will try to improve its accuracy.

Conclusions

In this project, we build a semantic browser to search on more than 50k COVID-19 papers. The original project in which I have been working with students from the Universita of Trieste is here. A live demo is available here

You can also play around with the command line app, you need to follow the instruction from here.

Acknowledgment

I would like to thank Gabriele Sarti for helping me in the writing of this article, Marco Franzon and Tommaso Rodani for supporting me in the elastic search implementation.

Thank you for reading

Be safe,

Francesco Saverio Zuppichini