Update build_rag.py

build_rag.py

@@ -1,5 +1,3 @@
-# build_rag.py (Updated with Normalization and Cosine Distance)
-
 import json
 import os
 import pandas as pd
@@ -15,12 +13,14 @@ import traceback
 # --- Configuration ---
 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 TO AVOID CONFUSION ***
+DATASET_REPO = "broadfield-dev/bible-chromadb-mpnet"
 STATUS_FILE = "build_status.log"
 JSON_DIRECTORY = 'bible_json'
 CHUNK_SIZE = 3
-EMBEDDING_BATCH_SIZE = 16
+EMBEDDING_BATCH_SIZE = 16 # Adjust based on available VRAM
 # (BOOK_ID_TO_NAME dictionary remains the same)
 BOOK_ID_TO_NAME = {
     1: "Genesis", 2: "Exodus", 3: "Leviticus", 4: "Numbers", 5: "Deuteronomy",
@@ -44,6 +44,12 @@ def update_status(message):
     with open(STATUS_FILE, "w") as f:
         f.write(message)
 
+# 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 process_bible_json_files(directory_path: str, chunk_size: int) -> pd.DataFrame:
     # (This function is unchanged)
     all_verses = []
@@ -84,7 +90,6 @@ def main():
         shutil.rmtree(CHROMA_PATH)
     client = chromadb.PersistentClient(path=CHROMA_PATH)
     
-    # *** FIX 1: SET THE DISTANCE FUNCTION FOR THE COLLECTION ***
     collection = client.create_collection(
         name=COLLECTION_NAME,
         metadata={"hnsw:space": "cosine"} # Use cosine distance
@@ -99,16 +104,18 @@ def main():
         batch_df = bible_chunks_df.iloc[i:i+EMBEDDING_BATCH_SIZE]
         texts = batch_df['text'].tolist()
         
-        inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=512).to(model.device)
+        # *** CHANGE 3: USE THE CORRECT POOLING STRATEGY FOR SBERT MODELS ***
+        encoded_input = tokenizer(texts, padding=True, truncation=True, return_tensors='pt').to(model.device)
         with torch.no_grad():
-            outputs = model(**inputs)
+            model_output = model(**encoded_input)
         
-        # *** FIX 2: NORMALIZE THE EMBEDDINGS ***
-        embeddings = F.normalize(outputs.last_hidden_state.mean(dim=1), p=2, dim=1)
+        # Perform pooling and normalization
+        sentence_embeddings = mean_pooling(model_output, encoded_input['attention_mask'])
+        normalized_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)
         
         collection.add(
             ids=[str(j) for j in range(i, i + len(batch_df))],
-            embeddings=embeddings.cpu().tolist(), # Convert to list after normalization
+            embeddings=normalized_embeddings.cpu().tolist(), # Convert to list
             documents=texts,
             metadatas=batch_df[['reference', 'version']].to_dict('records')
         )