add progress streaming function

handler.py

@@ -1,5 +1,6 @@
 from minicheck_web.minicheck import MiniCheck
 from web_retrieval import *
+from flask import Response
 
 
 def sort_chunks_single_doc_claim(used_chunk, support_prob_per_chunk):
@@ -39,32 +40,31 @@ class EndpointHandler():
             assert len(data['inputs']['claims']) == 1, "Only one claim is allowed for web retrieval for the current version."
 
             claim = data['inputs']['claims'][0]
-            ranked_docs, scores, ranked_urls = self.search_relevant_docs(claim)
-
-            outputs = {
-                'ranked_docs': ranked_docs,
-                'scores': scores,
-                'ranked_urls': ranked_urls
-            }
+            progress_stream = self.search_relevant_docs(claim)
+            outputs = Response(progress_stream, mimetype='text/event-stream')
             
         return outputs
     
     
     def search_relevant_docs(self, claim, timeout=10, max_search_results_per_query=5, allow_duplicated_urls=False):
 
+        def progress(message):
+            yield f"data: {message}\n\n"
+
         search_results = search_google(claim, timeout=timeout)
 
-        print('Searching webpages...')
+        yield from progress('Searching webpages...')
         start = time()
         with concurrent.futures.ThreadPoolExecutor() as e:
             scraped_results = e.map(scrape_url, search_results, itertools.repeat(timeout))
         end = time()
-        print(f"Finished searching in {round((end - start), 1)} seconds.\n")
+        yield from progress(f"Finished searching in {round((end - start), 1)} seconds.")
+
         scraped_results = [(r[0][:50000], r[1]) for r in scraped_results if r[0] and '��' not in r[0] and ".pdf" not in r[1]]
 
         retrieved_docs, urls = zip(*scraped_results[:max_search_results_per_query])
 
-        print('Scoring webpages...')
+        yield from progress('Scoring webpages...')
         start = time()
         retrieved_data = {
             'inputs': {
@@ -75,7 +75,7 @@ class EndpointHandler():
         _, _, used_chunk, support_prob_per_chunk = self.scorer.score(data=retrieved_data)
         end = time()
         num_chunks = len([item for items in used_chunk for item in items])
-        print(f'Finished {num_chunks} entailment checks in {round((end - start), 1)} seconds ({round(num_chunks / (end - start) * 60)} Doc./min).')
+        yield from progress(f'Finished {num_chunks} entailment checks in {round((end - start), 1)} seconds ({round(num_chunks / (end - start) * 60)} Doc./min).')
 
         ranked_docs, scores, ranked_urls = order_doc_score_url(used_chunk, support_prob_per_chunk, urls, allow_duplicated_urls=allow_duplicated_urls)
 

requirements.txt

@@ -4,4 +4,5 @@ nltk==3.8.1
 pandas==2.2.1
 numpy==1.26.2
 tqdm
-bs4
+bs4
+flask