add span highlight

handler.py

@@ -1,5 +1,7 @@
 from minicheck_web.minicheck import MiniCheck
 from web_retrieval import *
+from nltk.tokenize import sent_tokenize
+import evaluate
 
 
 def sort_chunks_single_doc_claim(used_chunk, support_prob_per_chunk):
@@ -22,23 +24,37 @@ def sort_chunks_single_doc_claim(used_chunk, support_prob_per_chunk):
 class EndpointHandler():
     def __init__(self, path="./"):
         self.scorer = MiniCheck(path=path)
+        self.rouge = evaluate.load('rouge')
+
 
     def __call__(self, data):
 
+        claim = data['inputs']['claims'][0]
+
         # Using user-provided document to do fact-checking
         if len(data['inputs']['docs']) == 1 and data['inputs']['docs'][0] != '':
             _, _, used_chunk, support_prob_per_chunk = self.scorer.score(data=data)
             ranked_docs, scores = sort_chunks_single_doc_claim(used_chunk, support_prob_per_chunk)
 
+            span_to_highlight = []
+            for doc_chunk, score in zip(ranked_docs, scores):
+                # If the chunk can support the claim, find the sentence with the highest rouge score
+                if score > 0.5:
+                    highest_score_sent, _ = self.chunk_and_highest_rouge_score(doc_chunk, claim)
+                    span_to_highlight.append(highest_score_sent)
+                else:
+                    span_to_highlight.append("")
+            
             outputs = {
                 'ranked_docs': ranked_docs,
-                'scores': scores
-                }
+                'scores': scores,
+                'span_to_highlight': span_to_highlight
+            }
             
         else:
             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 = {
@@ -60,7 +76,7 @@ class EndpointHandler():
             scraped_results = e.map(scrape_url, search_results, itertools.repeat(timeout))
         end = time()
         print(f"Finished searching in {round((end - start), 1)} seconds.\n")
-        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]]
+        scraped_results = [(r[0][:20000], 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])
 
@@ -79,4 +95,28 @@ class EndpointHandler():
 
         ranked_docs, scores, ranked_urls = order_doc_score_url(used_chunk, support_prob_per_chunk, urls, allow_duplicated_urls=allow_duplicated_urls)
 
-        return ranked_docs, scores, ranked_urls
+        return ranked_docs, scores, ranked_urls
+    
+
+    def chunk_and_highest_rouge_score(self, doc, claim):
+
+        '''
+        Given a document and a claim, return the sentence with the highest rouge score and the score
+        '''
+
+        doc_sentences = sent_tokenize(doc)
+        claims = [claim] * len(doc_sentences)
+
+        results = self.rouge.compute(
+            predictions=doc_sentences,                    
+            references=claims, 
+            use_aggregator=False)
+
+        highest_score = 0
+        highest_score_sent = ""
+        for i in range(len(doc_sentences)):
+            if results['rouge1'][i] > highest_score:
+                highest_score = results['rouge1'][i]
+                highest_score_sent = doc_sentences[i]
+        
+        return highest_score_sent, highest_score

requirements.txt

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