AabhavAmitabh/document-intelligence-chatbot

1# embedder.py
2# This file has ONE job: take text chunks and store them in ChromaDB
3# as embeddings so we can search them by meaning later.
4
5from sentence_transformers import SentenceTransformer
6import chromadb
7from ingestor import ingest_pdf
8
9
10# We use a small but powerful open-source embedding model.
11# It runs entirely on your machine — no API key needed for this step.
12# "all-MiniLM-L6-v2" is fast, accurate, and used in real production systems.
13EMBEDDING_MODEL = "all-MiniLM-L6-v2"
14
15# This is the folder where ChromaDB will save your vector database.
16# It will be created automatically the first time you run this.
17CHROMA_PATH = "chroma_db"
18
19# We give our collection of document chunks a name inside ChromaDB.
20COLLECTION_NAME = "documents"
21
22
23def get_embedding_model():
24    """
25    Loads the embedding model.
26    The first time this runs it downloads ~90MB from the internet.
27    After that it's cached locally and loads in seconds.
28    """
29    print("Loading embedding model...")
30    model = SentenceTransformer(EMBEDDING_MODEL)
31    print("Model ready.")
32    return model
33
34
35def embed_and_store(chunks: list, doc_name: str = "document") -> chromadb.Collection:
36    """
37    Takes a list of text chunks, converts each one to an embedding,
38    and stores everything in ChromaDB.
39
40    doc_name: a label so you know which document these chunks came from.
41               Useful later when you support multiple documents.
42    """
43    model = get_embedding_model()
44
45    # Create (or connect to) the ChromaDB database on disk
46    client = chromadb.PersistentClient(path=CHROMA_PATH)
47
48    # Delete the collection if it already exists so we start fresh.
49    # This prevents duplicate chunks if you run the script twice.
50    existing = [c.name for c in client.list_collections()]
51    if COLLECTION_NAME in existing:
52        client.delete_collection(COLLECTION_NAME)
53        print("Cleared existing collection.")
54
55    collection = client.create_collection(COLLECTION_NAME)
56
57    print(f"Embedding {len(chunks)} chunks — this may take 30-60 seconds...")
58    if not chunks:
59        raise ValueError(
60            "No text could be extracted from this PDF. "
61            "It may be a scanned image. Try a text-based PDF instead."
62        )
63
64    # Convert all chunks to embeddings in one batch (faster than one at a time)
65    embeddings = model.encode(chunks, show_progress_bar=True)
66
67    # Store each chunk in ChromaDB with:
68    # - a unique ID
69    # - the embedding (the numbers)
70    # - the original text (so we can retrieve it later)
71    # - metadata (which document it came from)
72    collection.add(
73        ids=[f"{doc_name}_chunk_{i}" for i in range(len(chunks))],
74        embeddings=embeddings.tolist(),
75        documents=chunks,
76        metadatas=[{"source": doc_name, "chunk_index": i} for i in range(len(chunks))]
77    )
78
79    print(f"Stored {collection.count()} chunks in ChromaDB.")
80    return collection
81
82
83def load_collection() -> chromadb.Collection:
84    """
85    Connects to an existing ChromaDB database and returns the collection.
86    This is what the rest of the app calls when answering questions —
87    it doesn't re-embed, it just opens what's already saved on disk.
88    """
89    client = chromadb.PersistentClient(path=CHROMA_PATH)
90    collection = client.get_collection(COLLECTION_NAME)
91    return collection
92
93
94# Test it directly when you run this file
95if __name__ == "__main__":
96    import sys
97
98    if len(sys.argv) < 2:
99        print("Usage: python embedder.py path/to/your/file.pdf")
100        sys.exit(1)
101
102    pdf_path = sys.argv[1]
103
104    # Get a clean name for the document (just the filename, no path)
105    import os
106    doc_name = os.path.splitext(os.path.basename(pdf_path))[0]
107
108    # Step 1: ingest the PDF into chunks
109    chunks = ingest_pdf(pdf_path)
110
111    # Step 2: embed and store
112    collection = embed_and_store(chunks, doc_name=doc_name)
113
114    # Step 3: do a quick test search to prove it's working
115    print("\n--- TEST SEARCH ---")
116    test_query = "Who is this document about?"
117    print(f"Query: '{test_query}'")
118
119    model = SentenceTransformer(EMBEDDING_MODEL)
120    query_embedding = model.encode([test_query]).tolist()
121
122    results = collection.query(
123        query_embeddings=query_embedding,
124        n_results=2
125    )
126
127    print("\nTop 2 most relevant chunks found:")
128    for i, doc in enumerate(results["documents"][0]):
129        print(f"\nResult {i+1}:\n{doc[:300]}...")
130        print("-" * 40)
131
132    print("\nPhase 3 complete. ChromaDB is populated and searchable.")
133