text_plain_output_5.png text_plain_output_3.png text_plain_output_1.png
from sklearn.feature_extraction.text import TfidfVectorizer

from sklearn.metrics.pairwise import cosine_similarity

from transformers import BertTokenizer, BertForQuestionAnswering

from wasabi import msg

import json

import os

import pandas as pd

import time

import torch

import pandas as pd
import json
import os
meta_path = '/kaggle/input/CORD-19-research-challenge/metadata.csv'
def get_data_texts():
    meta_data = pd.read_csv(meta_path, low_memory=True)
    paperID2year = {}
    for _, meta_row in meta_data.iterrows():
        if pd.notnull(meta_row['pmcid']):
            paperID2year[meta_row['pmcid']] = meta_row['publish_time']
        if pd.notnull(meta_row['sha']):
            paper_ids = meta_row['sha'].split('; ')
            for paper_id in paper_ids:
                paperID2year[paper_id] = meta_row['publish_time']
    data_text = {}
    index2paperID = {}
    index2paperPath = {}
    i = 0
    for dirname, _, filenames in os.walk('/kaggle/input/CORD-19-research-challenge'):
        for filename in filenames:
            paper_path = os.path.join(dirname, filename)
            if paper_path[-4:] != 'json':
                continue
            with open(paper_path) as json_file:
                article_data = json.load(json_file)
                paper_date = paperID2year.get(article_data['paper_id'], None)
                if paper_date:
                    if paper_date[0:4] == '2020':
                        data_text[article_data['paper_id']] = ' '.join([d['text'] for d in article_data['body_text']])
                        index2paperID[i] = article_data['paper_id']
                        index2paperPath[i] = paper_path
                        i += 1
    return (data_text, index2paperID, index2paperPath)
data_text, index2paperID, index2paperPath = get_data_texts()

from transformers import BertTokenizer, BertForQuestionAnswering
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import torch
from wasabi import msg
import time
class QuestionCovid:
    def __init__(self, TOKENIZER, MODEL, index2paperID, index2paperPath):
        self.TOKENIZER = TOKENIZER
        self.MODEL = MODEL
        self.index2paperID = index2paperID
        self.index2paperPath = index2paperPath
    def fit(self, data_text):
        self.TFIDF_VECTORIZER = TfidfVectorizer()
        with msg.loading('   Fitting TFIDF'):
            start = time.time()
            self.TFIDF_VECTORIZER.fit(data_text.values())
        msg.good('   TFIDF fitted - Took {:.2f}s'.format(time.time() - start))
        with msg.loading('   Creating Articles matrix'):
            start = time.time()
            self.ARTICLES_MATRIX = self.TFIDF_VECTORIZER.transform(data_text.values())
        msg.good('   Article matrix created - Took {:.2f}s'.format(time.time() - start))
    def get_answer(self, text, question):
        input_text = '[CLS] ' + question + ' [SEP] ' + text + ' [SEP]'
        input_ids = self.TOKENIZER.encode(input_text)
        token_type_ids = [0 if i <= input_ids.index(102) else 1 for i in range(len(input_ids))]
        start_scores, end_scores = self.MODEL(torch.tensor([input_ids]), token_type_ids=torch.tensor([token_type_ids]))
        all_tokens = self.TOKENIZER.convert_ids_to_tokens(input_ids)
        answer = ' '.join(all_tokens[torch.argmax(start_scores):torch.argmax(end_scores) + 1])
        score = round(start_scores.max().item(), 2)
        return (answer, score)
    def predict(self, question):
        query = self.TFIDF_VECTORIZER.transform([question + ' covid'])
        best_matches = sorted([(i, c) for i, c in enumerate(cosine_similarity(query, self.ARTICLES_MATRIX).ravel())], key=lambda x: x[1], reverse=True)
        for i, tfidf_score in best_matches[:5]:
            best_score = 0
            best_answer = 'No answer'
            best_text = 'No snippet'
            paper_path = self.index2paperPath[i]
            with open(paper_path) as json_file:
                article_data = json.load(json_file)
                text = ' '.join([d['text'] for d in article_data['body_text']])
            sentences = text.split('.')
            n = 3
            sentences_grouped = ['.'.join(sentences[i:i + n]) for i in range(0, len(sentences), n)]
            for subtext in sentences_grouped:
                answer, score = self.get_answer(subtext, question)
                if score > best_score:
                    best_score = score
                    best_answer = answer
                    best_text = subtext
            yield (self.index2paperID[i], best_answer, best_score, best_text, tfidf_score)

TOKENIZER = BertTokenizer.from_pretrained('bert-base-uncased')
MODEL = BertForQuestionAnswering.from_pretrained('bert-large-uncased-whole-word-masking-finetuned-squad')

covid_q = QuestionCovid(TOKENIZER, MODEL, index2paperID, index2paperPath)
covid_q.fit(data_text)

question = 'How long is the incubation period for the virus?'

for i, (paper_id, answer, score, snippet, tfidf_score) in enumerate(covid_q.predict(question)):
    print(f'Answer {i}: {answer}')
    print(f'Text segment: {snippet}')
    print(f'Paper id: {paper_id}')