Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,5 +1,5 @@
 # ======================================================
-# 📚 Handbook Assistant — Handbook-only, source-cited answers
+# 📘 Handbook Assistant (FAST OPTIMIZED VERSION)
 # ======================================================
 # Requirements:
 # pip install streamlit python-dotenv PyPDF2 numpy faiss-cpu scikit-learn huggingface-hub streamlit-chat sentence-transformers
@@ -15,305 +15,238 @@ import numpy as np
 import streamlit as st
 from dotenv import load_dotenv
 import PyPDF2
-from huggingface_hub import InferenceClient, login
 from streamlit_chat import message as st_message
-from sentence_transformers import SentenceTransformer  # ⚡ NEW: local embedder
+
+# Optional fast embedding model
+from sentence_transformers import SentenceTransformer
+
+# Try FAISS
+try:
+    import faiss
+except Exception:
+    faiss = None
 
 # ======================================================
-# ⚙️ PAGE CONFIG — must be first Streamlit call
+# ⚙️ CONFIGURATION
 # ======================================================
 st.set_page_config(page_title="📚 Handbook Assistant", page_icon="📘", layout="wide")
-st.title("📚 Handbook Assistant — Handbook-only, source-cited answers")
-st.caption("Place your handbook PDF(s) in the same folder as this app (e.g., handbook.pdf).")
+st.title("📚 Handbook Assistant — Fast Local Version")
+st.caption("Place your handbook PDF (e.g., handbook.pdf) beside this script or upload below.")
 
-# ======================================================
-# 🔐 ENVIRONMENT SETUP
-# ======================================================
 load_dotenv()
-HF_TOKEN = os.getenv("HF_TOKEN")
 
-if HF_TOKEN:
-    try:
-        login(HF_TOKEN)
-    except Exception:
-        pass
-else:
-    st.info("HF_TOKEN not found. Hugging Face model calls will be limited. Local embeddings are used instead.")
-
-# ======================================================
-# 📦 MODEL SETTINGS
-# ======================================================
-DEFAULT_OPEN_SOURCE_MODEL = "mistralai/Mistral-7B-Instruct-v0.3"
-AVAILABLE_MODELS = {
-    "Mistral 7B Instruct (Open Source – recommended)": DEFAULT_OPEN_SOURCE_MODEL,
-}
+# File names for saving
+HAND_INDEX_FN = "handbook_faiss.index"
+HAND_META_FN = "handbook_metadata.json"
+HAND_EMB_DIM_FN = "handbook_emb_dim.json"
 
 # ======================================================
-# 🧭 SIDEBAR
+# ⚙️ SIDEBAR SETTINGS
 # ======================================================
 with st.sidebar:
     st.header("⚙️ Settings")
-    selected_model_name = st.selectbox("Model", options=list(AVAILABLE_MODELS.keys()), index=0)
-    selected_model_id = AVAILABLE_MODELS[selected_model_name]
+
     similarity_threshold = st.slider("Similarity threshold", 0.3, 0.95, 0.62, 0.01)
-    top_k = st.slider("Number of retrieved chunks (top k)", 1, 10, 4)
-    chunk_size_chars = st.number_input("Chunk size (characters)", 400, 2500, 1200, 100)
-    chunk_overlap = st.number_input("Chunk overlap (chars)", 20, 600, 150, 10)
-    regenerate_index = st.button("🔁 Rebuild handbook index")
-    st.markdown("**Storage:** FAISS index + metadata saved to disk for faster restarts.")
+    top_k = st.slider("Top chunks retrieved", 1, 10, 4)
+    chunk_size_chars = st.number_input("Chunk size (chars)", min_value=400, max_value=3000, value=2000, step=100)
+    chunk_overlap = st.number_input("Chunk overlap (chars)", min_value=20, max_value=600, value=100, step=10)
+    regenerate_index = st.button("🔁 Rebuild handbook index (force re-embed)")
+    st.markdown("**Storage:** Cached FAISS index + metadata for fast restarts.")
 
