Update app.py

app.py

@@ -1,48 +1,44 @@
 import os
 from sentence_transformers import SentenceTransformer, util
 from datasets import load_dataset
-from datasets import load_dataset
 from transformers import pipeline
 import streamlit as st
 
-
-
-dataset_id = "sentence-transformers/natural-questions"
-dataset_file = load_dataset(dataset_id, split="train")
-
-# Use the allenai-specter model with SentenceTransformers
-model = SentenceTransformer('allenai-specter')
-
-# Prepare paper texts by combining query and answer fields
-paper_texts = [
-    record['query'] + '[SEP]' + record['answer'] for record in dataset_file.select(range(32))
-]
-
-# Compute embeddings for all paper texts
-corpus_embeddings = model.encode(paper_texts, convert_to_tensor=True, show_progress_bar=True)
-
-# Function to search for answers given a query
-def search_papers(query):
-    # Encode the query
-    query_embedding = model.encode(query, convert_to_tensor=True)
-
-    # Perform semantic search
-    search_hits = util.semantic_search(query_embedding, corpus_embeddings)
-    search_hits = search_hits[0]  # Get the hits for the first query
-
-    print("\n\nQuery:", query)
-    print("Most similar answers:")
-    for hit in search_hits[:5]:  # Limit to top 5 results for clarity
-        related_text = dataset_file[int(hit['corpus_id'])]  # Access related record
-        print("{:.2f}\tAnswer: {}".format(
-            hit['score'], related_text['answer']
-        ))
-
-
-# Summarization pipeline
-summarizer = pipeline("summarization")
-
-# Collect the relevant answers from the search function
+# Cache dataset loading
+@st.cache_data
+def load_data(dataset_id="sentence-transformers/natural-questions", split="train"):
+    return load_dataset(dataset_id, split=split)
+
+# Cache model loading
+@st.cache_resource
+def load_model():
+    return SentenceTransformer('allenai-specter')
+
+# Cache corpus embedding generation
+@st.cache_data
+def generate_embeddings(model, dataset_file, sample_size=32):
+    # Prepare paper texts by combining query and answer fields
+    paper_texts = [
+        record['query'] + '[SEP]' + record['answer'] for record in dataset_file.select(range(sample_size))
+    ]
+    # Compute embeddings for all paper texts
+    return paper_texts, model.encode(paper_texts, convert_to_tensor=True, show_progress_bar=True)
+
+# Cache summarization pipeline
+@st.cache_resource
+def load_summarizer():
+    return pipeline("summarization")
+
+# Streamlit app
+st.title("Semantic Search with Summarization")
+
+# Load resources
+dataset_file = load_data()
+model = load_model()
+paper_texts, corpus_embeddings = generate_embeddings(model, dataset_file)
+summarizer = load_summarizer()
+
+# Function to search and summarize
 def search_papers_and_summarize(query, max_summary_length=45):
     # Encode the query
     query_embedding = model.encode(query, convert_to_tensor=True)
@@ -62,11 +58,12 @@ def search_papers_and_summarize(query, max_summary_length=45):
 
     # Summarize the combined text
     summary = summarizer(combined_text, max_length=max_summary_length, clean_up_tokenization_spaces=True)
-    print("Summary:")
-    print(summary[0]['summary_text'])
-
-
-title = st.text_input("Ask a question", "What is Wimpy Kid")
-new_preds = search_papers_and_summarize(title)
-st.write("The Answer is", new_preds)
-
+    return summary[0]['summary_text']
+
+# Streamlit input
+query = st.text_input("Enter your query:", "")
+if query:
+    st.write("Searching for relevant answers...")
+    summary = search_papers_and_summarize(query)
+    st.subheader("Summary")
+    st.write(summary)