Spaces:
Runtime error
Runtime error
| 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 | |