Spaces:
Sleeping
Sleeping
File size: 3,697 Bytes
2db8ee1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 | from vectorstore.document_index import DocumentIndex
from vectorstore.chunk_index import ChunkIndex
from vectorstore.image_store import ImageVectorStore
from vectorstore.index_manager import IndexManager
from rag.generator import generate_answer, decompose_query
from rag.reranker import rerank
from ingest.pdf_ingest import ingest_pdf
from utils.text_utils import chunk_by_slide
import os
import glob
doc_index = DocumentIndex()
chunk_index = ChunkIndex()
image_store = ImageVectorStore()
index_manager = IndexManager(doc_index, chunk_index)
DATA_DIR = "../data"
INDICES_DIR = os.path.join(DATA_DIR, "indices")
CHUNK_INDEX_PATH = os.path.join(INDICES_DIR, "chunks")
DOC_INDEX_PATH = os.path.join(INDICES_DIR, "docs")
IMAGE_INDEX_PATH = os.path.join(INDICES_DIR, "images")
if os.path.exists(CHUNK_INDEX_PATH) and os.path.exists(DOC_INDEX_PATH) and os.path.exists(IMAGE_INDEX_PATH):
print("Loading indices from disk (skipping ingestion)...")
chunk_index.load_local(CHUNK_INDEX_PATH)
doc_index.load_local(DOC_INDEX_PATH)
image_store.load_local(IMAGE_INDEX_PATH)
index_manager = IndexManager(doc_index, chunk_index)
else:
print("Indices not found. Starting fresh ingesetion...")
pdf_list = [
"../data/raw/os.pdf",
"../data/raw/DEVOPS.pdf",
"../data/raw/DBMS_Notes.pdf"
]
text_chunks, _ = ingest_pdf(pdf_list, "../data/processed/images")
from collections import defaultdict
chunks_by_source = defaultdict(list)
for chunk in text_chunks:
chunks_by_source[chunk["source"]].append(chunk["text"])
for source, chunks in chunks_by_source.items():
slide_chunks = []
full_text = ""
for text in chunks:
text = text.strip()
if len(text) > 50:
full_text += text + "\n"
slide_chunks.append(text)
if(slide_chunks):
doc_index.add_document(full_text, source)
chunk_index.add_chunks(source, slide_chunks)
image_paths = glob.glob("../data/processed/images/*.png")
image_metadata = []
for p in image_paths:
try:
parts = p.split("_page_")
page_num = int(parts[1].split("_img_")[0]) if len(parts) > 1 else 0
except:
page_num = 0
image_metadata.append({"image_path": p, "page": page_num})
image_store.add_images(image_paths[:20], image_metadata[:20])
print("Saving indices to disk...")
chunk_index.save_local(CHUNK_INDEX_PATH)
doc_index.save_local(DOC_INDEX_PATH)
image_store.save_local(IMAGE_INDEX_PATH)
print("System ready!")
query = input("Enter your question: ")
# Optional: decompose
# Optional: decompose
# Skip decomposition for speed
subqueries = [query]
all_chunks = []
for sq in subqueries:
retrieved_chunks = index_manager.retrieve(sq)
all_chunks.extend(retrieved_chunks)
# Deduplicate
# Deduplicate
unique = []
seen = set()
for r in all_chunks:
content_text = r["content"] # this should be string
if content_text not in seen:
unique.append(r)
seen.add(content_text)
print(f"\nFound {len(unique)} unique chunks.")
for i, r in enumerate(unique[:3]):
print(f"Chunk {i}: {r['content'][:100]}...")
text_results = rerank(query, unique, top_k=6)
print(f"Reranked to {len(text_results)} chunks.")
# Image retrieval stays global (can improve later)
image_results = image_store.search(query, k=6)
try:
answer = generate_answer(query, text_results, image_results)
print("\n=== FINAL ANSWER ===\n")
print(answer)
except Exception as e:
print(f"\nGeneration failed: {e}")
# Don't fail the script, just report error
|