Update app.py

app.py

@@ -1,104 +1,78 @@
+# app.py (Updated with Normalization for the query)
+
 import sys
 import subprocess
 from flask import Flask, render_template, request, flash, redirect, url_for, jsonify
 import torch
+import torch.nn.functional as F # Import the functional module
 from transformers import AutoTokenizer, AutoModel
 import os
 import chromadb
 from huggingface_hub import snapshot_download
 
-# --- 1. Initialize Flask App ---
+# (App setup and load_resources function are unchanged)
 app = Flask(__name__)
 app.secret_key = os.urandom(24)
 
-# --- 2. Configuration & Resource Loading ---
-print("Starting application...")
-
-# --- Configuration (Must match build_rag.py) ---
 CHROMA_PATH = "chroma_db"
 COLLECTION_NAME = "bible_verses"
 MODEL_NAME = "google/embeddinggemma-300m"
 DATASET_REPO = "broadfield-dev/bible-chromadb-gemma"
-STATUS_FILE = "build_status.log"  # File to track build status
+STATUS_FILE = "build_status.log"
 
-# --- Global variables for resources ---
 chroma_collection = None
 tokenizer = None
 embedding_model = None
 
 def load_resources():
-    """Downloads the DB from the Hub if not present, then loads it and the model."""
+    # (This function is unchanged)
     global chroma_collection, tokenizer, embedding_model
-    if chroma_collection and embedding_model:
-        return True
-    
+    if chroma_collection and embedding_model: return True
     print("Attempting to load resources...")
     try:
         if not os.path.exists(CHROMA_PATH) or not os.listdir(CHROMA_PATH):
             print(f"Local DB not found. Downloading from '{DATASET_REPO}'...")
-            snapshot_download(
-                repo_id=DATASET_REPO,
-                repo_type="dataset",
-                local_dir=CHROMA_PATH,
-                local_dir_use_symlinks=False
-            )
+            snapshot_download(repo_id=DATASET_REPO, repo_type="dataset", local_dir=CHROMA_PATH, local_dir_use_symlinks=False)
             print("Database files downloaded.")
         else:
             print("Local database files found.")
-
         client = chromadb.PersistentClient(path=CHROMA_PATH)
         collection = client.get_collection(name=COLLECTION_NAME)
         if collection.count() == 0:
             print("Warning: Database collection is empty.")
             return False
-            
         chroma_collection = collection
         print(f"Successfully connected to DB with {collection.count()} items.")
-        
         tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
         embedding_model = AutoModel.from_pretrained(MODEL_NAME)
         print(f"Embedding model '{MODEL_NAME}' loaded successfully.")
-        
         return True
     except Exception as e:
-        print(f"Could not load resources. The database may not be built yet or the repo is empty.")
+        print(f"Could not load resources. DB may not be built or repo is empty.")
         print(f"Error: {e}")
         return False
 
 resources_loaded = load_resources()
 
-# --- 3. Define App Routes ---
-
+# (home, build_rag_route, and status routes are unchanged)
 @app.route('/')
 def home():
     return render_template('index.html')
 
 @app.route('/build-rag', methods=['POST'])
 def build_rag_route():
-    """Triggers the build script and immediately responds."""
     try:
-        # Clear old status and set to "In Progress"
-        with open(STATUS_FILE, "w") as f:
-            f.write("IN_PROGRESS: Starting build process...")
-        
-        # Start the build process in the background
+        with open(STATUS_FILE, "w") as f: f.write("IN_PROGRESS: Starting build process...")
         subprocess.Popen([sys.executable, "build_rag.py"])
-        
         return jsonify({"status": "started"})
     except Exception as e:
-        with open(STATUS_FILE, "w") as f:
-            f.write(f"FAILED: Could not start process - {e}")
+        with open(STATUS_FILE, "w") as f: f.write(f"FAILED: Could not start process - {e}")
         return jsonify({"status": "error", "message": str(e)}), 500
 
 @app.route('/status')
 def status():
-    """Endpoint for the frontend to poll for build status."""
-    if not os.path.exists(STATUS_FILE):
-        return jsonify({"status": "NOT_STARTED"})
-    
-    with open(STATUS_FILE, "r") as f:
-        status_line = f.read().strip()
-    
+    if not os.path.exists(STATUS_FILE): return jsonify({"status": "NOT_STARTED"})
+    with open(STATUS_FILE, "r") as f: status_line = f.read().strip()
     status_code, _, message = status_line.partition(': ')
     return jsonify({"status": status_code, "message": message})
     
@@ -118,10 +92,12 @@ def search():
     inputs = tokenizer(user_query, return_tensors="pt")
     with torch.no_grad():
         outputs = embedding_model(**inputs)
-    query_embedding = outputs.last_hidden_state.mean(dim=1).cpu().numpy()
+    
+    # *** FIX: NORMALIZE THE QUERY EMBEDDING ***
+    query_embedding = F.normalize(outputs.last_hidden_state.mean(dim=1), p=2, dim=1)
 
     search_results = chroma_collection.query(
-        query_embeddings=query_embedding.tolist(),
+        query_embeddings=query_embedding.cpu().tolist(),
         n_results=10
     )
 
@@ -138,6 +114,5 @@ def search():
     
     return render_template('index.html', results=results_list, query=user_query)
 
-# --- 4. Run the App ---
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=7860)