Update app.py

app.py

@@ -1,103 +1,51 @@
 import os
-import numpy as np
-import faiss
+import urllib.request
+import zipfile
 import streamlit as st
-from sentence_transformers import SentenceTransformer
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
-from beir import util
-from beir.datasets.data_loader import GenericDataLoader
-#from beir import EvaluateRetrieval
 
+def download_and_extract_dataset():
+    dataset_url = "https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/nq.zip"
+    dataset_zip_path = "nq.zip"
+    data_path = "./datasets/nq"
 
-# Function to load the dataset
-def load_dataset():
-    dataset_name = "nq"
-    data_path = f"datasets/{dataset_name}.zip"
-    if not os.path.exists(data_path):
-        url = f"https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/{dataset_name}.zip"
-        util.download_and_unzip(url, "datasets/")
-    
-    corpus, queries, qrels = GenericDataLoader(data_path).load(split="test")
-    return corpus, queries, qrels
-
-# Function for candidate retrieval
-def candidate_retrieval(corpus, queries):
-    embed_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
-    corpus_ids = list(corpus.keys())
-    corpus_texts = [corpus[pid]["text"] for pid in corpus_ids]
-    corpus_embeddings = embed_model.encode(corpus_texts, convert_to_numpy=True)
+    # Download the dataset if not already downloaded
+    if not os.path.exists(dataset_zip_path):
+        st.write("Downloading the dataset... This may take a few minutes.")
+        urllib.request.urlretrieve(dataset_url, dataset_zip_path)
+        st.write("Download complete!")
 
-    index = faiss.IndexFlatL2(corpus_embeddings.shape[1])
-    index.add(corpus_embeddings)
+    # Unzip the dataset if not already unzipped
+    if not os.path.exists(data_path):
+        st.write("Unzipping the dataset...")
+        with zipfile.ZipFile(dataset_zip_path, 'r') as zip_ref:
+            zip_ref.extractall("./datasets")
+        st.write("Dataset unzipped!")
 
-    query_texts = [queries[qid] for qid in queries.keys()]
-    query_embeddings = embed_model.encode(query_texts, convert_to_numpy=True)
+    return data_path
 
-    _, retrieved_indices = index.search(query_embeddings, 10)
-    return retrieved_indices, corpus_ids
+# Function to load the dataset
+def load_dataset():
+    from beir.datasets.data_loader import GenericDataLoader
+    
+    data_path = download_and_extract_dataset()
 
-# Function for reranking
-def rerank_passages(retrieved_indices, corpus, queries):
-    cross_encoder_model = AutoModelForSequenceClassification.from_pretrained("cross-encoder/ms-marco-MiniLM-L-12-v2")
-    tokenizer = AutoTokenizer.from_pretrained("cross-encoder/ms-marco-MiniLM-L-12-v2")
+    # 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)}")
 
-    reranked_passages = []
-    for i, query in enumerate(queries.values()):
-        query_passage_pairs = [(query, corpus[corpus_ids[idx]]["text"]) for idx in retrieved_indices[i]]
-        inputs = tokenizer(query_passage_pairs, padding=True, truncation=True, return_tensors="pt")
-        scores = cross_encoder_model(**inputs).logits.squeeze(-1)
-        
-        top_reranked_passages = [passage for _, passage in sorted(zip(scores, query_passage_pairs), key=lambda x: x[0], reverse=True)]
-        reranked_passages.append(top_reranked_passages)
-    
-    return reranked_passages
+    return corpus, queries, qrels
 
-# Function for evaluation
-""""
-def evaluate(qrels, retrieved_indices, reranked_passages, queries):
-    evaluator = EvaluateRetrieval()
-    
-    results_stage1 = {}
-    for i, query_id in enumerate(queries.keys()):
-        results_stage1[query_id] = {corpus_ids[idx]: 1 for idx in retrieved_indices[i]}
-    
-    ndcg_score_stage1 = evaluator.evaluate(qrels, results_stage1, [10])['NDCG@10']
-    
-    results_stage2 = {}
-    for i, query_id in enumerate(queries.keys()):
-        results_stage2[query_id] = {}
-        for passage in reranked_passages[i]:
-            for pid, doc in corpus.items():
-                if doc["text"] == passage[1]:
-                    results_stage2[query_id][pid] = 1
-                    break
-    
-    ndcg_score_stage2 = evaluator.evaluate(qrels, results_stage2, [10])['NDCG@10']
-    return ndcg_score_stage1, ndcg_score_stage2
-    """
-# Streamlit app
+# Streamlit main execution
 def main():
-    st.title("Multi-Stage Text Retrieval Pipeline")
-    
-    if st.button("Load Dataset"):
-        corpus, queries, qrels = load_dataset()
-        st.success("Dataset loaded successfully!")
-    
-    if st.button("Run Candidate Retrieval"):
-        retrieved_indices, corpus_ids = candidate_retrieval(corpus, queries)
-        st.success("Candidate retrieval completed!")
-        st.write("Retrieved indices:", retrieved_indices)
+    st.title("Multi-Stage Retrieval Pipeline")
 
-    if st.button("Run Reranking"):
-        reranked_passages = rerank_passages(retrieved_indices, corpus, queries)
-        st.success("Reranking completed!")
-        st.write("Reranked passages:", reranked_passages)
+    # Load the dataset
+    corpus, queries, qrels = load_dataset()
 
-    """if st.button("Evaluate"):
-        ndcg_score_stage1, ndcg_score_stage2 = evaluate(qrels, retrieved_indices, reranked_passages, queries)
-        st.write(f"NDCG@10 for Stage 1 (Candidate Retrieval): {ndcg_score_stage1}")
-        st.write(f"NDCG@10 for Stage 2 (Reranking): {ndcg_score_stage2}")
-        """
+    st.write("Dataset loaded successfully!")
 
 if __name__ == "__main__":
     main()