Update streamlit_app.py

streamlit_app.py

@@ -5,36 +5,40 @@ import numpy as np
 from sentence_transformers import SentenceTransformer
 from transformers import pipeline
 
-# Since quote_embeddings.pkl is in the same directory as this script
 EMBEDDING_PATH = "quote_embeddings.pkl"
 
-# Load saved quotes and embeddings
-with open(EMBEDDING_PATH, "rb") as f:
-    quotes, embeddings = pickle.load(f)
+@st.cache_resource(show_spinner=False)
+def load_data_and_models():
+    # Load quotes and embeddings
+    with open(EMBEDDING_PATH, "rb") as f:
+        quotes, embeddings = pickle.load(f)
 
-# Initialize embedder and FAISS index
-embedder = SentenceTransformer('all-MiniLM-L6-v2')
+    # Initialize embedder
+    embedder = SentenceTransformer('all-MiniLM-L6-v2')
 
-embeddings = embeddings.astype('float32')
-index = faiss.IndexFlatIP(embeddings.shape[1])
-faiss.normalize_L2(embeddings)
-index.add(embeddings)
+    # Prepare FAISS index
+    embeddings_np = embeddings.astype('float32')
+    index = faiss.IndexFlatIP(embeddings_np.shape[1])
+    faiss.normalize_L2(embeddings_np)
+    index.add(embeddings_np)
 
-# Initialize text generator pipeline
-generator = pipeline('text-generation', model='distilgpt2')
+    # Initialize text generation pipeline with smaller max length for speed
+    generator = pipeline('text-generation', model='distilgpt2')
+
+    return quotes, index, embedder, generator
+
+quotes, index, embedder, generator = load_data_and_models()
 
-# Define RAG search function
 def rag_search(query, top_k=3):
     q_emb = embedder.encode([query]).astype('float32')
     faiss.normalize_L2(q_emb)
     scores, indices = index.search(q_emb, top_k)
     context = "\n".join([f"{quotes[i]['quote']} — {quotes[i].get('author','Unknown')}" for i in indices[0]])
     prompt = f"Answer using these quotes:\n{context}\nQuestion: {query}\nAnswer:"
-    outputs = generator(prompt, max_length=100, num_return_sequences=1)
+    outputs = generator(prompt, max_length=80, num_return_sequences=1, do_sample=False)
     answer = outputs[0]['generated_text'].split('Answer:')[-1].strip()
     return answer
 
-# Streamlit UI
 st.title("🧠 RAG Quote-Based Q&A App")
 user_query = st.text_input("💬 Ask something related to quotes:")