Upload 6 files

app.py

@@ -0,0 +1,70 @@
+import streamlit as st
+from data_preparation import load_dataset
+from retrieval import load_embedding_model, retrieve_top_k
+from reranking import load_ranking_model, rerank
+from evaluation import evaluate_ndcg
+
+# Set up the Streamlit interface
+st.title("Multi-Stage Text Retrieval Pipeline for QA")
+
+# Query Input
+query = st.text_input("Enter a question:", "What is the capital of France?")
+
+# Embedding model selection
+embedding_model = st.selectbox(
+    "Select Embedding Model for Candidate Retrieval",
+    ["sentence-transformers/all-MiniLM-L6-v2", "nvidia/nv-embedqa-e5-v5"]
+)
+
+# Ranking model selection
+ranking_model = st.selectbox(
+    "Select Ranking Model for Re-Ranking",
+    ["cross-encoder/ms-marco-MiniLM-L-12-v2", "nvidia/nv-rerankqa-mistral-4b-v3"]
+)
+
+# Run retrieval pipeline on button click
+if st.button("Run Retrieval"):
+    # Load dataset
+    st.write("Loading dataset...")
+    corpus, queries, qrels = load_dataset("nq")
+
+    # Load selected embedding model
+    st.write(f"Loading embedding model: {embedding_model}...")
+    embed_model = load_embedding_model(embedding_model)
+
+    # Retrieve top-k passages using embedding model
+    st.write("Retrieving top-k passages...")
+    top_k_passages = retrieve_top_k(embed_model, query, corpus, k=10)
+    
+    # Display retrieved passages
+    st.write("Top-k passages before reranking:")
+    for i, (passage, score) in enumerate(top_k_passages):
+        st.write(f"{i+1}. Passage: {passage}, Score: {score:.4f}")
+
+    # Load selected ranking model
+    st.write(f"Loading ranking model: {ranking_model}...")
+    rank_model, rank_tokenizer = load_ranking_model(ranking_model)
+
+    # Rerank the retrieved passages
+    st.write("Reranking passages...")
+    ranked_passages = rerank(rank_model, rank_tokenizer, query, top_k_passages)
+    
+    # Display reranked passages
+    st.write("Top-k passages after reranking:")
+    for i, (passage, score) in enumerate(ranked_passages):
+        st.write(f"{i+1}. Passage: {passage}, Score: {score:.4f}")
+
+    # Evaluate using NDCG@10
+    st.write("Evaluating NDCG@10...")
+    query_id = list(queries.keys())[0]  # Assuming we are using the first query for evaluation
+    ndcg_score = evaluate_ndcg(ranked_passages, qrels[query_id])
+    st.write(f"NDCG@10: {ndcg_score:.4f}")
+
+# Sidebar with instructions
+st.sidebar.title("Instructions")
+st.sidebar.write("""
+1. Enter a question in the text input.
+2. Select the embedding model for candidate retrieval.
+3. Select the ranking model for reranking the retrieved passages.
+4. Click 'Run Retrieval' to start the pipeline and display the results.
+""")

data_preparation.py

@@ -0,0 +1,7 @@
+from beir import util
+from beir.datasets.data_loader import GenericDataLoader
+
+def load_dataset(dataset_name="nq"):
+    data_path = util.download_and_unzip(f"https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/{dataset_name}.zip", dataset_name)
+    corpus, queries, qrels = GenericDataLoader(data_folder=data_path).load(split="test")
+    return corpus, queries, qrels

evaluation.py

@@ -0,0 +1,5 @@
+from sklearn.metrics import ndcg_score
+
+def evaluate_ndcg(top_k_passages, qrels):
+    relevance_scores = [1 if doc in qrels else 0 for doc, _ in top_k_passages]
+    return ndcg_score([relevance_scores], [[1]*len(relevance_scores)], k=10)

requirements.txt

@@ -0,0 +1,6 @@
+streamlit
+beir
+sentence-transformers
+transformers
+torch
+scikit-learn

reranking.py

@@ -0,0 +1,16 @@
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+import torch
+
+def load_ranking_model(model_name):
+    model = AutoModelForSequenceClassification.from_pretrained(model_name)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    return model, tokenizer
+
+def rerank(model, tokenizer, query, top_k_passages):
+    inputs = tokenizer([f"{query} [SEP] {passage}" for passage, _ in top_k_passages], return_tensors="pt", truncation=True, padding=True)
+    with torch.no_grad():
+        outputs = model(**inputs).logits
+    scores = outputs.squeeze(-1)
+    
+    ranked_passages = sorted(zip(top_k_passages, scores), key=lambda x: x[1], reverse=True)
+    return [(passage, score.item()) for (passage, _), score in ranked_passages]

retrieval.py

@@ -0,0 +1,13 @@
+from sentence_transformers import SentenceTransformer, util
+
+def load_embedding_model(model_name):
+    return SentenceTransformer(model_name)
+
+def retrieve_top_k(model, query, corpus, k=10):
+    query_embedding = model.encode(query, convert_to_tensor=True)
+    corpus_embeddings = model.encode([corpus[doc_id]["text"] for doc_id in corpus], convert_to_tensor=True)
+    
+    hits = util.semantic_search(query_embedding, corpus_embeddings, top_k=k)[0]
+    top_k_passages = [(corpus[list(corpus.keys())[hit['corpus_id']]]["text"], hit['score']) for hit in hits]
+    
+    return top_k_passages