Update app.py

app.py

@@ -11,7 +11,7 @@ import torch
 from transformers import pipeline
 import pdfminer
 from pdfminer.high_level import extract_text
-from termcolor import colored
+#from termcolor import colored
 
 def read_pdf(file):
   text = extract_text(file.name)
@@ -61,8 +61,12 @@ def qa_ranker(query, docs_, top_k_ranker):
         ans.append(answer)
     return sorted(ans, key=lambda x: x['score'], reverse=True)[:top_k_ranker]
 
+def cstr(s, color='black'):
+    return "<text style=color:{}>{}</text>".format(color, s)
+
 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:])
+    a = cstr(' '.join([text[:start_idx], cstr(text[start_idx:end_idx], color='red'), text[end_idx:]]), color='black')
+    #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):
@@ -78,7 +82,8 @@ def final_qa_pipeline(file, query):
 
     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), fnl_rank[0]['doc'])
+        #return (fnl_rank[0]["answer"], str(np.round(100*fnl_rank[0]["score"],2))+"%" , fnl_rank[0]['doc'])
+        return (fnl_rank[0]["answer"], str(np.round(100*fnl_rank[0]["score"],2))+"%" , print_colored(fnl_rank[0]['doc'], fnl_rank[0]['start_idx'], fnl_rank[0]['end_idx']))
         #for fnl_ in fnl_rank:
         #    print("\n")
         #    print_colored(fnl_['doc'], fnl_['start'], fnl_['end'])
@@ -87,17 +92,18 @@ def final_qa_pipeline(file, query):
         return ("No match", 0)
         
 examples = [
+    [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("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")],
+   outputs = [gr.outputs.Textbox(label="Answer"), gr.outputs.Textbox(label="Confidence"), gr.outputs.HTML(label="Reference text")],
    examples=examples,
    title = "Question Answering on company annual reports",
    description = "Simply upload any annual report pdf you are interested in and ask model a question OR load an example from below."