Update app.py

app.py

@@ -1,5 +1,3 @@
-# app.py (Updated with Normalization for the query)
-
 import sys
 import subprocess
 from flask import Flask, render_template, request, flash, redirect, url_for, jsonify
@@ -16,14 +14,22 @@ app.secret_key = os.urandom(24)
 
 CHROMA_PATH = "chroma_db"
 COLLECTION_NAME = "bible_verses"
-MODEL_NAME = "google/embeddinggemma-300m"
-DATASET_REPO = "broadfield-dev/bible-chromadb-gemma"
+# *** CHANGE 1: UPDATE THE MODEL NAME ***
+MODEL_NAME = "sentence-transformers/all-mpnet-base-v2"
+# *** CHANGE 2: UPDATE THE DATASET REPO NAME ***
+DATASET_REPO = "broadfield-dev/bible-chromadb-mpnet"
 STATUS_FILE = "build_status.log"
 
 chroma_collection = None
 tokenizer = None
 embedding_model = None
 
+# Mean Pooling Function - Take attention mask into account for correct averaging
+def mean_pooling(model_output, attention_mask):
+    token_embeddings = model_output[0] #First element of model_output contains all token embeddings
+    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
+
 def load_resources():
     # (This function is unchanged)
     global chroma_collection, tokenizer, embedding_model
@@ -89,12 +95,13 @@ def search():
     if not user_query:
         return render_template('index.html', results=[])
 
-    inputs = tokenizer(user_query, return_tensors="pt")
+    # *** CHANGE 3: USE THE CORRECT POOLING STRATEGY FOR SBERT MODELS ***
+    encoded_input = tokenizer([user_query], padding=True, truncation=True, return_tensors='pt')
     with torch.no_grad():
-        outputs = embedding_model(**inputs)
+        model_output = embedding_model(**encoded_input)
     
-    # *** FIX: NORMALIZE THE QUERY EMBEDDING ***
-    query_embedding = F.normalize(outputs.last_hidden_state.mean(dim=1), p=2, dim=1)
+    sentence_embeddings = mean_pooling(model_output, encoded_input['attention_mask'])
+    query_embedding = F.normalize(sentence_embeddings, p=2, dim=1)
 
     search_results = chroma_collection.query(
         query_embeddings=query_embedding.cpu().tolist(),