Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import os
 import numpy as np
-os.system("pip install pdfminer.six rank_bm25 torch transformers")
+os.system("pip install pdfminer.six rank_bm25 torch transformers termcolor")
 
 from gradio.mix import Series
 import re
@@ -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)
@@ -60,7 +60,13 @@ def qa_ranker(query, docs_, top_k_ranker):
         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 = []
@@ -74,7 +80,7 @@ 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"], 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'])
@@ -93,7 +99,7 @@ examples = [
 iface = gr.Interface(
    fn = final_qa_pipeline,
    inputs = [gr.inputs.File(label="input pdf file"), gr.inputs.Textbox(label="Question:")],
-   outputs = [gr.outputs.HTML(label="Answer"), gr.outputs.HTML(label="Score"), gr.outputs.HTML(label="Reference text")],
+   outputs = [gr.outputs.TextBox(label="Answer"), gr.outputs.TextBox(label="Score"), gr.outputs.HTML(label="Reference text")],
    examples=examples,
    )
 iface.launch()