-# ======================================================
-# 🔗 HUGGING FACE CLIENT
-# ======================================================
-hf_client = InferenceClient(token=HF_TOKEN) if HF_TOKEN else None
+    uploaded_pdf = st.file_uploader("📄 Upload handbook PDF", type=["pdf"])
+    if uploaded_pdf:
+        temp_path = os.path.join(os.path.dirname(__file__), uploaded_pdf.name)
+        with open(temp_path, "wb") as f:
+            f.write(uploaded_pdf.getbuffer())
+        st.session_state.uploaded_pdf_path = temp_path
+        st.success(f"✅ Uploaded and saved: {uploaded_pdf.name}")
 
 # ======================================================
-# 📘 FILE UTILITIES
+# 🧩 UTILITIES
 # ======================================================
-HAND_INDEX_FN = "handbook_faiss.index"
-HAND_META_FN = "handbook_metadata.json"
-HAND_EMB_DIM_FN = "handbook_emb_dim.json"
-
-try:
-    import faiss
-except Exception:
-    faiss = None
+@st.cache_resource(show_spinner=False)
+def get_local_embedder():
+    """Load MiniLM model (only once)."""
+    return SentenceTransformer("all-MiniLM-L6-v2")
 
 def find_pdfs(patterns=["handbook*.pdf", "*.pdf"]) -> List[str]:
-    """Find handbook PDFs robustly."""
+    """Find handbook PDFs in script folder or uploaded ones."""
     base_dir = os.path.dirname(os.path.abspath(__file__))
     files = []
     for patt in patterns:
-        matched = glob.glob(os.path.join(base_dir, patt))
-        if matched:
-            files = matched
-            break
-    if not files:
-        for patt in patterns:
-            matched = glob.glob(patt)
-            if matched:
-                files = matched
-                break
-    if "uploaded_pdf_path" in st.session_state and os.path.exists(st.session_state.uploaded_pdf_path):
+        files += glob.glob(os.path.join(base_dir, patt))
+    if not files and "uploaded_pdf_path" in st.session_state:
         files = [st.session_state.uploaded_pdf_path]
-    return sorted(files)
+    return sorted(list(set(files)))
 
 def load_pdf_texts_with_page_info(pdf_paths: List[str]) -> List[Dict[str, Any]]:
-    pages = []
+    """Extract text from each page with filename and page number."""
+    all_pages = []
     for p in pdf_paths:
         try:
             with open(p, "rb") as f:
                 reader = PyPDF2.PdfReader(f)
                 for i, page in enumerate(reader.pages):
-                    text = page.extract_text() or ""
+                    try:
+                        text = page.extract_text() or ""
+                    except Exception:
+                        text = ""
                     if text.strip():
-                        pages.append({"filename": os.path.basename(p), "page": i + 1, "text": text})
+                        all_pages.append({"filename": os.path.basename(p), "page": i + 1, "text": text})
         except Exception as e:
-            st.warning(f"Failed reading {p}: {e}")
-    return pages
+            st.warning(f"⚠️ Failed to read {p}: {e}")
+    return all_pages
 
-def chunk_pages_into_segments(pages: List[Dict[str, Any]], chunk_size: int, overlap: int):
+def chunk_pages_into_segments(pages: List[Dict[str, Any]], chunk_size: int, overlap: int) -> List[Dict[str, Any]]:
+    """Split long page text into overlapping chunks."""
     chunks = []
     for pg in pages:
-        text, filename, page_no = pg["text"], pg["filename"], pg["page"]
+        text = pg["text"]
+        filename, page_no = pg["filename"], pg["page"]
         start, chunk_id = 0, 0
         while start < len(text):
             end = min(start + chunk_size, len(text))
             seg = text[start:end].strip()
-            if len(seg) >= 30:
+            if len(seg) > 50:
                 chunks.append({
                     "filename": filename,
                     "page": page_no,
                     "chunk_id": f"{filename}_p{page_no}_c{chunk_id}",
                     "text": seg
                 })
-                chunk_id += 1
+            chunk_id += 1
             start = end - overlap
             if start < 0:
                 start = 0
     return chunks
 
