# -*- coding: utf-8 -*- """ build_vector_db.py — Build the local ChromaDB knowledge base for Serinity. Run once before starting the app: python scripts/build_vector_db.py IMPORTANT: Ollama must be running and nomic-embed-text must be pulled: ollama pull nomic-embed-text """ import os import json from collections import Counter from datasets import load_dataset from langchain_core.documents import Document from langchain_chroma import Chroma from langchain_ollama import OllamaEmbeddings from tqdm import tqdm from dotenv import load_dotenv load_dotenv() # Configuration from .env OLLAMA_HOST = os.getenv("OLLAMA_HOST", "http://localhost:11434") EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "nomic-embed-text") CHROMA_PERSIST_DIR = os.getenv("CHROMA_PERSIST_DIR", "./chroma_db") CHROMA_COLLECTION = os.getenv("CHROMA_COLLECTION_NAME","mhcva-knowledge") DATASET_NAME = "Compumacy/Psych_data" BATCH_SIZE = 100 # Generic wellness filter markers GENERIC_WELLNESS_MARKERS = [ "general wellness", "lifestyle advice", "wellness discussion", "healthy habits", "nutrition", "exercise tips", "sleep hygiene tips", "mindfulness general", "self-care tips", ] def _is_generic_wellness(medical_context: str) -> bool: """Return True if the row looks like off-topic generic wellness content.""" if not medical_context: return False lower = medical_context.lower() return any(marker in lower for marker in GENERIC_WELLNESS_MARKERS) def download_dataset(): """Download the specified dataset from HuggingFace.""" print("Downloading dataset from HuggingFace...") ds = load_dataset(DATASET_NAME) print(f"Downloaded {len(ds['train']):,} rows.") return ds def filter_and_audit(data): """ Sample medical_context values, drop generic wellness rows, and print a summary for sanity checking before committing. """ print("\nRunning corpus quality filter...") context_counts = Counter() dropped_indices = set() dropped_samples = [] for idx, item in enumerate(tqdm(data, desc="Auditing medical_context")): metadata = item.get("metadata", {}) if isinstance(metadata, str): try: metadata = json.loads(metadata) except Exception: metadata = {} ctx = metadata.get("medical_context", "") or "" context_counts[ctx] += 1 if _is_generic_wellness(ctx): dropped_indices.add(idx) if len(dropped_samples) < 10: dropped_samples.append(ctx) print(f"\n{'='*60}") print("CORPUS FILTER SUMMARY") print(f"{'='*60}") print(f" Total rows: {len(data):,}") print(f" Rows dropped: {len(dropped_indices):,} ({100*len(dropped_indices)/len(data):.1f}%)") print(f" Rows kept: {len(data) - len(dropped_indices):,}") print("\n Top 15 medical_context values (before filtering):") for ctx, count in context_counts.most_common(15): label = ctx[:60] + "..." if len(ctx) > 60 else ctx flag = " [DROPPED]" if _is_generic_wellness(ctx) else "" print(f" [{count:>6}] {label}{flag}") if dropped_samples: print("\n Sample dropped contexts:") for s in dropped_samples: print(f" - {s[:80]}") print(f"{'='*60}\n") return dropped_indices def create_documents(data, dropped_indices: set, mode: str = "assistant_only"): """ Build Document objects from the dataset. Modes: assistant_only | qa_pairs | both_separate """ print(f"Creating documents (mode: {mode}, skipping {len(dropped_indices):,} filtered rows)...") documents = [] for idx, item in enumerate(tqdm(data, desc="Building documents")): if idx in dropped_indices: continue user_msg = item.get("user_message", "") or "" assistant_msg = item.get("assistant_message","") or "" metadata = item.get("metadata", {}) or {} if isinstance(metadata, str): try: metadata = json.loads(metadata) except Exception: metadata = {} doc_metadata = { "chunk_id": idx, "source_pdf": metadata.get("source_pdf", "unknown"), "page_number": metadata.get("page_number", -1), "confidence_score":metadata.get("confidence_score", 0.0), "medical_context": metadata.get("medical_context", ""), } if mode == "assistant_only": if assistant_msg.strip(): documents.append(Document(page_content=assistant_msg.strip(), metadata=doc_metadata)) elif mode == "qa_pairs": if user_msg and assistant_msg: combined = f"Question: {user_msg.strip()}\n\nAnswer: {assistant_msg.strip()}" documents.append(Document(page_content=combined, metadata=doc_metadata)) elif mode == "both_separate": if user_msg.strip(): q_meta = {**doc_metadata, "type": "question"} documents.append(Document(page_content=f"Question: {user_msg.strip()}", metadata=q_meta)) if assistant_msg.strip(): a_meta = {**doc_metadata, "type": "answer"} documents.append(Document(page_content=assistant_msg.strip(), metadata=a_meta)) print(f"Created {len(documents):,} document objects.") return documents def build_chroma_db(documents: list): """ Push documents to local ChromaDB using Ollama embeddings. Checks how many docs are already stored and resumes from the last batch. """ if not documents: raise ValueError("No documents to index — check filter settings.") print(f"\nInitializing OllamaEmbeddings ({EMBEDDING_MODEL})...") embeddings = OllamaEmbeddings( model=EMBEDDING_MODEL, base_url=OLLAMA_HOST, keep_alive=-1, ) print("Verifying Ollama connection...") test_vec = embeddings.embed_query("test") print(f"Embedding model ready (dimension: {len(test_vec)}).") os.makedirs(CHROMA_PERSIST_DIR, exist_ok=True) vectorstore = Chroma( collection_name=CHROMA_COLLECTION, embedding_function=embeddings, persist_directory=CHROMA_PERSIST_DIR, ) already_stored = vectorstore._collection.count() start_batch = (already_stored // BATCH_SIZE) start_doc = start_batch * BATCH_SIZE if already_stored > 0: print(f"\nResuming from batch {start_batch} ({already_stored:,} docs already stored).") print(f"Skipping first {start_doc:,} documents.") else: print(f"\nPushing to ChromaDB at {CHROMA_PERSIST_DIR} ...") print(f"Collection: {CHROMA_COLLECTION}") print(f"Total docs: {len(documents):,}") print(f"Remaining docs: {len(documents) - start_doc:,}") print(f"Batch size: {BATCH_SIZE}") print("This will take a while (embedding locally)...\n") remaining_docs = documents[start_doc:] total_batches = (len(remaining_docs) + BATCH_SIZE - 1) // BATCH_SIZE for i in tqdm(range(0, len(remaining_docs), BATCH_SIZE), total=total_batches, desc="Embedding batches"): batch = remaining_docs[i : i + BATCH_SIZE] vectorstore.add_documents(batch) final_count = vectorstore._collection.count() print(f"\nChromaDB build complete! {final_count:,} vectors stored.") return vectorstore def main(): import sys resuming = "--resume" in sys.argv print("=" * 60) print(" Serinity — LOCAL KNOWLEDGE BASE BUILDER") print(" Target: ChromaDB + nomic-embed-text (Ollama)") if resuming: print(" MODE: RESUME (skipping already-embedded batches)") print("=" * 60) print("\nChoose indexing mode:") print(" 1. Assistant messages only (default — cleanest clinical answers)") print(" 2. Q&A pairs (combined question + answer)") print(" 3. Both separate (individual question and answer docs)") choice = input("\nEnter choice (1/2/3) [1]: ").strip() or "1" mode_map = {"1": "assistant_only", "2": "qa_pairs", "3": "both_separate"} mode = mode_map.get(choice, "assistant_only") try: dataset = download_dataset() data = dataset["train"] dropped_idx = filter_and_audit(data) if not resuming: print("\nProceed with the filter above? (Ctrl+C to abort, Enter to continue)") input() else: print("\n(Resume mode — skipping filter confirmation, using same filter as before.)") documents = create_documents(data, dropped_idx, mode=mode) build_chroma_db(documents) print("\nKnowledge base is ready. Start the app with: uvicorn main:app --reload") except KeyboardInterrupt: print("\nAborted by user. Run with --resume to continue from where you stopped.") except Exception as e: print(f"\nERROR: {e}") raise if __name__ == "__main__": main()