|
import pandas as pd |
|
import numpy as np |
|
import datetime, time |
|
import pickle |
|
import glob |
|
import json |
|
from pandas.io.json import json_normalize |
|
from nltk.tokenize import sent_tokenize |
|
import nltk |
|
import scipy.spatial |
|
from transformers import AutoTokenizer, AutoModel, pipeline, AutoModelForQuestionAnswering |
|
from sentence_transformers import models, SentenceTransformer |
|
import torch |
|
import spacy |
|
import subprocess |
|
import streamlit as st |
|
from utils import * |
|
|
|
|
|
@st.cache(allow_output_mutation=True) |
|
def load_spacy_model(): |
|
subprocess.call(['python', '-m','spacy', 'download', 'en_core_web_sm']) |
|
|
|
@st.cache(allow_output_mutation=True) |
|
def load_prep_data(): |
|
with open('listfile_3.data', 'rb') as filehandle: |
|
articles = pickle.load(filehandle) |
|
|
|
for article in range(len(articles)): |
|
if articles[article][1] != []: |
|
articles[article][1] = sent_tokenize(articles[article][1]) |
|
|
|
return articles |
|
|
|
@st.cache(allow_output_mutation=True) |
|
def build_sent_trans_model(): |
|
|
|
word_embedding_model = models.BERT('./') |
|
|
|
|
|
pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension(), |
|
pooling_mode_mean_tokens=True, |
|
pooling_mode_cls_token=False, |
|
pooling_mode_max_tokens=False) |
|
|
|
model = SentenceTransformer(modules=[word_embedding_model, pooling_model]) |
|
return model |
|
|
|
@st.cache(allow_output_mutation=True) |
|
def load_embedded_articles(): |
|
with open('list_of_articles.pkl', 'rb') as f: |
|
list_of_articles = pickle.load(f) |
|
|
|
return list_of_articles |
|
|
|
@st.cache(allow_output_mutation=True) |
|
def load_comprehension_model(): |
|
|
|
comprehension_model = pipeline("question-answering", |
|
model=AutoModelForQuestionAnswering.\ |
|
from_pretrained("graviraja/covidbert_squad"), |
|
tokenizer=AutoTokenizer.\ |
|
from_pretrained("graviraja/covidbert_squad"), |
|
device=-1) |
|
|
|
return comprehension_model |
|
|
|
|
|
|
|
def main(): |
|
|
|
nltk.download('punkt') |
|
|
|
load_spacy_model() |
|
|
|
spacy_nlp = spacy.load('en_core_web_sm') |
|
|
|
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') |
|
|
|
|
|
|
|
embeddings = load_prep_data() |
|
|
|
model = build_sent_trans_model() |
|
model.to(device) |
|
|
|
list_of_articles = load_embedded_articles() |
|
|
|
comprehension_model = load_comprehension_model() |
|
|
|
st.title('Co-Search') |
|
|
|
query = st.text_input("Enter Query",'What are the corona viruses?', key="query") |
|
|
|
st.write('Using device type: {}'.format(device)) |
|
|
|
with st.spinner('Please wait...'): |
|
dt1 = datetime.datetime.now() |
|
|
|
query_embedding, results1 = fetch_stage1(query, model, list_of_articles) |
|
results2 = fetch_stage2(results1, model, embeddings, query_embedding) |
|
results3 = fetch_stage3(results2, query, embeddings, comprehension_model, spacy_nlp) |
|
|
|
dt2 = datetime.datetime.now() |
|
tdelta = dt2-dt1 |
|
|
|
st.write('Time taken in minutes: %.2f' % (tdelta.seconds/60)) |
|
|
|
if results3: |
|
count = 1 |
|
|
|
for res in results3: |
|
st.write('{}> {}'.format(count, res[2])) |
|
st.write('Score: %.4f' % (res[1])) |
|
st.write("From the article with title:\n{}".format(embeddings[res[0]][0])) |
|
st.write("\n") |
|
|
|
if count > 3: |
|
break |
|
count += 1 |
|
else: |
|
st.info("There isn't any answer") |
|
|
|
st.success('Done!') |
|
|
|
|
|
if __name__ == '__main__': |
|
main() |
|
|
