visual-rag-toolkit 0.1.1__py3-none-any.whl

visual_rag/config.py

@@ -0,0 +1,230 @@
+"""
+Configuration utilities for Visual RAG Toolkit.
+
+Provides:
+- YAML configuration loading with caching
+- Environment variable overrides
+- Convenience getters for common settings
+"""
+
+import copy
+import logging
+import os
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+logger = logging.getLogger(__name__)
+
+# Global config cache (raw YAML only; env overrides applied on demand)
+_raw_config_cache: Optional[Dict[str, Any]] = None
+_raw_config_cache_path: Optional[str] = None
+
+
+def _env_qdrant_url() -> Optional[str]:
+    return os.getenv("SIGIR_QDRANT_URL") or os.getenv("DEST_QDRANT_URL") or os.getenv("QDRANT_URL")
+
+
+def _env_qdrant_api_key() -> Optional[str]:
+    return (
+        os.getenv("SIGIR_QDRANT_KEY")
+        or os.getenv("SIGIR_QDRANT_API_KEY")
+        or os.getenv("DEST_QDRANT_API_KEY")
+        or os.getenv("QDRANT_API_KEY")
+    )
+
+
+def load_config(
+    config_path: Optional[str] = None,
+    force_reload: bool = False,
+    apply_env_overrides: bool = True,
+) -> Dict[str, Any]:
+    """
+    Load configuration from YAML file.
+
+    Uses caching to avoid repeated file I/O.
+    Environment variables can override config values.
+
+    Args:
+        config_path: Path to config file (auto-detected if None)
+        force_reload: Bypass cache and reload from file
+
+    Returns:
+        Configuration dictionary
+    """
+    global _raw_config_cache, _raw_config_cache_path
+
+    # Determine the effective config path (used for caching)
+    effective_path: Optional[str] = None
+
+    # Find config file
+    if config_path is None:
+        config_path = os.getenv("VISUALRAG_CONFIG")
+
+        if config_path is None:
+            # Check common locations
+            search_paths = [
+                Path.cwd() / "config.yaml",
+                Path.cwd() / "visual_rag.yaml",
+                Path.home() / ".visual_rag" / "config.yaml",
+            ]
+
+            for path in search_paths:
+                if path.exists():
+                    config_path = str(path)
+                    break
+    effective_path = str(config_path) if config_path else None
+
+    # Return cached raw config if available.
+    # - If caller doesn't specify a path (effective_path is None), use whatever was
+    #   loaded most recently (common pattern in apps).
+    # - If a path is specified, only reuse cache when it matches.
+    if (
+        _raw_config_cache is not None
+        and not force_reload
+        and (effective_path is None or _raw_config_cache_path == effective_path)
+    ):
+        cfg = copy.deepcopy(_raw_config_cache)
+        return _apply_env_overrides(cfg) if apply_env_overrides else cfg
+
+    # Load YAML if file exists
+    config = {}
+    if config_path and Path(config_path).exists():
+        try:
+            import yaml
+
+            with open(config_path, "r") as f:
+                config = yaml.safe_load(f) or {}
+
+            logger.info(f"Loaded config from: {config_path}")
+        except ImportError:
+            logger.warning("PyYAML not installed, using environment variables only")
+        except Exception as e:
+            logger.warning(f"Could not load config file: {e}")
+
+    # Cache RAW config (no env overrides)
+    _raw_config_cache = copy.deepcopy(config)
+    _raw_config_cache_path = effective_path
+
+    # Return resolved or raw depending on caller preference
+    cfg = copy.deepcopy(config)
+    return _apply_env_overrides(cfg) if apply_env_overrides else cfg
+
+
+def _apply_env_overrides(config: Dict[str, Any]) -> Dict[str, Any]:
+    """Apply environment variable overrides."""
+
+    env_mappings = {
+        # Qdrant
+        "QDRANT_URL": ["qdrant", "url"],
+        "QDRANT_API_KEY": ["qdrant", "api_key"],
+        "QDRANT_COLLECTION": ["qdrant", "collection"],
+        # Model
+        "VISUALRAG_MODEL": ["model", "name"],
+        "COLPALI_MODEL_NAME": ["model", "name"],  # Alias
+        "EMBEDDING_BATCH_SIZE": ["model", "batch_size"],
+        # Cloudinary
+        "CLOUDINARY_CLOUD_NAME": ["cloudinary", "cloud_name"],
+        "CLOUDINARY_API_KEY": ["cloudinary", "api_key"],
+        "CLOUDINARY_API_SECRET": ["cloudinary", "api_secret"],
+        # Processing
+        "PDF_DPI": ["processing", "dpi"],
+        "JPEG_QUALITY": ["processing", "jpeg_quality"],
+        # Search
+        "SEARCH_STRATEGY": ["search", "strategy"],
+        "PREFETCH_K": ["search", "prefetch_k"],
+        # Special token handling
+        "VISUALRAG_INCLUDE_SPECIAL_TOKENS": ["embedding", "include_special_tokens"],
+    }
+
+    for env_var, path in env_mappings.items():
+        value = os.getenv(env_var)
+        if value is not None:
+            # Navigate to the right place in config
+            current = config
+            for key in path[:-1]:
+                if key not in current:
+                    current[key] = {}
+                current = current[key]
+
+            # Convert value to appropriate type
+            final_key = path[-1]
+            if final_key in current:
+                existing_type = type(current[final_key])
+                # Use `is` for type comparisons (Ruff E721).
+                if existing_type is bool:
+                    value = value.lower() in ("true", "1", "yes", "on")
+                elif existing_type is int:
+                    value = int(value)
+                elif existing_type is float:
+                    value = float(value)
+
+            current[final_key] = value
+            logger.debug(f"Config override: {'.'.join(path)} = {value}")
+
+    return config
+
+
+def get(key: str, default: Any = None) -> Any:
+    """
+    Get a configuration value by dot-notation path.
+
+    Examples:
+        >>> get("qdrant.url")
+        >>> get("model.name", "vidore/colSmol-500M")
+        >>> get("search.strategy", "multi_vector")
+    """
+    config = load_config(apply_env_overrides=True)
+
+    keys = key.split(".")
+    current = config
+
+    for k in keys:
+        if isinstance(current, dict) and k in current:
+            current = current[k]
+        else:
+            return default
+
+    return current
+
+
+def get_section(section: str, *, apply_env_overrides: bool = True) -> Dict[str, Any]:
+    """Get an entire configuration section."""
+    config = load_config(apply_env_overrides=apply_env_overrides)
+    return config.get(section, {})
+
+
+# Convenience getters
+def get_qdrant_config() -> Dict[str, Any]:
+    """Get Qdrant configuration with defaults."""
+    return {
+        "url": get("qdrant.url", _env_qdrant_url()),
+        "api_key": get("qdrant.api_key", _env_qdrant_api_key()),
+        "collection": get("qdrant.collection", "visual_documents"),
+    }
+
+
+def get_model_config() -> Dict[str, Any]:
+    """Get model configuration with defaults."""
+    return {
+        "name": get("model.name", "vidore/colSmol-500M"),
+        "batch_size": get("model.batch_size", 4),
+        "device": get("model.device", "auto"),
+    }
+
+
+def get_processing_config() -> Dict[str, Any]:
+    """Get processing configuration with defaults."""
+    return {
+        "dpi": get("processing.dpi", 140),
+        "jpeg_quality": get("processing.jpeg_quality", 95),
+        "page_batch_size": get("processing.page_batch_size", 50),
+    }
+
+
+def get_search_config() -> Dict[str, Any]:
+    """Get search configuration with defaults."""
+    return {
+        "strategy": get("search.strategy", "multi_vector"),
+        "prefetch_k": get("search.prefetch_k", 200),
+        "top_k": get("search.top_k", 10),
+    }

