Spaces:
Runtime error
Runtime error
File size: 3,283 Bytes
9f23e0b 4e8a231 9f23e0b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 |
import numpy as np
import nltk
nltk.download('stopwords')
from nltk.corpus import stopwords
from nltk.tokenize import RegexpTokenizer
class QueryProcessor:
def __init__(self, question, section_texts, N, avg_doc_len):
self.section_texts = section_texts
self.N = N
self.avg_doc_len = avg_doc_len
# self.bm25_scores = {}
self.query_items = self.set_query(question)
self.section_document_idx = None
def set_query(self, question):
punct_regex = RegexpTokenizer(r'\w+')
return [q for q in punct_regex.tokenize(question.lower()) if q not in stopwords.words('english')]
def get_query(self):
return self.query_items
def bm25(self, word, paragraph, k=1.2, b=0.75):
# frequency of word (word) in doc (paragraph)
freq = paragraph.split().count(word)
# term frequency
tf = (freq * (k+1)) / (freq + k * (1 - b + b * len(paragraph.split()) / self.avg_doc_len))
# number of docs that contain the word
N_q = sum([1 for _, docs in self.section_texts.items() for doc in docs if word in doc.split()])
# inverse document frequency
idf = np.log(((self.N - N_q + 0.5) / (N_q + 0.5)) + 1)
return round(tf*idf, 4)
def get_bm25_scores(self):
bm25_scores = {}
for query in self.query_items:
bm25_scores[query] = {}
for section, docs in self.section_texts.items():
bm25_scores[query][section] = {}
for doc_index in range(len(docs)):
score = self.bm25(query, docs[doc_index])
if score > 0.0:
bm25_scores[query][section][doc_index] = score
if len(bm25_scores[query][section]) <= 0:
del bm25_scores[query][section]
return bm25_scores
def filter_bad_documents(self, bm25_scores):
section_document_idx = {}
for sec_docs in bm25_scores.values():
for sec, doc_scores in sec_docs.items():
if sec not in section_document_idx:
section_document_idx[sec] = []
for doc_idx, score in doc_scores.items():
if score > 0.5 and doc_idx not in section_document_idx[sec]:
section_document_idx[sec].append(doc_idx)
if len(section_document_idx[sec]) <= 0:
del section_document_idx[sec]
return section_document_idx
def get_context(self):
bm25_scores = self.get_bm25_scores()
self.section_document_idx = self.filter_bad_documents(bm25_scores)
# print(bm25_scores)
context = ' '.join([self.section_texts[section][d_id] for section, doc_ids in self.section_document_idx.items() for d_id in doc_ids])
# print(section_document_idx)
return context
def match_section_with_answer_text(self, text):
# print(text)
sections = []
for sec, doc_ids in self.section_document_idx.items():
for d_id in doc_ids:
if self.section_texts[sec][d_id].find(text) > -1:
sections.append(sec)
return sections
|