Update app.py

app.py

@@ -24,26 +24,22 @@ if st.button('Run semantic question answering'):
         except Exception as e:
             qa_result = str(e)
 
-        # top_5_hits = kws_result['hits']['hits'][:5] # print("First 5 results:")
         top_10_hits = kws_result['hits']['hits'][:10] # print("First 10 results:")
         top_5_text = [{'text': hit['_source']['content'][:500],
                        'confidence': hit['_score']} for hit in top_10_hits[:5] ]
-        top_5_para = [hit['_source']['content'] for hit in top_10_hits[:5]]
-        # top_3_para = [hit['_source']['content'] for hit in top_10_hits[:3]]
-        # top_5_para = [hit['_source']['content'][:5000] for hit in top_5_hits]
+        top_3_para = [hit['_source']['content'][:5000] for hit in top_10_hits[:3]]
+        # TODO: split + re-rank
 
         DPR_MODEL = "deepset/roberta-base-squad2" #, model="distilbert-base-cased-distilled-squad"
         pipe_exqa = pipeline("question-answering", model=DPR_MODEL) 
-        qa_results = [pipe_exqa(question=question, context=paragraph) for paragraph in top_5_para]
-        # qa_results = [pipe_exqa(question=question, context=paragraph) for paragraph in top_3_para]
+        qa_results = [pipe_exqa(question=question, context=paragraph) for paragraph in top_3_para]
         qa_results = sorted(qa_results, key=lambda x: x['score'], reverse=True)
         
         for i, qa_result in enumerate(qa_results):
             if "answer" in qa_result.keys(): # and qa_result["answer"] is not ""
                 answer_span, answer_score = qa_result["answer"], qa_result["score"]
                 st.write(f'Answer: **{answer_span}**')
-                # paragraph = top_3_para[i]
-                paragraph = top_5_para[i]
+                paragraph = top_3_para[i]
                 start_par, stop_para = max(0, qa_result["start"]-86), min(qa_result["end"]+90, len(paragraph))
                 answer_context = paragraph[start_par:stop_para].replace(answer_span, f'**{answer_span}**')
                 qa_result.update({'context': answer_context, 'paragraph': paragraph})