-# ======================================================
-# ⚡ LOCAL EMBEDDING SETUP
-# ======================================================
-@st.cache_resource(show_spinner=False)
-def get_local_embedder():
-    """Load and cache the fast local sentence transformer."""
-    return SentenceTransformer("all-MiniLM-L6-v2")
-
-def local_embeddings_for_texts(texts: List[str]) -> List[np.ndarray]:
-    """Generate normalized MiniLM embeddings quickly."""
+def embed_texts(texts: List[str], batch_size: int = 16) -> np.ndarray:
+    """Fast local embedding using MiniLM in batches."""
     model = get_local_embedder()
-    embs = model.encode(texts, convert_to_numpy=True, normalize_embeddings=True)
-    return [np.array(e, dtype=np.float32) for e in embs]
-
-def fallback_vectorize(texts: List[str]):
-    """Lightweight TF-IDF-like fallback (no external model)."""
-    token_doc_freq, token_lists = {}, []
-    for t in texts:
-        tokens = [w.lower().strip(".,:;()[]{}\"'") for w in t.split() if len(w) > 2]
-        token_lists.append(tokens)
-        for tok in set(tokens):
-            token_doc_freq[tok] = token_doc_freq.get(tok, 0) + 1
-    token_index = {tok: i for i, tok in enumerate(token_doc_freq.keys())}
-    n_docs = len(texts)
-    vecs = []
-    for tokens in token_lists:
-        vec = np.zeros(len(token_index), dtype=np.float32)
-        for tok in tokens:
-            vec[token_index[tok]] += 1.0
-        for tok, idx in token_index.items():
-            df = token_doc_freq[tok]
-            if df > 0:
-                vec[idx] *= math.log((1 + n_docs) / (1 + df))
-        norm = np.linalg.norm(vec)
-        if norm > 0:
-            vec /= norm
-        vecs.append(vec)
-    return vecs
-
-# ======================================================
-# 🧠 FAISS HELPERS
-# ======================================================
-def build_faiss_index(embeddings: List[np.ndarray]):
+    all_embeddings = []
+    for i in range(0, len(texts), batch_size):
+        batch = texts[i:i + batch_size]
+        emb = model.encode(batch, convert_to_numpy=True, normalize_embeddings=True, show_progress_bar=False)
+        all_embeddings.append(emb)
+    return np.vstack(all_embeddings)
+
+def build_faiss_index(embeddings: np.ndarray):
+    """Build FAISS cosine index."""
     if faiss is None:
-        raise RuntimeError("faiss not installed. Run: pip install faiss-cpu")
-    arr = np.vstack(embeddings).astype("float32")
-    arr /= np.linalg.norm(arr, axis=1, keepdims=True)
-    index = faiss.IndexFlatIP(arr.shape[1])
-    index.add(arr)
-    return index, arr.shape[1]
-
-def save_index_and_metadata(index, metadata, dim):
+        raise RuntimeError("❌ FAISS not installed (pip install faiss-cpu)")
+    d = embeddings.shape[1]
+    index = faiss.IndexFlatIP(d)
+    index.add(embeddings)
+    return index, d
+
+def save_index_and_metadata(index, metadata, emb_dim: int):
     faiss.write_index(index, HAND_INDEX_FN)
-    json.dump(metadata, open(HAND_META_FN, "w", encoding="utf-8"), indent=2)
-    json.dump({"dim": dim}, open(HAND_EMB_DIM_FN, "w", encoding="utf-8"))
+    with open(HAND_META_FN, "w", encoding="utf-8") as f:
+        json.dump(metadata, f, indent=2)
+    with open(HAND_EMB_DIM_FN, "w") as f:
+        json.dump({"dim": emb_dim}, f)
 
 def load_index_and_metadata():
     if not (os.path.exists(HAND_INDEX_FN) and os.path.exists(HAND_META_FN)):
         return None, None
     index = faiss.read_index(HAND_INDEX_FN)
-    metadata = json.load(open(HAND_META_FN, "r", encoding="utf-8"))
-    return index, metadata
+    with open(HAND_META_FN, "r", encoding="utf-8") as f:
+        meta = json.load(f)
+    with open(HAND_EMB_DIM_FN, "r") as f:
+        emb_dim = json.load(f)["dim"]
+    return index, meta
 
 # ======================================================
