Update app.py

refactored for document matching versus chunk matching

app.py

@@ -8,83 +8,37 @@ from huggingface_hub import HfApi, hf_hub_download
 from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError
 import pypdf
 import docx
-import time
 
 # --- CONFIGURATION ---
-DATASET_REPO_ID = "NavyDevilDoc/navy-policy-index" # Your Dataset
+DATASET_REPO_ID = "NavyDevilDoc/navy-policy-index" 
 HF_TOKEN = os.environ.get("HF_TOKEN")
-
-# File paths for local storage
 INDEX_FILE = "navy_index.faiss"
 META_FILE = "navy_metadata.pkl"
 
-st.set_page_config(page_title="Navy Search (FAISS)", layout="wide")
+st.set_page_config(page_title="Document Finder", layout="wide")
 
-# --- PERSISTENCE MANAGER ---
+# --- PERSISTENCE (SAME AS BEFORE) ---
 class IndexManager:
-    """Manages loading/saving the FAISS index and Metadata from Hugging Face"""
-    
     @staticmethod
     def load_from_hub():
-        """Downloads the index files from HF Dataset"""
-        if not HF_TOKEN:
-            st.warning("HF_TOKEN missing. Running in local-only mode.")
-            return False
-            
+        if not HF_TOKEN: return False
         try:
-            with st.spinner("Downloading Knowledge Base..."):
-                # Download Vector Index
-                hf_hub_download(
-                    repo_id=DATASET_REPO_ID,
-                    filename=INDEX_FILE,
-                    repo_type="dataset",
-                    local_dir=".",
-                    token=HF_TOKEN
-                )
-                # Download Metadata
-                hf_hub_download(
-                    repo_id=DATASET_REPO_ID,
-                    filename=META_FILE,
-                    repo_type="dataset",
-                    local_dir=".",
-                    token=HF_TOKEN
-                )
+            hf_hub_download(repo_id=DATASET_REPO_ID, filename=INDEX_FILE, local_dir=".", token=HF_TOKEN)
+            hf_hub_download(repo_id=DATASET_REPO_ID, filename=META_FILE, local_dir=".", token=HF_TOKEN)
             return True
-        except (EntryNotFoundError, RepositoryNotFoundError):
-            st.toast("No existing index found in Cloud. Starting fresh.", icon="🆕")
-            return False
-        except Exception as e:
-            st.error(f"Sync Error: {e}")
-            return False
+        except: return False
 
     @staticmethod
     def save_to_hub():
-        """Uploads the local files to HF Dataset"""
-        if not HF_TOKEN:
-            return
-            
+        if not HF_TOKEN: return
         api = HfApi(token=HF_TOKEN)
         try:
-            st.toast("Syncing to Cloud...", icon="☁️")
-            api.upload_file(
-                path_or_fileobj=INDEX_FILE,
-                path_in_repo=INDEX_FILE,
-                repo_id=DATASET_REPO_ID,
-                repo_type="dataset",
-                commit_message="Update FAISS Index"
-            )
-            api.upload_file(
-                path_or_fileobj=META_FILE,
-                path_in_repo=META_FILE,
-                repo_id=DATASET_REPO_ID,
-                repo_type="dataset",
-                commit_message="Update Metadata"
-            )
-            st.success("Knowledge Base Saved!")
-        except Exception as e:
-            st.error(f"Upload failed: {e}")
-
-# --- HELPER FUNCTIONS ---
+            api.upload_file(path_or_fileobj=INDEX_FILE, path_in_repo=INDEX_FILE, repo_id=DATASET_REPO_ID, repo_type="dataset")
+            api.upload_file(path_or_fileobj=META_FILE, path_in_repo=META_FILE, repo_id=DATASET_REPO_ID, repo_type="dataset")
+            st.toast("Database Synced!", icon="☁️")
+        except Exception as e: st.error(f"Sync Error: {e}")
+
+# --- PARSING & CHUNKING (SAME AS BEFORE) ---
 def parse_file(uploaded_file):
     text = ""
     filename = uploaded_file.name
@@ -92,159 +46,174 @@ def parse_file(uploaded_file):
         if filename.endswith(".pdf"):
             reader = pypdf.PdfReader(uploaded_file)
             for i, page in enumerate(reader.pages):
-                page_text = page.extract_text()
-                if page_text:
-                    text += f"\n[PAGE {i+1}] {page_text}"
+                if page.extract_text(): text += f"\n[PAGE {i+1}] {page.extract_text()}"
         elif filename.endswith(".docx"):
             doc = docx.Document(uploaded_file)
             text = "\n".join([para.text for para in doc.paragraphs])
         elif filename.endswith(".txt"):
             text = uploaded_file.read().decode("utf-8")
-    except Exception as e:
-        st.error(f"Error parsing {filename}: {e}")
+    except: pass
     return text, filename
 
 def recursive_chunking(text, source, chunk_size=500, overlap=100):
     words = text.split()
     chunks = []
     for i in range(0, len(words), chunk_size - overlap):
