Spaces:
Sleeping
Sleeping
| """ | |
| PDF Ingestion & Vector Store Builder for Study Ocean RAG. | |
| This module: | |
| 1. Recursively scans the PDF directory for all .pdf files | |
| 2. Extracts text from each PDF using PyMuPDF (faster & more reliable) | |
| 3. Splits text into overlapping chunks | |
| 4. Generates embeddings using a local sentence-transformer model | |
| 5. Stores embeddings in a FAISS vector store for fast retrieval | |
| """ | |
| import os | |
| import sys | |
| import time | |
| import fitz # PyMuPDF | |
| from langchain.text_splitter import RecursiveCharacterTextSplitter | |
| from langchain_community.vectorstores import FAISS | |
| from langchain_huggingface import HuggingFaceEmbeddings | |
| import config | |
| def collect_pdf_paths(pdf_dir: str) -> list[str]: | |
| """Recursively collect all PDF file paths from the given directory.""" | |
| pdf_paths = [] | |
| for root, _, files in os.walk(pdf_dir): | |
| for file in files: | |
| if file.lower().endswith(".pdf"): | |
| pdf_paths.append(os.path.join(root, file)) | |
| return sorted(pdf_paths) | |
| def extract_text_from_pdf(pdf_path: str) -> str: | |
| """Extract all text from a PDF using PyMuPDF.""" | |
| text = "" | |
| try: | |
| doc = fitz.open(pdf_path) | |
| for page_num in range(len(doc)): | |
| page = doc[page_num] | |
| text += page.get_text() | |
| doc.close() | |
| except Exception as e: | |
| print(f" ⚠ Error reading {os.path.basename(pdf_path)}: {e}") | |
| return text | |
| def build_documents(pdf_paths: list[str]) -> list: | |
| """ | |
| Extract text from all PDFs and split into chunks. | |
| Each chunk carries metadata about its source PDF and page info. | |
| """ | |
| text_splitter = RecursiveCharacterTextSplitter( | |
| chunk_size=config.CHUNK_SIZE, | |
| chunk_overlap=config.CHUNK_OVERLAP, | |
| separators=["\n\n", "\n", ". ", " ", ""], | |
| length_function=len, | |
| ) | |
| all_documents = [] | |
| for i, pdf_path in enumerate(pdf_paths, 1): | |
| filename = os.path.basename(pdf_path) | |
| relative_path = os.path.relpath(pdf_path, config.PDF_DIR) | |
| # Derive subject from folder structure | |
| parts = relative_path.replace("\\", "/").split("/") | |
| subject = parts[0] if len(parts) > 1 else "General" | |
| print(f" [{i}/{len(pdf_paths)}] Processing: {relative_path}") | |
| text = extract_text_from_pdf(pdf_path) | |
| if not text.strip(): | |
| print(f" ⚠ No text extracted (might be scanned/image PDF)") | |
| continue | |
| # Create chunks with rich metadata | |
| chunks = text_splitter.create_documents( | |
| texts=[text], | |
| metadatas=[{ | |
| "source": relative_path, | |
| "filename": filename, | |
| "subject": subject, | |
| }] | |
| ) | |
| all_documents.extend(chunks) | |
| return all_documents | |
| def create_vector_store(documents: list, embeddings) -> FAISS: | |
| """Create a FAISS vector store from documents.""" | |
| print(f"\n📊 Creating vector store with {len(documents)} chunks...") | |
| # Process in batches to avoid memory issues | |
| batch_size = 500 | |
| vector_store = None | |
| for i in range(0, len(documents), batch_size): | |
| batch = documents[i:i + batch_size] | |
| batch_num = (i // batch_size) + 1 | |
| total_batches = (len(documents) + batch_size - 1) // batch_size | |
| print(f" Embedding batch {batch_num}/{total_batches} ({len(batch)} chunks)...") | |
| if vector_store is None: | |
| vector_store = FAISS.from_documents(batch, embeddings) | |
| else: | |
| batch_store = FAISS.from_documents(batch, embeddings) | |
| vector_store.merge_from(batch_store) | |
| return vector_store | |
| def run_ingestion(): | |
| """Main ingestion pipeline.""" | |
| print("=" * 60) | |
| print("🎓 Study Ocean RAG — PDF Ingestion Pipeline") | |
| print("=" * 60) | |
| # 1. Collect PDFs | |
| print(f"\n📁 Scanning for PDFs in: {config.PDF_DIR}") | |
| pdf_paths = collect_pdf_paths(config.PDF_DIR) | |
| if not pdf_paths: | |
| print("❌ No PDF files found! Make sure your PDFs are in the data/pdfs directory.") | |
| print(" Extract StudyOceanPDFs.zip into data/pdfs/ first.") | |
| sys.exit(1) | |
| print(f" Found {len(pdf_paths)} PDF files\n") | |
| # 2. Extract and chunk | |
| print("📄 Extracting text and creating chunks...") | |
| start_time = time.time() | |
| documents = build_documents(pdf_paths) | |
| extract_time = time.time() - start_time | |
| print(f"\n ✅ Created {len(documents)} text chunks in {extract_time:.1f}s") | |
| if not documents: | |
| print("❌ No text could be extracted from any PDF. Exiting.") | |
| sys.exit(1) | |
| # 3. Initialize embeddings | |
| print(f"\n🧠 Loading embedding model: {config.EMBEDDING_MODEL}") | |
| embeddings = HuggingFaceEmbeddings( | |
| model_name=config.EMBEDDING_MODEL, | |
| model_kwargs={"device": "cpu"}, | |
| encode_kwargs={"normalize_embeddings": True, "batch_size": 64}, | |
| ) | |
| # 4. Build vector store | |
| start_time = time.time() | |
| vector_store = create_vector_store(documents, embeddings) | |
| embed_time = time.time() - start_time | |
| print(f" ✅ Vector store created in {embed_time:.1f}s") | |
| # 5. Save to disk | |
| os.makedirs(config.VECTOR_STORE_PATH, exist_ok=True) | |
| print(f"\n💾 Saving vector store to: {config.VECTOR_STORE_PATH}") | |
| vector_store.save_local(config.VECTOR_STORE_PATH) | |
| print(" ✅ Vector store saved successfully!") | |
| # Summary | |
| print("\n" + "=" * 60) | |
| print("✅ INGESTION COMPLETE") | |
| print(f" PDFs processed: {len(pdf_paths)}") | |
| print(f" Chunks created: {len(documents)}") | |
| print(f" Extract time: {extract_time:.1f}s") | |
| print(f" Embedding time: {embed_time:.1f}s") | |
| print(f" Store location: {config.VECTOR_STORE_PATH}") | |
| print("=" * 60) | |
| if __name__ == "__main__": | |
| run_ingestion() | |