Update app.py

app.py

@@ -11,10 +11,14 @@ import torch
 from transformers import pipeline
 import pdfminer
 from pdfminer.high_level import extract_text
-#from termcolor import colored
 
-len_doc = 400
-overlap = 50 
+len_doc = 500
+overlap = 30 
+
+qa_model = pipeline("question-answering", 
+                    #model = "deepset/minilm-uncased-squad2")
+                    model = "deepset/roberta-base-squad2")
+                    
 
 def read_pdf(file):
   text = extract_text(file.name)
@@ -50,10 +54,6 @@ def retrieval(query, top_k_retriver, docs, bm25_):
     
     return bm25_hits
 
-qa_model = pipeline("question-answering", 
-                    model = "deepset/minilm-uncased-squad2")
-                    #model = "deepset/roberta-base-squad2")
-                    
 def qa_ranker(query, docs_, top_k_ranker):
     ans = []
     for doc in docs_:
@@ -86,33 +86,27 @@ def final_qa_pipeline(file, query):
     
     bm25 = BM25Okapi(tokenized_corpus)
     
-    top_k_retriver, top_k_ranker = 30,3
+    top_k_retriver, top_k_ranker = 20,3
     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"], str(np.round(100*fnl_rank[0]["score"],2))+"%" , fnl_rank[0]['doc'])
-        #return (print_colored(fnl_rank[0]['doc'], fnl_rank[0]['start'], fnl_rank[0]['end']), str(np.round(100*fnl_rank[0]["score"],2))+"%"
         top1 = print_colored(fnl_rank[0]['doc'], fnl_rank[0]['start'], fnl_rank[0]['end'], str(np.round(100*fnl_rank[0]["score"],1))+"%")  
         if len(lvl1)>1:
             top2 = print_colored(fnl_rank[1]['doc'], fnl_rank[1]['start'], fnl_rank[1]['end'], str(np.round(100*fnl_rank[1]["score"],1))+"%")
         else:
             top2 = "None"
         return (top1, top2)
-        #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","No match")
         
 examples = [
+    [os.path.abspath("dbs-annual-report-2020.pdf"), "how many times has DBS won Best bank in the world ?"],
+    [os.path.abspath("dbs-annual-report-2020.pdf"), "what is PURE ?"],
     [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 ?"],
     [os.path.abspath("dbs-annual-report-2020.pdf"), "what is the sustainability focus ?"],
-    [os.path.abspath("dbs-annual-report-2020.pdf"), "what is PURE ?"],
     [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("NASDAQ_AAPL_2020.pdf"), "who is the chief executive officer ?"],
 ]