Update app.py

app.py

@@ -63,11 +63,19 @@ def rerank(retrieved_docs, query, corpus, top_k=5):
 
     reranked_indices = np.argsort(scores)[::-1][:top_k]
     reranked_docs = [retrieved_docs[idx] for idx in reranked_indices]
-    return reranked_docs
+    return reranked_docs, scores
+
+# Function to evaluate using NDCG@10
+def evaluate_ndcg(reranked_docs, qrels, query_id, k=10):
+    true_relevance = [qrels.get((query_id, doc_id), 0) for doc_id in reranked_docs]
+    ideal_relevance = sorted(true_relevance, reverse=True)
+    
+    # NDCG expects input as 2D arrays
+    return ndcg_score([ideal_relevance], [true_relevance], k=k)
 
 # Streamlit main function
 def main():
-    st.title("Multi-Stage Retrieval Pipeline")
+    st.title("Multi-Stage Retrieval Pipeline with Evaluation")
 
     st.write("Loading the dataset...")
     corpus, queries, qrels = load_dataset()
@@ -83,13 +91,21 @@ def main():
         retrieved_docs = candidate_retrieval(user_query, corpus, top_k=10)
         
         st.write("Running Reranking...")
-        reranked_docs = rerank(retrieved_docs, user_query, corpus, top_k=5)
+        reranked_docs, rerank_scores = rerank(retrieved_docs, user_query, corpus, top_k=5)
         
         st.write("Top Reranked Documents:")
         for doc_id in reranked_docs:
             st.write(f"Document ID: {doc_id}")
             st.write(f"Document Text: {corpus[doc_id]['text'][:500]}...")  # Show the first 500 characters of the document
         
+        # Evaluation if the user query exists in the qrels (ground truth relevance labels)
+        query_id = list(queries.keys())[0]  # Dummy query ID for now
+        if query_id in queries:
+            ndcg_score_value = evaluate_ndcg(reranked_docs, qrels, query_id, k=10)
+            st.write(f"NDCG@10 Score: {ndcg_score_value}")
+        else:
+            st.write("No ground truth available for this query.")
+        
         st.write("Query executed successfully!")
 
 if __name__ == "__main__":