Update app.py

app.py

@@ -1,66 +1,100 @@
-import streamlit as st
-import pandas as pd
+import os
 import numpy as np
+import faiss
+import streamlit as st
 from sentence_transformers import SentenceTransformer
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-import torch
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+from beir import util
+from beir.datasets.data_loader import GenericDataLoader
+from beir.evaluation.evaluator import EvaluateRetrieval
 
-# Load dataset
-def load_qa_dataset():
-    # Example using BEIR's NQ dataset
-    dataset = load_dataset('beir/nq')
-    return dataset['train']  # Use the training set for simplicity
+# 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
 
-# Load embedding models
-small_embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
-large_embedding_model = SentenceTransformer('sentence-transformers/paraphrase-mpnet-base-v2')
+# 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)
 
-# Load ranking models
-small_ranking_model = AutoModelForSequenceClassification.from_pretrained("cross-encoder/ms-marco-MiniLM-L-12-v2")
-small_ranking_tokenizer = AutoTokenizer.from_pretrained("cross-encoder/ms-marco-MiniLM-L-12-v2")
+    index = faiss.IndexFlatL2(corpus_embeddings.shape[1])
+    index.add(corpus_embeddings)
 
-large_ranking_model = AutoModelForSequenceClassification.from_pretrained("nvidia/nv-rerankqa-mistral-4b-v3")
-large_ranking_tokenizer = AutoTokenizer.from_pretrained("nvidia/nv-rerankqa-mistral-4b-v3")
+    query_texts = [queries[qid] for qid in queries.keys()]
+    query_embeddings = embed_model.encode(query_texts, convert_to_numpy=True)
 
-# Streamlit UI
-st.title("Multi-Stage Text Retrieval Pipeline for QA")
+    _, retrieved_indices = index.search(query_embeddings, 10)
+    return retrieved_indices, corpus_ids
 
-# Load dataset
-dataset = load_qa_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")
 
-# Input question
-question = st.text_input("Enter a question:")
-if question:
-    passages = dataset['text'][:100]  # Limit to 100 passages for demo purposes
+    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)
     
-    # Stage 1: Candidate Retrieval using Embedding Models
-    st.write("**Stage 1: Candidate Retrieval**")
-    question_embedding = small_embedding_model.encode(question)
-    passage_embeddings = small_embedding_model.encode(passages)
+    return reranked_passages
+
+# Function for evaluation
+def evaluate(qrels, retrieved_indices, reranked_passages, queries):
+    evaluator = EvaluateRetrieval()
     
-    # Find top-k similar passages using cosine similarity
-    top_k = 5
-    similarities = np.inner(question_embedding, passage_embeddings)
-    top_k_indices = np.argsort(similarities)[-top_k:][::-1]
+    results_stage1 = {}
+    for i, query_id in enumerate(queries.keys()):
+        results_stage1[query_id] = {corpus_ids[idx]: 1 for idx in retrieved_indices[i]}
     
-    retrieved_passages = [passages[i] for i in top_k_indices]
-    st.write("Top-k retrieved passages:")
-    for passage in retrieved_passages:
-        st.write(passage)
+    ndcg_score_stage1 = evaluator.evaluate(qrels, results_stage1, [10])['NDCG@10']
     
-    # Stage 2: Reranking with Ranking Models
-    st.write("**Stage 2: Reranking**")
-    inputs = [small_ranking_tokenizer(question, passage, return_tensors='pt', truncation=True, padding=True) for passage in retrieved_passages]
-    reranked_scores = []
+    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
     
-    for input_pair in inputs:
-        with torch.no_grad():
-            ranking_outputs = small_ranking_model(**input_pair)
-            score = ranking_outputs.logits.softmax(dim=1)[:, 1].item()  # Score for positive class
-            reranked_scores.append(score)
+    ndcg_score_stage2 = evaluator.evaluate(qrels, results_stage2, [10])['NDCG@10']
+    return ndcg_score_stage1, ndcg_score_stage2
+
+# Streamlit app
+def main():
+    st.title("Multi-Stage Text Retrieval Pipeline")
     
-    # Sort passages by ranking scores
-    ranked_passages = sorted(zip(retrieved_passages, reranked_scores), key=lambda x: x[1], reverse=True)
-    st.write("Ranked passages by relevance score:")
-    for passage, score in ranked_passages:
-        st.write(f"{passage} (Score: {score:.2f})")
+    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)
+
+    if st.button("Run Reranking"):
+        reranked_passages = rerank_passages(retrieved_indices, corpus, queries)
+        st.success("Reranking completed!")
+        st.write("Reranked passages:", reranked_passages)
+
+    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}")
+
+if __name__ == "__main__":
+    main()