return ranked docs and scores

handler.py

@@ -1,11 +1,40 @@
 from minicheck_web.minicheck import MiniCheck
 
+
+def sort_chunks_single_doc_claim(used_chunk, support_prob_per_chunk):
+    '''
+    Sort the chunks in a single document based on the probability of "supported" in descending order.
+    This function is used when a user document is provided.
+    '''
+
+    flattened_docs = [doc for chunk in used_chunk for doc in chunk]
+    flattened_scores = [score for chunk in support_prob_per_chunk for score in chunk]
+
+    doc_score = list(zip(flattened_docs, flattened_scores))
+    ranked_doc_score = sorted(doc_score, key=lambda x: x[1], reverse=True)
+
+    ranked_docs, scores = zip(*ranked_doc_score)
+
+    return ranked_docs, scores
+            
+
 class EndpointHandler():
     def __init__(self, path="./"):
         self.scorer = MiniCheck(path=path)
 
     def __call__(self, data):
 
-        _, raw_prob, _, _ = self.scorer.score(data=data)
+        _, _, used_chunk, support_prob_per_chunk = self.scorer.score(data=data)
+
+        # Using user-provided document to do fact-checking
+        if data['inputs']['docs'] != "":
+            ranked_docs, scores = sort_chunks_single_doc_claim(used_chunk, support_prob_per_chunk)
+        else:
+            raise NotImplementedError("Currently, only user-provided document is supported.")
             
-        return raw_prob
+        outputs = {
+            'ranked_docs': ranked_docs,
+            'scores': scores
+        }
+
+        return outputs