# build_rag.py import json import os import pandas as pd import torch from transformers import AutoTokenizer, AutoModel import chromadb import sys from tqdm import tqdm from huggingface_hub import HfApi, create_repo import traceback # --- Configuration --- CHROMA_PATH = "chroma_db" COLLECTION_NAME = "bible_verses" MODEL_NAME = "sentence-transformers/multi-qa-mpnet-base-dot-v1" DATASET_REPO = "broadfield-dev/bible-chromadb-multi-qa-mpnet" # This can remain the same STATUS_FILE = "build_status.log" JSON_DIRECTORY = 'bible_json' CHUNK_SIZE = 3 EMBEDDING_BATCH_SIZE = 16 # (BOOK_ID_TO_NAME dictionary remains the same) BOOK_ID_TO_NAME = { 1: "Genesis", 2: "Exodus", 3: "Leviticus", 4: "Numbers", 5: "Deuteronomy", 6: "Joshua", 7: "Judges", 8: "Ruth", 9: "1 Samuel", 10: "2 Samuel", 11: "1 Kings", 12: "2 Kings", 13: "1 Chronicles", 14: "2 Chronicles", 15: "Ezra", 16: "Nehemiah", 17: "Esther", 18: "Job", 19: "Psalms", 20: "Proverbs", 21: "Ecclesiastes", 22: "Song of Solomon", 23: "Isaiah", 24: "Jeremiah", 25: "Lamentations", 26: "Ezekiel", 27: "Daniel", 28: "Hosea", 29: "Joel", 30: "Amos", 31: "Obadiah", 32: "Jonah", 33: "Micah", 34: "Nahum", 35: "Habakkuk", 36: "Zephaniah", 37: "Haggai", 38: "Zechariah", 39: "Malachi", 40: "Matthew", 41: "Mark", 42: "Luke", 43: "John", 44: "Acts", 45: "Romans", 46: "1 Corinthians", 47: "2 Corinthians", 48: "Galatians", 49: "Ephesians", 50: "Philippians", 51: "Colossians", 52: "1 Thessalonians", 53: "2 Thessalonians", 54: "1 Timothy", 55: "2 Timothy", 56: "Titus", 57: "Philemon", 58: "Hebrews", 59: "James", 60: "1 Peter", 61: "2 Peter", 62: "1 John", 63: "2 John", 64: "3 John", 65: "Jude", 66: "Revelation" } def update_status(message): print(message) with open(STATUS_FILE, "w") as f: f.write(message) # Mean Pooling Function def mean_pooling(model_output, attention_mask): token_embeddings = model_output[0] 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: all_verses = [] if not os.path.exists(directory_path) or not os.listdir(directory_path): raise FileNotFoundError(f"Directory '{directory_path}' is empty or does not exist.") for filename in os.listdir(directory_path): if filename.endswith('.json'): version_name = os.path.splitext(filename)[0].split('_')[-1].upper() file_path = os.path.join(directory_path, filename) with open(file_path, 'r') as f: data = json.load(f) rows = data.get("resultset", {}).get("row", []) for row in rows: field = row.get("field", []) if len(field) == 5: _id, book_id, chapter, verse, text = field book_name = BOOK_ID_TO_NAME.get(book_id, "Unknown Book") all_verses.append({'version': version_name, 'book_name': book_name, 'chapter': chapter, 'verse': verse, 'text': text.strip()}) if not all_verses: raise ValueError("No verses were processed.") df = pd.DataFrame(all_verses) all_chunks = [] for (version, book_name, chapter), group in df.groupby(['version', 'book_name', 'chapter']): group = group.sort_values('verse').reset_index(drop=True) for i in range(0, len(group), chunk_size): chunk_df = group.iloc[i:i+chunk_size] combined_text = " ".join(chunk_df['text']) start_verse, end_verse = chunk_df.iloc[0]['verse'], chunk_df.iloc[-1]['verse'] reference = f"{book_name} {chapter}:{start_verse}" if start_verse == end_verse else f"{book_name} {chapter}:{start_verse}-{end_verse}" # *** CHANGE 1: ADD MORE METADATA TO EACH CHUNK *** all_chunks.append({ 'text': combined_text, 'reference': reference, 'version': version, 'book_name': book_name, 'chapter': chapter }) return pd.DataFrame(all_chunks) def main(): update_status("IN_PROGRESS: Step 1/5 - Processing JSON files...") bible_chunks_df = process_bible_json_files(JSON_DIRECTORY, chunk_size=CHUNK_SIZE) update_status("IN_PROGRESS: Step 2/5 - Setting up local ChromaDB...") if os.path.exists(CHROMA_PATH): import shutil shutil.rmtree(CHROMA_PATH) client = chromadb.PersistentClient(path=CHROMA_PATH) collection = client.create_collection(name=COLLECTION_NAME, metadata={"hnsw:space": "cosine"}) update_status(f"IN_PROGRESS: Step 3/5 - Loading embedding model '{MODEL_NAME}'...") tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) model = AutoModel.from_pretrained(MODEL_NAME, device_map="auto") update_status("IN_PROGRESS: Step 4/5 - Generating embeddings...") for i in tqdm(range(0, len(bible_chunks_df), EMBEDDING_BATCH_SIZE), desc="Embedding Chunks"): batch_df = bible_chunks_df.iloc[i:i+EMBEDDING_BATCH_SIZE] texts = batch_df['text'].tolist() encoded_input = tokenizer(texts, padding=True, truncation=True, return_tensors='pt').to(model.device) with torch.no_grad(): model_output = model(**encoded_input) embeddings = mean_pooling(model_output, encoded_input['attention_mask']) collection.add( ids=[str(j) for j in range(i, i + len(batch_df))], embeddings=embeddings.cpu().tolist(), documents=texts, # *** CHANGE 2: SAVE THE NEW METADATA FIELDS TO THE DATABASE *** metadatas=batch_df[['reference', 'version', 'book_name', 'chapter']].to_dict('records') ) update_status(f"IN_PROGRESS: Step 5/5 - Pushing database to Hugging Face Hub '{DATASET_REPO}'...") create_repo(repo_id=DATASET_REPO, repo_type="dataset", exist_ok=True) api = HfApi() api.upload_folder(folder_path=CHROMA_PATH, repo_id=DATASET_REPO, repo_type="dataset") update_status("SUCCESS: Build complete! The application is ready.") if __name__ == "__main__": try: main() except Exception as e: error_message = traceback.format_exc() if "401" in str(e) or "Unauthorized" in str(e): update_status("FAILED: Hugging Face authentication error. Ensure your HF_TOKEN secret has WRITE permissions.") else: update_status(f"FAILED: An unexpected error occurred. Check Space logs. Error: {e}") print(error_message, file=sys.stderr)