Update app.py

app.py

@@ -1,9 +1,16 @@
-import os
-import urllib.request
-import zipfile
 import streamlit as st
+from sentence_transformers import SentenceTransformer, util
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
+import numpy as np
+from sklearn.metrics import ndcg_score
 
+# Helper function to load the dataset
 def download_and_extract_dataset():
+    import urllib.request
+    import zipfile
+    import os
+
     dataset_url = "https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/nq.zip"
     dataset_zip_path = "nq.zip"
     data_path = "./datasets/nq"
@@ -23,29 +30,67 @@ def download_and_extract_dataset():
 
     return data_path
 
-# Function to load the dataset
+# Function to load corpus, queries, and qrels
 def load_dataset():
     from beir.datasets.data_loader import GenericDataLoader
-    
     data_path = download_and_extract_dataset()
-
-    # Load dataset using GenericDataLoader
-    st.write("Loading the dataset...")
     corpus, queries, qrels = GenericDataLoader(data_path).load(split="test")
-    st.write(f"Corpus Size: {len(corpus)}")
-    st.write(f"Queries Size: {len(queries)}")
-    st.write(f"Qrels Size: {len(qrels)}")
-
     return corpus, queries, qrels
 
-# Streamlit main execution
+# Stage 1: Candidate retrieval using Sentence Transformer
+def candidate_retrieval(query, corpus, top_k=10):
+    model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+    corpus_ids = list(corpus.keys())
+    corpus_embeddings = model.encode([corpus[doc_id]['text'] for doc_id in corpus_ids], convert_to_tensor=True)
+
+    query_embedding = model.encode(query, convert_to_tensor=True)
+    hits = util.semantic_search(query_embedding, corpus_embeddings, top_k=top_k)[0]
+    
+    retrieved_docs = [corpus_ids[hit['corpus_id']] for hit in hits]
+    return retrieved_docs
+
+# Stage 2: Reranking using cross-encoder
+def rerank(retrieved_docs, query, corpus, top_k=5):
+    tokenizer = AutoTokenizer.from_pretrained("cross-encoder/ms-marco-MiniLM-L-12-v2")
+    model = AutoModelForSequenceClassification.from_pretrained("cross-encoder/ms-marco-MiniLM-L-12-v2")
+
+    scores = []
+    for doc_id in retrieved_docs:
+        text = corpus[doc_id]['text']
+        inputs = tokenizer(query, text, return_tensors="pt", truncation=True, padding=True)
+        outputs = model(**inputs)
+        scores.append(outputs.logits.item())
+
+    reranked_indices = np.argsort(scores)[::-1][:top_k]
+    reranked_docs = [retrieved_docs[idx] for idx in reranked_indices]
+    return reranked_docs
+
+# Streamlit main function
 def main():
     st.title("Multi-Stage Retrieval Pipeline")
 
-    # Load the dataset
+    st.write("Loading the dataset...")
     corpus, queries, qrels = load_dataset()
-
-    st.write("Dataset loaded successfully!")
+    st.write(f"Corpus Size: {len(corpus)}")
+    
+    # User input for asking a question
+    user_query = st.text_input("Ask a question:")
+    
+    if user_query:
+        st.write(f"Your query: {user_query}")
+        
+        st.write("Running Candidate Retrieval...")
+        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)
+        
+        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
+        
+        st.write("Query executed successfully!")
 
 if __name__ == "__main__":
     main()