-        chunk_words = words[i:i + chunk_size]
-        chunk_text = " ".join(chunk_words)
-        
-        # Simple Page Extraction
-        page_num = "Unknown"
-        if "[PAGE" in chunk_text:
-            try:
-                start = chunk_text.rfind("[PAGE") + 6
-                end = chunk_text.find("]", start)
-                page_num = chunk_text[start:end]
-            except: pass
-
+        chunk_text = " ".join(words[i:i + chunk_size])
         if len(chunk_text) > 50:
-            chunks.append({
-                "text": chunk_text,
-                "source": source,
-                "page": page_num
-            })
+            chunks.append({"text": chunk_text, "source": source})
     return chunks
 
-# --- CORE SEARCH ENGINE (FAISS VERSION) ---
-class RobustSearchEngine:
+# --- CORE SEARCH ENGINE (UPDATED FOR DOC LEVEL) ---
+class DocSearchEngine:
     def __init__(self):
-        # Load Models (Force CPU to avoid meta tensor errors)
-        self.bi_encoder = SentenceTransformer('all-MiniLM-L6-v2', device="cpu")
+        self.bi_encoder = SentenceTransformer('all-mpnet-base-v2', device="cpu")
         self.cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2', device="cpu", automodel_args={"low_cpu_mem_usage": False})
-        
         self.index = None
-        self.metadata = [] # List of dicts matching index order
+        self.metadata = []
         
-        # Try to load existing index from disk
         if os.path.exists(INDEX_FILE) and os.path.exists(META_FILE):
             self.index = faiss.read_index(INDEX_FILE)
-            with open(META_FILE, "rb") as f:
-                self.metadata = pickle.load(f)
-        else:
-            # Initialize new index
-            self.index = None # Will init on first add
-            self.metadata = []
+            with open(META_FILE, "rb") as f: self.metadata = pickle.load(f)
 
     def add_documents(self, chunks):
-        # 1. Encode
         texts = [c["text"] for c in chunks]
         embeddings = self.bi_encoder.encode(texts)
-        faiss.normalize_L2(embeddings) # Normalize for Cosine Sim
+        faiss.normalize_L2(embeddings)
         
-        # 2. Init Index if needed
         if self.index is None:
-            dimension = embeddings.shape[1]
-            self.index = faiss.IndexFlatIP(dimension) # Inner Product = Cosine
+            self.index = faiss.IndexFlatIP(embeddings.shape[1])
             
-        # 3. Add to Index
         self.index.add(embeddings)
         self.metadata.extend(chunks)
         
-        # 4. Save to Disk
         faiss.write_index(self.index, INDEX_FILE)
-        with open(META_FILE, "wb") as f:
-            pickle.dump(self.metadata, f)
-            
+        with open(META_FILE, "wb") as f: pickle.dump(self.metadata, f)
         return len(texts)
 
-    def search(self, query, top_k=5):
-        if not self.index or self.index.ntotal == 0:
-            return []
+    def search_documents(self, query, top_k=5):
+        if not self.index or self.index.ntotal == 0: return []
             
-        # 1. Retrieval
-        candidate_k = top_k * 3
+        # 1. Retrieve MANY chunks (to ensure we find diverse documents)
+        # If we only get top 5 chunks, they might all be from the same document.
+        candidate_k = top_k * 10 
+        
         q_vec = self.bi_encoder.encode([query])
         faiss.normalize_L2(q_vec)
         
         scores, indices = self.index.search(q_vec, min(self.index.ntotal, candidate_k))
         
-        candidates = []
+        # 2. Extract Raw Candidates
+        raw_candidates = []
         for i, idx in enumerate(indices[0]):
             if idx != -1:
