Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import os | |
| import numpy as np | |
| os.system("pip install pdfminer.six rank_bm25 torch transformers termcolor") | |
| from gradio.mix import Series | |
| import re | |
| from rank_bm25 import BM25Okapi | |
| import string | |
| import torch | |
| from transformers import pipeline | |
| import pdfminer | |
| from pdfminer.high_level import extract_text | |
| from termcolor import colored | |
| def read_pdf(file): | |
| text = extract_text(file.name) | |
| # Split text into smaller docs | |
| len_doc = 400 | |
| overlap = 50 | |
| docs = [] | |
| i = 0 | |
| while i < len(text): | |
| docs.append(text[i:i+len_doc]) | |
| i = i + len_doc - overlap | |
| return docs | |
| # We use BM25 as retriver which will do 1st round of candidate filtering based on word based matching | |
| def bm25_tokenizer(text): | |
| stop_w = ['a', 'the', 'am', 'is' , 'are', 'who', 'how', 'where', 'when', 'why', 'what'] | |
| tokenized_doc = [] | |
| for token in text.lower().split(): | |
| token = token.strip(string.punctuation) | |
| if len(token) > 0 and token not in stop_w: | |
| tokenized_doc.append(token) | |
| return tokenized_doc | |
| def retrieval(query, top_k_retriver, docs, bm25_): | |
| bm25_scores = bm25_.get_scores(bm25_tokenizer(query)) | |
| top_n = np.argsort(bm25_scores)[::-1][:top_k_retriver] | |
| bm25_hits = [{'corpus_id': idx, | |
| 'score': bm25_scores[idx], | |
| 'docs':docs[idx]} for idx in top_n if bm25_scores[idx] > 0] | |
| bm25_hits = sorted(bm25_hits, key=lambda x: x['score'], reverse=True) | |
| return bm25_hits | |
| qa_model = pipeline("question-answering", | |
| model = "deepset/roberta-base-squad2") | |
| def qa_ranker(query, docs_, top_k_ranker): | |
| ans = [] | |
| for doc in docs_: | |
| answer = qa_model(question = query, | |
| context = doc) | |
| answer['doc'] = doc | |
| ans.append(answer) | |
| return sorted(ans, key=lambda x: x['score'], reverse=True)[:top_k_ranker] | |
| def print_colored(text, start_idx, end_idx): | |
| a = colored(text[:start_idx]) + \ | |
| colored(text[start_idx:end_idx], 'red', 'on_yellow') + \ | |
| colored(text[end_idx:])) | |
| return a | |
| def final_qa_pipeline(file, query): | |
| docs = read_pdf(file) | |
| tokenized_corpus = [] | |
| for doc in docs: | |
| tokenized_corpus.append(bm25_tokenizer(doc)) | |
| bm25 = BM25Okapi(tokenized_corpus) | |
| top_k_retriver, top_k_ranker = 10,1 | |
| lvl1 = retrieval(query, top_k_retriver, docs, bm25) | |
| if len(lvl1) > 0: | |
| fnl_rank = qa_ranker(query, [l["docs"] for l in lvl1], top_k_ranker) | |
| return (fnl_rank[0]["answer"], np.round(fnl_rank[0]["score"],3), print_colored(fnl_rank[0]['doc'], fnl_rank[0]['start'], fnl_rank[0]['end'])) | |
| #for fnl_ in fnl_rank: | |
| # print("\n") | |
| # print_colored(fnl_['doc'], fnl_['start'], fnl_['end']) | |
| # print(colored("Confidence score of ") + colored(str(fnl_['score'])[:4], attrs=['bold'])) | |
| else: | |
| return ("No match", 0) | |
| examples = [ | |
| [os.path.abspath("NASDAQ_AAPL_2020.pdf"), "how much are the outstanding shares ?"], | |
| [os.path.abspath("NASDAQ_AAPL_2020.pdf"), "How high is shareholders equity ?"], | |
| [os.path.abspath("NASDAQ_AAPL_2020.pdf"), "what is competitors strategy ?"], | |
| [os.path.abspath("dbs-annual-report-2020.pdf"), "how much dividend was paid to shareholders ?"], | |
| [os.path.abspath("dbs-annual-report-2020.pdf"), "what are the key risks ?"], | |
| ] | |
| iface = gr.Interface( | |
| fn = final_qa_pipeline, | |
| inputs = [gr.inputs.File(label="input pdf file"), gr.inputs.Textbox(label="Question:")], | |
| outputs = [gr.outputs.TextBox(label="Answer"), gr.outputs.TextBox(label="Score"), gr.outputs.HTML(label="Reference text")], | |
| examples=examples, | |
| ) | |
| iface.launch() |