add TFIDF ranking

handler.py

@@ -3,6 +3,9 @@ from web_retrieval import *
 from nltk.tokenize import sent_tokenize
 import evaluate
 
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.metrics.pairwise import cosine_similarity
+
 
 def sort_chunks_single_doc_claim(used_chunk, support_prob_per_chunk):
     '''
@@ -21,6 +24,29 @@ def sort_chunks_single_doc_claim(used_chunk, support_prob_per_chunk):
     return ranked_docs, scores
 
 
+def rank_documents_TFIDF(claim, scraped_results):
+
+    """
+    each element in scraped_results is a tuple of (document, URL)
+    """
+
+    documents = [result[0] for result in scraped_results]
+    corpus = [claim] + documents
+    
+    vectorizer = TfidfVectorizer()
+    tfidf_matrix = vectorizer.fit_transform(corpus)
+    
+    claim_vector = tfidf_matrix[0]
+    similarity_scores = cosine_similarity(claim_vector, tfidf_matrix[1:])
+
+    ranked_results = [(scraped_results[i][0], scraped_results[i][1], score)
+                      for i, score in enumerate(similarity_scores[0])]
+    ranked_results.sort(key=lambda x: x[2], reverse=True)
+    ranked_documents = [(result[0], result[1]) for result in ranked_results]
+
+    return ranked_documents
+
+
 class EndpointHandler():
     def __init__(self, path="./"):
         self.scorer = MiniCheck(path=path)
@@ -78,7 +104,7 @@ class EndpointHandler():
         return outputs
     
     
-    def search_relevant_docs(self, claim, timeout=10, max_search_results_per_query=10, allow_duplicated_urls=False):
+    def search_relevant_docs(self, claim, timeout=10, max_search_results_per_query=5, allow_duplicated_urls=False):
 
         search_results = search_google(claim, timeout=timeout)
 
@@ -91,6 +117,7 @@ class EndpointHandler():
         print(f"Finished searching in {round((end - start), 1)} seconds.\n")
         scraped_results = [(r[0][:20000], r[1]) for r in scraped_results if r[0] and '��' not in r[0]]   # those can be ranked based on TF-IDF to be more efficient
 
+        scraped_results = rank_documents_TFIDF(claim, scraped_results)
         retrieved_docs, urls = zip(*scraped_results[:max_search_results_per_query])
 
         print('Scoring webpages...')