-                candidates.append({
+                raw_candidates.append({
                     "text": self.metadata[idx]["text"],
                     "source": self.metadata[idx]["source"],
-                    "page": self.metadata[idx]["page"],
-                    "base_score": scores[0][i]
+                    "bi_score": scores[0][i]
                 })
-                
-        # 2. Re-Ranking
-        pairs = [[query, c["text"]] for c in candidates]
-        cross_scores = self.cross_encoder.predict(pairs)
+
+        # 3. Aggregation: Find the BEST chunk for each document
+        # We group by 'source' and keep the max score
+        doc_map = {} # {filename: {best_score, best_snippet}}
+        
+        for cand in raw_candidates:
+            source = cand['source']
+            score = cand['bi_score']
+            
+            # Initialization
+            if source not in doc_map:
+                doc_map[source] = {"score": score, "snippet": cand['text']}
+            else:
+                # Update if we found a better chunk in the same doc
+                if score > doc_map[source]["score"]:
+                    doc_map[source]["score"] = score
+                    doc_map[source]["snippet"] = cand['text']
+        
+        # 4. Sort Documents by their Best Chunk Score
+        ranked_docs = sorted(doc_map.items(), key=lambda item: item[1]['score'], reverse=True)
         
-        for i, c in enumerate(candidates):
-            c["score"] = cross_scores[i]
+        # 5. Cross-Encoder Verification (Optional but recommended)
+        # We verify the "Best Snippet" to ensure it's not a hallucination
+        final_results = []
+        top_docs = ranked_docs[:top_k] # Only re-rank the top contenders
+        
+        if top_docs:
+            pairs = [[query, doc[1]['snippet']] for doc in top_docs]
+            cross_scores = self.cross_encoder.predict(pairs)
             
-        # Sort
-        final_results = sorted(candidates, key=lambda x: x["score"], reverse=True)
-        return final_results[:top_k]
+            for i, (source, data) in enumerate(top_docs):
+                final_results.append({
+                    "source": source,
+                    "score": cross_scores[i], # High accuracy score
+                    "snippet": data['snippet']
+                })
+                
+            # Final Sort after Cross-Encoder
+            final_results = sorted(final_results, key=lambda x: x["score"], reverse=True)
+            
+        return final_results
 
 # --- UI LOGIC ---
 if 'engine' not in st.session_state:
-    # 1. Try cloud sync first
     IndexManager.load_from_hub()
-    # 2. Start engine
-    st.session_state.engine = RobustSearchEngine()
+    st.session_state.engine = DocSearchEngine()
 
 with st.sidebar:
-    st.header("🗄️ Knowledge Base")
-    uploaded_files = st.file_uploader("Ingest Documents", accept_multiple_files=True)
-    
-    if uploaded_files and st.button("Index Documents"):
-        with st.spinner("Processing..."):
+    st.header("🗄️ Upload Documents")
+    uploaded_files = st.file_uploader("Upload Files", accept_multiple_files=True)
+    if uploaded_files and st.button("Index"):
+        with st.spinner("Indexing..."):
             new_chunks = []
             for f in uploaded_files:
                 txt, fname = parse_file(f)
-                chunks = recursive_chunking(txt, fname)
-                new_chunks.extend(chunks)
-            
+                new_chunks.extend(recursive_chunking(txt, fname))
             if new_chunks:
-                count = st.session_state.engine.add_documents(new_chunks)
+                st.session_state.engine.add_documents(new_chunks)
                 IndexManager.save_to_hub()
-                st.success(f"Added {count} chunks!")
+                st.success("Indexed!")
+
+st.title("⚓ Document Finder")
+st.caption("Locates the specific Instruction or NAVADMIN relevant to your query.")
 
-st.title("⚓ Navy Search (FAISS Architecture)")
-query = st.text_input("Enter Query:")
+query = st.text_input("What are you looking for?", placeholder="e.g. 'FY25 Retention Bonuses'")
 
 if query:
-    results = st.session_state.engine.search(query)
+    results = st.session_state.engine.search_documents(query, top_k=5)
     
-    st.markdown("### 🔍 Results")
-    context_text = ""
+    st.subheader("Top Relevant Documents")
+    
+    if not results:
+        st.info("No documents found.")
+        
     for res in results:
-        context_text += f"Source: {res['source']}\n{res['text']}\n\n"
-        with st.expander(f"{res['source']} (Pg {res['page']}) - Score {res['score']:.2f}", expanded=True):
-            st.markdown(res['text'])
-
-    if st.button("Generate Summary"):
-        from huggingface_hub import InferenceClient
-        client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.3", token=HF_TOKEN)
-        prompt = f"Context:\n{context_text}\n\nUser: {query}\nAnswer:"
-        with st.spinner("Thinking..."):
-            try:
-                st.write(client.text_generation(prompt, max_new_tokens=400))
-            except Exception as e:
-                st.error(f"LLM Error: {e}")
+        score = res['score']
+        
+        # Color coding the confidence
+        if score > 2:
+            border_color = "#09ab3b" # Green
+            confidence = "High Match"
+        elif score > 0:
+            border_color = "#ffbd45" # Orange
+            confidence = "Possible Match"
+        else:
+            border_color = "#ff4b4b" # Red
+            confidence = "Low Match"
+
+        # --- DOCUMENT CARD UI ---
+        with st.container():
+            st.markdown(f"""
+            <div style="
+                border: 1px solid #ddd;
+                border-left: 5px solid {border_color};
+                padding: 15px;
+                border-radius: 5px;
+                margin-bottom: 10px;
+            ">
+                <h3 style="margin:0; padding:0;">📄 {res['source']}</h3>
+                <small style="color: gray;">Confidence: {confidence} ({score:.2f})</small>
+            </div>
+            """, unsafe_allow_html=True)
+            
+            with st.expander("View matching excerpt"):
+                st.markdown(f"**...{res['snippet']}...**")