Initial commit

examples/document_search.py

@@ -0,0 +1,146 @@
+#!/usr/bin/env python3
+"""
+Document search demo with recompute mode
+"""
+
+import os
+from pathlib import Path
+import shutil
+import time
+
+# Import backend packages to trigger plugin registration
+try:
+    import leann_backend_diskann
+    import leann_backend_hnsw
+    print("INFO: Backend packages imported successfully.")
+except ImportError as e:
+    print(f"WARNING: Could not import backend packages. Error: {e}")
+
+# Import upper-level API from leann-core
+from leann.api import LeannBuilder, LeannSearcher, LeannChat
+
+
+def load_sample_documents():
+    """Create sample documents for demonstration"""
+    docs = [
+        {"title": "Intro to Python", "content": "Python is a high-level, interpreted language known for simplicity."},
+        {"title": "ML Basics", "content": "Machine learning builds systems that learn from data."},
+        {"title": "Data Structures", "content": "Data structures like arrays, lists, and graphs organize data."},
+    ]
+    return docs
+
+def main():
+    print("==========================================================")
+    print("=== Leann Document Search Demo (DiskANN + Recompute) ===")
+    print("==========================================================")
+    
+    INDEX_DIR = Path("./test_indices")
+    INDEX_PATH = str(INDEX_DIR / "documents.diskann")
+    BACKEND_TO_TEST = "diskann"
+
+    if INDEX_DIR.exists():
+        print(f"--- Cleaning up old index directory: {INDEX_DIR} ---")
+        shutil.rmtree(INDEX_DIR)
+
+    # --- 1. Build index ---
+    print(f"\n[PHASE 1] Building index using '{BACKEND_TO_TEST}' backend...")
+    
+    builder = LeannBuilder(
+        backend_name=BACKEND_TO_TEST, 
+        graph_degree=32, 
+        complexity=64
+    )
+    
+    documents = load_sample_documents()
+    print(f"Loaded {len(documents)} sample documents.")
+    for doc in documents:
+        builder.add_text(doc["content"], metadata={"title": doc["title"]})
+        
+    builder.build_index(INDEX_PATH)
+    print(f"\nIndex built!")
+
+    # --- 2. Basic search demo ---
+    print(f"\n[PHASE 2] Basic search using '{BACKEND_TO_TEST}' backend...")
+    searcher = LeannSearcher(index_path=INDEX_PATH)
+    
+    query = "What is machine learning?"
+    print(f"\nQuery: '{query}'")
+    
+    print("\n--- Basic search mode (PQ computation) ---")
+    start_time = time.time()
+    results = searcher.search(query, top_k=2)
+    basic_time = time.time() - start_time
+    
+    print(f"⏱️  Basic search time: {basic_time:.3f} seconds")
+    print(">>> Basic search results <<<")
+    for i, res in enumerate(results, 1):
+        print(f"  {i}. ID: {res['id']}, Score: {res['score']:.4f}, Text: '{res['text']}', Metadata: {res['metadata']}")
+
+    # --- 3. Recompute search demo ---
+    print(f"\n[PHASE 3] Recompute search using embedding server...")
+    
+    print("\n--- Recompute search mode (get real embeddings via network) ---")
+    
+    # Configure recompute parameters
+    recompute_params = {
+        "recompute_beighbor_embeddings": True,  # Enable network recomputation
+        "USE_DEFERRED_FETCH": False,           # Don't use deferred fetch
+        "skip_search_reorder": True,           # Skip search reordering
+        "dedup_node_dis": True,               # Enable node distance deduplication
+        "prune_ratio": 0.1,                   # Pruning ratio 10%
+        "batch_recompute": False,             # Don't use batch recomputation
+        "global_pruning": False,              # Don't use global pruning
+        "zmq_port": 5555,                     # ZMQ port
+        "embedding_model": "sentence-transformers/all-mpnet-base-v2"
+    }
+    
+    print("Recompute parameter configuration:")
+    for key, value in recompute_params.items():
+        print(f"  {key}: {value}")
+    
+    print(f"\n🔄 Executing Recompute search...")
+    try:
+        start_time = time.time()
+        recompute_results = searcher.search(query, top_k=2, **recompute_params)
+        recompute_time = time.time() - start_time
+        
+        print(f"⏱️  Recompute search time: {recompute_time:.3f} seconds")
+        print(">>> Recompute search results <<<")
+        for i, res in enumerate(recompute_results, 1):
+            print(f"  {i}. ID: {res['id']}, Score: {res['score']:.4f}, Text: '{res['text']}', Metadata: {res['metadata']}")
+        
+        # Compare results
+        print(f"\n--- Result comparison ---")
+        print(f"Basic search time: {basic_time:.3f} seconds")
+        print(f"Recompute time: {recompute_time:.3f} seconds")
+        
+        print("\nBasic search vs Recompute results:")
+        for i in range(min(len(results), len(recompute_results))):
+            basic_score = results[i]['score']
+            recompute_score = recompute_results[i]['score']
+            score_diff = abs(basic_score - recompute_score)
+            print(f"  Position {i+1}: PQ={basic_score:.4f}, Recompute={recompute_score:.4f}, Difference={score_diff:.4f}")
+        
+        if recompute_time > basic_time:
+            print(f"✅ Recompute mode working correctly (more accurate but slower)")
+        else:
+            print(f"ℹ️  Recompute time is unusually fast, network recomputation may not be enabled")
+            
+    except Exception as e:
+        print(f"❌ Recompute search failed: {e}")
+        print("This usually indicates an embedding server connection issue")
+
+    # --- 4. Chat demo ---
+    print(f"\n[PHASE 4] Starting chat session...")
+    chat = LeannChat(index_path=INDEX_PATH)
+    chat_response = chat.ask(query)
+    print(f"You: {query}")
+    print(f"Leann: {chat_response}")
+
+    print("\n==========================================================")
+    print("✅ Demo finished successfully!")
+    print("==========================================================")
+
+
+if __name__ == "__main__":
+    main()