feat: implement hospital-specific customization pipeline with two-stage ANNOY retrieval

- Restructure pdf-version to customization/ for hospital-specific deployment
- Add customization_pipeline.py with two-stage retrieval (tag -> document -> chunk)
- Implement ANNOY indices for fast medical concept and chunk similarity search
- Add generate_embeddings.py for building hospital-specific embeddings
- Create test suite validating end-to-end pipeline functionality
- Add customization_requirements.txt with all necessary dependencies
- Update .gitignore to exclude rag_env/ virtual environment

This enables hospitals to deploy their own customized medical RAG systems
with private documents while maintaining the base medical knowledge framework.

.gitignore

@@ -1,6 +1,7 @@
 # 🧠 Virtual environments
 genAIvenv/
 .final_project_env/
+rag_env/
 .env
 .venv
 env/

customization/customization_pipeline.py

@@ -0,0 +1,156 @@
+#!/usr/bin/env python3
+"""Customization Pipeline - Hospital-Specific Document Retrieval
+
+This module provides the interface for hospital-specific document processing and retrieval.
+"""
+
+import sys
+from pathlib import Path
+from typing import List, Dict
+
+# Add src directory to Python path
+sys.path.insert(0, str(Path(__file__).parent / 'src'))
+
+# Import necessary modules
+from models.embedding_models import load_biomedbert_model
+from data.loaders import load_annotations
+from indexing.document_indexer import build_document_index
+from indexing.embedding_creator import create_tag_embeddings, create_chunk_embeddings
+from indexing.storage import save_document_system, load_document_system_with_annoy
+from retrieval.document_retriever import create_document_tag_mapping
+from retrieval.chunk_retriever import find_relevant_chunks_with_fallback
+
+
+def build_customization_embeddings():
+    """Build embeddings for the hospital-specific documents in the docs folder."""
+    print("🏥 Building hospital-specific embeddings...")
+    
+    # Paths
+    base_path = Path(__file__).parent
+    docs_path = base_path / "docs"
+    processing_path = base_path / "processing"
+    
+    # Load model and annotations
+    embedding_model = load_biomedbert_model()
+    annotations = load_annotations(file_path=str(processing_path / "mapping.json"))
+    
+    if not annotations:
+        print("❌ Unable to load annotation data")
+        return False
+    
+    # Build document index with chunks
+    print("📄 Processing documents...")
+    document_index = build_document_index(
+        annotations, 
+        assets_dir=str(docs_path),
+        chunk_size=256, 
+        chunk_overlap=25
+    )
+    
+    # Create embeddings
+    print("🔢 Creating embeddings...")
+    tag_embeddings = create_tag_embeddings(embedding_model, document_index)
+    doc_tag_mapping = create_document_tag_mapping(document_index, tag_embeddings)
+    chunk_embeddings = create_chunk_embeddings(embedding_model, document_index)
+    
+    # Save everything
+    print("💾 Saving to processing folder...")
+    save_document_system(
+        document_index,
+        tag_embeddings,
+        doc_tag_mapping,
+        chunk_embeddings,
+        output_dir=str(processing_path / "embeddings"),
+        build_annoy_indices=True
+    )
+    
+    print("✅ Embeddings built successfully!")
+    return True
+
+
+def retrieve_document_chunks(query: str, top_k: int = 5) -> List[Dict]:
+    """Retrieve relevant document chunks using two-stage ANNOY retrieval.
+    
+    Stage 1: Find relevant documents using tag embeddings (medical concepts)
+    Stage 2: Find relevant chunks within those documents using chunk embeddings
+    
+    Args:
+        query: The search query
+        top_k: Number of chunks to retrieve
+        
+    Returns:
+        List of dictionaries containing chunk information
+    """
+    # Load model and existing embeddings
+    embedding_model = load_biomedbert_model()
+    
+    # Load from processing folder
+    processing_path = Path(__file__).parent / "processing"
+    
+    # Load the saved system with ANNOY indices
+    document_index, tag_embeddings, doc_tag_mapping, chunk_embeddings, annoy_manager = \
+        load_document_system_with_annoy(
+            input_dir=str(processing_path / "embeddings"),
+            annoy_dir=str(processing_path / "indices")
+        )
+    
+    if annoy_manager is None:
+        print("❌ Failed to load ANNOY manager")
+        return []
+    
+    # Create query embedding
+    query_embedding = embedding_model.encode(query)
+    
+    # Stage 1: Find relevant documents using tag ANNOY index
+    print(f"🔍 Stage 1: Finding relevant documents for query: '{query}'")
+    relevant_tags, tag_distances = annoy_manager.search_tags(
+        query_embedding=query_embedding,
+        n_neighbors=20,  # Get more tags to find diverse documents
+        include_distances=True
+    )
+    
+    # Get documents that contain these relevant tags
+    relevant_docs = set()
+    for tag in relevant_tags[:10]:  # Use top 10 tags
+        for doc_name, doc_info in doc_tag_mapping.items():
+            if tag in doc_info['tags']:
+                relevant_docs.add(doc_name)
+    
+    relevant_docs = list(relevant_docs)
+    print(f"✅ Found {len(relevant_docs)} relevant documents based on medical tags")
+    
+    if not relevant_docs:
+        print("❌ No relevant documents found")
+        return []
+    
+    # Stage 2: Find relevant chunks within these documents using chunk ANNOY index
+    print(f"🔍 Stage 2: Finding relevant chunks within {len(relevant_docs)} documents")
+    chunks, chunk_distances = annoy_manager.search_chunks_in_documents(
+        query_embedding=query_embedding,
+        document_names=relevant_docs,
+        n_neighbors=top_k,
+        include_distances=True
+    )
+    
+    # Convert ANNOY distances to cosine similarities 
+    from indexing.annoy_manager import convert_angular_distance_to_cosine_similarity
+    
+    # Format results
+    results = []
+    for chunk, distance in zip(chunks, chunk_distances):
+        # Convert angular distance to cosine similarity
+        similarity = convert_angular_distance_to_cosine_similarity(distance)
+        
+        results.append({
+            'document': chunk['document'],
+            'chunk_text': chunk['text'],
+            'score': similarity,
+            'metadata': {
+                'chunk_id': chunk['chunk_id'],
+                'start_char': chunk.get('start_char', 0),
+                'end_char': chunk.get('end_char', 0)
+            }
+        })
+    
+    print(f"✅ Retrieved {len(results)} relevant chunks")
+    return results

{src/pdf-version → customization}/generate_embeddings.py

@@ -1,18 +1,12 @@
 #!/usr/bin/env python3
 """
-Quick script to generate new embeddings with sentence-based chunking
+Generate embeddings for hospital-specific documents
 """
 
-import sys
-from pathlib import Path
-
-# Add pdf-version directory to Python path
-sys.path.insert(0, str(Path(__file__).parent))
-
-from demos.demo_runner import build_medical_rag_system
+from customization_pipeline import build_customization_embeddings
 
 def main():
-    print("🚀 Starting to build medical RAG system with new sentence-based chunking...")
+    print("🚀 Starting to build hospital-specific embeddings...")
     print("📋 Configuration:")
     print("   - Chunk size: 256 tokens")
     print("   - Chunk overlap: 25 tokens (10%)")
@@ -22,17 +16,19 @@ def main():
     print("")
     
     try:
-        result = build_medical_rag_system(enable_chunk_embeddings=True)
+        success = build_customization_embeddings()
         
-        if result[0] is not None:
-            print("✅ Successfully built medical RAG system!")
-            print("📁 Generated files:")
-            print("   - document_index.json")
-            print("   - tag_embeddings.json") 
-            print("   - document_tag_mapping.json")
-            print("   - chunk_embeddings.json")
+        if success:
+            print("\n✅ Successfully built embeddings!")
+            print("📁 Generated files in processing folder:")
+            print("   - embeddings/document_index.json")
+            print("   - embeddings/tag_embeddings.json") 
+            print("   - embeddings/document_tag_mapping.json")
+            print("   - embeddings/chunk_embeddings.json")
+            print("   - indices/annoy_metadata.json")
+            print("   - indices/*.ann files")
         else:
-            print("❌ Failed to build system")
+            print("\n❌ Failed to build embeddings")
             
     except KeyboardInterrupt:
         print("\n⚠️ Process interrupted by user")

customization/src/data/__init__.py

@@ -0,0 +1,23 @@
+"""Data loading and PDF processing."""
+
+from .loaders import load_annotations, filter_pdf_files
+
+# Try to import PDF processing functions, but handle missing dependencies gracefully
+try:
+    from .pdf_processing import (
+        extract_pdf_text, 
+        extract_tables_from_pdf,
+        extract_images_ocr_from_pdf,
+        extract_pdf_content_enhanced
+    )
+    PDF_PROCESSING_AVAILABLE = True
+    __all__ = [
+        'load_annotations', 'filter_pdf_files',
+        'extract_pdf_text', 'extract_tables_from_pdf', 
+        'extract_images_ocr_from_pdf', 'extract_pdf_content_enhanced'
+    ]
+except ImportError as e:
+    print(f"⚠️ PDF processing not available: {e}")
+    print("📝 Only working with existing embeddings")
+    PDF_PROCESSING_AVAILABLE = False
+    __all__ = ['load_annotations', 'filter_pdf_files']

