generate_emb.py

import os
import argparse
import uuid
import time
import chromadb
from tqdm import tqdm
from langchain_text_splitters import MarkdownHeaderTextSplitter, RecursiveCharacterTextSplitter

from gemini_api import NewGeminiEmbeddings


def generate_embeddings(input_file, db_path):
    # Load the Markdown File
    try:
        with open(input_file, "r", encoding="utf-8") as file:
            markdown_content = file.read()
        print(f"Loaded Markdown file: {len(markdown_content)} characters.")
    except FileNotFoundError:
        print(f"Error: Could not find {input_file}. Please check the path.")
        return
    
    # Pass 1: Semantic Chunking (By Markdown Headers)
    # This ensures that sections like "Methodofor _ in chunkslogy" or "Conclusion" stay logically grouped.
    headers_to_split_on = [
        ("#", "Header 1"),
        ("##", "Header 2"),
        ("###", "Header 3"),
        ("####", "Header 4"),
    ]

    markdown_splitter = MarkdownHeaderTextSplitter(
        headers_to_split_on=headers_to_split_on,
        strip_headers=False # Keep the headers in the text for context
    )

    semantic_chunks = markdown_splitter.split_text(markdown_content)
    print(f"Pass 1 Complete: Split into {len(semantic_chunks)} semantic sections based on headers.")

    # Pass 2: Size Chunking (Recursive Character Splitter)
    # Markdown sections can still be too large for an embedding model. 
    # We break them down to 1000 characters with a 200 character overlap.
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=4000,
        chunk_overlap=200,
        # These separators ensure we don't split in the middle of a paragraph or table row if possible
        separators=["\n\n", "\n", "(?<=\. )", " ", ""] 
    )

    final_chunks = text_splitter.split_documents(semantic_chunks)
    print(f"Pass 2 Complete: Split into {len(final_chunks)} perfectly sized chunks for embedding.")

    # Add the source document name to the metadata of each chunk
    for chunk in final_chunks:
        chunk.metadata["source"] = os.path.basename(input_file)

    # Generate Embeddings
    texts_to_embed = [chunk.page_content for chunk in final_chunks]
    metadatas = [chunk.metadata for chunk in final_chunks]
    ids = [str(uuid.uuid4()) for chunk in final_chunks]
    print(f"Generating embeddings for {len(texts_to_embed)} chunks...")
    embeddings = []
    batch_size = 32  # Adjust based on your needs and API limits
    embedding_model = NewGeminiEmbeddings(model="gemini-embedding-001")
    for i in tqdm(range(0, len(texts_to_embed), batch_size)):
        batch = texts_to_embed[i:i+batch_size]
        batch_embeddings = embedding_model.embed_documents(batch)
        embeddings.extend(batch_embeddings)
        time.sleep(60)  # Sleep to respect API rate limits
    print(f"Embeddings generated for all {len(embeddings)} chunks.")

    # Store in ChromaDB
    print(f"Storing embeddings in ChromaDB at {db_path}...")
    collection_name = "all_arxiv_papers" 
    chroma_client = chromadb.PersistentClient(path=db_path)
    collection = chroma_client.get_or_create_collection(name=collection_name)
    insert_batch_size = 32
    for i in range(0, len(texts_to_embed), insert_batch_size):
        collection.upsert(
            ids=ids[i:i+insert_batch_size],
            embeddings=embeddings[i:i+insert_batch_size],
            metadatas=metadatas[i:i+insert_batch_size],
            documents=texts_to_embed[i:i+insert_batch_size]
        )
    print("Done!")
    
    print("Success! Your Markdown document has been chunked, embedded, and stored.")


def main(input_dir, db_path):
    # Process each Markdown file in the input directory
    for filename in os.listdir(input_dir):
        if filename.endswith(".md"):
            input_file = os.path.join(input_dir, filename)
            print(f"Processing {input_file}...")
            generate_embeddings(input_file, db_path)
            # break


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Convert Markdown to Embeddings")
    parser.add_argument("--input_dir", type=str, required=True, help="Directory containing Markdown files")
    parser.add_argument("--db_path", type=str, required=True, help="Directory to save ChromaDB")
    args = parser.parse_args()
    os.makedirs(args.db_path, exist_ok=True)
    main(args.input_dir, args.db_path)