fullstack/pylate-document-search
0
1#!/usr/bin/env python32"""3PyLate ZeroGPU Document Search with Runtime Package Installation4Complete version that installs all dependencies at startup if needed.5"""6 7import subprocess8import sys9import os10import time11 12print("🚀 Starting PyLate ZeroGPU Document Search...")13print("🔧 Checking and installing required packages...")14 15# ===== RUNTIME PACKAGE INSTALLATION =====16def install_package(package, quiet=True):17 """Install a package at runtime."""18 try:19 if quiet:20 subprocess.check_call([21 sys.executable, '-m', 'pip', 'install', package, 22 '--quiet', '--disable-pip-version-check'23 ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)24 else:25 subprocess.check_call([sys.executable, '-m', 'pip', 'install', package])26 return True27 except Exception as e:28 print(f"⚠️ Failed to install {package}: {e}")29 return False30 31def check_and_install_packages():32 """Check and install all required packages."""33 34 # Define packages with their import names and pip names35 packages_to_check = [36 # (import_name, pip_package, test_import)37 ('gradio', 'gradio==4.44.0', lambda: __import__('gradio')),38 ('spaces', 'spaces', lambda: __import__('spaces')),39 ('sentence_transformers', 'sentence-transformers', lambda: __import__('sentence_transformers')),40 ('docx', 'python-docx', lambda: __import__('docx')),41 ('fitz', 'pymupdf', lambda: __import__('fitz')),42 ('unstructured', 'unstructured', lambda: __import__('unstructured')),43 ('pandas', 'pandas', lambda: __import__('pandas')),44 ('numpy', 'numpy', lambda: __import__('numpy')),45 ('huggingface_hub', 'huggingface_hub', lambda: __import__('huggingface_hub')),46 ('accelerate', 'accelerate', lambda: __import__('accelerate')),47 ('pylate', 'pylate==1.2.0', lambda: __import__('pylate')),48 ]49 50 installed_count = 051 failed_packages = []52 53 for import_name, pip_package, test_func in packages_to_check:54 try:55 test_func()56 print(f"✅ {import_name} - already installed")57 installed_count += 158 except ImportError:59 print(f"📦 Installing {pip_package}...")60 success = install_package(pip_package, quiet=False)61 if success:62 try:63 # Test import after installation64 test_func()65 print(f"✅ {import_name} - installed successfully")66 installed_count += 167 except ImportError:68 print(f"❌ {import_name} - installation failed (import test failed)")69 failed_packages.append(import_name)70 else:71 print(f"❌ {import_name} - installation failed")72 failed_packages.append(import_name)73 74 print(f"\n📊 Installation Summary:")75 print(f" ✅ Successfully installed/verified: {installed_count}/{len(packages_to_check)}")76 77 if failed_packages:78 print(f" ❌ Failed packages: {', '.join(failed_packages)}")79 print(f" ⚠️ App may not work correctly with missing packages")80 else:81 print(f" 🎉 All packages ready!")82 83 return len(failed_packages) == 084 85# Install packages before importing anything else86installation_success = check_and_install_packages()87 88# Now import everything89print("\n🔄 Loading modules...")90 91try:92 import gradio as gr93 import spaces94 import torch95 import tempfile96 import sqlite397 import json98 import hashlib99 from pathlib import Path100 from typing import List, Dict, Any, Tuple101 print("✅ Core modules loaded")102except ImportError as e:103 print(f"❌ Failed to import core modules: {e}")104 sys.exit(1)105 106# Import document processing modules with fallbacks107try:108 import docx109 print("✅ python-docx loaded")110except ImportError:111 print("⚠️ python-docx not available - DOCX processing will be disabled")112 docx = None113 114try:115 import fitz # pymupdf116 print("✅ PyMuPDF loaded")117except ImportError:118 print("⚠️ PyMuPDF not available - PDF processing will be limited")119 fitz = None120 121try:122 from unstructured.partition.auto import partition123 print("✅ Unstructured loaded")124except ImportError:125 print("⚠️ Unstructured not available - fallback text extraction disabled")126 partition = None127 128try:129 from pylate import models, indexes, retrieve130 print("✅ PyLate loaded")131except ImportError as e:132 print(f"❌ PyLate failed to load: {e}")133 print("🔄 Attempting to install PyLate...")134 install_package('pylate==1.2.0', quiet=False)135 try:136 from pylate import models, indexes, retrieve137 print("✅ PyLate loaded after installation")138 except ImportError:139 print("❌ PyLate installation failed - core functionality unavailable")140 sys.exit(1)141 142# Set environment variables143os.environ["TRITON_CACHE_DIR"] = "/tmp/triton_cache"144os.environ["TORCH_COMPILE_DISABLE"] = "1"145 146print("🎯 All modules loaded successfully!\n")147 148# Global variables for PyLate components149model = None150index = None151retriever = None152metadata_db = None153 154# ===== DOCUMENT PROCESSING FUNCTIONS =====155 156def extract_text_from_pdf(file_path: str) -> str:157 """Extract text from PDF file using PyMuPDF and unstructured as fallback."""158 text = ""159 160 if not fitz:161 return "Error: PyMuPDF not available for PDF processing"162 163 try:164 # Use PyMuPDF (fitz) - more reliable than PyPDF2165 doc = fitz.open(file_path)166 for page in doc:167 text += page.get_text() + "\n"168 doc.close()169 170 # If no text extracted, try unstructured171 if not text.strip() and partition:172 elements = partition(filename=file_path)173 text = "\n".join([str(element) for element in elements])174 175 except Exception as e:176 # Final fallback to unstructured177 if partition:178 try:179 elements = partition(filename=file_path)180 text = "\n".join([str(element) for element in elements])181 except:182 text = f"Error: Could not extract text from PDF: {str(e)}"183 else:184 text = f"Error: Could not extract text from PDF: {str(e)}"185 186 return text.strip()187 188def extract_text_from_docx(file_path: str) -> str:189 """Extract text from DOCX file."""190 if not docx:191 return "Error: python-docx not available for DOCX processing"192 193 try:194 doc = docx.Document(file_path)195 text = ""196 for paragraph in doc.paragraphs:197 text += paragraph.text + "\n"198 return text.strip()199 except Exception as e:200 return f"Error: Could not extract text from DOCX: {str(e)}"201 202def extract_text_from_txt(file_path: str) -> str:203 """Extract text from TXT file."""204 try:205 with open(file_path, 'r', encoding='utf-8') as file:206 return file.read().strip()207 except UnicodeDecodeError:208 try:209 with open(file_path, 'r', encoding='latin1') as file:210 return file.read().strip()211 except Exception as e:212 return f"Error: Could not read text file: {str(e)}"213 except Exception as e:214 return f"Error: Could not read text file: {str(e)}"215 216def chunk_text(text: str, chunk_size: int = 1000, overlap: int = 100) -> List[Dict[str, Any]]:217 """Chunk text with overlap and return metadata."""218 chunks = []219 start = 0220 chunk_index = 0221 222 while start < len(text):223 end = start + chunk_size224 chunk_text = text[start:end]225 226 # Try to break at sentence boundary227 if end < len(text):228 last_period = chunk_text.rfind('.')229 last_newline = chunk_text.rfind('\n')230 break_point = max(last_period, last_newline)231 232 if break_point > chunk_size * 0.7:233 chunk_text = chunk_text[:break_point + 1]234 end = start + break_point + 1235 236 if chunk_text.strip():237 chunks.append({238 'text': chunk_text.strip(),239 'start': start,240 'end': end,241 'index': chunk_index,242 'length': len(chunk_text.strip())243 })244 chunk_index += 1245 246 start = max(start + 1, end - overlap)247 248 return chunks249 250# ===== METADATA DATABASE =====251 252def init_metadata_db():253 """Initialize SQLite database for metadata."""254 global metadata_db255 256 db_path = "metadata.db"257 metadata_db = sqlite3.connect(db_path, check_same_thread=False)258 259 metadata_db.execute("""260 CREATE TABLE IF NOT EXISTS documents (261 doc_id TEXT PRIMARY KEY,262 filename TEXT NOT NULL,263 file_hash TEXT NOT NULL,264 original_text TEXT NOT NULL,265 chunk_index INTEGER NOT NULL,266 total_chunks INTEGER NOT NULL,267 chunk_start INTEGER NOT NULL,268 chunk_end INTEGER NOT NULL,269 chunk_size INTEGER NOT NULL,270 created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP271 )272 """)273 274 metadata_db.execute("""275 CREATE INDEX IF NOT EXISTS idx_filename ON documents(filename);276 """)277 278 metadata_db.commit()279 280def add_document_metadata(doc_id: str, filename: str, file_hash: str,281 original_text: str, chunk_info: Dict[str, Any], total_chunks: int):282 """Add document metadata to database."""283 global metadata_db284 285 metadata_db.execute("""286 INSERT OR REPLACE INTO documents 287 (doc_id, filename, file_hash, original_text, chunk_index, total_chunks,288 chunk_start, chunk_end, chunk_size)289 VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)290 """, (291 doc_id, filename, file_hash, original_text,292 chunk_info['index'], total_chunks,293 chunk_info['start'], chunk_info['end'], chunk_info['length']294 ))295 metadata_db.commit()296 297def get_document_metadata(doc_id: str) -> Dict[str, Any]:298 """Get document metadata by ID."""299 global metadata_db300 301 cursor = metadata_db.execute(302 "SELECT * FROM documents WHERE doc_id = ?", (doc_id,)303 )304 row = cursor.fetchone()305 306 if row:307 columns = [desc[0] for desc in cursor.description]308 return dict(zip(columns, row))309 return {}310 311# ===== PYLATE INITIALIZATION =====312 313@spaces.GPU(duration=120) # Allow 2 minutes for initialization314def initialize_pylate(model_name: str = "colbert-ir/colbertv2.0") -> str:315 """Initialize PyLate components on ZeroGPU H200."""316 global model, index, retriever317 318 try:319 # Initialize metadata database320 init_metadata_db()321 322 # Load ColBERT model323 model = models.ColBERT(model_name_or_path=model_name)324 325 # Move to GPU - ZeroGPU provides CUDA access326 device_info = "CPU"327 if torch.cuda.is_available():328 model = model.to('cuda')329 device_name = torch.cuda.get_device_name()330 device_info = f"GPU: {device_name}"331 332 # Initialize PLAID index with optimized settings for ZeroGPU333 index = indexes.PLAID(334 index_folder="./pylate_index",335 index_name="documents",336 override=True,337 kmeans_niters=1, # Reduce k-means iterations for faster setup338 nbits=2 # Optimized for memory efficiency339 )340 341 # Initialize retriever342 retriever = retrieve.ColBERT(index=index)343 344 return f"✅ PyLate initialized successfully on ZeroGPU!\n🔥 Model: {model_name}\n🎯 Device: {device_info}\n💾 VRAM: ~70GB available\n🚀 Ready for document processing!"345 346 except Exception as e:347 return f"❌ Error initializing PyLate: {str(e)}\n\nPlease check the logs for more details."348 349# ===== DOCUMENT PROCESSING =====350 351@spaces.GPU(duration=300) # Allow 5 minutes for processing352def process_documents(files, chunk_size: int = 1000, overlap: int = 100) -> str:353 """Process uploaded documents and add to index using ZeroGPU."""354 global model, index, metadata_db355 356 if not model or not index:357 return "❌ Please initialize PyLate first!"358 359 if not files:360 return "❌ No files uploaded!"361 362 try:363 all_documents = []364 all_doc_ids = []365 processed_files = []366 skipped_files = []367 368 for file in files:369 # Get file info370 filename = Path(file.name).name371 file_path = file.name372 373 # Calculate file hash374 with open(file_path, 'rb') as f:375 file_hash = hashlib.md5(f.read()).hexdigest()376 377 # Extract text based on file type378 text = ""379 if filename.lower().endswith('.pdf'):380 if fitz:381 text = extract_text_from_pdf(file_path)382 else:383 skipped_files.append(f"{filename}: PDF processing not available")384 continue385 elif filename.lower().endswith('.docx'):386 if docx:387 text = extract_text_from_docx(file_path)388 else:389 skipped_files.append(f"{filename}: DOCX processing not available")390 continue391 elif filename.lower().endswith('.txt'):392 text = extract_text_from_txt(file_path)393 else:394 skipped_files.append(f"{filename}: Unsupported file type")395 continue396 397 if not text or text.startswith("Error:"):398 skipped_files.append(f"{filename}: Failed to extract text")399 continue400 401 # Chunk the text402 chunks = chunk_text(text, chunk_size, overlap)403 404 if not chunks:405 skipped_files.append(f"{filename}: No valid chunks created")406 continue407 408 # Process each chunk409 for chunk in chunks:410 doc_id = f"{filename}_chunk_{chunk['index']}"411 all_documents.append(chunk['text'])412 all_doc_ids.append(doc_id)413 414 # Store metadata415 add_document_metadata(416 doc_id=doc_id,417 filename=filename,418 file_hash=file_hash,419 original_text=chunk['text'],420 chunk_info=chunk,421 total_chunks=len(chunks)422 )423 424 processed_files.append(f"{filename}: {len(chunks)} chunks")425 426 if not all_documents:427 return "❌ No text could be extracted from uploaded files!\n" + "\n".join(skipped_files)428 429 # Encode documents with PyLate on H200 GPU430 document_embeddings = model.encode(431 all_documents,432 batch_size=32, # Optimized batch size for H200's 70GB VRAM433 is_query=False,434 show_progress_bar=True435 )436 437 # Add to PLAID index438 index.add_documents(439 documents_ids=all_doc_ids,440 documents_embeddings=document_embeddings441 )442 443 result = f"✅ Successfully processed {len([f for f in files if not any(f.name in skip for skip in skipped_files)])} files on ZeroGPU H200:\n"444 result += f"📄 Total chunks indexed: {len(all_documents)}\n"445 result += f"🔍 Documents processed:\n"446 for file_info in processed_files:447 result += f" • {file_info}\n"448 449 if skipped_files:450 result += f"\n⚠️ Skipped files:\n"451 for skip_info in skipped_files:452 result += f" • {skip_info}\n"453 454 result += f"\n🎉 Document index ready for search!"455 return result456 457 except Exception as e:458 return f"❌ Error processing documents: {str(e)}\n\nPlease check your files and try again."459 460# ===== SEARCH FUNCTION =====461 462@spaces.GPU(duration=60) # 1 minute for search463def search_documents(query: str, k: int = 5, show_chunks: bool = True) -> str:464 """Search documents using PyLate on ZeroGPU."""465 global model, retriever, metadata_db466 467 if not model or not retriever:468 return "❌ Please initialize PyLate and process documents first!"469 470 if not query.strip():471 return "❌ Please enter a search query!"472 473 try:474 # Encode query on GPU475 query_embedding = model.encode([query], is_query=True)476 477 # Search478 results = retriever.retrieve(query_embedding, k=k)[0]479 480 if not results:481 return "🔍 No results found for your query.\n\nTry:\n• Different keywords\n• Broader search terms\n• Check if documents were processed correctly"482 483 # Format results with metadata484 formatted_results = [f"🔍 **Search Results for:** '{query}' (powered by ZeroGPU H200)\n"]485 486 for i, result in enumerate(results):487 doc_id = result['id']488 score = result['score']489 490 # Get metadata491 metadata = get_document_metadata(doc_id)492 493 formatted_results.append(f"## Result {i+1} (Relevance: {score:.3f})")494 formatted_results.append(495 f"**📄 File:** {metadata.get('filename', 'Unknown')}")496 formatted_results.append(497 f"**📑 Chunk:** {metadata.get('chunk_index', 0) + 1}/{metadata.get('total_chunks', 1)}")498 499 if show_chunks:500 text = metadata.get('original_text', '')501 if len(text) > 400:502 preview = text[:400] + "..."503 else:504 preview = text505 formatted_results.append(f"**💬 Text:** {preview}")506 507 formatted_results.append("---")508 509 formatted_results.append(f"\n🎯 Found {len(results)} relevant results using ColBERT semantic search")510 return "\n".join(formatted_results)511 512 except Exception as e:513 return f"❌ Error searching: {str(e)}\n\nPlease try again or check if PyLate is properly initialized."514 515# ===== GRADIO INTERFACE =====516 517def create_interface():518 """Create the Gradio interface for ZeroGPU."""519 520 with gr.Blocks(title="PyLate ZeroGPU Document Search", theme=gr.themes.Soft()) as demo:521 gr.Markdown("""522 # 🚀 PyLate ZeroGPU Document Search523 ### Powered by ColBERT and NVIDIA H200 (70GB VRAM)524 525 Upload documents, process them with PyLate on ZeroGPU, and perform lightning-fast semantic search!526 527 **🔥 ZeroGPU Features:**528 - 🎯 NVIDIA H200 GPU with 70GB VRAM529 - ⚡ Dynamic GPU allocation (only when needed)530 - 🆓 Free for HF Pro subscribers531 - 🚀 Optimized for PyTorch/ColBERT workloads532 - 🔄 Automatic package installation533 """)534 535 # Status indicator536 with gr.Row():537 gr.Markdown(f"""538 **📊 System Status:**539 - ✅ PyLate: Ready540 - ✅ Document Processing: {"PDF ✅" if fitz else "PDF ❌"} | {"DOCX ✅" if docx else "DOCX ❌"} | TXT ✅541 - ✅ ZeroGPU: Available542 """)543 544 with gr.Tab("🚀 Setup"):545 gr.Markdown("### Initialize PyLate System on ZeroGPU H200")546 547 model_choice = gr.Dropdown(548 choices=[549 "colbert-ir/colbertv2.0",550 "sentence-transformers/all-MiniLM-L6-v2"551 ],552 value="colbert-ir/colbertv2.0",553 label="Select ColBERT Model",554 info="ColBERT v2.0 is recommended for best performance"555 )556 557 init_btn = gr.Button("🚀 Initialize PyLate on ZeroGPU", variant="primary", size="lg")558 init_status = gr.Textbox(label="Initialization Status", lines=6, max_lines=10)559 560 init_btn.click(561 initialize_pylate,562 inputs=model_choice,563 outputs=init_status564 )565 566 with gr.Tab("📄 Document Upload"):567 gr.Markdown("### Upload and Process Documents on H200 GPU")568 569 with gr.Row():570 with gr.Column():571 file_upload = gr.File(572 file_count="multiple",573 file_types=[".pdf", ".docx", ".txt"],574 label="Upload Documents",575 info="Supported: PDF, DOCX, TXT files"576 )577 578 with gr.Row():579 chunk_size = gr.Slider(580 minimum=500,581 maximum=3000,582 value=1000,583 step=100,584 label="Chunk Size (characters)",585 info="Larger chunks = more context, smaller chunks = more precise"586 )587 588 overlap = gr.Slider(589 minimum=0,590 maximum=500,591 value=100,592 step=50,593 label="Chunk Overlap (characters)",594 info="Overlap helps maintain context between chunks"595 )596 597 process_btn = gr.Button(598 "⚡ Process Documents on ZeroGPU", variant="primary", size="lg")599 600 with gr.Column():601 process_status = gr.Textbox(602 label="Processing Status",603 lines=15,604 max_lines=20,605 info="Processing status and results will appear here"606 )607 608 process_btn.click(609 process_documents,610 inputs=[file_upload, chunk_size, overlap],611 outputs=process_status612 )613 614 with gr.Tab("🔍 Search"):615 gr.Markdown("### Search Your Documents with H200 Power")616 617 with gr.Row():618 with gr.Column():619 search_query = gr.Textbox(620 label="Search Query",621 placeholder="Enter your search query... (e.g., 'machine learning algorithms', 'financial projections')",622 lines=2,623 info="Use natural language - ColBERT understands semantic meaning"624 )625 626 with gr.Row():627 num_results = gr.Slider(628 minimum=1,629 maximum=20,630 value=5,631 step=1,632 label="Number of Results",633 info="How many relevant chunks to return"634 )635 636 show_chunks = gr.Checkbox(637 value=True,638 label="Show Text Chunks",639 info="Display the actual text content"640 )641 642 search_btn = gr.Button("🔍 Search with ZeroGPU", variant="primary", size="lg")643 644 with gr.Column():645 search_results = gr.Textbox(646 label="Search Results",647 lines=18,648 max_lines=25,649 info="Semantic search results will appear here"650 )651 652 search_btn.click(653 search_documents,654 inputs=[search_query, num_results, show_chunks],655 outputs=search_results656 )657 658 with gr.Tab("ℹ️ ZeroGPU Info"):659 gr.Markdown("""660 ### About ZeroGPU PyLate Search661 662 **🔥 Powered by NVIDIA H200 Tensor Core GPU**663 664 #### 🚀 ZeroGPU Features:665 - **70GB HBM3 Memory** - Massive capacity for large document collections666 - **Dynamic Allocation** - GPU assigned only when functions need it667 - **Optimized for PyTorch** - Perfect for ColBERT/PyLate workloads668 - **Free for Pro Users** - No additional charges beyond HF Pro669 - **Auto Scaling** - Efficient resource usage and queue management670 671 #### 🧠 How ColBERT Works:672 1. **Late Interaction** - Processes queries and documents separately673 2. **Token-level Matching** - Fine-grained semantic understanding674 3. **Efficient Retrieval** - Fast search with high-quality results675 4. **GPU Acceleration** - Leverages H200 for rapid inference676 677 #### 📊 Performance Benefits:678 - **10-100x faster** than CPU-based search679 - **Large batch processing** - 32+ documents simultaneously680 - **Real-time search** - Sub-second query responses681 - **Massive scale** - 70GB VRAM handles huge document sets682 683 #### 🛠️ Technical Details:684 - **Runtime Package Installation** - Automatically installs dependencies685 - **Gradio SDK Required** - ZeroGPU doesn't support Docker686 - **Smart Chunking** - Intelligent text segmentation with overlap687 - **Metadata Tracking** - SQLite database for chunk information688 689 #### 🎯 Usage Tips:690 1. **Initialize first** - Required before processing documents691 2. **Natural language queries** - ColBERT understands meaning, not just keywords692 3. **Adjust chunk size** - Larger for context, smaller for precision693 4. **Multiple file types** - Mix PDFs, DOCX, and TXT files694 5. **Semantic search** - Try "concepts similar to X" type queries695 696 #### 🔒 Privacy & Security:697 - Documents processed in-memory only698 - No permanent storage of your content699 - Processing happens on HF infrastructure700 - Automatic cleanup after session ends701 702 ---703 704 **Built with ❤️ using:**705 - 🤖 PyLate & ColBERT for semantic search706 - ⚡ ZeroGPU H200 for GPU acceleration 707 - 🎨 Gradio for the interface708 - 🐍 Python ecosystem for document processing709 """)710 711 return demo712 713# ===== MAIN =====714 715if __name__ == "__main__":716 print("🎉 Launching PyLate ZeroGPU Document Search interface...")717 718 # Check if running on ZeroGPU719 if torch.cuda.is_available():720 print(f"🔥 GPU detected: {torch.cuda.get_device_name()}")721 else:722 print("💻 Running on CPU (GPU will be allocated when @spaces.GPU functions are called)")723 724 demo = create_interface()725 demo.launch(726 share=False,727 server_name="0.0.0.0",728 server_port=7860,729 show_error=True730 )