{src/pdf-version → customization/src}/demos/demo_runner.py

@@ -6,9 +6,15 @@ from models.embedding_models import load_biomedbert_model
 from data.loaders import load_annotations
 from indexing.document_indexer import build_document_index
 from indexing.embedding_creator import create_tag_embeddings, create_chunk_embeddings
-from indexing.storage import save_document_system, load_document_system
-from retrieval.document_retriever import create_document_tag_mapping, find_relevant_documents
-from retrieval.chunk_retriever import find_relevant_chunks, get_documents_for_rag, get_chunks_for_rag
+from indexing.storage import save_document_system, load_document_system, load_document_system_with_annoy
+from retrieval.document_retriever import (
+    create_document_tag_mapping, find_relevant_documents, 
+    find_relevant_documents_with_fallback
+)
+from retrieval.chunk_retriever import (
+    find_relevant_chunks, get_documents_for_rag, get_chunks_for_rag,
+    find_relevant_chunks_with_fallback
+)
 
 
 def build_medical_rag_system(enable_chunk_embeddings: bool = True):
@@ -135,4 +141,151 @@ def demo_all_strategies(query: str = "chest pain and shortness of breath"):
     for strategy, docs in results.items():
         print(f"{strategy:>10}: {len(docs)} documents selected")
     
-    return results
+    return results
+
+
+def demo_rag_query_with_annoy(query: str = "chest pain and shortness of breath", 
+                             strategy: str = "top_p", use_chunks: bool = True, **kwargs):
+    """Demo RAG query functionality with ANNOY acceleration."""
+    print(f"\n🚀 Demo ANNOY Query: '{query}' (Strategy: {strategy}, Use chunks: {use_chunks})")
+    print("=" * 80)
+    
+    # Try to load existing system with ANNOY
+    document_index, tag_embeddings, doc_tag_mapping, chunk_embeddings, annoy_manager = load_document_system_with_annoy()
+    
+    if document_index is None:
+        print("📦 No saved system found, building new one...")
+        build_result = build_medical_rag_system(enable_chunk_embeddings=use_chunks)
+        if build_result[0] is None:
+            return
+        embedding_model, document_index, tag_embeddings, doc_tag_mapping, chunk_embeddings = build_result
+        
+        # Try to load ANNOY manager after building
+        from indexing.storage import load_annoy_manager
+        annoy_manager = load_annoy_manager()
+    else:
+        embedding_model = load_biomedbert_model()
+    
+    print(f"🔧 ANNOY Status: {'Available' if annoy_manager else 'Not available (using fallback)'}")
+    
+    # Find relevant documents using ANNOY-accelerated method with fallback
+    print(f"\n🔍 Finding relevant documents...")
+    import time
+    start_time = time.time()
+    
+    relevant_docs = find_relevant_documents_with_fallback(
+        query, embedding_model, tag_embeddings, doc_tag_mapping,
+        annoy_manager=annoy_manager, strategy=strategy, **kwargs
+    )
+    
+    doc_search_time = time.time() - start_time
+    print(f"⏱️ Document search completed in {doc_search_time:.4f}s")
+    
+    if use_chunks and chunk_embeddings:
+        # Find relevant chunks using ANNOY-accelerated method with fallback
+        print(f"\n🔍 Finding relevant chunks within selected documents...")
+        start_time = time.time()
+        
+        relevant_chunks = find_relevant_chunks_with_fallback(
+            query, embedding_model, relevant_docs, chunk_embeddings,
+            annoy_manager=annoy_manager, strategy=strategy, 
+            top_chunks_per_doc=3, **kwargs
+        )
+        
+        chunk_search_time = time.time() - start_time
+        print(f"⏱️ Chunk search completed in {chunk_search_time:.4f}s")
+        
+        # Get chunks for RAG
+        rag_content = get_chunks_for_rag(relevant_chunks, max_chunks=10)
+        print(f"\n📋 Ready for RAG with {len(rag_content)} chunks")
+        
+        total_time = doc_search_time + chunk_search_time
+        print(f"🏁 Total search time: {total_time:.4f}s")
+        
+    else:
+        # Get full documents for RAG
+        rag_content = get_documents_for_rag(relevant_docs, document_index)
+        print(f"\n📋 Ready for RAG with {len(rag_content)} full documents")
+        print(f"🏁 Total search time: {doc_search_time:.4f}s")
+    
+    return rag_content
+
+
+def demo_performance_comparison(query: str = "chest pain and shortness of breath"):
+    """Demo performance comparison between original and ANNOY methods."""
+    print(f"\n⚡ Performance Comparison Demo")
+    print("=" * 80)
+    print(f"Query: '{query}'")
+    
+    # Load system with ANNOY
+    document_index, tag_embeddings, doc_tag_mapping, chunk_embeddings, annoy_manager = load_document_system_with_annoy()
+    
+    if document_index is None:
+        print("❌ No saved system found")
+        return
+    
+    embedding_model = load_biomedbert_model()
+    strategy = "top_p"
+    strategy_params = {"top_p": 0.8, "min_similarity": 0.3}
+    
+    print(f"\n📊 Testing document retrieval performance...")
+    
+    # Test original method
+    import time
+    start_time = time.time()
+    original_docs = find_relevant_documents(
+        query, embedding_model, tag_embeddings, doc_tag_mapping, 
+        strategy=strategy, **strategy_params
+    )
+    original_time = time.time() - start_time
+    
+    # Test ANNOY method (with fallback)
+    start_time = time.time()
+    annoy_docs = find_relevant_documents_with_fallback(
+        query, embedding_model, tag_embeddings, doc_tag_mapping,
+        annoy_manager=annoy_manager, strategy=strategy, **strategy_params
+    )
+    annoy_time = time.time() - start_time
+    
+    # Results
+    print(f"🔍 Original method: {len(original_docs)} docs in {original_time:.4f}s")
+    print(f"🚀 ANNOY method: {len(annoy_docs)} docs in {annoy_time:.4f}s")
+    
+    if annoy_time > 0:
+        speedup = original_time / annoy_time
+        print(f"⚡ Speedup: {speedup:.2f}x")
+    
+    # Check result similarity
+    if original_docs and annoy_docs:
+        overlap = set(original_docs) & set(annoy_docs)
+        print(f"📊 Result overlap: {len(overlap)}/{len(original_docs)} documents")
+    
+    # Test chunk retrieval if available
+    if chunk_embeddings and len(original_docs) > 0:
+        print(f"\n📊 Testing chunk retrieval performance...")
+        relevant_docs = original_docs[:2]  # Test with first 2 documents
+        
+        # Original method
+        start_time = time.time()
+        original_chunks = find_relevant_chunks(
+            query, embedding_model, relevant_docs, chunk_embeddings, 
+            strategy=strategy, **strategy_params
+        )
+        original_chunk_time = time.time() - start_time
+        
+        # ANNOY method (with fallback)
+        start_time = time.time()
+        annoy_chunks = find_relevant_chunks_with_fallback(
+            query, embedding_model, relevant_docs, chunk_embeddings,
+            annoy_manager=annoy_manager, strategy=strategy, **strategy_params
+        )
+        annoy_chunk_time = time.time() - start_time
+        
+        print(f"🔍 Original chunks: {len(original_chunks)} chunks in {original_chunk_time:.4f}s")
+        print(f"🚀 ANNOY chunks: {len(annoy_chunks)} chunks in {annoy_chunk_time:.4f}s")
+        
+        if annoy_chunk_time > 0:
+            chunk_speedup = original_chunk_time / annoy_chunk_time
+            print(f"⚡ Chunk speedup: {chunk_speedup:.2f}x")
+    
+    print(f"\n✅ Performance comparison completed!")

customization/src/indexing/annoy_manager.py