visual_rag/demo_runner.py

@@ -0,0 +1,90 @@
+"""
+Launch the Streamlit demo from an installed package.
+
+Why:
+- After `pip install visual-rag-toolkit`, the repo layout isn't present.
+- We package the `demo/` module and expose `visual_rag.demo()` + `visual-rag-demo`.
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import os
+import subprocess
+import sys
+from pathlib import Path
+from typing import Optional
+
+
+def demo(
+    *,
+    host: str = "0.0.0.0",
+    port: int = 7860,
+    headless: bool = True,
+    open_browser: bool = False,
+    extra_args: Optional[list[str]] = None,
+) -> int:
+    """
+    Launch the Streamlit demo UI.
+
+    Requirements:
+    - `visual-rag-toolkit[ui,qdrant,embedding,pdf]` (or `visual-rag-toolkit[all]`)
+
+    Returns:
+        Streamlit process exit code.
+    """
+    try:
+        import streamlit  # noqa: F401
+    except Exception as e:  # pragma: no cover
+        raise RuntimeError(
+            "Streamlit is not installed. Install with:\n"
+            '  pip install "visual-rag-toolkit[ui,qdrant,embedding,pdf]"'
+        ) from e
+
+    # Resolve the installed demo entrypoint path.
+    mod = importlib.import_module("demo.app")
+    app_path = Path(getattr(mod, "__file__", "")).resolve()
+    if not app_path.exists():  # pragma: no cover
+        raise RuntimeError("Could not locate installed demo app (demo.app).")
+
+    # Build a stable Streamlit invocation.
+    cmd = [sys.executable, "-m", "streamlit", "run", str(app_path)]
+    cmd += ["--server.address", str(host)]
+    cmd += ["--server.port", str(int(port))]
+    cmd += ["--server.headless", "true" if headless else "false"]
+    cmd += ["--browser.gatherUsageStats", "false"]
+    cmd += ["--server.runOnSave", "false"]
+    cmd += ["--browser.serverAddress", str(host)]
+    if not open_browser:
+        cmd += ["--browser.serverPort", str(int(port))]
+        cmd += ["--browser.open", "false"]
+
+    if extra_args:
+        cmd += list(extra_args)
+
+    env = os.environ.copy()
+    # Make sure the demo doesn't spam internal Streamlit warnings in logs.
+    env.setdefault("STREAMLIT_BROWSER_GATHER_USAGE_STATS", "false")
+
+    return subprocess.call(cmd, env=env)
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description="Launch the Visual RAG Toolkit Streamlit demo.")
+    p.add_argument("--host", default="0.0.0.0")
+    p.add_argument("--port", type=int, default=7860)
+    p.add_argument(
+        "--no-headless", action="store_true", help="Run with a browser window (not headless)."
+    )
+    p.add_argument("--open", action="store_true", help="Open browser automatically.")
+    args, unknown = p.parse_known_args()
+
+    rc = demo(
+        host=args.host,
+        port=args.port,
+        headless=(not args.no_headless),
+        open_browser=bool(args.open),
+        extra_args=unknown,
+    )
+    raise SystemExit(rc)

visual_rag/embedding/__init__.py

@@ -0,0 +1,26 @@
+"""
+Embedding module - Visual and text embedding generation.
+
+Provides:
+- VisualEmbedder: Backend-agnostic visual embedder (ColPali, etc.)
+- Pooling utilities: tile-level, global, MaxSim scoring
+"""
+
+from visual_rag.embedding.pooling import (
+    compute_maxsim_batch,
+    compute_maxsim_score,
+    global_mean_pooling,
+    tile_level_mean_pooling,
+)
+from visual_rag.embedding.visual_embedder import ColPaliEmbedder, VisualEmbedder
+
+__all__ = [
+    # Main embedder
+    "VisualEmbedder",
+    "ColPaliEmbedder",  # Backward compatibility alias
+    # Pooling functions
+    "tile_level_mean_pooling",
+    "global_mean_pooling",
+    "compute_maxsim_score",
+    "compute_maxsim_batch",
+]

visual_rag/embedding/pooling.py

@@ -0,0 +1,343 @@
+"""
+Pooling strategies for multi-vector embeddings.
+
+Provides:
+- Tile-level mean pooling: Preserves spatial structure (num_tiles × dim)
+- Global mean pooling: Single vector (1 × dim)
+- MaxSim scoring for ColBERT-style late interaction
+"""
+
+import logging
+from typing import Optional, Union
+
+import numpy as np
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+def _infer_output_dtype(
+    embedding: Union[torch.Tensor, np.ndarray],
+    output_dtype: Optional[np.dtype] = None,
+) -> np.dtype:
+    """Infer output dtype: use provided, else match input (fp16→fp16, bf16→fp32, fp32→fp32)."""
+    if output_dtype is not None:
+        return output_dtype
+    if isinstance(embedding, torch.Tensor):
+        if embedding.dtype == torch.float16:
+            return np.float16
+        return np.float32
+    if isinstance(embedding, np.ndarray) and embedding.dtype == np.float16:
+        return np.float16
+    return np.float32
+
+
+def tile_level_mean_pooling(
+    embedding: Union[torch.Tensor, np.ndarray],
+    num_tiles: int,
+    patches_per_tile: int = 64,
+    output_dtype: Optional[np.dtype] = None,
+) -> np.ndarray:
+    """
+    Compute tile-level mean pooling for multi-vector embeddings.
+
+    Instead of collapsing to 1×dim (global pooling), this preserves spatial
+    structure by computing mean per tile → num_tiles × dim.
+
+    This is our NOVEL contribution for scalable visual retrieval:
+    - Faster than full MaxSim (fewer vectors to compare)
+    - More accurate than global pooling (preserves spatial info)
+    - Ideal for two-stage retrieval (prefetch with pooled, rerank with full)
+
+    Args:
+        embedding: Visual token embeddings [num_visual_tokens, dim]
+        num_tiles: Number of tiles (including global tile)
+        patches_per_tile: Patches per tile (64 for ColSmol)
+        output_dtype: Output dtype (default: infer from input, fp16→fp16, bf16→fp32)
+
+    Returns:
+        Tile-level pooled embeddings [num_tiles, dim]
+
+    Example:
+        >>> # Image with 4×3 tiles + 1 global = 13 tiles
+        >>> # Each tile has 64 patches → 832 visual tokens
+        >>> pooled = tile_level_mean_pooling(embedding, num_tiles=13)
+        >>> print(pooled.shape)  # (13, 128)
+    """
+    out_dtype = _infer_output_dtype(embedding, output_dtype)
+    if isinstance(embedding, torch.Tensor):
+        if embedding.dtype == torch.bfloat16:
+            emb_np = embedding.cpu().float().numpy()
+        else:
+            emb_np = embedding.cpu().numpy().astype(np.float32)
+    else:
+        emb_np = np.array(embedding, dtype=np.float32)
+
+    num_visual_tokens = emb_np.shape[0]
+    expected_tokens = num_tiles * patches_per_tile
+
+    if num_visual_tokens != expected_tokens:
+        logger.debug(f"Token count mismatch: {num_visual_tokens} vs expected {expected_tokens}")
+        actual_tiles = num_visual_tokens // patches_per_tile
+        if actual_tiles * patches_per_tile != num_visual_tokens:
+            actual_tiles += 1
+        num_tiles = actual_tiles
+
+    tile_embeddings = []
+    for tile_idx in range(num_tiles):
+        start_idx = tile_idx * patches_per_tile
+        end_idx = min(start_idx + patches_per_tile, num_visual_tokens)
+
+        if start_idx >= num_visual_tokens:
+            break
+
+        tile_patches = emb_np[start_idx:end_idx]
+        tile_mean = tile_patches.mean(axis=0)
+        tile_embeddings.append(tile_mean)
+
+    return np.array(tile_embeddings, dtype=out_dtype)
+
+
+def colpali_row_mean_pooling(
+    embedding: Union[torch.Tensor, np.ndarray],
+    grid_size: int = 32,
+    output_dtype: Optional[np.dtype] = None,
+) -> np.ndarray:
+    out_dtype = _infer_output_dtype(embedding, output_dtype)
+    if isinstance(embedding, torch.Tensor):
+        if embedding.dtype == torch.bfloat16:
+            emb_np = embedding.cpu().float().numpy()
+        else:
+            emb_np = embedding.cpu().numpy().astype(np.float32)
+    else:
+        emb_np = np.array(embedding, dtype=np.float32)
+
+    num_tokens, dim = emb_np.shape
+    expected = int(grid_size) * int(grid_size)
+    if num_tokens != expected:
+        raise ValueError(
+            f"Expected {expected} visual tokens for grid_size={grid_size}, got {num_tokens}"
+        )
+
+    grid = emb_np.reshape(int(grid_size), int(grid_size), int(dim))
+    pooled = grid.mean(axis=1)
+    return pooled.astype(out_dtype)
+
+
+def colsmol_experimental_pooling(
+    embedding: Union[torch.Tensor, np.ndarray],
+    num_tiles: int,
+    patches_per_tile: int = 64,
+    output_dtype: Optional[np.dtype] = None,
+) -> np.ndarray:
+    out_dtype = _infer_output_dtype(embedding, output_dtype)
+    if isinstance(embedding, torch.Tensor):
+        if embedding.dtype == torch.bfloat16:
+            emb_np = embedding.cpu().float().numpy()
+        else:
+            emb_np = embedding.cpu().numpy().astype(np.float32)
+    else:
+        emb_np = np.array(embedding, dtype=np.float32)
+
+    num_visual_tokens, dim = emb_np.shape
+    if num_tiles <= 0:
+        raise ValueError("num_tiles must be > 0")
+    if patches_per_tile <= 0:
+        raise ValueError("patches_per_tile must be > 0")
+
+    last_tile_start = (int(num_tiles) - 1) * int(patches_per_tile)
+    if last_tile_start >= num_visual_tokens:
+        actual_tiles = int(num_visual_tokens) // int(patches_per_tile)
+        if actual_tiles * int(patches_per_tile) != int(num_visual_tokens):
+            actual_tiles += 1
+        if actual_tiles <= 0:
+            raise ValueError(
+                f"Not enough tokens for num_tiles={num_tiles}, patches_per_tile={patches_per_tile}: got {num_visual_tokens}"
+            )
+        num_tiles = actual_tiles
+        last_tile_start = (int(num_tiles) - 1) * int(patches_per_tile)
+
+    prefix = emb_np[:last_tile_start]
+    last_tile = emb_np[
+        last_tile_start : min(last_tile_start + int(patches_per_tile), num_visual_tokens)
+    ]
+
+    if prefix.size:
+        prefix_tiles = prefix.reshape(-1, int(patches_per_tile), int(dim))
+        prefix_means = prefix_tiles.mean(axis=1)
+    else:
+        prefix_means = np.zeros((0, int(dim)), dtype=out_dtype)
+
+    return np.concatenate([prefix_means.astype(out_dtype), last_tile.astype(out_dtype)], axis=0)
+
+
+def colpali_experimental_pooling_from_rows(
+    row_vectors: Union[torch.Tensor, np.ndarray],
+    output_dtype: Optional[np.dtype] = None,
+) -> np.ndarray:
+    """
+    Experimental "convolution-style" pooling with window size 3.
+
+    For N input rows, produces N + 2 output vectors:
+    - Position 0: row[0] alone (1 row)
+    - Position 1: mean(rows[0:2]) (2 rows)
+    - Position 2: mean(rows[0:3]) (3 rows)
+    - Positions 3 to N-1: sliding window of 3 (rows[i-2:i+1])
+    - Position N: mean(rows[N-2:N]) (last 2 rows)
+    - Position N+1: row[N-1] alone (last row)
+
+    For N=32 rows: produces 34 vectors.
+    """
+    out_dtype = _infer_output_dtype(row_vectors, output_dtype)
+    if isinstance(row_vectors, torch.Tensor):
+        if row_vectors.dtype == torch.bfloat16:
+            rows = row_vectors.cpu().float().numpy()
+        else:
+            rows = row_vectors.cpu().numpy().astype(np.float32)
+    else:
+        rows = np.array(row_vectors, dtype=np.float32)
+
+    n, dim = rows.shape
+    if n < 1:
+        raise ValueError("row_vectors must be non-empty")
+    if n == 1:
+        return rows.astype(out_dtype)
+    if n == 2:
+        return np.stack([rows[0], rows[:2].mean(axis=0), rows[1]], axis=0).astype(out_dtype)
+    if n == 3:
+        return np.stack(
+            [
+                rows[0],
+                rows[:2].mean(axis=0),
+                rows[:3].mean(axis=0),
+                rows[1:3].mean(axis=0),
+                rows[2],
+            ],
+            axis=0,
+        ).astype(out_dtype)
+
+    out = np.zeros((n + 2, dim), dtype=np.float32)
+    out[0] = rows[0]
+    out[1] = rows[:2].mean(axis=0)
+    out[2] = rows[:3].mean(axis=0)
+    for i in range(3, n):
+        out[i] = rows[i - 2 : i + 1].mean(axis=0)
+    out[n] = rows[n - 2 : n].mean(axis=0)
+    out[n + 1] = rows[n - 1]
+    return out.astype(out_dtype)
+
+
+def global_mean_pooling(
+    embedding: Union[torch.Tensor, np.ndarray],
+    output_dtype: Optional[np.dtype] = None,
+) -> np.ndarray:
+    """
+    Compute global mean pooling → single vector.
+
+    This is the simplest pooling but loses all spatial information.
+    Use for fastest retrieval when accuracy can be sacrificed.
+
+    Args:
+        embedding: Multi-vector embeddings [num_tokens, dim]
+        output_dtype: Output dtype (default: infer from input, fp16→fp16, bf16→fp32)
+
+    Returns:
+        Pooled vector [dim]
+    """
+    out_dtype = _infer_output_dtype(embedding, output_dtype)
+    if isinstance(embedding, torch.Tensor):
+        if embedding.dtype == torch.bfloat16:
+            emb_np = embedding.cpu().float().numpy()
+        else:
+            emb_np = embedding.cpu().numpy()
+    else:
+        emb_np = np.array(embedding)
+
+    return emb_np.mean(axis=0).astype(out_dtype)
+
+
+def compute_maxsim_score(
+    query_embedding: np.ndarray,
+    doc_embedding: np.ndarray,
+    normalize: bool = True,
+) -> float:
+    """
+    Compute ColBERT-style MaxSim late interaction score.
+
+    For each query token, finds max similarity with any document token,
+    then sums across query tokens.
+
+    This is the standard scoring for ColBERT/ColPali:
+    score = Σ_q max_d (sim(q, d))
+
+    Args:
+        query_embedding: Query embeddings [num_query_tokens, dim]
+        doc_embedding: Document embeddings [num_doc_tokens, dim]
+        normalize: L2 normalize embeddings before scoring (recommended)
+
+    Returns:
+        MaxSim score (higher is better)
+
+    Example:
+        >>> query = embedder.embed_query("budget allocation")
+        >>> doc = embeddings[0]  # From embed_images
+        >>> score = compute_maxsim_score(query, doc)
+    """
+    if normalize:
+        # L2 normalize
+        query_norm = query_embedding / (
+            np.linalg.norm(query_embedding, axis=1, keepdims=True) + 1e-8
+        )
+        doc_norm = doc_embedding / (np.linalg.norm(doc_embedding, axis=1, keepdims=True) + 1e-8)
+    else:
+        query_norm = query_embedding
+        doc_norm = doc_embedding
+
+    # Compute similarity matrix: [num_query, num_doc]
+    similarity_matrix = np.dot(query_norm, doc_norm.T)
+
+    # MaxSim: For each query token, take max similarity with any doc token
+    max_similarities = similarity_matrix.max(axis=1)
+
+    # Sum across query tokens
+    score = float(max_similarities.sum())
+
+    return score
+
+
+def compute_maxsim_batch(
+    query_embedding: np.ndarray,
+    doc_embeddings: list,
+    normalize: bool = True,
+) -> list:
+    """
+    Compute MaxSim scores for multiple documents efficiently.
+
+    Args:
+        query_embedding: Query embeddings [num_query_tokens, dim]
+        doc_embeddings: List of document embeddings
+        normalize: L2 normalize embeddings
+
+    Returns:
+        List of MaxSim scores
+    """
+    # Pre-normalize query once
+    if normalize:
+        query_norm = query_embedding / (
+            np.linalg.norm(query_embedding, axis=1, keepdims=True) + 1e-8
+        )
+    else:
+        query_norm = query_embedding
+
+    scores = []
+    for doc_emb in doc_embeddings:
+        if normalize:
+            doc_norm = doc_emb / (np.linalg.norm(doc_emb, axis=1, keepdims=True) + 1e-8)
+        else:
+            doc_norm = doc_emb
+
+        sim_matrix = np.dot(query_norm, doc_norm.T)
+        max_sims = sim_matrix.max(axis=1)
+        scores.append(float(max_sims.sum()))
+
+    return scores