|
|
| import pandas as pd |
| import numpy as np |
| 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 |
|
|
|
|
| def get_full_sentence(spacy_nlp, para_text, start_index, end_index): |
| """ |
| Returns the relative sentence of original text, |
| given a specific paragraph (body text). |
| """ |
| sent_start = 0 |
| sent_end = len(para_text) |
| for sent in spacy_nlp(para_text).sents: |
| if (sent.start_char <= start_index) and (sent.end_char >= start_index): |
| sent_start = sent.start_char |
| if (sent.start_char <= end_index) and (sent.end_char >= end_index): |
| sent_end = sent.end_char |
| sentence = para_text[sent_start:sent_end + 1] |
| return sentence |
|
|
|
|
| def fetch_stage1(query, model, list_of_articles): |
| """ |
| Compare all the articles' abstract content with each query |
| """ |
|
|
| |
| query_embedding = model.encode([query])[0] |
|
|
| |
| all_abs_distances = [] |
|
|
| for idx_of_article,article in enumerate(list_of_articles): |
| if article: |
| distances = [] |
| cdists = scipy.spatial.distance.cdist([query_embedding], np.vstack(article), "cosine").reshape(-1,1) |
| for idx,sentence in enumerate(article): |
| distances.append((idx, 1 - cdists[idx][0])) |
|
|
| results = sorted(distances, key=lambda x: x[1], reverse=True) |
| if results: |
| all_abs_distances.append((idx_of_article, results[0][0], results[0][1])) |
|
|
| results = sorted(all_abs_distances, key=lambda x: x[2], reverse=True) |
|
|
| return query_embedding, results |
|
|
|
|
| def fetch_stage2(results, model, embeddings, query_embedding): |
| """ |
| Take the 20 most similar articles, based on the relevant abstracts and |
| compare all the body texts content to the query |
| """ |
| |
| all_text_distances = [] |
| for top in results[0:20]: |
| article_idx = top[0] |
|
|
| body_texts = [text[0] for text in embeddings[article_idx][2]] |
| body_text_embeddings = model.encode(body_texts, show_progress_bar=False) |
|
|
| |
| |
|
|
| qbody = scipy.spatial.distance.cdist([query_embedding], |
| np.vstack(body_text_embeddings), |
| "cosine").reshape(-1,1) |
| |
| body_text_distances = [(idx, 1-dist[0]) for idx,dist in enumerate(qbody)] |
| |
| |
| |
| |
|
|
| |
| |
| |
|
|
| results = sorted(body_text_distances, key=lambda x: x[1], reverse=True) |
|
|
| if results: |
| all_text_distances.append((article_idx, results[0][0], results[0][1])) |
|
|
| results = sorted(all_text_distances, key=lambda x: x[2], reverse=True) |
|
|
| return results |
|
|
|
|
| def fetch_stage3(results, query, embeddings, comprehension_model, spacy_nlp): |
| """ |
| For the top 20 retrieved paragraphs in the document, |
| answer will be comprehended on each paragraph using the model. |
| """ |
| |
| answers = [] |
|
|
| |
|
|
| for top_text in results[0:20]: |
| article_idx = top_text[0] |
| body_text_idx = top_text[1] |
|
|
| query_ = {"context": embeddings[article_idx][2][body_text_idx][0], "question": query} |
| pred = comprehension_model(query_, topk=1, show_progress_bar=False) |
|
|
| |
| if pred["answer"] and round(pred["score"], 4) > 0: |
| |
| sent = get_full_sentence(spacy_nlp, query_['context'], pred["start"], pred["end"]) |
| answers.append((article_idx, round(pred["score"], 4), sent)) |
|
|
| results = sorted(answers, key=lambda x: x[1], reverse=True) |
|
|
| return results |
|
|