-# 🔍 INDEX BUILDER
+# 🧠 INDEX BUILDER
 # ======================================================
 def ensure_handbook_index(rebuild=False):
-    if st.session_state.get("handbook_ready") and not rebuild:
+    """Build or load handbook FAISS index efficiently."""
+    if "handbook_ready" in st.session_state and st.session_state.handbook_ready and not rebuild:
         return
+
     pdfs = find_pdfs()
     if not pdfs:
-        st.error("No handbook PDF found.")
+        st.error("❌ No handbook PDF found.")
+        st.session_state.handbook_ready = False
         return
-    if not rebuild and os.path.exists(HAND_INDEX_FN):
-        try:
-            index, meta = load_index_and_metadata()
+
+    # Try loading cached index
+    if os.path.exists(HAND_INDEX_FN) and not rebuild:
+        index, metadata = load_index_and_metadata()
+        if index is not None:
             st.session_state.faiss_index = index
-            st.session_state.metadata = meta
+            st.session_state.metadata = metadata
             st.session_state.handbook_ready = True
-            st.success(f"Loaded existing FAISS index ({len(meta)} chunks).")
+            st.success(f"✅ Loaded FAISS index with {len(metadata)} chunks.")
             return
-        except Exception as e:
-            st.warning(f"Reload failed: {e}. Rebuilding…")
 
-    with st.spinner("⚙️ Building FAISS index locally with MiniLM…"):
-        pages = load_pdf_texts_with_page_info(pdfs)
-        chunks = chunk_pages_into_segments(pages, chunk_size_chars, chunk_overlap)
-        texts = [c["text"] for c in chunks]
-        try:
-            embs = local_embeddings_for_texts(texts)
-        except Exception as e:
-            st.warning(f"Local MiniLM failed ({e}); using fallback.")
-            embs = fallback_vectorize(texts)
-        index, dim = build_faiss_index(embs)
-        save_index_and_metadata(index, chunks, dim)
-        st.session_state.faiss_index = index
-        st.session_state.metadata = chunks
-        st.session_state.handbook_ready = True
-        st.success(f"✅ Indexed {len(chunks)} chunks.")
-
-if regenerate_index:
-    ensure_handbook_index(True)
-elif "handbook_ready" not in st.session_state:
-    ensure_handbook_index(False)
+    st.info("⚙️ Building FAISS index locally with MiniLM… this may take 30–60 seconds.")
+    start_time = time.time()
+
+    pages = load_pdf_texts_with_page_info(pdfs)
+    chunks = chunk_pages_into_segments(pages, int(chunk_size_chars), int(chunk_overlap))
+    if not chunks:
+        st.error("❌ No readable text found in the handbook.")
+        return
+
+    texts = [c["text"] for c in chunks]
+    embeddings = embed_texts(texts, batch_size=16)
+
+    index, emb_dim = build_faiss_index(embeddings)
+    save_index_and_metadata(index, chunks, emb_dim)
+
+    st.session_state.faiss_index = index
+    st.session_state.metadata = chunks
+    st.session_state.handbook_ready = True
+
+    elapsed = time.time() - start_time
+    st.success(f"✅ Handbook indexed in {elapsed:.1f} seconds ({len(chunks)} chunks).")
 
 # ======================================================
-# 🔎 RETRIEVAL + CHAT
+# 🔍 RETRIEVAL
 # ======================================================
-def embed_query(query: str):
-    try:
-        emb = local_embeddings_for_texts([query])[0]
-    except Exception:
-        emb = fallback_vectorize([query])[0]
-    return emb / (np.linalg.norm(emb) or 1)
+def embed_query(query: str) -> np.ndarray:
+    model = get_local_embedder()
+    emb = model.encode([query], convert_to_numpy=True, normalize_embeddings=True)[0]
+    return emb.astype("float32")
 
 def retrieve_top_chunks(query: str, k: int):
     index = st.session_state.get("faiss_index")
     metadata = st.session_state.get("metadata", [])