@@ -0,0 +1,392 @@
+"""ANNOY index management for PDF-based RAG system."""
+
+import os
+import json
+import numpy as np
+from typing import Dict, List, Optional, Tuple, Union
+from pathlib import Path
+import logging
+
+try:
+    from annoy import AnnoyIndex
+except ImportError:
+    raise ImportError("annoy package is required. Install with: pip install annoy")
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class AnnoyIndexManager:
+    """Manages ANNOY indices for fast vector similarity search."""
+    
+    def __init__(self, embedding_dim: int = 1024, metric: str = 'angular'):
+        """
+        Initialize ANNOY index manager.
+        
+        Args:
+            embedding_dim: Dimension of embeddings (1024 for BGE Large Medical)
+            metric: Distance metric ('angular' for cosine similarity, 'euclidean', 'manhattan', 'hamming', 'dot')
+        """
+        self.embedding_dim = embedding_dim
+        self.metric = metric
+        self.tag_index = None
+        self.chunk_index = None
+        self.tag_to_id_mapping = {}
+        self.id_to_tag_mapping = {}
+        self.chunk_to_id_mapping = {}
+        self.id_to_chunk_mapping = {}
+        
+        logger.info(f"Initialized AnnoyIndexManager: dim={embedding_dim}, metric={metric}")
+    
+    def build_tag_index(self, tag_embeddings: Dict[str, np.ndarray], n_trees: int = 50) -> AnnoyIndex:
+        """
+        Build ANNOY index for tag embeddings.
+        
+        Args:
+            tag_embeddings: Dictionary mapping tags to their embeddings
+            n_trees: Number of trees (more trees = better precision, slower build)
+            
+        Returns:
+            Built ANNOY index
+        """
+        logger.info(f"Building tag ANNOY index with {len(tag_embeddings)} tags...")
+        
+        # Create index
+        self.tag_index = AnnoyIndex(self.embedding_dim, self.metric)
+        
+        # Create mappings
+        self.tag_to_id_mapping = {}
+        self.id_to_tag_mapping = {}
+        
+        # Add embeddings to index
+        for tag_id, (tag, embedding) in enumerate(tag_embeddings.items()):
+            self.tag_index.add_item(tag_id, embedding)
+            self.tag_to_id_mapping[tag] = tag_id
+            self.id_to_tag_mapping[tag_id] = tag
+        
+        # Build index
+        logger.info(f"Building index with {n_trees} trees...")
+        self.tag_index.build(n_trees)
+        
+        logger.info(f"✅ Tag ANNOY index built successfully: {len(tag_embeddings)} tags")
+        return self.tag_index
+    
+    def build_chunk_index(self, chunk_embeddings: Dict[str, List[Dict]], n_trees: int = 50) -> AnnoyIndex:
+        """
+        Build ANNOY index for chunk embeddings.
+        
+        Args:
+            chunk_embeddings: Dictionary mapping document names to lists of chunk dictionaries
+            n_trees: Number of trees
+            
+        Returns:
+            Built ANNOY index
+        """
+        # Count total chunks
+        total_chunks = sum(len(chunks) for chunks in chunk_embeddings.values())
+        logger.info(f"Building chunk ANNOY index with {total_chunks} chunks...")
+        
+        # Create index
+        self.chunk_index = AnnoyIndex(self.embedding_dim, self.metric)
+        
+        # Create mappings
+        self.chunk_to_id_mapping = {}
+        self.id_to_chunk_mapping = {}
+        
+        chunk_id = 0
+        for doc_name, chunks in chunk_embeddings.items():
+            for chunk in chunks:
+                # Create unique chunk identifier
+                chunk_key = f"{doc_name}#{chunk['chunk_id']}"
+                
+                # Add to index
+                self.chunk_index.add_item(chunk_id, chunk['embedding'])
+                
+                # Create mappings
+                self.chunk_to_id_mapping[chunk_key] = chunk_id
+                self.id_to_chunk_mapping[chunk_id] = {
+                    'document': doc_name,
+                    'chunk_id': chunk['chunk_id'],
+                    'text': chunk['text'],
+                    'start_char': chunk.get('start_char', 0),
+                    'end_char': chunk.get('end_char', len(chunk['text'])),
+                    'token_count': chunk.get('token_count', len(chunk['text'].split())),
+                    'chunk_key': chunk_key
+                }
+                
+                chunk_id += 1
+        
+        # Build index
+        logger.info(f"Building chunk index with {n_trees} trees...")
+        self.chunk_index.build(n_trees)
+        
+        logger.info(f"✅ Chunk ANNOY index built successfully: {total_chunks} chunks")
+        return self.chunk_index
+    
+    def save_indices(self, output_dir: Union[str, Path]):
+        """
+        Save ANNOY indices and mappings to disk.
+        
+        Args:
+            output_dir: Directory to save indices
+        """
+        output_dir = Path(output_dir)
+        # Save indices at the same level as embeddings, not inside embeddings
+        indices_dir = output_dir.parent / 'indices'
+        indices_dir.mkdir(exist_ok=True)
+        
+        # Save tag index
+        if self.tag_index is not None:
+            tag_index_path = indices_dir / 'tag_embeddings.ann'
+            self.tag_index.save(str(tag_index_path))
+            
+            # Save tag mappings
+            tag_mappings_path = indices_dir / 'tag_mappings.json'
+            with open(tag_mappings_path, 'w', encoding='utf-8') as f:
+                json.dump({
+                    'tag_to_id': self.tag_to_id_mapping,
+                    'id_to_tag': self.id_to_tag_mapping
+                }, f, indent=2, ensure_ascii=False)
+            
+            logger.info(f"✅ Tag index saved: {tag_index_path}")
+        
+        # Save chunk index
+        if self.chunk_index is not None:
+            chunk_index_path = indices_dir / 'chunk_embeddings.ann'
+            self.chunk_index.save(str(chunk_index_path))
+            
+            # Save chunk mappings
+            chunk_mappings_path = indices_dir / 'chunk_mappings.json'
+            with open(chunk_mappings_path, 'w', encoding='utf-8') as f:
+                json.dump({
+                    'chunk_to_id': self.chunk_to_id_mapping,
+                    'id_to_chunk': self.id_to_chunk_mapping
+                }, f, indent=2, ensure_ascii=False)
+            
+            logger.info(f"✅ Chunk index saved: {chunk_index_path}")
+        
+        # Save index metadata
+        metadata_path = indices_dir / 'annoy_metadata.json'
+        with open(metadata_path, 'w', encoding='utf-8') as f:
+            json.dump({
+                'embedding_dim': self.embedding_dim,
+                'metric': self.metric,
+                'tag_index_exists': self.tag_index is not None,
+                'chunk_index_exists': self.chunk_index is not None,
+                'num_tags': len(self.tag_to_id_mapping),
+                'num_chunks': len(self.chunk_to_id_mapping)
+            }, f, indent=2)
+        
+        logger.info(f"✅ ANNOY indices saved to: {indices_dir}")
+    
+    def load_indices(self, input_dir: Union[str, Path]) -> bool:
+        """
+        Load ANNOY indices and mappings from disk.
+        
+        Args:
+            input_dir: Directory containing saved indices
+            
+        Returns:
+            True if successfully loaded, False otherwise
+        """
+        input_dir = Path(input_dir)
+        # Load indices from the same level as embeddings, not inside embeddings
+        indices_dir = input_dir.parent / 'indices'
+        
+        if not indices_dir.exists():
+            logger.warning(f"Indices directory not found: {indices_dir}")
+            return False
+        
+        try:
+            # Load metadata
+            metadata_path = indices_dir / 'annoy_metadata.json'
+            if metadata_path.exists():
+                with open(metadata_path, 'r', encoding='utf-8') as f:
+                    metadata = json.load(f)
+                self.embedding_dim = metadata['embedding_dim']
+                self.metric = metadata['metric']
+                logger.info(f"Loaded metadata: dim={self.embedding_dim}, metric={self.metric}")
+            
+            # Load tag index
+            tag_index_path = indices_dir / 'tag_embeddings.ann'
+            tag_mappings_path = indices_dir / 'tag_mappings.json'
+            
+            if tag_index_path.exists() and tag_mappings_path.exists():
+                self.tag_index = AnnoyIndex(self.embedding_dim, self.metric)
+                self.tag_index.load(str(tag_index_path))
+                
+                with open(tag_mappings_path, 'r', encoding='utf-8') as f:
+                    mappings = json.load(f)
+                    self.tag_to_id_mapping = mappings['tag_to_id']
+                    self.id_to_tag_mapping = {int(k): v for k, v in mappings['id_to_tag'].items()}
+                
+                logger.info(f"✅ Tag index loaded: {len(self.tag_to_id_mapping)} tags")
+            
+            # Load chunk index
+            chunk_index_path = indices_dir / 'chunk_embeddings.ann'
+            chunk_mappings_path = indices_dir / 'chunk_mappings.json'
+            
+            if chunk_index_path.exists() and chunk_mappings_path.exists():
+                self.chunk_index = AnnoyIndex(self.embedding_dim, self.metric)
+                self.chunk_index.load(str(chunk_index_path))
+                
+                with open(chunk_mappings_path, 'r', encoding='utf-8') as f:
+                    mappings = json.load(f)
+                    self.chunk_to_id_mapping = mappings['chunk_to_id']
+                    self.id_to_chunk_mapping = {int(k): v for k, v in mappings['id_to_chunk'].items()}
+                
+                logger.info(f"✅ Chunk index loaded: {len(self.chunk_to_id_mapping)} chunks")
+            
+            return True
+            
+        except Exception as e:
+            logger.error(f"Failed to load ANNOY indices: {e}")
+            return False
+    
+    def search_tags(self, query_embedding: np.ndarray, n_neighbors: int = 10, 
+                   include_distances: bool = True) -> Union[List[str], Tuple[List[str], List[float]]]:
+        """
+        Search for similar tags using ANNOY index.
+        
+        Args:
+            query_embedding: Query embedding vector
+            n_neighbors: Number of nearest neighbors to return
+            include_distances: Whether to return distances
+            
+        Returns:
+            List of tag names, or tuple of (tag_names, distances)
+        """
+        if self.tag_index is None:
+            raise ValueError("Tag index not built or loaded")
+        
+        # Search using ANNOY
+        if include_distances:
+            neighbor_ids, distances = self.tag_index.get_nns_by_vector(
+                query_embedding, n_neighbors, include_distances=True
+            )
+        else:
+            neighbor_ids = self.tag_index.get_nns_by_vector(
+                query_embedding, n_neighbors, include_distances=False
+            )
+        
+        # Convert IDs to tag names
+        tag_names = [self.id_to_tag_mapping[neighbor_id] for neighbor_id in neighbor_ids]
+        
+        if include_distances:
+            return tag_names, distances
+        else:
+            return tag_names
+    
+    def search_chunks(self, query_embedding: np.ndarray, n_neighbors: int = 10,
+                     include_distances: bool = True) -> Union[List[Dict], Tuple[List[Dict], List[float]]]:
+        """
+        Search for similar chunks using ANNOY index.
+        
+        Args:
+            query_embedding: Query embedding vector
+            n_neighbors: Number of nearest neighbors to return
+            include_distances: Whether to return distances
+            
+        Returns:
+            List of chunk dictionaries, or tuple of (chunks, distances)
+        """
+        if self.chunk_index is None:
+            raise ValueError("Chunk index not built or loaded")
+        
+        # Search using ANNOY
+        if include_distances:
+            neighbor_ids, distances = self.chunk_index.get_nns_by_vector(
+                query_embedding, n_neighbors, include_distances=True
+            )
+        else:
+            neighbor_ids = self.chunk_index.get_nns_by_vector(
+                query_embedding, n_neighbors, include_distances=False
+            )
+        
+        # Convert IDs to chunk info
+        chunks = [self.id_to_chunk_mapping[neighbor_id] for neighbor_id in neighbor_ids]
+        
+        if include_distances:
+            return chunks, distances
+        else:
+            return chunks
+    
+    def search_chunks_in_documents(self, query_embedding: np.ndarray, 
+                                  document_names: List[str], n_neighbors: int = 10,
+                                  include_distances: bool = True) -> Union[List[Dict], Tuple[List[Dict], List[float]]]:
+        """
+        Search for similar chunks within specific documents.
+        
+        Args:
+            query_embedding: Query embedding vector
+            document_names: List of document names to search within
+            n_neighbors: Number of nearest neighbors to return
+            include_distances: Whether to return distances
+            
+        Returns:
+            List of chunk dictionaries, or tuple of (chunks, distances)
+        """
+        if self.chunk_index is None:
+            raise ValueError("Chunk index not built or loaded")
+        
+        # Get more candidates than needed since we'll filter by document
+        search_candidates = min(n_neighbors * 5, len(self.id_to_chunk_mapping))
+        
+        # Search using ANNOY
+        if include_distances:
+            candidate_ids, distances = self.chunk_index.get_nns_by_vector(
+                query_embedding, search_candidates, include_distances=True
+            )
+        else:
+            candidate_ids = self.chunk_index.get_nns_by_vector(
+                query_embedding, search_candidates, include_distances=False
+            )
+        
+        # Filter by document names and take top n_neighbors
+        filtered_chunks = []
+        filtered_distances = [] if include_distances else None
+        
+        for i, candidate_id in enumerate(candidate_ids):
+            chunk_info = self.id_to_chunk_mapping[candidate_id]
+            if chunk_info['document'] in document_names:
+                filtered_chunks.append(chunk_info)
+                if include_distances:
+                    filtered_distances.append(distances[i])
+                
+                if len(filtered_chunks) >= n_neighbors:
+                    break
+        
+        if include_distances:
+            return filtered_chunks, filtered_distances
+        else:
+            return filtered_chunks
+    
+    def get_index_stats(self) -> Dict:
+        """Get statistics about the loaded indices."""
+        stats = {
+            'embedding_dim': self.embedding_dim,
+            'metric': self.metric,
+            'tag_index_loaded': self.tag_index is not None,
+            'chunk_index_loaded': self.chunk_index is not None,
+            'num_tags': len(self.tag_to_id_mapping) if self.tag_index else 0,
+            'num_chunks': len(self.chunk_to_id_mapping) if self.chunk_index else 0
+        }
+        return stats
+
+
+def convert_angular_distance_to_cosine_similarity(angular_distance: float) -> float:
+    """
+    Convert ANNOY angular distance to cosine similarity.
+    
+    Args:
+        angular_distance: Angular distance from ANNOY
+        
+    Returns:
+        Cosine similarity (0 to 1)
+    """
+    # Angular distance is related to cosine similarity by:
+    # angular_distance = 2 * arccos(cosine_similarity) / π
+    # Therefore: cosine_similarity = cos(angular_distance * π / 2)
+    import math
+    return math.cos(angular_distance * math.pi / 2)

