aiagents4pharma 1.40.1__py3-none-any.whl → 1.41.0__py3-none-any.whl

aiagents4pharma/talk2scholars/tools/pdf/utils/gpu_detection.py

@@ -0,0 +1,154 @@
+"""
+GPU Detection Utility for Milvus Index Selection
+Handle COSINE -> IP conversion for GPU indexes
+"""
+
+import logging
+import subprocess
+from typing import Dict, Any, Tuple
+
+logger = logging.getLogger(__name__)
+
+
+def detect_nvidia_gpu(config=None) -> bool:
+    """
+    Detect if NVIDIA GPU is available and should be used.
+
+    Args:
+        config: Hydra config object that may contain force_cpu_mode flag
+
+    Returns:
+        bool: True if GPU should be used, False if CPU should be used
+    """
+
+    # Check for force CPU mode in config
+    if config and hasattr(config, "gpu_detection"):
+        force_cpu = getattr(config.gpu_detection, "force_cpu_mode", False)
+        if force_cpu:
+            logger.info(
+                "Force CPU mode enabled in config - using CPU even though GPU may be available"
+            )
+            return False
+
+    # Normal GPU detection logic
+    try:
+        result = subprocess.run(
+            ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
+            capture_output=True,
+            text=True,
+            timeout=10,
+            check=False,
+        )
+
+        if result.returncode == 0 and result.stdout.strip():
+            gpu_names = result.stdout.strip().split("\n")
+            logger.info("Detected NVIDIA GPU(s): %s", gpu_names)
+            logger.info("To force CPU mode, set 'force_cpu_mode: true' in config")
+            return True
+
+        logger.info("nvidia-smi command failed or no GPUs detected")
+        return False
+
+    except (subprocess.TimeoutExpired, FileNotFoundError) as e:
+        logger.info("NVIDIA GPU detection failed: %s", e)
+        return False
+
+
+def get_optimal_index_config(
+    has_gpu: bool, embedding_dim: int = 768, use_cosine: bool = True
+) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+    """
+    Get optimal index and search parameters based on GPU availability.
+
+    IMPORTANT: GPU indexes don't support COSINE distance. When using GPU with COSINE,
+    vectors must be normalized and IP distance used instead.
+
+    Args:
+        has_gpu (bool): Whether NVIDIA GPU is available
+        embedding_dim (int): Dimension of embeddings
+        use_cosine (bool): Whether to use cosine similarity (will be converted to IP for GPU)
+
+    Returns:
+        Tuple[Dict[str, Any], Dict[str, Any]]: (index_params, search_params)
+    """
+    if has_gpu:
+        logger.info("Configuring GPU_CAGRA index for NVIDIA GPU")
+
+        # For GPU: COSINE is not supported, must use IP with normalized vectors
+        if use_cosine:
+            logger.warning(
+                "GPU indexes don't support COSINE distance. "
+                "Vectors will be normalized and IP distance will be used instead."
+            )
+            metric_type = (
+                "IP"  # Inner Product for normalized vectors = cosine similarity
+            )
+        else:
+            metric_type = "IP"  # Default to IP for GPU
+
+        # GPU_CAGRA index parameters - optimized for performance
+        index_params = {
+            "index_type": "GPU_CAGRA",
+            "metric_type": metric_type,
+            "params": {
+                "intermediate_graph_degree": 64,  # Higher for better recall
+                "graph_degree": 32,  # Balanced performance/recall
+                "build_algo": "IVF_PQ",  # Higher quality build
+                "cache_dataset_on_device": "true",  # Cache for better recall
+                "adapt_for_cpu": "false",  # Pure GPU mode
+            },
+        }
+
+        # GPU_CAGRA search parameters
+        search_params = {
+            "metric_type": metric_type,
+            "params": {
+                "itopk_size": 128,  # Power of 2, good for intermediate results
+                "search_width": 16,  # Balanced entry points
+                "team_size": 16,  # Optimize for typical vector dimensions
+            },
+        }
+
+    else:
+        logger.info("Configuring CPU index (IVF_FLAT) - no NVIDIA GPU detected")
+
+        # CPU supports COSINE directly
+        metric_type = "COSINE" if use_cosine else "IP"
+
+        # CPU IVF_FLAT index parameters
+        index_params = {
+            "index_type": "IVF_FLAT",
+            "metric_type": metric_type,
+            "params": {
+                "nlist": min(
+                    1024, max(64, embedding_dim // 8)
+                )  # Dynamic nlist based on dimension
+            },
+        }
+
+        # CPU search parameters
+        search_params = {
+            "metric_type": metric_type,
+            "params": {"nprobe": 16},  # Slightly higher than original for better recall
+        }
+
+    return index_params, search_params
+
+
+def log_index_configuration(
+    index_params: Dict[str, Any], search_params: Dict[str, Any], use_cosine: bool = True
+) -> None:
+    """Log the selected index configuration for debugging."""
+    index_type = index_params.get("index_type", "Unknown")
+    metric_type = index_params.get("metric_type", "Unknown")
+
+    logger.info("=== Milvus Index Configuration ===")
+    logger.info("Index Type: %s", index_type)
+    logger.info("Metric Type: %s", metric_type)
+
+    if index_type == "GPU_CAGRA" and use_cosine and metric_type == "IP":
+        logger.info("NOTE: Using IP with normalized vectors to simulate COSINE for GPU")
+
+    logger.info("Index Params: %s", index_params.get("params", {}))
+    logger.info("Search Params: %s", search_params.get("params", {}))
+    logger.info("===================================")

aiagents4pharma/talk2scholars/tools/pdf/utils/nvidia_nim_reranker.py

@@ -1,11 +1,10 @@
 """
-NVIDIA NIM Reranker Utility
+NVIDIA NIM Reranker Utility for Milvus Integration
+Rerank chunks instead of papers following traditional RAG pipeline
 """
 
 import logging
 import os
-
-
 from typing import Any, List
 
 from langchain_core.documents import Document
@@ -18,60 +17,81 @@ logger = logging.getLogger(__name__)
 logger.setLevel(getattr(logging, log_level))
 
 
-def rank_papers_by_query(self, query: str, config: Any, top_k: int = 5) -> List[str]:
+def rerank_chunks(
+    chunks: List[Document], query: str, config: Any, top_k: int = 25
+) -> List[Document]:
     """
-    Rank papers by relevance to the query using NVIDIA's off-the-shelf re-ranker.
+    Rerank chunks by relevance to the query using NVIDIA's reranker.
 
-    This function aggregates all chunks per paper, ranks them using the NVIDIA model,
-    and returns the top-k papers.
+    This follows the traditional RAG pipeline: first retrieve chunks, then rerank them.
 
     Args:
-        query (str): The query string.
-        config (Any): Configuration containing reranker settings (model, api_key).
-        top_k (int): Number of top papers to return.
+        chunks (List[Document]): List of chunks to rerank
+        query (str): The query string
+        config (Any): Configuration containing reranker settings
+        top_k (int): Number of top chunks to return after reranking
 
     Returns:
-        List of tuples (paper_id, dummy_score) sorted by relevance.
+        List[Document]: Reranked chunks (top_k most relevant)
     """
+    logger.info(
+        "Starting NVIDIA chunk reranker for query: '%s' with %d chunks, top_k=%d",
+        query[:50] + "..." if len(query) > 50 else query,
+        len(chunks),
+        top_k,
+    )
 
-    logger.info("Starting NVIDIA re-ranker for query: '%s' with top_k=%d", query, top_k)
-    # Aggregate all document chunks for each paper
-    paper_texts = {}
-    for doc in self.documents.values():
-        paper_id = doc.metadata["paper_id"]
-        paper_texts.setdefault(paper_id, []).append(doc.page_content)
-
-    aggregated_documents = []
-    for paper_id, texts in paper_texts.items():
-        aggregated_text = " ".join(texts)
-        aggregated_documents.append(
-            Document(page_content=aggregated_text, metadata={"paper_id": paper_id})
+    # If we have fewer chunks than top_k, just return all
+    if len(chunks) <= top_k:
+        logger.info(
+            "Number of chunks (%d) <= top_k (%d), returning all chunks without reranking",
+            len(chunks),
+            top_k,
         )
+        return chunks
 
-    logger.info(
-        "Aggregated %d papers into %d documents for reranking",
-        len(paper_texts),
-        len(aggregated_documents),
-    )
-    # Instantiate the NVIDIA re-ranker client using provided config
-    # Use NVIDIA API key from Hydra configuration (expected to be resolved via oc.env)
+    # Get API key from config
     api_key = config.reranker.api_key
     if not api_key:
         logger.error("No NVIDIA API key found in configuration for reranking")
         raise ValueError("Configuration 'reranker.api_key' must be set for reranking")
-    logger.info("Using NVIDIA API key from configuration for reranking")
-    # Truncate long inputs at the model-end to avoid exceeding max token size
-    logger.info("Setting NVIDIA reranker truncate mode to END to limit input length")
+
+    logger.info("Using NVIDIA reranker model: %s", config.reranker.model)
+
+    # Initialize reranker with truncation to handle long chunks
     reranker = NVIDIARerank(
         model=config.reranker.model,
         api_key=api_key,
-        truncate="END",
+        truncate="END",  # Truncate at the end if too long
+    )
+
+    # Log chunk metadata for debugging
+    logger.debug(
+        "Reranking chunks from papers: %s",
+        list(set(chunk.metadata.get("paper_id", "unknown") for chunk in chunks))[:5],
+    )
+
+    # Rerank the chunks
+    logger.info("Calling NVIDIA reranker API with %d chunks...", len(chunks))
+    reranked_chunks = reranker.compress_documents(query=query, documents=chunks)
+
+    for i, doc in enumerate(reranked_chunks[:top_k]):
+        score = doc.metadata.get("relevance_score", "N/A")
+        source = doc.metadata.get("paper_id", "unknown")
+        logger.info("Rank %d | Score: %.4f | Source: %s", i + 1, score, source)
+
+    logger.info(
+        "Successfully reranked chunks. Returning top %d chunks",
+        min(top_k, len(reranked_chunks)),
     )
 
-    # Get the ranked list of documents based on the query
-    response = reranker.compress_documents(query=query, documents=aggregated_documents)
-    logger.info("Received %d documents from NVIDIA reranker", len(response))
+    # Log which papers the top chunks come from
+    if reranked_chunks and logger.isEnabledFor(logging.DEBUG):
+        top_papers = {}
+        for chunk in reranked_chunks[:top_k]:
+            paper_id = chunk.metadata.get("paper_id", "unknown")
+            top_papers[paper_id] = top_papers.get(paper_id, 0) + 1
+        logger.debug("Top %d chunks distribution by paper: %s", top_k, top_papers)
 
-    ranked_papers = [doc.metadata["paper_id"] for doc in response[:top_k]]
-    logger.info("Top %d papers after reranking: %s", top_k, ranked_papers)
-    return ranked_papers
+    # Return only top_k chunks (convert to list to match return type)
+    return list(reranked_chunks[:top_k])

aiagents4pharma/talk2scholars/tools/pdf/utils/paper_loader.py

@@ -0,0 +1,123 @@
+"""
+Paper loading utilities for managing PDF documents in vector store.
+"""
+
+import logging
+from typing import Any, Dict
+
+from .batch_processor import add_papers_batch
+
+logger = logging.getLogger(__name__)
+
+
+def load_all_papers(
+    vector_store: Any,  # The Vectorstore instance
+    articles: Dict[str, Any],
+    call_id: str,
+    config: Any,
+    has_gpu: bool,
+) -> None:
+    """
+    Ensure all papers from article_data are loaded into the Milvus vector store.
+    Optimized for GPU/CPU processing.
+
+    Args:
+        vector_store: The Vectorstore instance
+        articles: Dictionary of article data
+        call_id: Call identifier for logging
+        config: Configuration object
+        has_gpu: Whether GPU is available
+    """
+    papers_to_load = []
+    skipped_papers = []
+    already_loaded = []
+
+    # Check which papers need to be loaded
+    for pid, article_info in articles.items():
+        if pid not in vector_store.loaded_papers:
+            pdf_url = article_info.get("pdf_url")
+            if pdf_url:
+                # Prepare tuple for batch loading
+                papers_to_load.append((pid, pdf_url, article_info))
+            else:
+                skipped_papers.append(pid)
+        else:
+            already_loaded.append(pid)
+
+    # Log summary of papers status with hardware info
+    hardware_info = f" (GPU acceleration: {'enabled' if has_gpu else 'disabled'})"
+    logger.info(
+        "%s: Paper loading summary%s - Total: %d, Already loaded: %d, To load: %d, No PDF: %d",
+        call_id,
+        hardware_info,
+        len(articles),
+        len(already_loaded),
+        len(papers_to_load),
+        len(skipped_papers),
+    )
+
+    if skipped_papers:
+        logger.warning(
+            "%s: Skipping %d papers without PDF URLs: %s%s",
+            call_id,
+            len(skipped_papers),
+            skipped_papers[:5],  # Show first 5
+            "..." if len(skipped_papers) > 5 else "",
+        )
+
+    if not papers_to_load:
+        logger.info("%s: All papers with PDFs are already loaded in Milvus", call_id)
+        return
+
+    # Use batch loading with parallel processing for ALL papers at once
+    # Adjust parameters based on hardware capabilities
+    if has_gpu:
+        # GPU can handle more parallel processing
+        max_workers = min(12, max(4, len(papers_to_load)))  # More workers for GPU
+        batch_size = config.get("embedding_batch_size", 2000)  # Larger batches for GPU
+        logger.info(
+            "%s: Using GPU-optimized loading parameters: %d workers, batch size %d",
+            call_id,
+            max_workers,
+            batch_size,
+        )
+    else:
+        # CPU - more conservative parameters
+        max_workers = min(8, max(3, len(papers_to_load)))  # Conservative for CPU
+        batch_size = config.get("embedding_batch_size", 1000)  # Smaller batches for CPU
+        logger.info(
+            "%s: Using CPU-optimized loading parameters: %d workers, batch size %d",
+            call_id,
+            max_workers,
+            batch_size,
+        )
+
+    logger.info(
+        "%s: Loading %d papers in ONE BATCH using %d parallel workers (batch size: %d, %s)",
+        call_id,
+        len(papers_to_load),
+        max_workers,
+        batch_size,
+        "GPU accelerated" if has_gpu else "CPU processing",
+    )
+
+    # This should process ALL papers at once with hardware optimization
+    add_papers_batch(
+        papers_to_add=papers_to_load,
+        vector_store=vector_store.vector_store,  # Pass the LangChain vector store
+        loaded_papers=vector_store.loaded_papers,
+        paper_metadata=vector_store.paper_metadata,
+        documents=vector_store.documents,
+        config=vector_store.config,
+        metadata_fields=vector_store.metadata_fields,
+        has_gpu=vector_store.has_gpu,
+        max_workers=max_workers,
+        batch_size=batch_size,
+    )
+
+    logger.info(
+        "%s: Successfully completed batch loading of all %d papers with %s",
+        call_id,
+        len(papers_to_load),
+        "GPU acceleration" if has_gpu else "CPU processing",
+    )

aiagents4pharma/talk2scholars/tools/pdf/utils/rag_pipeline.py

@@ -0,0 +1,122 @@
+"""
+RAG pipeline for retrieving and reranking chunks from a vector store.
+"""
+
+import logging
+from typing import Any, List
+
+
+# Import our GPU detection utility
+from .nvidia_nim_reranker import rerank_chunks
+from .retrieve_chunks import retrieve_relevant_chunks
+
+logger = logging.getLogger(__name__)
+
+
+def retrieve_and_rerank_chunks(
+    vector_store: Any, query: str, config: Any, call_id: str, has_gpu: bool
+) -> List[Any]:
+    """
+    Traditional RAG pipeline: retrieve chunks from all papers, then rerank.
+    Optimized for GPU/CPU hardware.
+
+    Args:
+        vs: Vector store instance
+        query: User query
+
+    Returns:
+        List of reranked chunks
+    """
+    hardware_mode = "GPU-accelerated" if has_gpu else "CPU-optimized"
+    logger.info(
+        "%s: Starting traditional RAG pipeline - retrieve then rerank (%s)",
+        call_id,
+        hardware_mode,
+    )
+
+    # Step 1: Retrieve chunks from ALL papers (cast wide net)
+    # Adjust initial retrieval count based on hardware
+    if has_gpu:
+        # GPU can handle larger initial retrieval efficiently
+        initial_chunks_count = config.get(
+            "initial_retrieval_k", 150
+        )  # Increased for GPU
+        mmr_diversity = config.get(
+            "mmr_diversity", 0.75
+        )  # Slightly more diverse for larger sets
+    else:
+        # CPU - use conservative settings
+        initial_chunks_count = config.get("initial_retrieval_k", 100)  # Original
+        mmr_diversity = config.get("mmr_diversity", 0.8)  # Original
+
+    logger.info(
+        "%s: Step 1 - Retrieving top %d chunks from ALL papers (%s mode)",
+        call_id,
+        initial_chunks_count,
+        hardware_mode,
+    )
+
+    retrieved_chunks = retrieve_relevant_chunks(
+        vector_store,
+        query=query,
+        paper_ids=None,  # No filter - retrieve from all papers
+        top_k=initial_chunks_count,
+        mmr_diversity=mmr_diversity,
+    )
+
+    if not retrieved_chunks:
+        logger.warning("%s: No chunks retrieved from vector store", call_id)
+        return []
+
+    logger.info(
+        "%s: Retrieved %d chunks from %d unique papers using %s",
+        call_id,
+        len(retrieved_chunks),
+        len(
+            set(chunk.metadata.get("paper_id", "unknown") for chunk in retrieved_chunks)
+        ),
+        hardware_mode,
+    )
+
+    # Step 2: Rerank the retrieved chunks
+    final_chunk_count = config.top_k_chunks
+    logger.info(
+        "%s: Step 2 - Reranking %d chunks to get top %d",
+        call_id,
+        len(retrieved_chunks),
+        final_chunk_count,
+    )
+
+    reranked_chunks = rerank_chunks(
+        chunks=retrieved_chunks,
+        query=query,
+        config=config,
+        top_k=final_chunk_count,
+    )
+
+    # Log final results with hardware info
+    final_papers = len(
+        set(chunk.metadata.get("paper_id", "unknown") for chunk in reranked_chunks)
+    )
+
+    logger.info(
+        "%s: Reranking complete using %s. Final %d chunks from %d unique papers",
+        call_id,
+        hardware_mode,
+        len(reranked_chunks),
+        final_papers,
+    )
+
+    # Log performance insights
+    if len(retrieved_chunks) > 0:
+        efficiency = len(reranked_chunks) / len(retrieved_chunks) * 100
+        logger.debug(
+            "%s: Pipeline efficiency: %.1f%% (%d final / %d initial chunks) - %s",
+            call_id,
+            efficiency,
+            len(reranked_chunks),
+            len(retrieved_chunks),
+            hardware_mode,
+        )
+
+    return reranked_chunks