smart-models/Placebo_AI

1import os2import json3from langchain_chroma import Chroma4from langchain_ollama import OllamaEmbeddings5from langchain_core.documents import Document6from langchain_text_splitters import RecursiveCharacterTextSplitter7from tqdm import tqdm8 9def run_incremental_indexing(textbooks, base_dir, persist_directory="vector_store", batch_size=50):10    print("=" * 70)11    print("STARTING INCREMENTAL INDEXING FOR MBBS TEXTBOOKS")12    print("=" * 70)13    14    # Initialize embeddings15    embeddings = OllamaEmbeddings(model="nomic-embed-text:latest")16    17    # Connect to the existing vector store18    if not os.path.exists(persist_directory):19        print(f"Error: Vector store not found at {persist_directory}. Please run a full setup first.")20        return21        22    print(f"Connecting to existing Chroma DB at {persist_directory}...")23    vectorstore = Chroma(24        persist_directory=persist_directory,25        embedding_function=embeddings26    )27    print(f"Connected. Current chunk count: {vectorstore._collection.count():,}")28    29    # Initialize Text Splitter30    text_splitter = RecursiveCharacterTextSplitter(31        chunk_size=500,32        chunk_overlap=150,33        length_function=len,34        is_separator_regex=False,35    )36    37    all_documents = []38    39    for textbook in textbooks:40        jsonl_path = os.path.join(base_dir, textbook)41        if not os.path.exists(jsonl_path):42            print(f"Skipping {textbook} (Not found)...")43            continue44 45        print(f"Reading data from {textbook}...")46        with open(jsonl_path, "r", encoding="utf-8") as f:47            for line in f:48                try:49                    data = json.loads(line)50                    content = data.get("content", "")51                    if not content:52                        continue53                    54                    base_metadata = {55                        "subject": data.get("subject"),56                        "book_name": data.get("book_name"),57                        "page_number": data.get("page_number"),58                        "category": data.get("category"),59                        "image_path": data.get("image_path"),60                        "track": "mbbs"61                    }62 63                    header = f"SUBJECT: {data.get('subject')} | BOOK: {data.get('book_name')} | PAGE: {data.get('page_number')}\n"64                    chunks = text_splitter.split_text(content)65                    66                    for chunk in chunks:67                        doc = Document(68                            page_content=header + chunk,69                            metadata=base_metadata70                        )71                        all_documents.append(doc)72                        73                except Exception as e:74                    print(f"Error parsing line: {e}")75 76    total_docs = len(all_documents)77    if total_docs == 0:78        print("No new documents to index.")79        return80 81    print(f"Adding {total_docs} new chunks to the vector database in batches of {batch_size}...")82    83    # Add documents in batches84    for i in tqdm(range(0, total_docs, batch_size), desc="Adding chunks"):85        batch = all_documents[i : i + batch_size]86        vectorstore.add_documents(batch)87    88    # Get updated count89    updated_count = vectorstore._collection.count()90    print(f"\nIncremental Indexing Complete!")91    print(f"Previous Chunk Count: {updated_count - total_docs:,}")92    print(f"Added Chunks:         {total_docs:,}")93    print(f"New Total Chunks:     {updated_count:,}")94    print("=" * 70)95 96if __name__ == "__main__":97    import os98    BASE_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "data", "processed")99    VECTOR_DB_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "db", "vector_store")100    101    # Find only the MBBS jsonl files in the BASE_DIR102    TEXTBOOKS = ["mbbs_ENT.jsonl", "mbbs_Embryology.jsonl", "mbbs_NEUROSCIENCE.jsonl"]103    104    if not TEXTBOOKS:105        print("No MBBS jsonl files found. Run process_mbbs_books.py first.")106    else:107        print(f"Found {len(TEXTBOOKS)} MBBS textbooks to index: {', '.join(TEXTBOOKS)}")108        run_incremental_indexing(TEXTBOOKS, BASE_DIR, VECTOR_DB_DIR, batch_size=50)109