{src/pdf-version → customization/src}/indexing/storage.py

@@ -2,13 +2,19 @@
 
 import json
 import os
+import logging
 from typing import Dict, Optional, Tuple
 import numpy as np
+from .annoy_manager import AnnoyIndexManager
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 
 def save_document_system(document_index: Dict, tag_embeddings: Dict, 
                         doc_tag_mapping: Dict, chunk_embeddings: Dict = None, 
-                        output_dir: str = None):
+                        output_dir: str = None, build_annoy_indices: bool = True):
     """Save the complete document indexing system.
     
     Args:
@@ -85,6 +91,31 @@ def save_document_system(document_index: Dict, tag_embeddings: Dict,
         with open(os.path.join(output_dir, 'chunk_embeddings.json'), 'w', encoding='utf-8') as f:
             json.dump(chunk_embeddings_serializable, f, indent=2, ensure_ascii=False)
     
+    # Build and save ANNOY indices if requested
+    if build_annoy_indices:
+        logger.info("🔧 Building ANNOY indices for fast retrieval...")
+        try:
+            # Initialize ANNOY manager (assuming BGE Large Medical embedding dimension)
+            annoy_manager = AnnoyIndexManager(embedding_dim=1024, metric='angular')
+            
+            # Build tag index
+            logger.info("Building tag ANNOY index...")
+            annoy_manager.build_tag_index(tag_embeddings, n_trees=50)
+            
+            # Build chunk index if chunk embeddings are provided
+            if chunk_embeddings:
+                logger.info("Building chunk ANNOY index...")
+                annoy_manager.build_chunk_index(chunk_embeddings, n_trees=50)
+            
+            # Save indices
+            logger.info("Saving ANNOY indices...")
+            annoy_manager.save_indices(output_dir)
+            
+            logger.info("✅ ANNOY indices built and saved successfully")
+        except Exception as e:
+            logger.error(f"❌ Failed to build ANNOY indices: {e}")
+            logger.warning("Continuing without ANNOY indices - will use original search methods")
+    
     print("✅ Document system saved to files")
 
 
@@ -161,4 +192,62 @@ def load_document_system(input_dir: str = None) -> Tuple[Optional[Dict], Optiona
         
     except Exception as e:
         print(f"❌ Failed to load document system: {e}")
-        return None, None, None, None
+        return None, None, None, None
+
+
+def load_annoy_manager(input_dir: str = None) -> Optional[AnnoyIndexManager]:
+    """
+    Load ANNOY index manager with pre-built indices.
+    
+    Args:
+        input_dir: Input directory containing saved indices
+        
+    Returns:
+        AnnoyIndexManager instance or None if loading fails
+    """
+    if input_dir is None:
+        # Get project root directory
+        from pathlib import Path
+        root_dir = Path(__file__).parent.parent.parent.parent
+        input_dir = root_dir / 'embeddings' / 'pdfembeddings'
+    
+    try:
+        # Initialize ANNOY manager
+        annoy_manager = AnnoyIndexManager(embedding_dim=1024, metric='angular')
+        
+        # Try to load indices
+        if annoy_manager.load_indices(input_dir):
+            logger.info("✅ ANNOY indices loaded successfully")
+            return annoy_manager
+        else:
+            logger.warning("⚠️ Failed to load ANNOY indices")
+            return None
+            
+    except Exception as e:
+        logger.error(f"❌ Failed to initialize ANNOY manager: {e}")
+        return None
+
+
+def load_document_system_with_annoy(input_dir: str = None, annoy_dir: str = None) -> Tuple[Optional[Dict], Optional[Dict], Optional[Dict], Optional[Dict], Optional[AnnoyIndexManager]]:
+    """
+    Load the complete document indexing system including ANNOY indices.
+    
+    Args:
+        input_dir: Input directory containing saved files
+        annoy_dir: Directory containing ANNOY indices (if different from input_dir)
+        
+    Returns:
+        Tuple of (document_index, tag_embeddings, doc_tag_mapping, chunk_embeddings, annoy_manager).
+        Returns all None values if loading fails.
+    """
+    # Load the standard document system
+    document_index, tag_embeddings, doc_tag_mapping, chunk_embeddings = load_document_system(input_dir)
+    
+    if document_index is None:
+        return None, None, None, None, None
+    
+    # Load ANNOY manager
+    # Use annoy_dir if provided, otherwise use input_dir
+    annoy_manager = load_annoy_manager(annoy_dir if annoy_dir else input_dir)
+    
+    return document_index, tag_embeddings, doc_tag_mapping, chunk_embeddings, annoy_manager

{src/pdf-version → customization/src}/retrieval/chunk_retriever.py

@@ -1,9 +1,15 @@
 """Chunk-level retrieval functionality."""
 
-from typing import List, Dict, Callable
+from typing import List, Dict, Callable, Optional
 import numpy as np
+import logging
 from sentence_transformers import SentenceTransformer
-from src.indexing.embedding_creator import create_text_embedding
+from indexing.embedding_creator import create_text_embedding
+from indexing.annoy_manager import AnnoyIndexManager, convert_angular_distance_to_cosine_similarity
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 
 def cosine_similarity(vec1: np.ndarray, vec2: np.ndarray) -> float:
@@ -190,4 +196,172 @@ def get_chunks_for_rag(relevant_chunks: List[Dict], max_chunks: int = 10) -> Lis
         rag_chunks.append(formatted_chunk)
     
     print(f"📄 Retrieved {len(rag_chunks)} chunks for RAG")
-    return rag_chunks
+    return rag_chunks
+
+
+# ANNOY-accelerated chunk retrieval functions
+
+def find_relevant_chunks_annoy_top_k(query: str, model: SentenceTransformer, 
+                                    relevant_docs: List[str], annoy_manager: AnnoyIndexManager,
+                                    top_chunks_per_doc: int = 3, 
+                                    similarity_metric: str = "angular") -> List[Dict]:
+    """Find most relevant chunks using ANNOY index and Top-K strategy."""
+    query_embedding = create_text_embedding(model, query)
+    
+    # Use ANNOY to search chunks in the relevant documents
+    all_chunks, distances = annoy_manager.search_chunks_in_documents(
+        query_embedding, relevant_docs, 
+        n_neighbors=len(relevant_docs) * top_chunks_per_doc,
+        include_distances=True
+    )
+    
+    # Convert distances to similarities and format results
+    all_relevant_chunks = []
+    for chunk, distance in zip(all_chunks, distances):
+        similarity = convert_angular_distance_to_cosine_similarity(distance)
+        
+        chunk_result = {
+            'document': chunk['document'],
+            'chunk_id': chunk['chunk_id'],
+            'text': chunk['text'],
+            'start_char': chunk.get('start_char', 0),
+            'end_char': chunk.get('end_char', len(chunk['text'])),
+            'token_count': chunk.get('token_count', len(chunk['text'].split())),
+            'similarity': similarity
+        }
+        all_relevant_chunks.append(chunk_result)
+    
+    # Group by document and take top chunks per document
+    doc_chunks = {}
+    for chunk in all_relevant_chunks:
+        doc_name = chunk['document']
+        if doc_name not in doc_chunks:
+            doc_chunks[doc_name] = []
+        doc_chunks[doc_name].append(chunk)
+    
+    # Take top chunks from each document
+    final_chunks = []
+    for doc_name in relevant_docs:
+        if doc_name in doc_chunks:
+            doc_chunks[doc_name].sort(key=lambda x: x['similarity'], reverse=True)
+            final_chunks.extend(doc_chunks[doc_name][:top_chunks_per_doc])
+    
+    # Sort all chunks by similarity
+    final_chunks.sort(key=lambda x: x['similarity'], reverse=True)
+    
+    logger.info(f"🚀 Found {len(final_chunks)} relevant chunks (ANNOY Top-K)")
+    for i, chunk in enumerate(final_chunks[:5]):  # Show top 5
+        logger.info(f"  {i+1}. {chunk['document']} (chunk {chunk['chunk_id']}, similarity: {chunk['similarity']:.3f})")
+        logger.info(f"     Preview: {chunk['text'][:100]}...")
+    
+    return final_chunks
+
+
+def find_relevant_chunks_annoy_top_p(query: str, model: SentenceTransformer,
+                                    relevant_docs: List[str], annoy_manager: AnnoyIndexManager,
+                                    top_p: float = 0.6, min_similarity: float = 0.3,
+                                    similarity_metric: str = "angular") -> List[Dict]:
+    """Find most relevant chunks using ANNOY index and Top-P strategy."""
+    query_embedding = create_text_embedding(model, query)
+    
+    # Search more chunks to ensure we have enough candidates for Top-P selection
+    search_candidates = min(len(relevant_docs) * 10, 100)  # Reasonable upper limit
+    
+    # Use ANNOY to search chunks in the relevant documents
+    all_chunks, distances = annoy_manager.search_chunks_in_documents(
+        query_embedding, relevant_docs, 
+        n_neighbors=search_candidates,
+        include_distances=True
+    )
+    
+    # Convert distances to similarities and filter by minimum similarity
+    filtered_chunks = []
+    for chunk, distance in zip(all_chunks, distances):
+        similarity = convert_angular_distance_to_cosine_similarity(distance)
+        
+        # Only include chunks above minimum similarity threshold
+        if similarity >= min_similarity:
+            chunk_result = {
+                'document': chunk['document'],
+                'chunk_id': chunk['chunk_id'],
+                'text': chunk['text'],
+                'start_char': chunk.get('start_char', 0),
+                'end_char': chunk.get('end_char', len(chunk['text'])),
+                'token_count': chunk.get('token_count', len(chunk['text'].split())),
+                'similarity': similarity
+            }
+            filtered_chunks.append(chunk_result)
+    
+    if not filtered_chunks:
+        logger.warning(f"⚠️ No chunks found above similarity threshold {min_similarity}")
+        return []
+    
+    # Sort by similarity
+    filtered_chunks.sort(key=lambda x: x['similarity'], reverse=True)
+    
+    # Apply Top-P selection
+    total_score = sum(chunk['similarity'] for chunk in filtered_chunks)
+    cumulative_prob = 0.0
+    selected_chunks = []
+    
+    for chunk in filtered_chunks:
+        prob = chunk['similarity'] / total_score
+        cumulative_prob += prob
+        selected_chunks.append(chunk)
+        
+        # Stop when we reach the Top-P threshold
+        if cumulative_prob >= top_p:
+            break
+    
+    logger.info(f"🚀 Found {len(selected_chunks)} relevant chunks (ANNOY Top-P={top_p})")
+    logger.info(f"📊 Filtered from {len(filtered_chunks)} chunks above threshold")
+    logger.info(f"📊 Cumulative probability: {cumulative_prob:.3f}")
+    
+    for i, chunk in enumerate(selected_chunks[:5]):  # Show top 5
+        logger.info(f"  {i+1}. {chunk['document']} (chunk {chunk['chunk_id']}, similarity: {chunk['similarity']:.3f})")
+        logger.info(f"     Preview: {chunk['text'][:100]}...")
+    
+    return selected_chunks
+
+
+def find_relevant_chunks_annoy(query: str, model: SentenceTransformer, 
+                              relevant_docs: List[str], annoy_manager: AnnoyIndexManager,
+                              strategy: str = "top_p", **kwargs) -> List[Dict]:
+    """Unified interface for ANNOY-accelerated chunk retrieval with different strategies."""
+    
+    similarity_metric = kwargs.get("similarity_metric", "angular")
+    
+    if strategy == "top_k":
+        top_chunks_per_doc = kwargs.get("top_chunks_per_doc", 3)
+        return find_relevant_chunks_annoy_top_k(query, model, relevant_docs, annoy_manager, 
+                                              top_chunks_per_doc, similarity_metric)
+    
+    elif strategy == "top_p":
+        top_p = kwargs.get("top_p", 0.6)
+        min_similarity = kwargs.get("min_similarity", 0.3)
+        return find_relevant_chunks_annoy_top_p(query, model, relevant_docs, annoy_manager, 
+                                              top_p, min_similarity, similarity_metric)
+    
+    else:
+        raise ValueError(f"Unknown strategy: {strategy}. Use 'top_k' or 'top_p'")
+
+
+def find_relevant_chunks_with_fallback(query: str, model: SentenceTransformer, 
+                                      relevant_docs: List[str], chunk_embeddings: Dict,
+                                      annoy_manager: Optional[AnnoyIndexManager] = None,
+                                      strategy: str = "top_p", **kwargs) -> List[Dict]:
+    """
+    Find relevant chunks with ANNOY acceleration and fallback to original method.
+    
+    This function automatically uses ANNOY if available, otherwise falls back to original search.
+    """
+    if annoy_manager is not None:
+        try:
+            logger.info("🚀 Using ANNOY-accelerated chunk retrieval")
+            return find_relevant_chunks_annoy(query, model, relevant_docs, annoy_manager, strategy, **kwargs)
+        except Exception as e:
+            logger.warning(f"⚠️ ANNOY chunk retrieval failed, falling back to original method: {e}")
+    
+    # Fallback to original method
+    logger.info("🔍 Using original chunk retrieval method")
+    return find_relevant_chunks(query, model, relevant_docs, chunk_embeddings, strategy, **kwargs)

{src/pdf-version → customization/src}/retrieval/document_retriever.py

@@ -1,9 +1,15 @@
 """Document retrieval strategies and functionality."""
 
-from typing import List, Dict
+from typing import List, Dict, Optional
 import numpy as np
+import logging
 from sentence_transformers import SentenceTransformer
-from src.indexing.embedding_creator import create_text_embedding
+from indexing.embedding_creator import create_text_embedding
+from indexing.annoy_manager import AnnoyIndexManager, convert_angular_distance_to_cosine_similarity
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 
 def find_relevant_documents_top_k(query: str, model: SentenceTransformer, 
@@ -189,4 +195,202 @@ def create_document_tag_mapping(document_index: Dict, tag_embeddings: Dict) -> D
             'treatments': doc_info.get('treatments', [])
         }
     
-    return doc_tag_mapping
+    return doc_tag_mapping
+
+
+# ANNOY-accelerated document retrieval functions
+
+def find_relevant_documents_annoy_top_k(query: str, model: SentenceTransformer, 
+                                       annoy_manager: AnnoyIndexManager, doc_tag_mapping: Dict, 
+                                       top_k: int = 3, search_neighbors: int = 20) -> List[str]:
+    """Find top-k most relevant documents using ANNOY index for fast tag search."""
+    query_embedding = create_text_embedding(model, query)
+    
+    # Use ANNOY to find similar tags quickly
+    similar_tags, distances = annoy_manager.search_tags(
+        query_embedding, n_neighbors=search_neighbors, include_distances=True
+    )
+    
+    # Convert angular distances to cosine similarities
+    tag_similarities = {}
+    for tag, distance in zip(similar_tags, distances):
+        similarity = convert_angular_distance_to_cosine_similarity(distance)
+        tag_similarities[tag] = similarity
+    
+    # Find documents that contain the most similar tags
+    doc_scores = {}
+    for pdf_name, doc_info in doc_tag_mapping.items():
+        doc_tags = doc_info['tags']
+        
+        # Calculate document score using max similarity for precise tag matching
+        if doc_tags:
+            similarities = [tag_similarities.get(tag, 0) for tag in doc_tags]
+            # Use max similarity to find documents with best tag matches
+            doc_score = max(similarities)
+            doc_scores[pdf_name] = doc_score
+    
+    # Sort and return top-k documents
+    sorted_docs = sorted(doc_scores.items(), key=lambda x: x[1], reverse=True)
+    relevant_docs = [doc_name for doc_name, score in sorted_docs[:top_k]]
+    
+    logger.info(f"🚀 Found {len(relevant_docs)} relevant documents for query: '{query}' (ANNOY TOP-K)")
+    for i, doc_name in enumerate(relevant_docs):
+        score = doc_scores[doc_name]
+        logger.info(f"  {i+1}. {doc_name} (similarity: {score:.3f})")
+    
+    return relevant_docs
+
+
+def find_relevant_documents_annoy_top_p(query: str, model: SentenceTransformer, 
+                                       annoy_manager: AnnoyIndexManager, doc_tag_mapping: Dict, 
+                                       top_p: float = 0.6, min_similarity: float = 0.5, 
+                                       search_neighbors: int = 30) -> List[str]:
+    """Find documents using TOP-P (nucleus sampling) approach with ANNOY acceleration."""
+    query_embedding = create_text_embedding(model, query)
+    
+    # Use ANNOY to find similar tags quickly
+    similar_tags, distances = annoy_manager.search_tags(
+        query_embedding, n_neighbors=search_neighbors, include_distances=True
+    )
+    
+    # Convert angular distances to cosine similarities
+    tag_similarities = {}
+    for tag, distance in zip(similar_tags, distances):
+        similarity = convert_angular_distance_to_cosine_similarity(distance)
+        tag_similarities[tag] = similarity
+    
+    # Find documents that contain the most similar tags
+    doc_scores = {}
+    for pdf_name, doc_info in doc_tag_mapping.items():
+        doc_tags = doc_info['tags']
+        
+        # Calculate document score using max similarity for precise tag matching
+        if doc_tags:
+            similarities = [tag_similarities.get(tag, 0) for tag in doc_tags]
+            # Use max similarity to find documents with best tag matches
+            doc_score = max(similarities)
+            doc_scores[pdf_name] = doc_score
+    
+    # Filter out documents below minimum similarity threshold
+    filtered_docs = {doc: score for doc, score in doc_scores.items() 
+                    if score >= min_similarity}
+    
+    if not filtered_docs:
+        logger.warning(f"⚠️ No documents found above similarity threshold {min_similarity}")
+        return []
+    
+    # Sort documents by similarity score
+    sorted_docs = sorted(filtered_docs.items(), key=lambda x: x[1], reverse=True)
+    
+    # Apply TOP-P selection
+    total_score = sum(score for _, score in sorted_docs)
+    cumulative_prob = 0.0
+    selected_docs = []
+    
+    for doc_name, score in sorted_docs:
+        prob = score / total_score
+        cumulative_prob += prob
+        selected_docs.append(doc_name)
+        
+        # Stop when we reach the TOP-P threshold
+        if cumulative_prob >= top_p:
+            break
+    
+    logger.info(f"🚀 Found {len(selected_docs)} relevant documents for query: '{query}' (ANNOY TOP-P={top_p})")
+    logger.info(f"📊 Cumulative probability: {cumulative_prob:.3f}")
+    
+    for i, doc_name in enumerate(selected_docs):
+        score = doc_scores[doc_name]
+        prob = score / total_score
+        logger.info(f"  {i+1}. {doc_name} (similarity: {score:.3f}, prob: {prob:.3f})")
+    
+    return selected_docs
+
+
+def find_relevant_documents_annoy_threshold(query: str, model: SentenceTransformer, 
+                                          annoy_manager: AnnoyIndexManager, doc_tag_mapping: Dict, 
+                                          similarity_threshold: float = 0.5, search_neighbors: int = 50) -> List[str]:
+    """Find all documents above a similarity threshold using ANNOY acceleration."""
+    query_embedding = create_text_embedding(model, query)
+    
+    # Use ANNOY to find similar tags quickly
+    similar_tags, distances = annoy_manager.search_tags(
+        query_embedding, n_neighbors=search_neighbors, include_distances=True
+    )
+    
+    # Convert angular distances to cosine similarities
+    tag_similarities = {}
+    for tag, distance in zip(similar_tags, distances):
+        similarity = convert_angular_distance_to_cosine_similarity(distance)
+        tag_similarities[tag] = similarity
+    
+    # Find documents that contain the most similar tags
+    doc_scores = {}
+    for pdf_name, doc_info in doc_tag_mapping.items():
+        doc_tags = doc_info['tags']
+        
+        # Calculate document score using weighted average
+        if doc_tags:
+            similarities = [tag_similarities.get(tag, 0) for tag in doc_tags]
+            avg_similarity = np.mean(similarities)
+            max_similarity = max(similarities)
+            # Weighted combination: 70% average (overall relevance) + 30% max (strongest match)
+            doc_score = avg_similarity * 0.7 + max_similarity * 0.3
+            if doc_score >= similarity_threshold:
+                doc_scores[pdf_name] = doc_score
+    
+    # Sort by similarity score
+    sorted_docs = sorted(doc_scores.items(), key=lambda x: x[1], reverse=True)
+    relevant_docs = [doc_name for doc_name, score in sorted_docs]
+    
+    logger.info(f"🚀 Found {len(relevant_docs)} relevant documents for query: '{query}' (ANNOY threshold={similarity_threshold})")
+    for i, doc_name in enumerate(relevant_docs):
+        score = doc_scores[doc_name]
+        logger.info(f"  {i+1}. {doc_name} (similarity: {score:.3f})")
+    
+    return relevant_docs
+
+
+def find_relevant_documents_annoy(query: str, model: SentenceTransformer, 
+                                 annoy_manager: AnnoyIndexManager, doc_tag_mapping: Dict, 
+                                 strategy: str = "top_k", **kwargs) -> List[str]:
+    """Unified interface for ANNOY-accelerated document retrieval with different strategies."""
+    if strategy == "top_k":
+        top_k = kwargs.get("top_k", 3)
+        search_neighbors = kwargs.get("search_neighbors", 20)
+        return find_relevant_documents_annoy_top_k(query, model, annoy_manager, doc_tag_mapping, top_k, search_neighbors)
+    
+    elif strategy == "top_p":
+        top_p = kwargs.get("top_p", 0.6)
+        min_similarity = kwargs.get("min_similarity", 0.5)
+        search_neighbors = kwargs.get("search_neighbors", 30)
+        return find_relevant_documents_annoy_top_p(query, model, annoy_manager, doc_tag_mapping, top_p, min_similarity, search_neighbors)
+    
+    elif strategy == "threshold":
+        similarity_threshold = kwargs.get("similarity_threshold", 0.5)
+        search_neighbors = kwargs.get("search_neighbors", 50)
+        return find_relevant_documents_annoy_threshold(query, model, annoy_manager, doc_tag_mapping, similarity_threshold, search_neighbors)
+    
+    else:
+        raise ValueError(f"Unknown strategy: {strategy}. Use 'top_k', 'top_p', or 'threshold'")
+
+
+def find_relevant_documents_with_fallback(query: str, model: SentenceTransformer, 
+                                         tag_embeddings: Dict, doc_tag_mapping: Dict,
+                                         annoy_manager: Optional[AnnoyIndexManager] = None,
+                                         strategy: str = "top_k", **kwargs) -> List[str]:
+    """
+    Find relevant documents with ANNOY acceleration and fallback to original method.
+    
+    This function automatically uses ANNOY if available, otherwise falls back to original search.
+    """
+    if annoy_manager is not None:
+        try:
+            logger.info("🚀 Using ANNOY-accelerated document retrieval")
+            return find_relevant_documents_annoy(query, model, annoy_manager, doc_tag_mapping, strategy, **kwargs)
+        except Exception as e:
+            logger.warning(f"⚠️ ANNOY retrieval failed, falling back to original method: {e}")
+    
+    # Fallback to original method
+    logger.info("🔍 Using original document retrieval method")
+    return find_relevant_documents(query, model, tag_embeddings, doc_tag_mapping, strategy, **kwargs)

customization/test/test_pipeline.py

@@ -0,0 +1,117 @@
+#!/usr/bin/env python3
+"""Test script to verify the customization pipeline with ANNOY indices."""
+
+import sys
+from pathlib import Path
+
+# Add parent directory to path
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from customization_pipeline import retrieve_document_chunks
+
+
+def test_pipeline():
+    """Test the complete pipeline with different queries."""
+    print("🧪 Testing Customization Pipeline with ANNOY Indices")
+    print("=" * 60)
+    
+    # Test queries
+    test_queries = [
+        "chest pain and shortness of breath",
+        "pregnancy bleeding emergency",
+        "atrial fibrillation treatment",
+        "fever of unknown origin",
+        "dizziness diagnostic approach"
+    ]
+    
+    for query in test_queries:
+        print(f"\n📋 Query: '{query}'")
+        print("-" * 60)
+        
+        try:
+            # Retrieve chunks
+            results = retrieve_document_chunks(query, top_k=3)
+            
+            if results:
+                print(f"✅ Found {len(results)} relevant chunks:\n")
+                
+                for i, result in enumerate(results, 1):
+                    print(f"Result {i}:")
+                    print(f"  📄 Document: {result['document']}")
+                    print(f"  📊 Score: {result['score']:.4f}")
+                    print(f"  📝 Chunk ID: {result['metadata']['chunk_id']}")
+                    print(f"  📖 Text Preview: {result['chunk_text'][:150]}...")
+                    print()
+            else:
+                print("❌ No results found")
+                
+        except Exception as e:
+            print(f"❌ Error processing query: {e}")
+            import traceback
+            traceback.print_exc()
+    
+    print("\n" + "=" * 60)
+    print("✅ Pipeline test completed!")
+
+
+def test_specific_medical_cases():
+    """Test specific medical scenarios."""
+    print("\n\n🏥 Testing Specific Medical Cases")
+    print("=" * 60)
+    
+    medical_cases = {
+        "Cardiac Emergency": "acute coronary syndrome ST elevation",
+        "Neurological": "stroke symptoms thrombolysis window",
+        "Respiratory": "pulmonary embolism Wells score",
+        "Obstetric Emergency": "eclampsia magnesium sulfate",
+        "Pediatric": "pediatric seizure management"
+    }
+    
+    for case_type, query in medical_cases.items():
+        print(f"\n🔍 {case_type}: '{query}'")
+        print("-" * 60)
+        
+        results = retrieve_document_chunks(query, top_k=2)
+        
+        if results:
+            for result in results:
+                print(f"📄 {result['document']}")
+                print(f"   Score: {result['score']:.4f}")
+                print(f"   Relevant content found in chunk {result['metadata']['chunk_id']}")
+        else:
+            print("   No specific guidance found")
+
+
+def test_performance():
+    """Test retrieval performance."""
+    import time
+    
+    print("\n\n⚡ Testing Retrieval Performance")
+    print("=" * 60)
+    
+    queries = [
+        "chest pain",
+        "headache emergency",
+        "fever neutropenia",
+        "pneumonia antibiotics",
+        "atrial fibrillation"
+    ]
+    
+    total_time = 0
+    for query in queries:
+        start_time = time.time()
+        results = retrieve_document_chunks(query, top_k=5)
+        elapsed = time.time() - start_time
+        total_time += elapsed
+        
+        print(f"Query: '{query}' - Retrieved {len(results)} chunks in {elapsed:.3f}s")
+    
+    avg_time = total_time / len(queries)
+    print(f"\n📊 Average retrieval time: {avg_time:.3f}s per query")
+
+
+if __name__ == "__main__":
+    # Run all tests
+    test_pipeline()
+    test_specific_medical_cases()
+    test_performance()

customization_requirements.txt

@@ -0,0 +1,188 @@
+# Customization Pipeline Requirements
+# Generated from rag_env environment for hospital-specific document processing
+# 
+# Key libraries:
+# - sentence-transformers: Medical domain embeddings (BGE-Large-Medical)
+# - torch: Deep learning framework
+# - annoy: Fast vector similarity search indices
+# - pdfplumber: PDF text and table extraction
+# - llama-index: Document chunking and processing
+# - transformers: Hugging Face model support
+# - openai: LLM integration (optional)
+#
+# Install with: pip install -r customization_requirements.txt
+#
+accelerate==1.9.0
+acres==0.5.0
+aiohappyeyeballs==2.6.1
+aiohttp==3.12.14
+aiosignal==1.4.0
+aiosqlite==0.21.0
+annotated-types==0.7.0
+annoy==1.17.3
+anyio==4.9.0
+appnope==0.1.4
+asttokens==3.0.0
+attrs==25.3.0
+banks==2.2.0
+beautifulsoup4==4.13.4
+bm25s==0.2.13
+certifi==2025.7.14
+cffi==1.17.1
+charset-normalizer==3.4.2
+ci-info==0.3.0
+click==8.2.1
+colorama==0.4.6
+comm==0.2.2
+configobj==5.0.9
+configparser==7.2.0
+cryptography==45.0.5
+dataclasses-json==0.6.7
+debugpy==1.8.15
+decorator==5.2.1
+defusedxml==0.7.1
+Deprecated==1.2.18
+dirtyjson==1.0.8
+distro==1.9.0
+easyocr==1.7.2
+etelemetry==0.3.1
+executing==2.2.0
+filelock==3.18.0
+filetype==1.2.0
+fitz==0.0.1.dev2
+frozenlist==1.7.0
+fsspec==2025.7.0
+greenlet==3.2.3
+griffe==1.7.3
+h11==0.16.0
+hf-xet==1.1.5
+httpcore==1.0.9
+httplib2==0.22.0
+httpx==0.28.1
+huggingface-hub==0.33.4
+idna==3.10
+imageio==2.37.0
+ipykernel==6.30.0
+ipython==9.4.0
+ipython_pygments_lexers==1.1.1
+jedi==0.19.2
+Jinja2==3.1.6
+jiter==0.10.0
+joblib==1.5.1
+jpype1==1.6.0
+jupyter_client==8.6.3
+jupyter_core==5.8.1
+lazy_loader==0.4
+llama-cloud==0.1.32
+llama-cloud-services==0.6.43
+llama-index==0.12.50
+llama-index-agent-openai==0.4.12
+llama-index-cli==0.4.4
+llama-index-core==0.12.50
+llama-index-embeddings-huggingface==0.5.5
+llama-index-embeddings-openai==0.3.1
+llama-index-indices-managed-llama-cloud==0.7.10
+llama-index-instrumentation==0.3.0
+llama-index-llms-huggingface==0.5.0
+llama-index-llms-openai==0.4.7
+llama-index-llms-openai-like==0.4.0
+llama-index-llms-openrouter==0.3.2
+llama-index-multi-modal-llms-openai==0.5.3
+llama-index-program-openai==0.3.2
+llama-index-question-gen-openai==0.3.1
+llama-index-readers-file==0.4.11
+llama-index-readers-llama-parse==0.4.0
+llama-index-retrievers-bm25==0.5.2
+llama-index-workflows==1.1.0
+llama-parse==0.6.43
+looseversion==1.3.0
+lxml==6.0.0
+MarkupSafe==3.0.2
+marshmallow==3.26.1
+matplotlib-inline==0.1.7
+mpmath==1.3.0
+multidict==6.6.3
+mypy_extensions==1.1.0
+nest-asyncio==1.6.0
+networkx==3.5
+nibabel==5.3.2
+ninja==1.11.1.4
+nipype==1.10.0
+nltk==3.9.1
+numpy==2.2.6
+openai==1.97.0
+opencv-python-headless==4.12.0.88
+packaging==25.0
+pandas==2.2.3
+parso==0.8.4
+pathlib==1.0.1
+pdfminer.six==20250506
+pdfplumber==0.11.7
+pexpect==4.9.0
+pillow==11.3.0
+platformdirs==4.3.8
+prompt_toolkit==3.0.51
+propcache==0.3.2
+prov==2.1.1
+psutil==7.0.0
+ptyprocess==0.7.0
+pure_eval==0.2.3
+puremagic==1.30
+pyclipper==1.3.0.post6
+pycparser==2.22
+pydantic==2.11.7
+pydantic_core==2.33.2
+pydot==4.0.1
+Pygments==2.19.2
+PyMuPDF==1.26.3
+pyparsing==3.2.3
+pypdf==5.8.0
+pypdfium2==4.30.0
+PyStemmer==2.2.0.3
+python-bidi==0.6.6
+python-dateutil==2.9.0.post0
+python-dotenv==1.1.1
+pytz==2025.2
+pyxnat==1.6.3
+PyYAML==6.0.2
+pyzmq==27.0.0
+rdflib==7.1.4
+regex==2024.11.6
+requests==2.32.4
+safetensors==0.5.3
+scikit-image==0.25.2
+scikit-learn==1.7.1
+scipy==1.16.0
+sentence-transformers==5.0.0
+setuptools==80.9.0
+shapely==2.1.1
+simplejson==3.20.1
+six==1.17.0
+sniffio==1.3.1
+soupsieve==2.7
+SQLAlchemy==2.0.41
+stack-data==0.6.3
+striprtf==0.0.26
+sympy==1.14.0
+tabula-py==2.10.0
+tabulate==0.9.0
+tenacity==9.1.2
+threadpoolctl==3.6.0
+tifffile==2025.6.11
+tiktoken==0.9.0
+tokenizers==0.21.2
+torch==2.7.1
+torchvision==0.22.1
+tornado==6.5.1
+tqdm==4.67.1
+traitlets==5.14.3
+traits==7.0.2
+transformers==4.53.2
+typing-inspect==0.9.0
+typing-inspection==0.4.1
+typing_extensions==4.14.1
+tzdata==2025.2
+urllib3==2.5.0
+wcwidth==0.2.13
+wrapt==1.17.2
+yarl==1.20.1

src/pdf-version/data/__init__.py

@@ -1,15 +0,0 @@
-"""Data loading and PDF processing."""
-
-from .loaders import load_annotations, filter_pdf_files
-from .pdf_processing import (
-    extract_pdf_text, 
-    extract_tables_from_pdf,
-    extract_images_ocr_from_pdf,
-    extract_pdf_content_enhanced
-)
-
-__all__ = [
-    'load_annotations', 'filter_pdf_files',
-    'extract_pdf_text', 'extract_tables_from_pdf', 
-    'extract_images_ocr_from_pdf', 'extract_pdf_content_enhanced'
-]

src/pdf-version/main.py

@@ -1,83 +0,0 @@
-#!/usr/bin/env python3
-"""OnCall AI - Medical RAG System
-
-Main entry point for the medical RAG system.
-"""
-
-import sys
-from pathlib import Path
-
-# Add pdf-version directory to Python path
-sys.path.insert(0, str(Path(__file__).parent))
-
-from demos.demo_runner import build_medical_rag_system, demo_rag_query, demo_all_strategies
-
-
-def main():
-    """Main program entry point."""
-    try:
-        # Build the system with chunk embeddings
-        build_medical_rag_system(enable_chunk_embeddings=True)
-        
-        # Demo chunk-based retrieval
-        print("\n" + "="*80)
-        print("🧩 CHUNK-BASED RETRIEVAL DEMO")
-        print("="*80)
-        demo_rag_query("chest pain and shortness of breath", 
-                      strategy="top_p", use_chunks=True, top_p=0.8)
-        
-    except KeyboardInterrupt:
-        print("\n\n👋 User interrupted, program exiting")
-    except Exception as e:
-        print(f"\n❌ Program execution error: {e}")
-        import traceback
-        traceback.print_exc()
-
-
-def interactive_demo():
-    """Interactive demo mode."""
-    print("🏥 OnCall AI - Interactive Demo Mode")
-    print("=" * 50)
-    
-    while True:
-        print("\nOptions:")
-        print("1. Build/rebuild system")
-        print("2. Query with TOP-P strategy")
-        print("3. Query with TOP-K strategy") 
-        print("4. Compare all strategies")
-        print("5. Custom query")
-        print("6. Exit")
-        
-        choice = input("\nSelect option (1-6): ").strip()
-        
-        if choice == "1":
-            build_medical_rag_system(enable_chunk_embeddings=True)
-        elif choice == "2":
-            query = input("Enter your query: ").strip()
-            if query:
-                demo_rag_query(query, strategy="top_p", use_chunks=True)
-        elif choice == "3":
-            query = input("Enter your query: ").strip()
-            if query:
-                demo_rag_query(query, strategy="top_k", use_chunks=True, top_k=3)
-        elif choice == "4":
-            query = input("Enter your query: ").strip()
-            if query:
-                demo_all_strategies(query)
-        elif choice == "5":
-            query = input("Enter your query: ").strip()
-            strategy = input("Enter strategy (top_k/top_p/threshold): ").strip()
-            if query and strategy:
-                demo_rag_query(query, strategy=strategy, use_chunks=True)
-        elif choice == "6":
-            print("👋 Goodbye!")
-            break
-        else:
-            print("❌ Invalid option. Please select 1-6.")
-
-
-if __name__ == "__main__":
-    if len(sys.argv) > 1 and sys.argv[1] == "--interactive":
-        interactive_demo()
-    else:
-        main()

src/pdf-version/oncall_ai.py

@@ -1,55 +0,0 @@
-#!/usr/bin/env python3
-"""OnCall AI - Medical RAG System (Backward Compatibility)
-
-This file provides backward compatibility with the original rag.py interface.
-Import everything from the new modular structure.
-"""
-
-import sys
-from pathlib import Path
-
-# Add pdf-version directory to Python path
-sys.path.insert(0, str(Path(__file__).parent))
-
-# Import all functions for backward compatibility
-from models.embedding_models import load_biomedbert_model, load_meditron_model
-from data.loaders import load_annotations, filter_pdf_files
-from data.pdf_processing import (
-    extract_pdf_text, extract_tables_from_pdf,
-    extract_images_ocr_from_pdf, extract_pdf_content_enhanced
-)
-from indexing.document_indexer import build_document_index, split_text_into_chunks
-from indexing.embedding_creator import create_text_embedding, create_tag_embeddings, create_chunk_embeddings
-from indexing.storage import save_document_system, load_document_system
-from retrieval.document_retriever import (
-    find_relevant_documents_top_k, find_relevant_documents_top_p,
-    find_relevant_documents_threshold, find_relevant_documents,
-    create_document_tag_mapping
-)
-from retrieval.chunk_retriever import find_relevant_chunks, get_documents_for_rag, get_chunks_for_rag
-from demos.demo_runner import build_medical_rag_system, demo_rag_query, demo_all_strategies
-
-# Main function for backward compatibility
-def main():
-    """Main program entry compatible with original rag.py."""
-    try:
-        # Build the system with chunk embeddings
-        build_medical_rag_system(enable_chunk_embeddings=True)
-        
-        # Demo chunk-based retrieval
-        print("\n" + "="*80)
-        print("🧩 CHUNK-BASED RETRIEVAL DEMO")
-        print("="*80)
-        demo_rag_query("chest pain and shortness of breath", 
-                      strategy="top_p", use_chunks=True, top_p=0.8)
-        
-    except KeyboardInterrupt:
-        print("\n\n👋 User interrupted, program exiting")
-    except Exception as e:
-        print(f"\n❌ Program execution error: {e}")
-        import traceback
-        traceback.print_exc()
-
-
-if __name__ == "__main__":
-    main()