Spaces:
Runtime error
Runtime error
# -*- coding: utf-8 -*- | |
""" | |
Created on Wed Jan 4 05:56:28 2023 | |
@author: dreji18 | |
""" | |
# loading the packages | |
from rake_nltk import Rake | |
import wikipedia | |
from rank_bm25 import BM25Okapi | |
import torch | |
from transformers import DistilBertTokenizer, DistilBertForQuestionAnswering | |
from fastapi import FastAPI | |
app = FastAPI() | |
def read_root(): | |
return {"Hello": "World"} | |
# keyword extraction function | |
def keyword_extractor(query): | |
""" | |
Rake has some features: | |
1. convert automatically to lower case | |
2. extract important key phrases | |
3. it will extract combine words also (eg. Deep Learning, Capital City) | |
""" | |
r = Rake() # Uses stopwords for english from NLTK, and all puntuation characters. | |
r.extract_keywords_from_text(query) | |
keywords = r.get_ranked_phrases() # To get keyword phrases ranked highest to lowest. | |
return keywords | |
# data collection using wikepedia | |
def data_collection(search_words): | |
"""wikipedia""" | |
search_query = ' '.join(search_words) | |
wiki_pages = wikipedia.search(search_query, results = 5) | |
information_list = [] | |
pages_list = [] | |
for i in wiki_pages: | |
try: | |
info = wikipedia.summary(i) | |
if any(word in info.lower() for word in search_words): | |
information_list.append(info) | |
pages_list.append(i) | |
except: | |
pass | |
original_info = information_list | |
information_list = [item[:1000] for item in information_list] # limiting the word len to 512 | |
return information_list, pages_list, original_info | |
# document ranking function | |
def document_ranking(documents, query, n): | |
"""BM25""" | |
try: | |
tokenized_corpus = [doc.split(" ") for doc in documents] | |
bm25 = BM25Okapi(tokenized_corpus) | |
tokenized_query = query.split(" ") | |
doc_scores = bm25.get_scores(tokenized_query) | |
datastore = bm25.get_top_n(tokenized_query, documents, n) | |
except: | |
pass | |
return datastore | |
def qna(context, question): | |
"""DistilBert""" | |
tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased',return_token_type_ids = True) | |
model = DistilBertForQuestionAnswering.from_pretrained('distilbert-base-uncased-distilled-squad', return_dict=False) | |
encoding = tokenizer.encode_plus(question, context) | |
input_ids, attention_mask = encoding["input_ids"], encoding["attention_mask"] | |
start_scores, end_scores = model(torch.tensor([input_ids]), attention_mask=torch.tensor([attention_mask])) | |
ans_tokens = input_ids[torch.argmax(start_scores) : torch.argmax(end_scores)+1] | |
answer_tokens = tokenizer.convert_ids_to_tokens(ans_tokens , skip_special_tokens=True) | |
answer_tokens_to_string = tokenizer.convert_tokens_to_string(answer_tokens) | |
return answer_tokens_to_string | |
def answergen(search_string: str): | |
try: | |
keyword_list = keyword_extractor(search_string) | |
information, pages, original_data = data_collection(keyword_list) | |
datastore = document_ranking(information, search_string, 3) | |
answers_list = [] | |
for i in range(len(datastore)): | |
result = qna(datastore[i], search_string) | |
answers_list.append(result) | |
return {"answer 1": answers_list[0], | |
"answer 2": answers_list[1], | |
"answer 3": answers_list[2]} | |
except: | |
return {"sorry couldn't process the request"} | |
#uvicorn app:app --port 8000 --reload | |
#%% | |