smart-models/Placebo_AI

1import json2import os3from langchain_chroma import Chroma4from langchain_ollama import OllamaEmbeddings5from langchain_core.documents import Document6from langchain_text_splitters import RecursiveCharacterTextSplitter7from tqdm import tqdm8 9def run_indexing(textbooks, base_dir, persist_directory="vector_store", batch_size=50):10    # Initialize embeddings11    embeddings = OllamaEmbeddings(model="nomic-embed-text:latest")12    13    # Initialize Text Splitter14    text_splitter = RecursiveCharacterTextSplitter(15        chunk_size=1000,16        chunk_overlap=150,17        length_function=len,18        is_separator_regex=False,19    )20    21    all_documents = []22    23    for textbook in textbooks:24        jsonl_path = os.path.join(base_dir, textbook)25        if not os.path.exists(jsonl_path):26            print(f"Skipping {textbook} (Not found)...")27            continue28 29        print(f"Loading data from {textbook}...")30        with open(jsonl_path, "r", encoding="utf-8") as f:31            for line in f:32                try:33                    data = json.loads(line)34                    content = data.get("content", "")35                    if not content:36                        continue37                    38                    base_metadata = {39                        "subject": data.get("subject"),40                        "book_name": data.get("book_name"),41                        "page_number": data.get("page_number"),42                        "category": data.get("category"),43                        "image_path": data.get("image_path"),44                        "track": "pharmacy"45                    }46 47                    header = f"SUBJECT: {data.get('subject')} | BOOK: {data.get('book_name')} | PAGE: {data.get('page_number')}\n"48                    chunks = text_splitter.split_text(content)49                    50                    for chunk in chunks:51                        doc = Document(52                            page_content=header + chunk,53                            metadata=base_metadata54                        )55                        all_documents.append(doc)56                        57                except Exception as e:58                    print(f"Error parsing line: {e}")59 60    total_docs = len(all_documents)61    if total_docs == 0:62        print("No documents to index.")63        return64 65    # Delete old vector store to ensure clean merge66    if os.path.exists(persist_directory):67        import shutil68        print(f"Cleaning old vector store at {persist_directory}...")69        try:70            shutil.rmtree(persist_directory)71        except Exception as e:72            print(f"Warning: Could not delete old vector store: {e}. Please ensure no other process is using it.")73            return74 75    print(f"Starting batch indexing: {total_docs} chunks created, batch size {batch_size}...")76    77    # Initialize the store with the first batch78    first_batch = all_documents[:batch_size]79    vectorstore = Chroma.from_documents(80        documents=first_batch,81        embedding=embeddings,82        persist_directory=persist_directory83    )84    85    # Add remaining batches86    for i in tqdm(range(batch_size, total_docs, batch_size), desc="Indexing progress"):87        batch = all_documents[i : i + batch_size]88        vectorstore.add_documents(batch)89    90    print(f"\nIndexing complete! {total_docs} chunks indexed into: {persist_directory}")91 92if __name__ == "__main__":93    import os94    BASE_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "data", "processed")95    VECTOR_DB_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "db", "vector_store")96    97    # Dynamically find all jsonl files in the BASE_DIR98    TEXTBOOKS = [f for f in os.listdir(BASE_DIR) if f.endswith(".jsonl")]99    100    print(f"Found {len(TEXTBOOKS)} textbooks for indexing: {', '.join(TEXTBOOKS)}")101    102    run_indexing(TEXTBOOKS, BASE_DIR, VECTOR_DB_DIR, batch_size=100)103