Update app.py

app.py

@@ -1,22 +1,141 @@
 import os
 import shutil
+import faiss
+import numpy as np
 import gradio as gr
+from sentence_transformers import SentenceTransformer
+from huggingface_hub import HfApi, hf_hub_download, login
 
-# Function to prepare FAISS files for download
+# 🔹 Hugging Face Repository Details
+HF_REPO_ID = "tstone87/repo"  # Your dataset repo
+HF_TOKEN = os.getenv("HF_TOKEN")  # Retrieve token securely
+
+if not HF_TOKEN:
+    raise ValueError("❌ ERROR: Hugging Face token not found. Add it as a secret in the Hugging Face Space settings.")
+
+# 🔹 Authenticate with Hugging Face
+login(token=HF_TOKEN)
+
+# 🔹 File Paths
+EMBEDDINGS_FILE = "policy_embeddings.npy"
+INDEX_FILE = "faiss_index.bin"
+TEXT_FILE = "combined_text_documents.txt"
+
+# 🔹 Load policy text from file
+if os.path.exists(TEXT_FILE):
+    with open(TEXT_FILE, "r", encoding="utf-8") as f:
+        POLICY_TEXT = f.read()
+    print("✅ Loaded policy text from combined_text_documents.txt")
+else:
+    print("❌ ERROR: combined_text_documents.txt not found! Ensure it's uploaded.")
+    POLICY_TEXT = ""
+
+# 🔹 Sentence Embedding Model (Optimized for Speed)
+model = SentenceTransformer("all-MiniLM-L6-v2")
+
+# 🔹 Split policy text into chunks for FAISS indexing
+chunk_size = 500
+chunks = [POLICY_TEXT[i:i+chunk_size] for i in range(0, len(POLICY_TEXT), chunk_size)] if POLICY_TEXT else []
+
+# 🔹 Function to Upload FAISS Files to Hugging Face Hub
+def upload_faiss_to_hf():
+    api = HfApi()
+
+    if os.path.exists(EMBEDDINGS_FILE):
+        print("📤 Uploading FAISS embeddings to Hugging Face...")
+        api.upload_file(
+            path_or_fileobj=EMBEDDINGS_FILE,
+            path_in_repo=EMBEDDINGS_FILE,
+            repo_id=HF_REPO_ID,
+            repo_type="dataset",
+            token=HF_TOKEN,
+        )
+
+    if os.path.exists(INDEX_FILE):
+        print("📤 Uploading FAISS index to Hugging Face...")
+        api.upload_file(
+            path_or_fileobj=INDEX_FILE,
+            path_in_repo=INDEX_FILE,
+            repo_id=HF_REPO_ID,
+            repo_type="dataset",
+            token=HF_TOKEN,
+        )
+
+    print("✅ FAISS files successfully uploaded to Hugging Face.")
+
+# 🔹 Function to Download FAISS Files from Hugging Face Hub if Missing
+def download_faiss_from_hf():
+    if not os.path.exists(EMBEDDINGS_FILE):
+        print("📥 Downloading FAISS embeddings from Hugging Face...")
+        hf_hub_download(repo_id=HF_REPO_ID, filename=EMBEDDINGS_FILE, local_dir=".", token=HF_TOKEN)
+
+    if not os.path.exists(INDEX_FILE):
+        print("📥 Downloading FAISS index from Hugging Face...")
+        hf_hub_download(repo_id=HF_REPO_ID, filename=INDEX_FILE, local_dir=".", token=HF_TOKEN)
+
+    print("✅ FAISS files downloaded from Hugging Face.")
+
+# 🔹 Check if FAISS Files Exist, Otherwise Download or Generate
+if os.path.exists(EMBEDDINGS_FILE) and os.path.exists(INDEX_FILE):
+    print("✅ FAISS files found locally. Loading from disk...")
+    embeddings = np.load(EMBEDDINGS_FILE)
+    index = faiss.read_index(INDEX_FILE)
+else:
+    print("🚀 FAISS files not found! Downloading from Hugging Face...")
+    download_faiss_from_hf()
+
+    if os.path.exists(EMBEDDINGS_FILE) and os.path.exists(INDEX_FILE):
+        embeddings = np.load(EMBEDDINGS_FILE)
+        index = faiss.read_index(INDEX_FILE)
+    else:
+        print("🚀 No FAISS files found. Recomputing...")
+        if chunks:
+            embeddings = np.array([model.encode(chunk) for chunk in chunks])
+
+            # Save embeddings for future use
+            np.save(EMBEDDINGS_FILE, embeddings)
+
+            # Use FAISS optimized index for faster lookup
+            d = embeddings.shape[1]
+            nlist = 10  # Number of clusters
+            index = faiss.IndexIVFFlat(faiss.IndexFlatL2(d), d, nlist)
+            index.train(embeddings)
+            index.add(embeddings)
+            index.nprobe = 2  # Speed optimization
+
+            # Save FAISS index
+            faiss.write_index(index, INDEX_FILE)
+            upload_faiss_to_hf()  # Upload FAISS files to Hugging Face
+            print("✅ FAISS index created and saved.")
+        else:
+            print("❌ ERROR: No text to index. Check combined_text_documents.txt.")
+            index = None
+
+# 🔹 Function to Search FAISS
+def search_policy(query, top_k=3):
+    if index is None:
+        return "Error: FAISS index is not available."
+
+    query_embedding = model.encode(query).reshape(1, -1)
+    distances, indices = index.search(query_embedding, top_k)
+
+    return "\n\n".join([chunks[i] for i in indices[0] if i < len(chunks)])
+
+# 🔹 Gradio UI to Download FAISS Files
 def prepare_faiss_files():
-    if os.path.exists("policy_embeddings.npy") and os.path.exists("faiss_index.bin"):
-        shutil.copy("policy_embeddings.npy", "/mnt/data/policy_embeddings.npy")
-        shutil.copy("faiss_index.bin", "/mnt/data/faiss_index.bin")
+    if os.path.exists(EMBEDDINGS_FILE) and os.path.exists(INDEX_FILE):
+        shutil.copy(EMBEDDINGS_FILE, "/mnt/data/policy_embeddings.npy")
+        shutil.copy(INDEX_FILE, "/mnt/data/faiss_index.bin")
         return "✅ FAISS files are ready for download. Go to the 'Files' tab in Hugging Face Space and download them."
     else:
         return "❌ FAISS files not found. Try running the chatbot first to generate them."
 
-# Gradio UI to trigger FAISS file preparation
 with gr.Blocks() as download_ui:
     gr.Markdown("### 🔽 Download FAISS Files")
     download_button = gr.Button("Prepare FAISS Files for Download")
     output_text = gr.Textbox()
     download_button.click(fn=prepare_faiss_files, outputs=output_text)
 
-# Launch the download interface
 download_ui.launch()
+
+print("✅ FAISS index successfully loaded.")