-    if index is None or not metadata:
+    if not index or not metadata:
         return [], []
     q_emb = embed_query(query).reshape(1, -1)
     D, I = index.search(q_emb, k)
     results = [metadata[i] for i in I[0] if i < len(metadata)]
     return results, D[0].tolist()
 
-PROMPT_TEMPLATE = """You are HandbookAssistant. Answer **only** from the handbook excerpts below.
-If the answer isn't found, say exactly:
-"Sorry, I can only answer questions based on the school's handbook."
-Always cite sources like (Source: <filename>, page <page>, chunk <chunk_id>)."""
-
-def build_prompt(chunks, question):
-    excerpts = "\n\n".join([f"--- {c['chunk_id']} ({c['filename']}, page {c['page']}) ---\n{c['text']}" for c in chunks])
-    return f"{PROMPT_TEMPLATE}\n\n{excerpts}\n\nUser: {question}\nAnswer:"
-
-def call_hf_model(prompt, model_id):
-    resp = hf_client.text_generation(model=model_id, inputs=prompt, max_new_tokens=512, temperature=0.2)
-    if isinstance(resp, dict) and "generated_text" in resp:
-        return resp["generated_text"]
-    if isinstance(resp, list) and resp and "generated_text" in resp[0]:
-        return resp[0]["generated_text"]
-    return str(resp)
-
 # ======================================================
-# 💬 CHAT INTERFACE
+# 🗣️ CHAT INTERFACE
 # ======================================================
-if "chat_history" not in st.session_state:
-    st.session_state.chat_history = []
+ensure_handbook_index(rebuild=regenerate_index)
 
 st.divider()
 st.subheader("💬 Ask the handbook")
 
-user_input = st.chat_input("Ask a question about the handbook…")
-
+user_input = st.chat_input("Ask a question about the handbook...")
 if user_input:
     st_message(user_input, is_user=True)
-    retrieved, scores = retrieve_top_chunks(user_input, int(top_k))
-    if not retrieved or (scores and max(scores) < similarity_threshold):
-        reply = "Sorry, I can only answer questions based on the school's handbook."
+
+    retrieved, scores = retrieve_top_chunks(user_input, top_k)
+    if not retrieved or max(scores) < similarity_threshold:
+        reply = "Sorry, I can only answer based on the handbook, and I couldn’t find relevant information."
+        st_message(reply, is_user=False)
     else:
-        prompt = build_prompt(retrieved, user_input)
-        try:
-            reply = call_hf_model(prompt, selected_model_id)
-        except Exception as e:
-            reply = f"⚠️ Model error: {e}"
-    st_message(reply, is_user=False)
-    st.session_state.chat_history.append({"role": "user", "content": user_input})
-    st.session_state.chat_history.append({"role": "assistant", "content": reply})
+        answer = "Based on the handbook:\n\n"
+        for r, s in zip(retrieved, scores):
+            short = (r["text"][:300] + "…") if len(r["text"]) > 300 else r["text"]
+            answer += f"📄 **{r['filename']}**, page {r['page']} — (score {s:.3f})\n> {short}\n\n"
+        st_message(answer.strip(), is_user=False)
 
 # ======================================================
-# 🗂️ CHAT HISTORY + TOOLS
+# 🧾 HISTORY & EXPORT
 # ======================================================
 st.divider()
 st.subheader("Conversation History")
-for i, msg in enumerate(st.session_state.chat_history):
-    st_message(msg["content"], is_user=(msg["role"] == "user"), key=f"hist_{i}")
-
-col1, col2 = st.columns([1, 1])
-with col1:
-    if st.button("🔄 Reset chat"):
-        st.session_state.chat_history = []
-        st.success("Chat reset.")
-with col2:
-    transcript = "\n\n".join([f"{m['role'].upper()}: {m['content']}" for m in st.session_state.chat_history])
-    st.download_button("📥 Download transcript", data=transcript, file_name="handbook_transcript.txt")
+if "chat_history" not in st.session_state:
+    st.session_state.chat_history = []