Update app.py

app.py

@@ -1,45 +1,46 @@
 import streamlit as st
-import pandas as pd
-from datasets import load_dataset
+import numpy as np
 from sentence_transformers import SentenceTransformer
-from transformers import pipeline
+from transformers import CrossEncoder
 
-# Load a subset of the Natural Questions dataset for testing
-def load_nq_dataset():
-    return load_dataset("nq", split='train[:1%]')  # Load a small subset for testing
+# Sample passages
+passages = [
+    "The sky is blue.",
+    "The grass is green.",
+    "The sun is bright.",
+    "Rain falls from the sky.",
+    "Flowers bloom in spring."
+]
 
 # Load models
-@st.cache_resource
-def load_models():
-    embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
-    ranking_model = pipeline("text-classification", model="cross-encoder/ms-marco-MiniLM-L-12-v2")
-    return embedding_model, ranking_model
+embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+ranking_model = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-12-v2')
+
+def get_relevant_passages(question, passages):
+    keywords = question.lower().split()
+    relevant_passages = [p for p in passages if any(keyword in p.lower() for keyword in keywords)]
+    return relevant_passages if relevant_passages else passages  # Return all if no match
 
-# Main function for the Streamlit app
 def main():
     st.title("Multi-Stage Text Retrieval Pipeline for QA")
-
-    # Load dataset
-    dataset = load_nq_dataset()
-
-    # User input
     question = st.text_input("Enter a question:")
     
     if question:
-        # Load models
-        embedding_model, ranking_model = load_models()
-
-        # Retrieve passages (mock implementation, replace with actual logic)
-        top_k_passages = dataset['context'][:5]  # Replace with actual retrieval logic
-        embeddings = embedding_model.encode(top_k_passages)
-
-        # Re-rank passages
-        ranked_passages = ranking_model([question + " " + passage for passage in top_k_passages])
-
-        # Display results
-        st.write("Top Retrieved Passages:")
-        for i, (passage, score) in enumerate(zip(top_k_passages, ranked_passages)):
-            st.write(f"{i + 1}: {passage} (Score: {score['score']:.2f})")
+        relevant_passages = get_relevant_passages(question, passages)
+        st.write("Relevant passages:")
+        for p in relevant_passages:
+            st.write(f"- {p}")
+        
+        # Embedding and ranking
+        if st.button("Retrieve Answers"):
+            passage_embeddings = embedding_model.encode(relevant_passages)
+            question_embedding = embedding_model.encode(question)
+            scores = np.dot(passage_embeddings, question_embedding.T)
+            ranked_indices = np.argsort(scores)[::-1]
+
+            st.write("Ranked passages:")
+            for idx in ranked_indices:
+                st.write(f"- {relevant_passages[idx]} (Score: {scores[idx]:.2f})")
 
 if __name__ == "__main__":
     main()