Initial deployment of ColPali Visual Retrieval backend

.gitignore

@@ -0,0 +1,15 @@
+.sesskey
+.venv/
+__pycache__/
+ipynb_checkpoints/
+.python-version
+.env
+template/
+*.json
+output/
+pdfs/
+colpalidemo/
+src/static/full_images/
+src/static/sim_maps/
+embeddings/
+hf_dataset/

Dockerfile

@@ -0,0 +1,35 @@
+FROM python:3.11-slim
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    git \
+    build-essential \
+    libglib2.0-0 \
+    libsm6 \
+    libxext6 \
+    libxrender-dev \
+    libgomp1 \
+    wget \
+    && rm -rf /var/lib/apt/lists/*
+
+# Create a non-root user
+RUN useradd -m -u 1000 user
+USER user
+ENV HOME=/home/user \
+    PATH=/home/user/.local/bin:$PATH
+
+# Set working directory
+WORKDIR $HOME/app
+
+# Copy files to container
+COPY --chown=user . $HOME/app
+
+# Install Python dependencies
+RUN pip install --no-cache-dir --upgrade pip
+RUN pip install --no-cache-dir -e .
+
+# Expose the port the app runs on
+EXPOSE 7860
+
+# Run the application
+CMD ["python", "main.py"]

README.md

@@ -0,0 +1,54 @@
+---
+title: ColPali Visual Retrieval
+emoji: 🔍
+colorFrom: green
+colorTo: blue
+sdk: docker
+sdk_version: "3.11"
+app_file: app.py
+pinned: false
+---
+
+# ColPali Visual Retrieval with Vespa
+
+A powerful visual document retrieval system that combines **ColPali** (Contextual Late Interaction with Patch-level Information) with **Vespa** for scalable, intelligent document search and question-answering.
+
+## 🌟 Features
+
+- **Visual Document Search**: Search through PDF documents using natural language queries
+- **Token-level Similarity Maps**: Visualize exactly which parts of documents match your query
+- **AI-Powered Chat**: Ask questions about retrieved documents using Google Gemini
+- **Multiple Ranking Methods**: Choose between ColPali, BM25, or Hybrid ranking
+
+## 🚀 Try It Out
+
+1. Enter a natural language query in the search box
+2. Select your preferred ranking method
+3. Click on token buttons to see visual attention maps
+4. Ask follow-up questions in the chat interface
+
+## 📄 Sample Queries
+
+- "Pie chart with model comparison"
+- "Speaker diarization evaluation"
+- "Results table from dense retrieval"
+- "Graph showing training loss"
+- "Architecture diagram with transformer"
+
+## 🛠️ Technology Stack
+
+- **ColPali**: Visual-language model for document understanding
+- **Vespa**: Distributed search engine for scalability
+- **FastHTML**: Modern web framework for the UI
+- **Google Gemini**: AI-powered question answering
+
+## 📊 About the Dataset
+
+This demo uses ~400 pages from AI-related research papers published in 2024. The documents are processed using ColPali to create visual embeddings that enable semantic search across document images.
+
+## 🔗 Links
+
+- [ColPali Paper](https://arxiv.org/abs/2404.09317)
+- [Vespa Documentation](https://docs.vespa.ai/)
+- [Blog Post](https://blog.vespa.ai/visual-retrieval-with-colpali-and-vespa/)
+- [GitHub Repository](https://github.com/vespa-engine/vespa/tree/master/examples/colpali-visual-retrieval)

README_HF.md

@@ -0,0 +1,54 @@
+---
+title: ColPali Visual Retrieval
+emoji: 🔍
+colorFrom: green
+colorTo: blue
+sdk: docker
+sdk_version: "3.11"
+app_file: app.py
+pinned: false
+---
+
+# ColPali Visual Retrieval with Vespa
+
+A powerful visual document retrieval system that combines **ColPali** (Contextual Late Interaction with Patch-level Information) with **Vespa** for scalable, intelligent document search and question-answering.
+
+## 🌟 Features
+
+- **Visual Document Search**: Search through PDF documents using natural language queries
+- **Token-level Similarity Maps**: Visualize exactly which parts of documents match your query
+- **AI-Powered Chat**: Ask questions about retrieved documents using Google Gemini
+- **Multiple Ranking Methods**: Choose between ColPali, BM25, or Hybrid ranking
+
+## 🚀 Try It Out
+
+1. Enter a natural language query in the search box
+2. Select your preferred ranking method
+3. Click on token buttons to see visual attention maps
+4. Ask follow-up questions in the chat interface
+
+## 📄 Sample Queries
+
+- "Pie chart with model comparison"
+- "Speaker diarization evaluation"
+- "Results table from dense retrieval"
+- "Graph showing training loss"
+- "Architecture diagram with transformer"
+
+## 🛠️ Technology Stack
+
+- **ColPali**: Visual-language model for document understanding
+- **Vespa**: Distributed search engine for scalability
+- **FastHTML**: Modern web framework for the UI
+- **Google Gemini**: AI-powered question answering
+
+## 📊 About the Dataset
+
+This demo uses ~400 pages from AI-related research papers published in 2024. The documents are processed using ColPali to create visual embeddings that enable semantic search across document images.
+
+## 🔗 Links
+
+- [ColPali Paper](https://arxiv.org/abs/2404.09317)
+- [Vespa Documentation](https://docs.vespa.ai/)
+- [Blog Post](https://blog.vespa.ai/visual-retrieval-with-colpali-and-vespa/)
+- [GitHub Repository](https://github.com/vespa-engine/vespa/tree/master/examples/colpali-visual-retrieval)

app.py

@@ -0,0 +1,17 @@
+#!/usr/bin/env python3
+# app.py - Entry point for Hugging Face Spaces
+from main import app, APP_DIR, setup_static_routes, IMG_DIR, SIM_MAP_DIR
+import os
+
+# Ensure directories exist
+os.makedirs(APP_DIR, exist_ok=True)
+os.makedirs(IMG_DIR, exist_ok=True)
+os.makedirs(SIM_MAP_DIR, exist_ok=True)
+
+# Set up static routes
+setup_static_routes(app)
+
+# For Hugging Face Spaces
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)

backend/about.md

@@ -0,0 +1,985 @@
+# ColPali 🤝 Vespa - Visual Retrieval System
+
+A powerful visual document retrieval system that combines **ColPali** (Contextual Late Interaction with Patch-level Information) with **Vespa** for scalable, intelligent document search and question-answering.
+
+## 🌟 Features
+
+### 🔍 **Visual Document Search**
+
+- **Multi-modal retrieval**: Search through PDF documents using natural language queries
+- **Visual understanding**: ColPali model processes document images and text simultaneously
+- **Token-level similarity maps**: Visualize exactly which parts of documents match your query
+- **Multiple ranking algorithms**: Choose between hybrid, semantic, and other ranking methods
+
+### 🧠 **AI-Powered Chat**
+
+- **Intelligent Q&A**: Ask questions about retrieved documents using Google Gemini 2.0
+- **Context-aware responses**: AI analyzes document images to provide accurate answers
+- **Real-time streaming**: Get responses as they're generated
+
+### ⚡ **Scalable Infrastructure**
+
+- **Vespa integration**: Enterprise-grade search platform for large document collections
+- **Real-time processing**: Instant search results and similarity map generation
+- **Cloud-ready**: Supports Vespa Cloud deployment with secure authentication
+
+## 🏗️ Architecture
+
+```
+┌─────────────────┐    ┌─────────────────┐    ┌─────────────────┐
+│   Frontend      │    │    Backend      │    │   Vespa Cloud   │
+│   (Browser)     │    │   (Your Local   │    │   (Remote)      │
+│                 │    │    Computer)    │    │                 │
+│ • Search UI     │◄──►│ • ColPali Model │◄──►│ • Document Store│
+│ • Similarity    │    │ • Query Proc.   │    │ • Vector Search │
+│   Maps          │    │ • Sim Map Gen.  │    │ • Ranking       │
+│ • Chat Interface│    │ • Gemini Int.   │    │                 │
+└─────────────────┘    └─────────────────┘    └─────────────────┘
+        ↑                        ↑                        ↑
+   Web Browser              LOCAL AI               REMOTE Storage
+```
+
+### 🏠 **LOCAL Processing (Your Computer)**
+
+**All AI model inference happens on YOUR local machine:**
+
+- **ColPali Model**: Runs locally on your GPU/CPU (~7GB model)
+- **Document Processing**: PDF → Images → Embeddings (local)
+- **Query Processing**: Text → Embeddings (local)
+- **Similarity Maps**: Visual attention generation (local)
+- **Gemini Chat**: Processes retrieved images locally
+
+**Device Detection:**
+
+```python
+device = get_torch_device("auto")  # Detects: CUDA, MPS (Apple), or CPU
+print(f"Using device: {device}")   # Shows YOUR hardware
+```
+
+### ☁️ **REMOTE Processing (Vespa Cloud)**
+
+**Only storage and search index operations happen remotely:**
+
+- **Document Storage**: Stores processed embeddings (not raw models)
+- **Vector Search**: Fast similarity search across document collection
+- **Query Routing**: Handles search requests and ranking
+- **Metadata Storage**: Document titles, URLs, page numbers
+
+### 🔄 **Complete Data Flow**
+
+#### **Document Upload Process:**
+
+1. **LOCAL**: Your computer downloads PDF from URL
+2. **LOCAL**: ColPali converts PDF pages to images
+3. **LOCAL**: ColPali generates visual embeddings (1024 patches × 128 dims)
+4. **LOCAL**: Embeddings converted to binary format for efficiency
+5. **REMOTE**: Binary embeddings uploaded to Vespa Cloud
+6. **REMOTE**: Vespa indexes embeddings for fast search
+
+#### **Search Query Process:**
+
+1. **LOCAL**: You enter search query in browser
+2. **LOCAL**: ColPali processes query → generates query embeddings
+3. **REMOTE**: Query embeddings sent to Vespa Cloud
+4. **REMOTE**: Vespa searches document index, returns matches
+5. **LOCAL**: ColPali generates similarity maps for results
+6. **BROWSER**: Results displayed with visual attention maps
+
+#### **AI Chat Process:**
+
+1. **LOCAL**: Retrieved document images processed by your machine
+2. **REMOTE**: Images + query sent to Google Gemini API
+3. **REMOTE**: Gemini generates response based on visual content
+4. **BROWSER**: Streaming response displayed in real-time
+
+### Core Components
+
+- **ColPali Model**: Visual-language model for document understanding (LOCAL)
+- **Vespa Search**: Distributed search and storage engine (REMOTE)
+- **FastHTML Frontend**: Modern, responsive web interface (BROWSER)
+- **Gemini Integration**: AI-powered question answering (REMOTE API)
+- **Similarity Map Generator**: Visual attention visualization (LOCAL)
+
+## 💻 **System Requirements**
+
+### **LOCAL Machine Requirements (For AI Processing)**
+
+**Minimum:**
+
+- **CPU**: Modern multi-core processor (Intel/AMD/Apple Silicon)
+- **RAM**: 8GB+ (16GB recommended)
+- **Storage**: 10GB free space (for model cache)
+- **Python**: 3.10+ (< 3.13)
+
+**Recommended:**
+
+- **GPU**: NVIDIA GPU with 8GB+ VRAM (RTX 3070/4060 or better)
+- **Apple**: M1/M2/M3 Mac (uses Metal Performance Shaders)
+- **RAM**: 16GB+ for smoother processing
+- **Storage**: SSD for faster model loading
+
+**Performance Examples:**
+
+- **RTX 4090**: ~1-2 seconds per query
+- **RTX 3070**: ~3-5 seconds per query
+- **Apple M2**: ~4-6 seconds per query
+- **CPU Only**: ~15-30 seconds per query
+
+### **REMOTE Requirements (Vespa Cloud)**
+
+**What you need:**
+
+- **Vespa Cloud account** (handles all remote processing)
+- **Internet connection** (for uploading embeddings and search queries)
+- **Authentication tokens** (provided by Vespa Cloud)
+
+**What Vespa Cloud provides:**
+
+- **Scalable storage** for any number of documents
+- **Sub-second search** across millions of embeddings
+- **High availability** with automatic failover
+- **Global CDN** for fast access worldwide
+
+## 💰 **Cost Breakdown**
+
+### **FREE Components**
+
+- **ColPali Model**: Open source, runs locally (no per-query costs)
+- **Python Application**: MIT/Apache licensed, completely free
+- **Local Processing**: Uses your own hardware (no cloud AI fees)
+
+### **PAID Components**
+
+- **Vespa Cloud**: Pay for storage and search operations
+  - ~$0.001 per 1000 searches
+  - ~$0.10 per GB storage per month
+- **Google Gemini API**: Optional, for chat features only
+  - ~$0.01 per 1000 image tokens
+  - Only used when you ask questions about documents
+
+### **Cost Examples (Monthly)**
+
+- **Personal Use** (100 documents, 1000 searches): ~$5-10/month
+- **Small Business** (1000 documents, 10k searches): ~$20-50/month
+- **Enterprise** (10k+ documents, 100k+ searches): $200+/month
+
+**💡 Cost Optimization Tips:**
+
+- Use local Vespa installation to avoid cloud costs
+- Disable Gemini chat if not needed (saves API costs)
+- Process documents in batches to minimize upload time
+
+## 🚀 Quick Start
+
+### Prerequisites
+
+- Python 3.10+ (< 3.13)
+- **8GB+ RAM** for ColPali model
+- **Vespa Cloud account** or local Vespa installation
+- **Google Gemini API key** (optional, for chat features)
+- **GPU recommended** but not required
+
+### 1. Installation
+
+```bash
+# Clone the repository
+git clone <repository-url>
+cd colpali-vespa-visual-retrieval
+
+# Install dependencies
+pip install -e .
+
+# For development
+pip install -e ".[dev]"
+
+# For document feeding capabilities
+pip install -e ".[feed]"
+```
+
+### 2. Environment Configuration
+
+Create a `.env` file with your configuration:
+
+```bash
+# Vespa Configuration
+VESPA_APP_TOKEN_URL=https://your-app.vespa-cloud.com
+VESPA_CLOUD_SECRET_TOKEN=your_secret_token
+
+# Alternative: mTLS Authentication
+USE_MTLS=false
+VESPA_APP_MTLS_URL=https://your-app.vespa-cloud.com
+VESPA_CLOUD_MTLS_KEY="-----BEGIN PRIVATE KEY-----..."
+VESPA_CLOUD_MTLS_CERT="-----BEGIN CERTIFICATE-----..."
+
+# Optional: Gemini AI (for chat features)
+GEMINI_API_KEY=your_gemini_api_key
+
+# Optional: Logging
+LOG_LEVEL=INFO
+HOT_RELOAD=false
+```
+
+### 3. Deploy Vespa Application
+
+```bash
+# Deploy the Vespa schema and configuration
+python deploy_vespa_app.py \
+  --tenant_name your_tenant \
+  --vespa_application_name colpalidemo \
+  --token_id_write colpalidemo_write \
+  --token_id_read colpalidemo_read
+```
+
+### 4. Run the Application
+
+```bash
+python main.py
+```
+
+The application will be available at `http://localhost:7860`
+
+## 📚 Document Management
+
+### Uploading Documents
+
+Use the feeding script to process and upload PDF documents:
+
+```bash
+python feed_vespa.py \
+  --application_name colpalidemo \
+  --vespa_schema_name pdf_page
+```
+
+**Document Processing Pipeline (LOCAL → REMOTE):**
+
+1. **PDF Download** (LOCAL): Your computer downloads PDFs from URLs
+2. **PDF Conversion** (LOCAL): PDFs converted to images (one per page)
+3. **ColPali Processing** (LOCAL): Each page processed by ColPali model on YOUR GPU/CPU
+4. **Embedding Generation** (LOCAL): Visual embeddings created (1024 patches × 128 dimensions)
+5. **Binary Encoding** (LOCAL): Embeddings converted to efficient binary format
+6. **Vespa Upload** (REMOTE): Binary embeddings uploaded to Vespa Cloud
+7. **Search Indexing** (REMOTE): Vespa indexes embeddings for fast retrieval
+
+**⚠️ Important Notes:**
+
+- **Processing Time**: Expect 5-30 seconds per page depending on your hardware
+- **Network Usage**: Only final embeddings uploaded (~1KB per page vs ~1MB original)
+- **Privacy**: Original PDFs and images stay on your local machine
+- **Storage**: Raw images cached locally for similarity map generation
+
+### Supported Operations
+
+- ✅ **Upload Documents**: Add new PDFs to the system
+- ✅ **Search Documents**: Query existing documents
+- ✅ **View Documents**: Browse stored documents
+- ❌ **Remove Documents**: _Not currently implemented_
+- ❌ **Update Documents**: _Not currently implemented_
+
+## 🔐 Authentication & Security
+
+### 🛡️ **Current Security Implementation**
+
+#### **SECURE Components:**
+
+**Vespa Authentication (REMOTE)**
+
+- **Token Authentication**: Bearer tokens for Vespa Cloud API access
+- **mTLS Certificates**: Mutual TLS for enterprise security
+- **Encrypted Communication**: HTTPS/TLS for all Vespa connections
+
+**API Key Management (LOCAL)**
+
+- **Environment Variables**: Sensitive keys stored in `.env` files
+- **API Key Rotation**: Google Gemini supports key rotation
+- **Local Storage**: Keys never transmitted except to authorized APIs
+
+#### **LIMITED Security Components:**
+
+**Session Management**
+
+```python
+# Basic UUID session tracking (FastHTML)
+session["session_id"] = str(uuid.uuid4())
+
+# HTTP-only cookies (Next.js)
+cookieStore.set(SESSION_KEY, newSessionId, {
+  httpOnly: true,
+  secure: process.env.NODE_ENV === "production",
+  sameSite: "lax",
+  maxAge: 60 * 60 * 24 * 30, // 30 days
+});
+```
+
+**Basic Request Validation**
+
+```python
+# HTMX request validation
+if "hx-request" not in request.headers:
+    return RedirectResponse("/search")
+
+# Parameter validation
+if not query:
+    return NextResponse.json({ error: "Query is required" }, { status: 400 });
+```
+
+### ⚠️ **Security Limitations & Risks**
+
+#### **MISSING Security Features:**
+
+**❌ No API Authentication**
+
+- Local API endpoints are **completely open**
+- No rate limiting or abuse protection
+- No user authentication or authorization
+- Anyone can access `/fetch_results`, `/get_sim_map` endpoints
+
+**❌ No Input Sanitization**
+
+```python
+# Raw user input passed directly to models
+query = searchParams.get("query")  # No validation/sanitization
+ranking = searchParams.get("ranking")  # No input filtering
+```
+
+**❌ No Security Headers**
+
+- No CORS configuration
+- No Content Security Policy (CSP)
+- No X-Frame-Options protection
+- No X-Content-Type-Options validation
+
+**❌ No Rate Limiting**
+
+- Unlimited API requests
+- No protection against DoS attacks
+- No query throttling or user limits
+
+**❌ No CSRF Protection**
+
+- No token validation for state-changing operations
+- Cross-site request forgery possible
+
+### 🎯 **Security Recommendations**
+
+#### **IMMEDIATE (High Priority)**
+
+**1. Add API Authentication**
+
+```typescript
+// middleware.ts - Add API key validation
+export function middleware(request: NextRequest) {
+  const apiKey = request.headers.get("X-API-Key");
+  if (!apiKey || apiKey !== process.env.COLPALI_API_KEY) {
+    return new Response("Unauthorized", { status: 401 });
+  }
+}
+```
+
+**2. Implement Rate Limiting**
+
+```typescript
+// Use next-rate-limit or similar
+import rateLimit from "@/lib/rate-limit";
+
+const limiter = rateLimit({
+  interval: 60 * 1000, // 1 minute
+  uniqueTokenPerInterval: 500, // Limit each IP to 100 requests per interval
+});
+
+await limiter.check(10, getClientIP(request)); // 10 requests per minute
+```
+
+**3. Add Security Headers**
+
+```typescript
+// next.config.js
+const securityHeaders = [
+  { key: "X-Frame-Options", value: "DENY" },
+  { key: "X-Content-Type-Options", value: "nosniff" },
+  { key: "Referrer-Policy", value: "strict-origin-when-cross-origin" },
+  {
+    key: "Content-Security-Policy",
+    value: "default-src 'self'; script-src 'self' 'unsafe-inline'",
+  },
+];
+```
+
+**4. Input Validation & Sanitization**
+
+```typescript
+import { z } from "zod";
+
+const SearchSchema = z.object({
+  query: z
+    .string()
+    .min(1)
+    .max(500)
+    .regex(/^[a-zA-Z0-9\s\.\?\!]*$/),
+  ranking: z.enum(["hybrid", "colpali", "bm25"]),
+});
+```
+
+#### **MEDIUM Priority**
+
+**5. CORS Configuration**
+
+```typescript
+// Restrict origins to known domains
+const corsHeaders = {
+  "Access-Control-Allow-Origin": "https://yourdomain.com",
+  "Access-Control-Allow-Methods": "GET, POST, OPTIONS",
+  "Access-Control-Allow-Headers": "Content-Type, Authorization",
+};
+```
+
+**6. Request Size Limits**
+
+```typescript
+// Limit request payload sizes
+export const config = {
+  api: {
+    bodyParser: {
+      sizeLimit: "1mb",
+    },
+  },
+};
+```
+
+**7. Audit Logging**
+
+```python
+# Log all API access with IP, timestamp, and queries
+logger.info(f"API_ACCESS: {client_ip} - {endpoint} - {query[:100]}")
+```
+
+#### **LONG-TERM (Production Ready)**
+
+**8. User Authentication (Optional)**
+
+```typescript
+// Add NextAuth.js or similar for user accounts
+// Implement role-based access control
+// Add document ownership and permissions
+```
+
+**9. Network Security**
+
+```bash
+# Deploy behind reverse proxy (nginx/cloudflare)
+# Enable DDoS protection
+# Use Web Application Firewall (WAF)
+```
+
+**10. Data Privacy Controls**
+
+```typescript
+// Implement data retention policies
+// Add user data deletion capabilities
+// GDPR compliance features
+```
+
+### 🔒 **Security Best Practices**
+
+#### **For LOCAL Development:**
+
+- **Never commit API keys** to version control
+- **Use strong environment variable names** (avoid `API_KEY`)
+- **Rotate API keys regularly** (monthly)
+- **Enable firewall** on development machines
+- **Use HTTPS even locally** for production testing
+
+#### **For PRODUCTION Deployment:**
+
+- **Deploy behind CDN/WAF** (Cloudflare, AWS Shield)
+- **Enable rate limiting** at infrastructure level
+- **Use container security scanning**
+- **Implement monitoring and alerting**
+- **Regular security audits and penetration testing**
+
+#### **For REMOTE Services:**
+
+- **Vespa Cloud**: Follows enterprise security standards
+- **Gemini API**: Google-managed security and compliance
+- **Environment Isolation**: Separate dev/staging/prod credentials
+
+### 🚨 **Current Risk Level: MEDIUM**
+
+**Suitable for:**
+
+- ✅ **Personal projects and demos**
+- ✅ **Internal company tools** (behind firewall)
+- ✅ **Research and development** environments
+
+**NOT suitable for:**
+
+- ❌ **Public internet deployment**
+- ❌ **Customer-facing applications**
+- ❌ **Production environments** with sensitive data
+- ❌ **Commercial applications** without security hardening
+
+## 🎯 Usage Guide
+
+### Basic Search
+
+1. Navigate to the homepage
+2. Enter your search query in natural language
+3. Select ranking method (hybrid, semantic, etc.)
+4. View results with similarity maps
+
+### Similarity Maps
+
+- Click on token buttons to see which parts of documents match specific query terms
+- Visual heatmaps show attention patterns
+- Reset button returns to original document view
+
+### AI Chat
+
+- Ask questions about retrieved documents
+- Chat responses are based on document content
+- Streaming responses for real-time interaction
+
+### Search Rankings
+
+- **Hybrid**: Combines multiple ranking signals
+- **Semantic**: Pure semantic similarity
+- **BM25**: Traditional text-based ranking
+- **ColPali**: Visual-first ranking
+
+## 🛠️ Development
+
+### Project Structure
+
+```
+├── main.py                 # Application entry point
+├── backend/
+│   ├── colpali.py         # ColPali model integration
+│   ├── vespa_app.py       # Vespa client and queries
+│   └── modelmanager.py    # Model management utilities
+├── frontend/
+│   ├── app.py             # UI components
+│   └── layout.py          # Layout templates
+├── feed_vespa.py          # Document upload script
+├── deploy_vespa_app.py    # Vespa deployment script
+├── colpali-with-snippets/ # Vespa schema definitions
+└── static/                # Static assets and generated files
+```
+
+### Running in Development
+
+```bash
+# Enable hot reload
+export HOT_RELOAD=true
+python main.py
+
+# Or set in .env
+echo "HOT_RELOAD=true" >> .env
+```
+
+### Code Quality
+
+```bash
+# Format code
+ruff format .
+
+# Lint code
+ruff check .
+```
+
+## 📊 API Endpoints
+
+### **Current API Routes (⚠️ UNSECURED)**
+
+| Endpoint         | Method | Description             | Security Status  |
+| ---------------- | ------ | ----------------------- | ---------------- |
+| `/`              | GET    | Homepage                | ✅ Public (safe) |
+| `/search`        | GET    | Search interface        | ✅ Public (safe) |
+| `/fetch_results` | GET    | Fetch search results    | ⚠️ **OPEN API**  |
+| `/get_sim_map`   | GET    | Get similarity maps     | ⚠️ **OPEN API**  |
+| `/get-message`   | GET    | Chat with AI (SSE)      | ⚠️ **OPEN API**  |
+| `/full_image`    | GET    | Get full document image | ⚠️ **OPEN API**  |
+| `/suggestions`   | GET    | Query autocomplete      | ⚠️ **OPEN API**  |
+| `/static/*`      | GET    | Static file serving     | ✅ Public (safe) |
+
+### **Security Analysis by Endpoint**
+
+#### **🔒 SECURE Endpoints**
+
+- **`/`** and **`/search`**: Static HTML pages, no sensitive data
+- **`/static/*`**: Public assets (CSS, JS, images)
+
+#### **⚠️ UNSECURED Endpoints (Risk)**
+
+**`/fetch_results`** - **HIGH RISK**
+
+```bash
+# Anyone can perform unlimited searches
+curl "http://localhost:7860/fetch_results?query=secret&ranking=hybrid"
+```
+
+- **Risks**: Resource abuse, server overload, competitive intelligence gathering
+- **Exposes**: Search capabilities, document metadata, processing times
+
+**`/get_sim_map`** - **MEDIUM RISK**
+
+```bash
+# Access similarity maps without authentication
+curl "http://localhost:7860/get_sim_map?query_id=123&idx=0&token=word&token_idx=5"
+```
+
+- **Risks**: Unauthorized access to visual analysis
+- **Exposes**: Document visual patterns, query insights
+
+**`/get-message`** - **HIGH RISK**
+
+```bash
+# Trigger AI processing without limits
+curl "http://localhost:7860/get-message?query_id=123&query=question&doc_ids=doc1,doc2"
+```
+
+- **Risks**: Gemini API abuse, cost exploitation, resource exhaustion
+- **Exposes**: AI-generated insights, document content analysis
+
+**`/full_image`** - **HIGH RISK**
+
+```bash
+# Download any document image
+curl "http://localhost:7860/full_image?doc_id=any_document_id"
+```
+
+- **Risks**: Unauthorized document access, data leakage
+- **Exposes**: Full document images, potentially sensitive content
+
+### **Immediate Security Fixes**
+
+#### **1. Add API Key Authentication**
+
+```python
+# Python FastHTML middleware
+@app.middleware("http")
+async def verify_api_key(request, call_next):
+    if request.url.path.startswith("/fetch_results"):
+        api_key = request.headers.get("X-API-Key")
+        if not api_key or api_key != os.getenv("COLPALI_API_KEY"):
+            return JSONResponse({"error": "Unauthorized"}, status_code=401)
+    return await call_next(request)
+```
+
+#### **2. Implement Rate Limiting**
+
+```python
+from slowapi import Limiter, _rate_limit_exceeded_handler
+from slowapi.util import get_remote_address
+
+limiter = Limiter(key_func=get_remote_address)
+
+@rt("/fetch_results")
+@limiter.limit("10/minute")  # 10 requests per minute per IP
+async def get_results(request, query: str, ranking: str):
+    # ... existing code
+```
+
+#### **3. Input Validation**
+
+```python
+from pydantic import BaseModel, validator
+
+class SearchRequest(BaseModel):
+    query: str
+    ranking: str
+
+    @validator('query')
+    def query_must_be_safe(cls, v):
+        if len(v) > 500:
+            raise ValueError('Query too long')
+        # Add sanitization logic
+        return v.strip()
+```
+
+#### **4. Request Origin Validation**
+
+```python
+ALLOWED_ORIGINS = ["http://localhost:3000", "https://yourdomain.com"]
+
+@app.middleware("http")
+async def cors_middleware(request, call_next):
+    origin = request.headers.get("origin")
+    if origin not in ALLOWED_ORIGINS:
+        return JSONResponse({"error": "Forbidden"}, status_code=403)
+    return await call_next(request)
+```
+
+### **📈 Recommended API Security Architecture**
+
+```
+┌─────────────────┐    ┌─────────────────┐    ┌─────────────────┐
+│   Frontend      │    │  Rate Limiter   │    │   Backend API   │
+│                 │    │                 │    │                 │
+│ • API Key       │◄──►│ • IP Limiting   │◄──►│ • Input Valid.  │
+│ • CORS Headers  │    │ • User Quotas   │    │ • Auth Checks   │
+│ • Request Valid.│    │ • DoS Protection│    │ • Audit Logs    │
+└─────────────────┘    └─────────────────┘    └─────────────────┘
+```
+
+**Benefits:**
+
+- **Layer 1**: Frontend validates requests before sending
+- **Layer 2**: Rate limiter prevents abuse and DoS attacks
+- **Layer 3**: Backend performs final validation and authorization
+
+### **🔒 Security Implementation Checklist**
+
+#### **Before Production Deployment:**
+
+**CRITICAL (Must Do):**
+
+- [ ] **Generate API Key**: Create strong API key for endpoint authentication
+- [ ] **Enable Rate Limiting**: Implement per-IP request limits
+- [ ] **Add Security Headers**: X-Frame-Options, CSP, X-Content-Type-Options
+- [ ] **Input Validation**: Sanitize all user inputs (query, ranking)
+- [ ] **CORS Configuration**: Restrict origins to known domains only
+- [ ] **Environment Security**: Never commit API keys, use secure .env
+- [ ] **HTTPS Only**: Force TLS in production (no HTTP)
+
+**HIGH Priority:**
+
+- [ ] **Audit Logging**: Log all API requests with IP and timestamp
+- [ ] **Request Size Limits**: Prevent large payload attacks
+- [ ] **Error Handling**: Don't expose stack traces or internal details
+- [ ] **Session Security**: HTTP-only, secure, SameSite cookies
+- [ ] **API Documentation**: Document authentication requirements
+
+**MEDIUM Priority:**
+
+- [ ] **User Authentication**: Consider adding user accounts for access control
+- [ ] **Request Timeout**: Prevent long-running request abuse
+- [ ] **Content Validation**: Verify response content types
+- [ ] **Monitoring**: Set up alerts for unusual API usage patterns
+- [ ] **Backup Strategy**: Secure backup of environment variables
+
+#### **Security Testing Commands:**
+
+**Test API Authentication:**
+
+```bash
+# Should fail without API key
+curl "http://localhost:7860/fetch_results?query=test&ranking=hybrid"
+
+# Should succeed with API key
+curl -H "X-API-Key: your_api_key" "http://localhost:7860/fetch_results?query=test&ranking=hybrid"
+```
+
+**Test Rate Limiting:**
+
+```bash
+# Run multiple requests to trigger rate limit
+for i in {1..15}; do
+  curl -H "X-API-Key: your_api_key" "http://localhost:7860/fetch_results?query=test$i&ranking=hybrid"
+  echo "Request $i"
+done
+```
+
+**Test Input Validation:**
+
+```bash
+# Should reject invalid/malicious inputs
+curl -H "X-API-Key: your_api_key" "http://localhost:7860/fetch_results?query=<script>alert('xss')</script>&ranking=invalid"
+```
+
+**Test Security Headers:**
+
+```bash
+# Check security headers in response
+curl -I "http://localhost:7860/"
+# Should see: X-Frame-Options, X-Content-Type-Options, etc.
+```
+
+#### **Security Monitoring:**
+
+**Log Analysis Queries:**
+
+```bash
+# Monitor API usage patterns
+grep "API_ACCESS" /var/log/colpali.log | tail -100
+
+# Detect potential abuse
+grep "RATE_LIMIT_EXCEEDED" /var/log/colpali.log
+
+# Check authentication failures
+grep "UNAUTHORIZED" /var/log/colpali.log
+```
+
+**Alerting Setup:**
+
+- **Rate Limit Violations**: Alert when >50 requests/minute from single IP
+- **Authentication Failures**: Alert on repeated unauthorized attempts
+- **Unusual Queries**: Alert on suspicious query patterns or injection attempts
+- **Resource Usage**: Alert on high CPU/memory usage (potential DoS)
+
+## 🧪 Models Used
+
+- **ColPali v1.2**: Visual document understanding
+- **ColPaliGemma 3B**: Base visual-language model
+- **Google Gemini 2.0**: AI chat and question answering
+
+## 🔧 Configuration Options
+
+### Environment Variables
+
+| Variable                   | Required | Description                                 | Security Impact                     |
+| -------------------------- | -------- | ------------------------------------------- | ----------------------------------- |
+| `VESPA_APP_TOKEN_URL`      | Yes\*    | Vespa application URL (token auth)          | **HIGH** - Remote access            |
+| `VESPA_CLOUD_SECRET_TOKEN` | Yes\*    | Vespa secret token                          | **CRITICAL** - Full database access |
+| `USE_MTLS`                 | No       | Use mTLS instead of token auth              | **MEDIUM** - Auth method            |
+| `VESPA_APP_MTLS_URL`       | Yes\*\*  | Vespa application URL (mTLS)                | **HIGH** - Remote access            |
+| `VESPA_CLOUD_MTLS_KEY`     | Yes\*\*  | mTLS private key                            | **CRITICAL** - TLS credentials      |
+| `VESPA_CLOUD_MTLS_CERT`    | Yes\*\*  | mTLS certificate                            | **HIGH** - TLS credentials          |
+| `GEMINI_API_KEY`           | No       | Google Gemini API key                       | **HIGH** - AI access/costs          |
+| `LOG_LEVEL`                | No       | Logging level (DEBUG, INFO, WARNING, ERROR) | **LOW** - Debug info                |
+| `HOT_RELOAD`               | No       | Enable hot reload in development            | **LOW** - Dev convenience           |
+
+#### **🔒 Security-Related Environment Variables (Recommended)**
+
+| Variable                   | Required  | Description                          | Default |
+| -------------------------- | --------- | ------------------------------------ | ------- |
+| `COLPALI_API_KEY`          | **YES\*** | API key for endpoint authentication  | None    |
+| `ALLOWED_ORIGINS`          | **YES\*** | Comma-separated allowed CORS origins | None    |
+| `RATE_LIMIT_REQUESTS`      | No        | Max requests per minute per IP       | `10`    |
+| `RATE_LIMIT_WINDOW`        | No        | Rate limit window in seconds         | `60`    |
+| `MAX_QUERY_LENGTH`         | No        | Maximum query string length          | `500`   |
+| `ENABLE_AUDIT_LOGGING`     | No        | Log all API requests for security    | `false` |
+| `SECURITY_HEADERS_ENABLED` | No        | Enable security headers              | `true`  |
+| `CSRF_SECRET`              | **YES\*** | Secret for CSRF token generation     | None    |
+
+**Example Security-Enhanced `.env`:**
+
+```bash
+# Existing configuration
+VESPA_APP_TOKEN_URL=https://your-app.vespa-cloud.com
+VESPA_CLOUD_SECRET_TOKEN=your_vespa_secret_token
+GEMINI_API_KEY=your_gemini_api_key
+
+# NEW: Security configuration
+COLPALI_API_KEY=your_strong_random_api_key_here
+ALLOWED_ORIGINS=http://localhost:3000,https://yourdomain.com
+RATE_LIMIT_REQUESTS=10
+RATE_LIMIT_WINDOW=60
+MAX_QUERY_LENGTH=500
+ENABLE_AUDIT_LOGGING=true
+SECURITY_HEADERS_ENABLED=true
+CSRF_SECRET=your_random_csrf_secret_here
+
+# Development vs Production
+NODE_ENV=production  # Enable secure cookies
+LOG_LEVEL=INFO       # Don't expose debug info in production
+```
+
+\*Required for token authentication  
+\*\*Required for mTLS authentication  
+\*\*\*Required for production security
+
+## 🚨 Troubleshooting
+
+### **LOCAL Processing Issues**
+
+**ColPali model fails to load:**
+
+```bash
+# Check GPU memory
+nvidia-smi  # For NVIDIA GPUs
+# or
+system_profiler SPDisplaysDataType  # For Apple Silicon
+
+# Clear model cache if corrupted
+rm -rf ~/.cache/huggingface/hub/models--vidore--colpali-v1.2
+```
+
+**Out of memory errors:**
+
+- Reduce batch size in `feed_vespa.py` (try `batch_size=1`)
+- Close other applications to free RAM/VRAM
+- Use CPU processing if GPU memory insufficient: `CUDA_VISIBLE_DEVICES="" python main.py`
+
+**Slow processing on CPU:**
+
+- Expected behavior - ColPali requires significant computation
+- Consider upgrading to GPU or Apple Silicon for 5-10x speedup
+- Process documents overnight for large collections
+
+### **REMOTE Processing Issues**
+
+**Connection to Vespa fails:**
+
+- Verify your Vespa URL and credentials in `.env`
+- Check if the Vespa application is deployed and running
+- Ensure network connectivity: `ping your-app.vespa-cloud.com`
+- Validate authentication tokens haven't expired
+
+**Document upload fails:**
+
+- Check Vespa Cloud storage quota and billing
+- Verify embedding format matches Vespa schema
+- Ensure stable internet connection for large uploads
+
+**Search returns no results:**
+
+- Confirm documents were successfully uploaded to Vespa
+- Check if embeddings were properly indexed
+- Verify query processing isn't failing locally
+
+### **MIXED (Local + Remote) Issues**
+
+**Chat features don't work:**
+
+- **LOCAL**: Verify document images are being generated locally
+- **REMOTE**: Check `GEMINI_API_KEY` is set correctly
+- **REMOTE**: Verify Gemini API quota and billing
+- **NETWORK**: Ensure images can be sent to Gemini API
+
+**Similarity maps missing:**
+
+- **LOCAL**: Confirm ColPali model loaded successfully
+- **LOCAL**: Check if similarity map generation completed
+- **REMOTE**: Verify Vespa returned similarity data
+- **BROWSER**: Clear browser cache for static files
+
+### Performance Tips
+
+**LOCAL Optimization:**
+
+- Use GPU acceleration for 5-10x faster model inference
+- Optimize batch sizes based on available memory
+- Use SSD storage for faster model loading
+- Consider quantized models for lower memory usage
+
+**REMOTE Optimization:**
+
+- Use Vespa's HNSW indexing for faster search
+- Optimize embedding dimensions vs accuracy tradeoff
+- Enable compression for faster network transfer
+- Use multiple Vespa instances for high availability
+
+**NETWORK Optimization:**
+
+- Process documents in batches to reduce upload overhead
+- Use compression for embedding transfer
+- Consider regional Vespa deployment for lower latency
+
+## 📄 License
+
+Apache-2.0
+
+## 🤝 Contributing
+
+1. Fork the repository
+2. Create a feature branch
+3. Make your changes
+4. Run tests and linting
+5. Submit a pull request
+
+## 📞 Support
+
+For issues and questions:
+
+- Check the troubleshooting section
+- Review Vespa and ColPali documentation
+- Open an issue on the repository

backend/cache.py

@@ -0,0 +1,26 @@
+from collections import OrderedDict
+
+
+# Initialize LRU Cache
+class LRUCache:
+    def __init__(self, max_size=20):
+        self.max_size = max_size
+        self.cache = OrderedDict()
+
+    def get(self, key):
+        if key in self.cache:
+            self.cache.move_to_end(key)
+            return self.cache[key]
+        return None
+
+    def set(self, key, value):
+        if key in self.cache:
+            self.cache.move_to_end(key)
+        else:
+            if len(self.cache) >= self.max_size:
+                self.cache.popitem(last=False)
+        self.cache[key] = value
+
+    def delete(self, key):
+        if key in self.cache:
+            del self.cache[key]

backend/colpali.py

@@ -0,0 +1,281 @@
+import torch
+from PIL import Image
+import numpy as np
+from typing import Generator, Tuple, List, Union, Dict
+from pathlib import Path
+import base64
+from io import BytesIO
+import re
+import io
+import matplotlib.cm as cm
+
+from colpali_engine.models import ColPali, ColPaliProcessor
+from colpali_engine.utils.torch_utils import get_torch_device
+from vidore_benchmark.interpretability.torch_utils import (
+    normalize_similarity_map_per_query_token,
+)
+from functools import lru_cache
+import logging
+
+
+class SimMapGenerator:
+    """
+    Generates similarity maps based on query embeddings and image patches using the ColPali model.
+    """
+
+    colormap = cm.get_cmap("viridis")  # Preload colormap for efficiency
+
+    def __init__(
+        self,
+        logger: logging.Logger,
+        model_name: str = "vidore/colpali-v1.2",
+        n_patch: int = 32,
+    ):
+        """
+        Initializes the SimMapGenerator class with a specified model and patch dimension.
+
+        Args:
+            model_name (str): The model name for loading the ColPali model.
+            n_patch (int): The number of patches per dimension.
+        """
+        self.model_name = model_name
+        self.n_patch = n_patch
+        self.device = get_torch_device("auto")
+        self.logger = logger
+        self.logger.info(f"Using device: {self.device}")
+        self.model, self.processor = self.load_model()
+
+    def load_model(self) -> Tuple[ColPali, ColPaliProcessor]:
+        """
+        Loads the ColPali model and processor.
+
+        Returns:
+            Tuple[ColPali, ColPaliProcessor]: Loaded model and processor.
+        """
+        model = ColPali.from_pretrained(
+            self.model_name,
+            torch_dtype=torch.bfloat16,  # Note that the embeddings created during feed were float32 -> binarized, yet setting this seem to produce the most similar results both locally (mps) and HF (Cuda)
+            device_map=self.device,
+        ).eval()
+
+        processor = ColPaliProcessor.from_pretrained(self.model_name)
+        return model, processor
+
+    def gen_similarity_maps(
+        self,
+        query: str,
+        query_embs: torch.Tensor,
+        token_idx_map: Dict[int, str],
+        images: List[Union[Path, str]],
+        vespa_sim_maps: List[Dict],
+    ) -> Generator[Tuple[int, str, str], None, None]:
+        """
+        Generates similarity maps for the provided images and query, and returns base64-encoded blended images.
+
+        Args:
+            query (str): The query string.
+            query_embs (torch.Tensor): Query embeddings tensor.
+            token_idx_map (dict): Mapping from indices to tokens.
+            images (List[Union[Path, str]]): List of image paths or base64-encoded strings.
+            vespa_sim_maps (List[Dict]): List of Vespa similarity maps.
+
+        Yields:
+            Tuple[int, str, str]: A tuple containing the image index, selected token, and base64-encoded image.
+        """
+        processed_images, original_images, original_sizes = [], [], []
+        for img in images:
+            img_pil = self._load_image(img)
+            original_images.append(img_pil.copy())
+            original_sizes.append(img_pil.size)
+            processed_images.append(img_pil)
+
+        vespa_sim_map_tensor = self._prepare_similarity_map_tensor(
+            query_embs, vespa_sim_maps
+        )
+        similarity_map_normalized = normalize_similarity_map_per_query_token(
+            vespa_sim_map_tensor
+        )
+
+        for idx, img in enumerate(original_images):
+            for token_idx, token in token_idx_map.items():
+                if self.should_filter_token(token):
+                    continue
+
+                sim_map = similarity_map_normalized[idx, token_idx, :, :]
+                blended_img_base64 = self._blend_image(
+                    img, sim_map, original_sizes[idx]
+                )
+                yield idx, token, token_idx, blended_img_base64
+
+    def _load_image(self, img: Union[Path, str]) -> Image:
+        """
+        Loads an image from a file path or a base64-encoded string.
+
+        Args:
+            img (Union[Path, str]): The image to load.
+
+        Returns:
+            Image: The loaded PIL image.
+        """
+        try:
+            if isinstance(img, Path):
+                return Image.open(img).convert("RGB")
+            elif isinstance(img, str):
+                return Image.open(BytesIO(base64.b64decode(img))).convert("RGB")
+        except Exception as e:
+            raise ValueError(f"Failed to load image: {e}")
+
+    def _prepare_similarity_map_tensor(
+        self, query_embs: torch.Tensor, vespa_sim_maps: List[Dict]
+    ) -> torch.Tensor:
+        """
+        Prepares a similarity map tensor from Vespa similarity maps.
+
+        Args:
+            query_embs (torch.Tensor): Query embeddings tensor.
+            vespa_sim_maps (List[Dict]): List of Vespa similarity maps.
+
+        Returns:
+            torch.Tensor: The prepared similarity map tensor.
+        """
+        vespa_sim_map_tensor = torch.zeros(
+            (len(vespa_sim_maps), query_embs.size(1), self.n_patch, self.n_patch)
+        )
+        for idx, vespa_sim_map in enumerate(vespa_sim_maps):
+            for cell in vespa_sim_map["quantized"]["cells"]:
+                patch = int(cell["address"]["patch"])
+                query_token = int(cell["address"]["querytoken"])
+                value = cell["value"]
+                if hasattr(self.processor, "image_seq_length"):
+                    image_seq_length = self.processor.image_seq_length
+                else:
+                    image_seq_length = 1024
+
+                if patch >= image_seq_length:
+                    continue
+                vespa_sim_map_tensor[
+                    idx,
+                    query_token,
+                    patch // self.n_patch,
+                    patch % self.n_patch,
+                ] = value
+        return vespa_sim_map_tensor
+
+    def _blend_image(
+        self, img: Image, sim_map: torch.Tensor, original_size: Tuple[int, int]
+    ) -> str:
+        """
+        Blends an image with a similarity map and encodes it to base64.
+
+        Args:
+            img (Image): The original image.
+            sim_map (torch.Tensor): The similarity map tensor.
+            original_size (Tuple[int, int]): The original size of the image.
+
+        Returns:
+            str: The base64-encoded blended image.
+        """
+        SCALING_FACTOR = 8
+        sim_map_resolution = (
+            max(32, int(original_size[0] / SCALING_FACTOR)),
+            max(32, int(original_size[1] / SCALING_FACTOR)),
+        )
+
+        sim_map_np = sim_map.cpu().float().numpy()
+        sim_map_img = Image.fromarray(sim_map_np).resize(
+            sim_map_resolution, resample=Image.BICUBIC
+        )
+        sim_map_resized_np = np.array(sim_map_img, dtype=np.float32)
+        sim_map_normalized = self._normalize_sim_map(sim_map_resized_np)
+
+        heatmap = self.colormap(sim_map_normalized)
+        heatmap_img = Image.fromarray((heatmap * 255).astype(np.uint8)).convert("RGBA")
+
+        buffer = io.BytesIO()
+        heatmap_img.save(buffer, format="PNG")
+        return base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+    @staticmethod
+    def _normalize_sim_map(sim_map: np.ndarray) -> np.ndarray:
+        """
+        Normalizes a similarity map to range [0, 1].
+
+        Args:
+            sim_map (np.ndarray): The similarity map.
+
+        Returns:
+            np.ndarray: The normalized similarity map.
+        """
+        sim_map_min, sim_map_max = sim_map.min(), sim_map.max()
+        if sim_map_max - sim_map_min > 1e-6:
+            return (sim_map - sim_map_min) / (sim_map_max - sim_map_min)
+        return np.zeros_like(sim_map)
+
+    @staticmethod
+    def should_filter_token(token: str) -> bool:
+        """
+        Determines if a token should be filtered out based on predefined patterns.
+
+        The function filters out tokens that:
+
+            - Start with '<' (e.g., '<bos>')
+            - Consist entirely of whitespace
+            - Are purely punctuation (excluding tokens that contain digits or start with '▁')
+            - Start with an underscore '_'
+            - Exactly match the word 'Question'
+            - Are exactly the single character '▁'
+
+        Output of test:
+            Token: '2'         | False
+            Token: '0'         | False
+            Token: '2'         | False
+            Token: '3'         | False
+            Token: '▁2'        | False
+            Token: '▁hi'       | False
+            Token: 'norwegian' | False
+            Token: 'unlisted'  | False
+            Token: '<bos>'     | True
+            Token: 'Question'  | True
+            Token: ':'         | True
+            Token: '<pad>'     | True
+            Token: '\n'        | True
+            Token: '▁'         | True
+            Token: '?'         | True
+            Token: ')'         | True
+            Token: '%'         | True
+            Token: '/)'        | True
+
+
+        Args:
+            token (str): The token to check.
+
+        Returns:
+            bool: True if the token should be filtered out, False otherwise.
+        """
+        pattern = re.compile(
+            r"^<.*$|^\s+$|^(?!.*\d)(?!▁)[^\w\s]+$|^_.*$|^Question$|^▁$"
+        )
+        return bool(pattern.match(token))
+
+    @lru_cache(maxsize=128)
+    def get_query_embeddings_and_token_map(
+        self, query: str
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Retrieves query embeddings and a token index map.
+
+        Args:
+            query (str): The query string.
+
+        Returns:
+            Tuple[torch.Tensor, dict]: Query embeddings and token index map.
+        """
+        inputs = self.processor.process_queries([query]).to(self.model.device)
+        with torch.no_grad():
+            q_emb = self.model(**inputs).to("cpu")[0]
+
+        query_tokens = self.processor.tokenizer.tokenize(
+            self.processor.decode(inputs.input_ids[0])
+        )
+        idx_to_token = {idx: token for idx, token in enumerate(query_tokens)}
+        return q_emb, idx_to_token

backend/modelmanager.py

@@ -0,0 +1,24 @@
+from .colpali import load_model
+
+
+class ModelManager:
+    _instance = None
+    model = None
+    processor = None
+    use_dummy_model = False
+
+    @staticmethod
+    def get_instance():
+        if ModelManager._instance is None:
+            ModelManager._instance = ModelManager()
+            if not ModelManager.use_dummy_model:
+                ModelManager._instance.initialize_model_and_processor()
+        return ModelManager._instance
+
+    def initialize_model_and_processor(self):
+        if self.model is None or self.processor is None:  # Ensure no reinitialization
+            self.model, self.processor, self.device = load_model()
+            if self.model is None or self.processor is None:
+                print("Failed to initialize model or processor at startup")
+            else:
+                print("Model and processor loaded at startup")

backend/stopwords.py

@@ -0,0 +1,18 @@
+import spacy
+import os
+
+# Download the model if it is not already present
+if not spacy.util.is_package("en_core_web_sm"):
+    spacy.cli.download("en_core_web_sm")
+nlp = spacy.load("en_core_web_sm")
+
+
+# It would be possible to remove bolding for stopwords without removing them from the query,
+# but that would require a java plugin which we didn't want to complicate this sample app with.
+def filter(text):
+    doc = nlp(text)
+    tokens = [token.text for token in doc if not token.is_stop]
+    if len(tokens) == 0:
+        # if we remove all the words we don't have a query at all, so use the original
+        return text
+    return " ".join(tokens)

backend/testquery.py

@@ -0,0 +1,3013 @@
+import torch
+
+token_to_idx = {
+    "<bos>": 0,
+    "Question": 1,
+    ":": 2,
+    "▁Percentage": 3,
+    "▁of": 4,
+    "▁non": 5,
+    "-": 6,
+    "fresh": 7,
+    "▁water": 8,
+    "▁as": 9,
+    "▁source": 10,
+    "?": 11,
+    "<pad>": 21,
+    "\n": 22,
+}
+idx_to_token = {v: k for k, v in token_to_idx.items()}
+q_embs = torch.tensor(
+    [
+        [
+            1.6547e-01,
+            -1.8838e-02,
+            1.0150e-01,
+            -2.2643e-02,
+            5.0652e-02,
+            3.2039e-02,
+            6.8322e-02,
+            2.7134e-02,
+            7.2197e-03,
+            -4.2341e-02,
+            -5.8006e-02,
+            -1.1389e-01,
+            7.0041e-02,
+            -6.8021e-02,
+            2.4681e-02,
+            5.3306e-02,
+            4.1714e-02,
+            6.2021e-02,
+            1.3488e-01,
+            3.4943e-02,
+            3.8032e-02,
+            -3.2724e-02,
+            -1.2960e-01,
+            1.1453e-02,
+            -2.6477e-02,
+            3.5219e-02,
+            -7.6606e-02,
+            2.2387e-01,
+            -3.4888e-02,
+            -4.0333e-02,
+            1.4128e-01,
+            4.2248e-02,
+            -1.2664e-01,
+            -7.8376e-02,
+            -2.0356e-02,
+            4.2198e-02,
+            -7.0776e-02,
+            1.3965e-02,
+            4.2442e-03,
+            7.1987e-02,
+            8.5172e-04,
+            -6.4878e-02,
+            -1.8954e-01,
+            -8.6171e-02,
+            9.1983e-02,
+            -9.3358e-02,
+            2.2704e-01,
+            1.3102e-02,
+            6.5327e-02,
+            2.4815e-02,
+            -1.4533e-01,
+            5.8823e-02,
+            -6.1434e-02,
+            5.2004e-02,
+            -8.4065e-02,
+            1.6298e-01,
+            8.1965e-02,
+            2.6553e-02,
+            -1.2377e-01,
+            -5.3495e-02,
+            -3.4537e-02,
+            5.1438e-02,
+            8.2665e-03,
+            7.9407e-02,
+            5.8799e-02,
+            -3.5538e-02,
+            1.9870e-01,
+            6.2459e-02,
+            1.6154e-01,
+            7.2921e-02,
+            -9.7275e-02,
+            3.0933e-02,
+            -1.0579e-02,
+            -1.4484e-01,
+            -4.8761e-02,
+            -5.3119e-02,
+            6.2644e-02,
+            2.2985e-02,
+            -2.1209e-01,
+            9.0963e-02,
+            -2.6955e-02,
+            -7.7520e-02,
+            1.2072e-01,
+            -1.9626e-02,
+            5.8813e-02,
+            -1.2730e-01,
+            1.5610e-01,
+            -1.6914e-01,
+            6.5033e-02,
+            8.5765e-02,
+            1.2701e-01,
+            6.5633e-02,
+            -1.0309e-01,
+            -8.0259e-02,
+            4.5913e-02,
+            -3.3277e-02,
+            1.9227e-01,
+            -2.3351e-02,
+            -8.0545e-02,
+            -9.8760e-03,
+            -4.1836e-02,
+            1.2041e-01,
+            8.1419e-02,
+            1.4848e-01,
+            7.2537e-02,
+            -4.7115e-03,
+            4.1489e-02,
+            4.3031e-02,
+            -1.3515e-01,
+            1.0383e-01,
+            -9.4411e-02,
+            2.8965e-02,
+            1.9185e-01,
+            -1.4600e-02,
+            -8.0910e-02,
+            -2.9022e-02,
+            -4.4347e-02,
+            -8.8980e-03,
+            2.8737e-03,
+            4.2124e-02,
+            -1.6609e-02,
+            4.2994e-02,
+            -7.2814e-02,
+            -2.9573e-02,
+            -1.2666e-01,
+            -4.3703e-02,
+            -7.2094e-02,
+            -2.7486e-02,
+        ],
+        [
+            -9.8454e-02,
+            -1.2954e-01,
+            6.5702e-02,
+            -9.0006e-03,
+            -1.0934e-01,
+            3.2155e-02,
+            -1.1444e-01,
+            -1.0309e-01,
+            5.7024e-02,
+            1.0124e-01,
+            2.0721e-02,
+            -5.2608e-03,
+            6.9916e-02,
+            1.8036e-02,
+            3.1653e-02,
+            2.1923e-02,
+            -9.2523e-02,
+            -1.8215e-02,
+            1.2974e-01,
+            -2.9632e-02,
+            -1.3854e-01,
+            2.5710e-02,
+            -1.1727e-03,
+            1.0245e-01,
+            -2.0731e-01,
+            4.3669e-02,
+            -7.0196e-02,
+            -1.6697e-01,
+            6.6050e-02,
+            9.9776e-02,
+            -1.2227e-01,
+            9.4000e-02,
+            1.1945e-01,
+            2.4611e-02,
+            -1.4073e-01,
+            9.3476e-02,
+            2.1170e-01,
+            -7.4522e-02,
+            -5.3362e-02,
+            -4.1198e-03,
+            8.3880e-02,
+            2.6590e-02,
+            -4.8489e-02,
+            -3.1279e-02,
+            -9.3401e-03,
+            -1.5945e-01,
+            -6.6368e-02,
+            7.5715e-02,
+            5.6884e-02,
+            1.2861e-01,
+            1.0073e-02,
+            -1.7185e-02,
+            1.1545e-01,
+            2.8725e-02,
+            -9.4969e-02,
+            -4.5517e-02,
+            3.1253e-02,
+            2.1135e-02,
+            -1.4505e-02,
+            9.0893e-02,
+            -6.2680e-02,
+            -7.2855e-02,
+            -1.1275e-01,
+            -1.8433e-01,
+            1.4693e-01,
+            4.0366e-02,
+            6.6879e-02,
+            -8.5653e-03,
+            1.0663e-01,
+            -1.2342e-01,
+            2.5350e-01,
+            -3.2227e-02,
+            9.9404e-02,
+            3.7340e-02,
+            4.5462e-04,
+            -2.3015e-01,
+            4.9006e-02,
+            1.0079e-01,
+            -4.7179e-02,
+            6.7642e-02,
+            -1.0833e-01,
+            1.0030e-01,
+            1.2838e-01,
+            -9.2911e-03,
+            1.2342e-01,
+            -4.8455e-02,
+            5.3904e-03,
+            4.5178e-02,
+            3.8961e-02,
+            1.3383e-01,
+            -1.2236e-02,
+            8.2026e-03,
+            4.3735e-03,
+            -7.1725e-02,
+            6.4360e-02,
+            1.0004e-01,
+            6.6840e-02,
+            6.5649e-02,
+            -1.3978e-02,
+            1.2810e-01,
+            4.4325e-03,
+            1.1136e-01,
+            -1.7329e-01,
+            -3.4472e-02,
+            -1.4066e-01,
+            1.5641e-02,
+            -3.3600e-02,
+            -7.6192e-02,
+            5.3085e-02,
+            -7.2859e-03,
+            2.8798e-02,
+            -3.3748e-02,
+            -7.7591e-02,
+            -4.0927e-02,
+            -3.6577e-02,
+            1.4012e-02,
+            -8.1780e-02,
+            -4.6315e-03,
+            -1.6508e-02,
+            -2.4506e-02,
+            1.5122e-01,
+            7.8270e-02,
+            8.6502e-02,
+            -2.4651e-02,
+            -1.0286e-01,
+            -1.2171e-02,
+            -7.9000e-02,
+            1.1161e-01,
+        ],
+        [
+            6.4320e-02,
+            -2.6815e-02,
+            1.2390e-01,
+            8.6642e-02,
+            5.3320e-02,
+            9.1074e-02,
+            -8.9753e-02,
+            1.6141e-02,
+            -5.0281e-02,
+            1.0177e-01,
+            4.5343e-02,
+            9.1281e-02,
+            -8.2592e-03,
+            -1.4100e-02,
+            -4.7048e-02,
+            -8.6034e-02,
+            -1.1608e-01,
+            -8.4754e-02,
+            1.0302e-01,
+            -1.2210e-02,
+            8.5147e-02,
+            1.3103e-01,
+            -3.3592e-03,
+            -1.4328e-02,
+            -4.6128e-03,
+            2.1401e-02,
+            -1.7464e-01,
+            -3.9111e-02,
+            2.0886e-02,
+            8.9284e-02,
+            1.3262e-01,
+            7.0918e-02,
+            -1.3693e-02,
+            -9.7673e-02,
+            -5.3411e-02,
+            -6.7563e-02,
+            2.3017e-02,
+            -3.4614e-02,
+            3.6464e-02,
+            -4.6408e-02,
+            1.4866e-01,
+            -2.1191e-01,
+            -6.4368e-02,
+            -3.0555e-02,
+            7.2177e-02,
+            2.4685e-02,
+            8.6115e-02,
+            -5.8688e-02,
+            -2.8230e-03,
+            1.1166e-01,
+            1.8380e-01,
+            3.6462e-02,
+            -1.4943e-02,
+            -1.4276e-01,
+            1.0795e-01,
+            1.9204e-02,
+            -6.5320e-02,
+            1.0561e-01,
+            -1.3642e-01,
+            -4.3444e-02,
+            -6.4725e-02,
+            -1.3094e-01,
+            1.9447e-02,
+            -1.3199e-01,
+            8.7880e-02,
+            5.8078e-02,
+            -2.5612e-02,
+            7.7620e-02,
+            -2.5044e-02,
+            1.0772e-02,
+            6.5417e-02,
+            8.2137e-02,
+            -2.8482e-02,
+            5.5003e-02,
+            -1.1163e-01,
+            1.7200e-03,
+            -1.0106e-01,
+            -1.8413e-02,
+            1.2838e-01,
+            -1.2991e-01,
+            -1.8546e-02,
+            1.0517e-01,
+            1.0279e-01,
+            -8.1887e-02,
+            1.0885e-01,
+            -1.0635e-01,
+            1.2035e-01,
+            1.1769e-01,
+            -2.8768e-02,
+            -3.3413e-02,
+            1.3779e-01,
+            1.4403e-02,
+            2.3429e-02,
+            -1.2761e-01,
+            7.2160e-02,
+            -1.0512e-01,
+            -1.7202e-02,
+            5.3549e-02,
+            5.5205e-02,
+            9.2863e-02,
+            -2.3728e-02,
+            -5.1368e-02,
+            -3.7719e-02,
+            -4.7308e-02,
+            -2.2489e-02,
+            4.5195e-02,
+            1.0398e-01,
+            -2.2197e-02,
+            -1.7208e-01,
+            7.4649e-02,
+            -7.7925e-02,
+            -6.4237e-02,
+            2.8195e-02,
+            2.2692e-01,
+            -9.7749e-02,
+            2.4283e-01,
+            -4.0124e-02,
+            1.8797e-02,
+            -6.1516e-02,
+            6.5331e-03,
+            1.3717e-01,
+            9.9761e-02,
+            5.4705e-02,
+            3.5325e-02,
+            1.9071e-01,
+            -6.1137e-02,
+            1.6656e-01,
+            3.4067e-02,
+        ],
+        [
+            1.0896e-01,
+            -9.4366e-03,
+            1.2956e-01,
+            7.8127e-02,
+            5.5422e-02,
+            3.9155e-02,
+            -5.8379e-03,
+            -4.4257e-02,
+            -7.5182e-02,
+            1.0452e-01,
+            4.4595e-02,
+            2.0972e-02,
+            1.1071e-01,
+            7.4710e-02,
+            -7.5500e-02,
+            -5.3393e-02,
+            -1.5478e-02,
+            -5.4455e-03,
+            5.6779e-02,
+            -6.3919e-02,
+            5.1792e-02,
+            1.2070e-01,
+            -5.3707e-02,
+            4.5715e-02,
+            -9.3062e-02,
+            3.0224e-02,
+            -1.5892e-01,
+            -2.1702e-02,
+            1.7942e-03,
+            1.4574e-01,
+            2.0721e-01,
+            -2.8224e-02,
+            -5.8104e-02,
+            -1.1645e-01,
+            -1.1515e-01,
+            -1.5202e-01,
+            -3.9751e-02,
+            6.0342e-02,
+            7.8182e-02,
+            -2.1132e-02,
+            9.6468e-02,
+            -5.3148e-02,
+            -3.0343e-02,
+            7.9363e-02,
+            1.0752e-01,
+            3.1086e-02,
+            1.9322e-02,
+            -1.1134e-01,
+            1.6342e-02,
+            3.0358e-02,
+            1.8543e-01,
+            5.5353e-02,
+            -8.6656e-02,
+            -1.5650e-01,
+            1.2087e-01,
+            -3.7852e-02,
+            -6.9116e-02,
+            5.9981e-03,
+            1.9205e-02,
+            -1.0314e-01,
+            -6.8082e-02,
+            -1.3078e-01,
+            3.8448e-02,
+            -9.2233e-02,
+            1.0965e-01,
+            6.6332e-02,
+            -2.5805e-02,
+            1.2299e-01,
+            2.8629e-02,
+            -4.4949e-02,
+            4.5560e-02,
+            1.0507e-01,
+            -1.0271e-01,
+            1.6237e-02,
+            -1.4555e-01,
+            -4.5335e-02,
+            -1.3477e-01,
+            1.0230e-02,
+            1.2380e-01,
+            -1.0681e-01,
+            1.4412e-02,
+            1.2396e-01,
+            8.5290e-02,
+            -2.5138e-02,
+            1.0191e-01,
+            -1.3413e-01,
+            8.5871e-02,
+            1.3389e-01,
+            -3.6357e-02,
+            -3.9740e-02,
+            2.1128e-01,
+            1.0263e-02,
+            2.5547e-02,
+            -7.0139e-02,
+            1.0178e-01,
+            -1.2729e-01,
+            -1.0717e-01,
+            -4.9394e-02,
+            7.7645e-02,
+            7.4589e-02,
+            -8.2835e-02,
+            -4.2227e-02,
+            -3.6417e-02,
+            -2.2900e-02,
+            -2.1010e-02,
+            2.7898e-02,
+            2.7314e-02,
+            2.2172e-02,
+            -7.1122e-02,
+            5.1570e-02,
+            2.1860e-02,
+            -3.5103e-03,
+            -5.4524e-02,
+            1.7485e-01,
+            -8.3810e-02,
+            2.3868e-01,
+            5.9468e-02,
+            -2.6706e-02,
+            -2.6617e-02,
+            3.3851e-02,
+            6.3651e-02,
+            1.0611e-01,
+            9.6252e-02,
+            -6.8701e-02,
+            1.8108e-01,
+            -1.0178e-01,
+            1.6935e-01,
+            5.9301e-02,
+        ],
+        [
+            1.3364e-01,
+            6.6797e-02,
+            1.0182e-01,
+            6.1569e-02,
+            -1.4169e-04,
+            1.1567e-01,
+            -3.1255e-03,
+            -8.9336e-02,
+            5.3206e-03,
+            1.7179e-01,
+            1.3974e-01,
+            -6.5797e-02,
+            1.2566e-01,
+            6.1290e-02,
+            -1.3671e-01,
+            -1.1113e-01,
+            1.6596e-01,
+            6.2694e-02,
+            -1.6573e-02,
+            -1.8648e-02,
+            2.1086e-02,
+            9.6320e-03,
+            -1.0017e-01,
+            2.3561e-02,
+            -6.2797e-02,
+            2.1202e-02,
+            -1.3333e-01,
+            1.4728e-01,
+            1.9824e-03,
+            1.6832e-01,
+            1.7138e-01,
+            4.6578e-02,
+            -1.2989e-01,
+            3.2167e-02,
+            -7.7753e-02,
+            -1.1548e-01,
+            4.3367e-02,
+            -4.5640e-02,
+            4.3583e-02,
+            -8.9393e-03,
+            5.3231e-02,
+            -4.8284e-02,
+            3.1841e-04,
+            6.6374e-02,
+            7.7363e-02,
+            4.0000e-02,
+            1.5414e-02,
+            -9.6156e-02,
+            1.1389e-01,
+            7.1118e-02,
+            3.0042e-02,
+            5.9752e-02,
+            -6.4565e-02,
+            -1.2914e-01,
+            1.1048e-01,
+            1.2409e-02,
+            -9.9385e-02,
+            1.8671e-02,
+            2.1383e-02,
+            3.7012e-03,
+            -1.1497e-01,
+            -6.8653e-02,
+            3.0582e-02,
+            -1.1567e-01,
+            1.4165e-01,
+            3.7493e-02,
+            -6.0779e-02,
+            1.0989e-01,
+            8.6001e-02,
+            -5.6139e-02,
+            2.0710e-02,
+            9.8577e-02,
+            -9.9427e-02,
+            5.8372e-02,
+            -1.3443e-01,
+            -1.3021e-02,
+            -1.3802e-01,
+            7.6053e-02,
+            1.2181e-01,
+            -8.7719e-02,
+            1.0967e-02,
+            1.3160e-01,
+            4.7032e-02,
+            -6.3351e-02,
+            7.1883e-02,
+            -9.7565e-02,
+            1.4424e-01,
+            1.2353e-01,
+            -6.1527e-02,
+            2.4263e-02,
+            2.9356e-01,
+            6.2813e-02,
+            -4.5265e-03,
+            -1.0213e-01,
+            1.4227e-02,
+            -7.9267e-02,
+            -1.0845e-01,
+            -2.0014e-02,
+            2.8542e-02,
+            9.7207e-02,
+            -2.5234e-02,
+            -7.3668e-02,
+            -3.0084e-02,
+            -1.2958e-02,
+            -3.9597e-02,
+            -7.2243e-02,
+            5.3054e-02,
+            3.1470e-03,
+            -1.9800e-02,
+            1.3476e-01,
+            -1.6873e-02,
+            5.2286e-02,
+            2.0254e-02,
+            1.0554e-01,
+            -3.0395e-02,
+            8.6349e-02,
+            7.6580e-02,
+            -3.0139e-02,
+            -4.8131e-02,
+            -4.5770e-02,
+            1.5154e-01,
+            1.1276e-01,
+            5.7244e-02,
+            8.0574e-02,
+            1.5610e-01,
+            -1.5523e-01,
+            1.0428e-01,
+            5.7947e-02,
+        ],
+        [
+            7.0440e-02,
+            1.4300e-01,
+            6.0559e-02,
+            1.9177e-02,
+            -9.0313e-02,
+            1.7104e-01,
+            -5.1137e-02,
+            -1.0229e-01,
+            -2.8831e-02,
+            1.5385e-01,
+            -8.4017e-03,
+            -4.9185e-03,
+            1.0820e-01,
+            1.0022e-01,
+            -1.9284e-01,
+            -2.7293e-02,
+            4.9526e-02,
+            5.5152e-02,
+            -6.7003e-02,
+            -7.0313e-03,
+            -3.3208e-02,
+            1.3815e-02,
+            -6.0694e-02,
+            6.2041e-02,
+            -3.2288e-02,
+            1.1629e-01,
+            -7.5270e-02,
+            2.1824e-01,
+            -2.5215e-03,
+            1.8179e-01,
+            1.5514e-01,
+            1.0494e-01,
+            -1.2390e-01,
+            -1.2241e-03,
+            6.1079e-04,
+            -5.8730e-02,
+            6.1313e-02,
+            -7.8853e-02,
+            6.0292e-02,
+            -9.1497e-04,
+            7.9087e-02,
+            -1.8246e-02,
+            5.0215e-03,
+            3.5083e-02,
+            5.9616e-02,
+            5.9520e-02,
+            6.0224e-02,
+            -1.3079e-01,
+            1.6500e-01,
+            4.4308e-03,
+            4.2712e-02,
+            5.5916e-02,
+            -5.4616e-02,
+            -8.5617e-02,
+            1.1235e-01,
+            6.5911e-03,
+            -6.1463e-02,
+            3.7832e-02,
+            3.4189e-02,
+            -1.1295e-02,
+            -8.0972e-02,
+            -1.0051e-02,
+            -2.6856e-02,
+            -7.9570e-02,
+            1.2776e-01,
+            6.5826e-02,
+            -3.1759e-02,
+            9.6016e-02,
+            6.7249e-02,
+            -4.5115e-02,
+            6.3695e-03,
+            1.2092e-01,
+            -1.3821e-01,
+            -9.7066e-02,
+            -1.5063e-02,
+            2.4618e-02,
+            -1.9589e-01,
+            5.8625e-02,
+            1.7886e-01,
+            -6.3740e-02,
+            -1.7241e-02,
+            7.3394e-02,
+            5.8903e-02,
+            -1.6557e-02,
+            -1.9226e-02,
+            -1.0912e-01,
+            8.7902e-02,
+            6.4426e-02,
+            4.8019e-02,
+            8.1223e-02,
+            2.9888e-01,
+            8.7458e-02,
+            4.4167e-02,
+            -1.3228e-01,
+            5.9629e-02,
+            -1.0696e-01,
+            -1.4102e-01,
+            -5.2509e-02,
+            -1.9981e-02,
+            1.6788e-01,
+            9.7499e-02,
+            -5.4125e-02,
+            -8.2383e-02,
+            -6.3908e-02,
+            -6.8830e-02,
+            -1.2622e-01,
+            3.1651e-02,
+            4.4592e-02,
+            -1.3325e-02,
+            1.1260e-01,
+            -3.9567e-02,
+            6.9631e-03,
+            1.4943e-01,
+            8.6930e-02,
+            -3.6171e-03,
+            -5.6886e-02,
+            -8.3102e-03,
+            -2.6001e-02,
+            -1.5187e-02,
+            -1.8835e-02,
+            2.3583e-02,
+            9.5520e-02,
+            -3.4944e-02,
+            4.5537e-02,
+            6.1444e-02,
+            -1.7165e-01,
+            1.0230e-01,
+            2.0319e-02,
+        ],
+        [
+            7.7065e-02,
+            1.4466e-01,
+            1.0796e-01,
+            9.7362e-03,
+            -9.3062e-02,
+            2.0065e-01,
+            -9.3982e-03,
+            -1.2871e-01,
+            -4.6724e-02,
+            1.6573e-01,
+            -1.7444e-02,
+            -3.1211e-02,
+            9.7404e-02,
+            9.6222e-02,
+            -1.4772e-01,
+            -5.6838e-02,
+            9.6276e-02,
+            7.7819e-02,
+            -1.1884e-01,
+            -4.5898e-02,
+            -7.3665e-02,
+            2.1005e-02,
+            -2.8032e-02,
+            7.0900e-02,
+            -7.7625e-02,
+            7.0269e-02,
+            -7.0747e-02,
+            2.5605e-01,
+            1.9427e-04,
+            1.4614e-01,
+            1.6723e-01,
+            1.0016e-01,
+            -7.8961e-02,
+            1.3307e-02,
+            -1.1207e-02,
+            -7.1492e-02,
+            7.5561e-02,
+            -1.3290e-02,
+            4.4527e-02,
+            -2.4224e-02,
+            8.7846e-02,
+            4.7492e-02,
+            -7.0398e-02,
+            -2.6663e-02,
+            4.9730e-02,
+            6.6743e-02,
+            3.1060e-02,
+            -7.8352e-02,
+            1.0100e-01,
+            6.4963e-02,
+            1.6746e-02,
+            3.0324e-02,
+            -2.6583e-02,
+            -7.8028e-02,
+            7.0180e-02,
+            3.8924e-02,
+            -8.5861e-02,
+            -3.3792e-02,
+            6.1542e-02,
+            1.6180e-02,
+            -7.7865e-02,
+            3.9551e-02,
+            2.9772e-02,
+            -7.0824e-02,
+            1.5235e-01,
+            4.8718e-02,
+            1.5973e-03,
+            8.7719e-02,
+            7.7414e-02,
+            -6.4385e-02,
+            -6.4330e-02,
+            1.3965e-01,
+            -1.6355e-01,
+            -6.5261e-02,
+            -6.2693e-02,
+            4.9435e-02,
+            -1.5245e-01,
+            6.6557e-02,
+            1.5213e-01,
+            -8.2073e-02,
+            1.4664e-02,
+            8.4507e-02,
+            3.0684e-02,
+            -8.7932e-02,
+            1.9927e-02,
+            -7.1788e-02,
+            9.4965e-02,
+            3.9220e-02,
+            4.5944e-02,
+            8.7249e-02,
+            3.3315e-01,
+            5.9872e-02,
+            2.6362e-02,
+            -1.7888e-01,
+            -1.6042e-02,
+            -9.4593e-02,
+            -1.7915e-01,
+            -2.9888e-02,
+            -4.3776e-02,
+            1.1388e-01,
+            2.3778e-02,
+            -4.0233e-02,
+            -4.7893e-02,
+            -4.4371e-02,
+            -2.7491e-02,
+            -9.6716e-02,
+            -2.0120e-02,
+            5.6864e-02,
+            1.8953e-02,
+            1.2741e-01,
+            -5.3045e-02,
+            3.2240e-02,
+            1.4479e-01,
+            9.5315e-02,
+            -2.7717e-02,
+            -5.7349e-02,
+            3.3824e-03,
+            -3.5642e-02,
+            -1.6905e-02,
+            -4.5765e-02,
+            1.1481e-02,
+            4.2545e-02,
+            1.2632e-02,
+            4.5401e-02,
+            7.7769e-02,
+            -1.7107e-01,
+            7.7265e-02,
+            1.3922e-02,
+        ],
+        [
+            -2.5666e-02,
+            9.8268e-02,
+            2.1571e-01,
+            6.8094e-02,
+            -1.0440e-01,
+            1.4299e-01,
+            6.5751e-02,
+            -5.0693e-02,
+            -5.3796e-02,
+            1.5239e-01,
+            2.9566e-02,
+            -7.9492e-02,
+            9.3274e-02,
+            7.6112e-02,
+            -1.8187e-02,
+            -1.1190e-01,
+            9.7962e-02,
+            -2.8204e-02,
+            -8.1216e-02,
+            -5.5618e-02,
+            -6.2378e-02,
+            4.4238e-02,
+            -1.6572e-02,
+            -3.4035e-02,
+            -8.8068e-02,
+            1.4164e-02,
+            -2.6908e-02,
+            1.9650e-01,
+            6.8845e-03,
+            8.7550e-02,
+            1.7410e-01,
+            1.0088e-01,
+            -1.1340e-02,
+            3.2057e-04,
+            -5.5130e-02,
+            -2.5234e-02,
+            8.2460e-02,
+            -6.0768e-02,
+            1.2448e-01,
+            6.8736e-02,
+            4.6176e-02,
+            1.0866e-01,
+            4.9560e-02,
+            -5.8322e-02,
+            4.4106e-02,
+            2.0739e-02,
+            -9.0032e-02,
+            -5.8815e-02,
+            7.8127e-03,
+            1.7999e-01,
+            1.2519e-01,
+            7.1377e-02,
+            9.3219e-02,
+            -1.3311e-01,
+            6.0305e-02,
+            5.9400e-02,
+            -1.6119e-01,
+            5.4173e-02,
+            -5.8663e-02,
+            -4.8149e-02,
+            1.6295e-02,
+            -6.9787e-02,
+            5.7512e-03,
+            -8.9745e-03,
+            1.3492e-01,
+            3.0659e-02,
+            -6.8611e-02,
+            2.1200e-02,
+            8.5522e-02,
+            -4.3482e-02,
+            -8.4601e-02,
+            1.4191e-01,
+            -1.5514e-01,
+            -5.8989e-03,
+            -4.5591e-02,
+            6.4905e-02,
+            -1.3198e-01,
+            1.3764e-01,
+            9.5549e-02,
+            -9.4689e-02,
+            -1.9705e-02,
+            2.1147e-01,
+            -9.8519e-03,
+            -7.7839e-02,
+            6.1447e-02,
+            -6.4708e-02,
+            3.1579e-03,
+            7.6588e-02,
+            -1.2452e-01,
+            -6.1076e-02,
+            2.5150e-01,
+            2.3101e-02,
+            -1.3632e-02,
+            -8.5695e-02,
+            -8.5841e-02,
+            -1.3152e-01,
+            -1.5294e-01,
+            -4.4509e-03,
+            8.6619e-02,
+            -1.0974e-02,
+            -3.9592e-02,
+            -3.0472e-02,
+            -1.4011e-01,
+            8.6485e-03,
+            -1.3633e-02,
+            -1.7940e-02,
+            3.3016e-02,
+            -5.7245e-02,
+            1.0200e-01,
+            1.2807e-01,
+            1.5249e-02,
+            -1.2197e-02,
+            -1.6867e-02,
+            3.4516e-02,
+            -9.0908e-02,
+            -2.6167e-02,
+            2.2975e-01,
+            4.2693e-02,
+            2.5415e-03,
+            7.1921e-03,
+            1.2855e-01,
+            5.5747e-03,
+            -2.7843e-02,
+            2.4283e-02,
+            -1.3484e-02,
+            -1.5948e-01,
+            2.1127e-02,
+            3.8017e-02,
+        ],
+        [
+            5.8410e-02,
+            3.7501e-02,
+            1.7189e-01,
+            4.5894e-02,
+            -6.7739e-02,
+            1.1421e-01,
+            4.6982e-02,
+            -1.2963e-01,
+            -4.2804e-02,
+            1.2137e-01,
+            1.0761e-01,
+            -7.2615e-02,
+            1.1811e-01,
+            1.1291e-01,
+            -1.6041e-01,
+            -6.6820e-02,
+            1.9071e-01,
+            -3.6201e-02,
+            -1.0659e-01,
+            -3.3226e-02,
+            -2.6535e-02,
+            7.8536e-02,
+            -3.2975e-02,
+            1.3015e-02,
+            -1.0903e-01,
+            -1.2502e-02,
+            -7.0142e-02,
+            1.6872e-01,
+            -3.7569e-02,
+            2.2090e-01,
+            1.5620e-01,
+            9.5620e-02,
+            -7.9306e-02,
+            7.8558e-02,
+            -6.6709e-02,
+            -6.6446e-02,
+            3.1198e-02,
+            -4.8808e-02,
+            6.7798e-02,
+            -4.9524e-02,
+            1.2496e-01,
+            3.1661e-02,
+            4.0690e-02,
+            -1.1341e-02,
+            2.9852e-03,
+            3.6166e-02,
+            -1.1723e-01,
+            -8.2746e-02,
+            1.2874e-01,
+            1.0894e-01,
+            7.6219e-02,
+            5.5564e-02,
+            4.3571e-02,
+            -1.4745e-01,
+            5.1090e-02,
+            -6.8233e-04,
+            -1.7107e-01,
+            -2.5612e-02,
+            1.9185e-02,
+            3.4466e-03,
+            -2.5362e-02,
+            -3.4807e-02,
+            8.1003e-02,
+            -2.1582e-02,
+            1.1581e-01,
+            -5.5891e-04,
+            4.2455e-02,
+            6.3277e-02,
+            7.9868e-02,
+            -2.0369e-02,
+            -7.6229e-02,
+            1.2525e-01,
+            -1.4795e-01,
+            5.8798e-03,
+            -2.8584e-02,
+            -3.2508e-04,
+            -5.5666e-02,
+            7.1217e-02,
+            1.3122e-01,
+            -1.8671e-02,
+            -6.6402e-02,
+            8.1978e-02,
+            -5.5194e-02,
+            -1.7259e-02,
+            7.1373e-02,
+            -2.4721e-02,
+            1.0176e-01,
+            1.7447e-02,
+            -8.2872e-02,
+            4.7707e-02,
+            2.9338e-01,
+            3.0481e-02,
+            1.3437e-02,
+            -1.1888e-01,
+            -6.1243e-02,
+            -8.7197e-02,
+            -1.1085e-01,
+            3.2847e-03,
+            1.0611e-02,
+            1.2908e-01,
+            -2.9064e-02,
+            4.5522e-03,
+            -7.8045e-02,
+            -1.4551e-02,
+            -3.3958e-02,
+            -1.3626e-01,
+            1.0859e-01,
+            -7.3438e-03,
+            9.5942e-02,
+            1.4176e-01,
+            7.3675e-02,
+            -2.3085e-03,
+            -1.1689e-02,
+            6.1119e-02,
+            -9.3085e-02,
+            -1.9003e-03,
+            1.9013e-01,
+            -2.5321e-02,
+            -7.0616e-02,
+            -3.7145e-02,
+            1.2449e-01,
+            5.3496e-02,
+            6.4021e-02,
+            2.2668e-02,
+            7.3348e-02,
+            -1.7717e-01,
+            8.7206e-02,
+            -9.0490e-03,
+        ],
+        [
+            1.0227e-01,
+            4.4796e-02,
+            4.0916e-02,
+            1.2557e-01,
+            6.2699e-02,
+            8.8194e-02,
+            1.0194e-02,
+            -1.8886e-01,
+            -2.6260e-02,
+            1.8107e-01,
+            1.1127e-01,
+            -4.9960e-03,
+            8.2291e-02,
+            6.0932e-02,
+            -1.3735e-01,
+            -8.1350e-02,
+            1.2948e-01,
+            1.0171e-01,
+            -7.5950e-02,
+            4.7169e-02,
+            3.5525e-03,
+            -1.4790e-02,
+            -3.9063e-02,
+            -1.9754e-02,
+            -4.6450e-02,
+            4.9071e-03,
+            -3.5914e-02,
+            1.7644e-01,
+            4.1860e-02,
+            1.0857e-01,
+            8.5257e-02,
+            1.0368e-01,
+            -1.7223e-01,
+            4.6316e-02,
+            -3.9480e-02,
+            -9.5817e-02,
+            1.5301e-03,
+            -4.5519e-02,
+            1.2380e-01,
+            -3.0402e-02,
+            4.9047e-02,
+            3.4955e-02,
+            -1.3018e-02,
+            7.4193e-02,
+            -3.0426e-02,
+            1.8414e-02,
+            -1.6606e-02,
+            -1.6301e-01,
+            1.5784e-01,
+            1.3759e-01,
+            3.1906e-02,
+            2.9389e-02,
+            -9.4501e-02,
+            -1.9930e-01,
+            1.3336e-01,
+            3.0685e-02,
+            -7.3809e-02,
+            6.2165e-02,
+            7.3050e-02,
+            3.2870e-02,
+            -1.2857e-01,
+            -7.5806e-04,
+            1.3985e-01,
+            -5.5108e-02,
+            9.2220e-02,
+            7.2490e-02,
+            -2.6251e-02,
+            3.4508e-02,
+            -3.1215e-02,
+            -8.1111e-02,
+            1.4316e-02,
+            1.2390e-01,
+            -9.6291e-03,
+            6.4214e-02,
+            -9.6013e-02,
+            7.5809e-02,
+            -1.5899e-01,
+            8.6961e-02,
+            1.5239e-03,
+            -5.6370e-02,
+            -4.3367e-02,
+            1.5813e-02,
+            9.7189e-04,
+            -5.2946e-02,
+            4.5950e-02,
+            -1.0028e-01,
+            1.1108e-01,
+            9.0491e-03,
+            -3.5540e-02,
+            -1.2020e-02,
+            2.2980e-01,
+            2.7125e-02,
+            -2.4191e-02,
+            -1.1363e-01,
+            -1.0109e-01,
+            -1.4781e-01,
+            -4.7656e-02,
+            3.9481e-02,
+            5.4198e-02,
+            8.2908e-02,
+            1.4034e-02,
+            -1.8492e-02,
+            -6.8612e-02,
+            -4.8741e-02,
+            1.1223e-02,
+            3.9220e-02,
+            5.4551e-04,
+            6.5554e-02,
+            4.3087e-02,
+            1.4678e-01,
+            -4.4496e-02,
+            6.2379e-02,
+            4.7876e-02,
+            7.0156e-02,
+            -6.4684e-02,
+            6.1076e-02,
+            1.4685e-01,
+            -6.3639e-02,
+            -8.7487e-02,
+            -1.3756e-02,
+            1.2724e-01,
+            1.7404e-01,
+            6.7980e-02,
+            6.8036e-02,
+            1.9786e-01,
+            -9.2910e-02,
+            1.9158e-01,
+            4.2686e-02,
+        ],
+        [
+            4.7626e-02,
+            9.3338e-02,
+            9.8020e-02,
+            9.2408e-02,
+            1.8267e-02,
+            1.9572e-02,
+            1.1056e-01,
+            -1.2639e-01,
+            -3.1999e-02,
+            1.3731e-01,
+            1.4826e-01,
+            -6.7136e-02,
+            8.9095e-02,
+            1.9683e-01,
+            -6.5839e-02,
+            -1.0708e-01,
+            1.4414e-01,
+            4.7389e-02,
+            -7.9921e-02,
+            2.7793e-02,
+            9.7240e-02,
+            3.8000e-03,
+            -4.2359e-02,
+            5.9870e-02,
+            -8.2957e-02,
+            -2.1429e-02,
+            -5.1302e-02,
+            7.6202e-02,
+            3.5426e-02,
+            8.1509e-02,
+            1.0849e-01,
+            2.1857e-01,
+            -9.1445e-02,
+            1.1869e-01,
+            -1.1225e-02,
+            -1.6530e-01,
+            2.9216e-02,
+            -7.2792e-02,
+            2.9544e-02,
+            5.3102e-02,
+            2.7746e-02,
+            1.8145e-01,
+            4.4050e-02,
+            8.2549e-02,
+            2.3603e-02,
+            1.1494e-02,
+            -9.1516e-02,
+            -6.2333e-02,
+            1.0970e-01,
+            1.4544e-01,
+            -2.3549e-02,
+            3.8348e-02,
+            -8.8504e-03,
+            -1.6986e-01,
+            5.3215e-02,
+            4.0085e-02,
+            -1.4667e-01,
+            8.6047e-02,
+            4.7751e-02,
+            -9.0753e-03,
+            -8.0938e-02,
+            7.1735e-02,
+            1.1711e-01,
+            -4.2525e-02,
+            9.4244e-02,
+            6.2260e-02,
+            -8.4798e-02,
+            -6.2766e-02,
+            6.1025e-02,
+            -7.7480e-02,
+            -5.3010e-02,
+            1.0182e-01,
+            -7.1081e-03,
+            1.4029e-01,
+            -8.2788e-02,
+            3.0182e-02,
+            -1.4279e-01,
+            4.7882e-02,
+            3.2736e-02,
+            1.0058e-02,
+            -1.2673e-02,
+            3.2723e-02,
+            5.7034e-02,
+            -2.0343e-02,
+            -6.3037e-03,
+            6.8914e-03,
+            2.8643e-02,
+            8.2739e-02,
+            -2.5124e-02,
+            -7.0611e-02,
+            2.7142e-01,
+            -1.6370e-02,
+            1.1240e-02,
+            -1.5776e-01,
+            -8.8755e-02,
+            -1.2389e-01,
+            -2.5785e-02,
+            -3.1258e-02,
+            2.2023e-02,
+            8.5497e-02,
+            6.8360e-02,
+            9.6662e-03,
+            -1.2843e-01,
+            -7.4122e-02,
+            7.1105e-02,
+            5.8334e-02,
+            -1.8609e-02,
+            7.2775e-02,
+            -1.1696e-03,
+            1.3165e-01,
+            -4.7326e-02,
+            6.8737e-02,
+            -2.3211e-02,
+            5.5425e-02,
+            -6.8077e-02,
+            -7.3743e-04,
+            1.6215e-01,
+            -4.9965e-02,
+            -6.8493e-02,
+            -6.8358e-03,
+            9.9518e-02,
+            1.5627e-01,
+            1.3100e-01,
+            9.0073e-03,
+            1.8023e-01,
+            -2.7992e-02,
+            1.3087e-01,
+            7.5672e-02,
+        ],
+        [
+            9.0390e-03,
+            1.3137e-01,
+            4.6586e-02,
+            1.1836e-01,
+            1.1406e-01,
+            -7.3655e-02,
+            -1.4340e-02,
+            -1.4285e-01,
+            2.0426e-02,
+            1.3532e-01,
+            1.3353e-01,
+            -4.7242e-02,
+            2.6449e-02,
+            2.0707e-02,
+            -7.5049e-02,
+            -4.4098e-02,
+            1.4599e-01,
+            1.1299e-01,
+            -2.2097e-02,
+            8.4679e-02,
+            5.8946e-02,
+            2.1213e-02,
+            -5.8343e-03,
+            -4.5902e-02,
+            -6.1018e-02,
+            4.4368e-02,
+            -5.7242e-02,
+            1.4734e-01,
+            -2.6191e-02,
+            8.6116e-02,
+            4.6004e-02,
+            9.4532e-02,
+            -1.7424e-01,
+            1.3697e-01,
+            4.6544e-02,
+            -1.1718e-01,
+            8.2567e-02,
+            -3.1494e-02,
+            6.2073e-02,
+            -4.6328e-02,
+            1.0782e-02,
+            6.6071e-02,
+            9.6953e-03,
+            2.0415e-02,
+            -7.0939e-02,
+            -3.1883e-02,
+            -4.9926e-02,
+            -1.2694e-01,
+            1.6351e-01,
+            1.1670e-01,
+            3.8149e-02,
+            7.2657e-02,
+            -1.3930e-01,
+            -1.6682e-01,
+            1.2471e-01,
+            1.9094e-02,
+            -6.1618e-02,
+            5.8628e-02,
+            1.8804e-03,
+            -9.8104e-03,
+            -1.8998e-01,
+            1.4179e-02,
+            1.4856e-01,
+            -4.2194e-02,
+            5.0934e-02,
+            7.4201e-02,
+            1.5546e-02,
+            2.3192e-02,
+            -5.5886e-02,
+            -4.3618e-02,
+            3.6677e-02,
+            1.2750e-01,
+            2.1240e-02,
+            9.5849e-02,
+            -1.2006e-01,
+            6.0424e-02,
+            -1.6516e-01,
+            7.8015e-02,
+            -7.6026e-02,
+            -1.6711e-02,
+            -3.3572e-02,
+            4.3704e-02,
+            1.2728e-02,
+            -6.9630e-02,
+            4.1065e-02,
+            -8.6589e-02,
+            1.1581e-01,
+            2.2666e-02,
+            -5.1476e-02,
+            2.2756e-02,
+            1.9309e-01,
+            7.0417e-02,
+            8.9253e-03,
+            -1.3546e-01,
+            -7.7783e-02,
+            -1.5058e-01,
+            -4.7874e-02,
+            3.3498e-02,
+            6.4585e-02,
+            7.8633e-02,
+            1.0162e-01,
+            4.5821e-02,
+            -2.8424e-02,
+            4.4524e-02,
+            4.6649e-02,
+            4.5805e-02,
+            8.9723e-02,
+            3.7597e-02,
+            -6.8016e-03,
+            1.1727e-01,
+            -6.6387e-02,
+            3.4958e-02,
+            4.0205e-02,
+            7.5854e-02,
+            -5.1322e-02,
+            6.0060e-02,
+            1.4631e-01,
+            -1.1550e-01,
+            -1.5388e-01,
+            6.0975e-03,
+            1.2013e-01,
+            2.1669e-01,
+            4.8245e-02,
+            8.4634e-02,
+            1.4934e-01,
+            -1.2099e-01,
+            1.6673e-01,
+            6.2539e-02,
+        ],
+        [
+            1.6222e-01,
+            7.5644e-03,
+            1.0667e-01,
+            8.3084e-02,
+            5.1888e-02,
+            8.7477e-02,
+            -1.1550e-01,
+            -4.7037e-02,
+            -3.9418e-02,
+            8.0987e-02,
+            1.2812e-01,
+            -6.6527e-02,
+            6.5405e-02,
+            5.2723e-03,
+            -6.2940e-02,
+            -5.8805e-02,
+            2.8526e-02,
+            4.0631e-02,
+            9.8513e-02,
+            -2.8448e-02,
+            9.2361e-02,
+            6.9724e-02,
+            -1.1792e-01,
+            -1.4806e-02,
+            -6.6172e-02,
+            5.9468e-02,
+            -1.8325e-01,
+            3.2851e-02,
+            -2.7042e-02,
+            1.7569e-01,
+            1.9061e-01,
+            4.1868e-02,
+            -1.4116e-01,
+            -1.6780e-02,
+            -6.2881e-02,
+            -9.4655e-02,
+            5.1488e-02,
+            -3.6501e-03,
+            5.6754e-02,
+            -6.3285e-03,
+            9.2234e-02,
+            -1.0761e-01,
+            -5.5392e-02,
+            7.5249e-02,
+            8.7539e-02,
+            4.3377e-02,
+            3.2473e-03,
+            -1.0990e-01,
+            5.2611e-02,
+            3.4955e-02,
+            1.1722e-01,
+            4.0339e-02,
+            -9.6605e-02,
+            -9.3660e-02,
+            9.6494e-02,
+            -1.4790e-02,
+            -5.7770e-02,
+            4.3192e-02,
+            -8.1741e-02,
+            -4.0909e-02,
+            -9.5212e-02,
+            -1.2459e-01,
+            3.6241e-02,
+            -1.2878e-01,
+            1.1544e-01,
+            3.6392e-03,
+            -7.7986e-03,
+            8.2339e-02,
+            1.4319e-01,
+            -5.8374e-02,
+            1.2483e-01,
+            5.6826e-02,
+            -6.2649e-02,
+            4.5884e-02,
+            -1.8261e-01,
+            -9.1633e-02,
+            -1.9414e-01,
+            2.8376e-02,
+            9.5654e-02,
+            -1.3532e-01,
+            2.1669e-03,
+            7.9960e-02,
+            5.0917e-02,
+            -5.4874e-02,
+            7.6823e-02,
+            -1.3932e-01,
+            1.2431e-01,
+            1.3001e-01,
+            -6.0069e-02,
+            2.6470e-02,
+            2.2378e-01,
+            2.7335e-02,
+            1.6487e-02,
+            -9.6473e-02,
+            8.3110e-02,
+            -8.8776e-02,
+            -6.2299e-02,
+            4.0821e-02,
+            6.5279e-02,
+            8.4030e-02,
+            -4.9298e-02,
+            -3.7441e-02,
+            -4.9233e-02,
+            -1.9428e-03,
+            -1.9763e-02,
+            -9.1249e-02,
+            1.1593e-01,
+            6.6291e-03,
+            -1.0686e-01,
+            5.9383e-02,
+            -3.1088e-02,
+            -1.9844e-03,
+            2.5891e-02,
+            1.3471e-01,
+            -3.5624e-02,
+            1.2647e-01,
+            6.7167e-02,
+            -1.8682e-02,
+            -2.4806e-02,
+            -1.1453e-02,
+            9.6992e-02,
+            1.2253e-01,
+            5.6707e-02,
+            1.0860e-02,
+            2.4775e-01,
+            -8.8641e-02,
+            1.1375e-01,
+            6.9383e-02,
+        ],
+        [
+            1.4724e-01,
+            3.7511e-02,
+            1.1238e-01,
+            6.3809e-02,
+            1.8174e-03,
+            8.4854e-02,
+            -7.8965e-02,
+            -3.1542e-02,
+            -3.7643e-02,
+            1.0136e-01,
+            1.5837e-01,
+            -7.5701e-02,
+            1.0313e-01,
+            2.2795e-02,
+            -1.1227e-01,
+            -4.3783e-02,
+            4.6561e-02,
+            4.1225e-02,
+            7.4273e-02,
+            -4.8200e-02,
+            2.7666e-02,
+            5.1417e-02,
+            -1.3928e-01,
+            2.0020e-02,
+            -9.3673e-02,
+            5.1440e-02,
+            -2.0076e-01,
+            1.0715e-01,
+            -3.0323e-02,
+            1.7241e-01,
+            1.9730e-01,
+            7.7685e-02,
+            -1.5949e-01,
+            4.8495e-02,
+            -8.0068e-02,
+            -7.1965e-02,
+            8.7222e-02,
+            -4.8383e-02,
+            7.0627e-02,
+            3.9791e-03,
+            7.0469e-02,
+            -8.6697e-02,
+            -3.7149e-02,
+            7.9965e-02,
+            8.8967e-02,
+            5.3325e-02,
+            -5.6073e-03,
+            -1.1336e-01,
+            1.0338e-01,
+            2.3824e-02,
+            6.9393e-02,
+            6.1018e-02,
+            -7.2880e-02,
+            -6.6328e-02,
+            8.9040e-02,
+            -8.8268e-03,
+            -7.5720e-02,
+            3.7544e-02,
+            -8.3740e-02,
+            -4.7977e-02,
+            -1.2425e-01,
+            -1.0304e-01,
+            -3.7850e-03,
+            -1.1540e-01,
+            1.0803e-01,
+            2.5964e-02,
+            -1.0915e-02,
+            1.1758e-01,
+            1.6561e-01,
+            -4.5171e-02,
+            1.1272e-01,
+            6.0721e-02,
+            -9.6027e-02,
+            6.0908e-02,
+            -1.6545e-01,
+            -6.5444e-02,
+            -1.6802e-01,
+            2.5850e-02,
+            1.2434e-01,
+            -1.2358e-01,
+            -5.2647e-04,
+            1.2329e-01,
+            6.3928e-02,
+            -5.5478e-02,
+            5.4642e-02,
+            -1.0372e-01,
+            1.0955e-01,
+            1.1810e-01,
+            -7.7028e-02,
+            5.0302e-02,
+            2.7576e-01,
+            6.4704e-02,
+            3.7824e-03,
+            -8.3016e-02,
+            6.7355e-02,
+            -4.9660e-02,
+            -1.2677e-01,
+            3.4498e-02,
+            2.7027e-02,
+            9.1476e-02,
+            -5.1170e-02,
+            -3.4786e-02,
+            -4.4356e-02,
+            6.2440e-03,
+            -3.5473e-02,
+            -1.1076e-01,
+            1.0014e-01,
+            -1.5927e-02,
+            -7.9539e-02,
+            1.1542e-01,
+            -3.6775e-03,
+            2.5437e-02,
+            2.1185e-02,
+            1.0472e-01,
+            -4.8937e-03,
+            4.5876e-02,
+            5.3122e-02,
+            -1.9815e-02,
+            -1.8126e-02,
+            -2.3358e-02,
+            1.2271e-01,
+            9.8860e-02,
+            4.5235e-02,
+            6.2841e-02,
+            1.6294e-01,
+            -1.0909e-01,
+            7.2012e-02,
+            5.3389e-02,
+        ],
+        [
+            1.2605e-01,
+            5.1723e-02,
+            1.2152e-01,
+            5.8291e-02,
+            -3.8019e-02,
+            8.8117e-02,
+            -4.1927e-02,
+            -2.3022e-02,
+            -3.1158e-02,
+            1.2518e-01,
+            1.4279e-01,
+            -7.1922e-02,
+            1.0867e-01,
+            3.3116e-02,
+            -1.3669e-01,
+            -3.2413e-02,
+            9.0301e-02,
+            5.9203e-02,
+            -7.5021e-03,
+            -6.9837e-02,
+            -1.2463e-03,
+            1.8891e-02,
+            -1.1878e-01,
+            6.1658e-02,
+            -1.2649e-01,
+            4.2450e-02,
+            -1.6298e-01,
+            1.5450e-01,
+            -2.6732e-02,
+            1.7891e-01,
+            2.1083e-01,
+            1.0432e-01,
+            -1.3396e-01,
+            6.9416e-02,
+            -7.3193e-02,
+            -6.8281e-02,
+            9.9454e-02,
+            -4.4528e-02,
+            7.7992e-02,
+            -2.6336e-03,
+            6.9279e-02,
+            -3.8836e-02,
+            -2.8872e-02,
+            5.6951e-02,
+            8.7154e-02,
+            6.1652e-02,
+            -1.8049e-02,
+            -1.2096e-01,
+            1.3820e-01,
+            5.1705e-02,
+            3.4052e-02,
+            1.1524e-01,
+            -4.1216e-02,
+            -8.5670e-02,
+            7.1571e-02,
+            -1.6001e-02,
+            -8.8427e-02,
+            3.1357e-02,
+            -5.4897e-02,
+            -3.1463e-02,
+            -1.2070e-01,
+            -5.7460e-02,
+            -2.1424e-02,
+            -1.2160e-01,
+            1.0026e-01,
+            4.7475e-02,
+            -1.4578e-02,
+            1.3247e-01,
+            1.5272e-01,
+            -2.9788e-02,
+            8.2947e-02,
+            7.9830e-02,
+            -1.1615e-01,
+            7.2949e-02,
+            -1.2386e-01,
+            -5.3969e-02,
+            -1.3578e-01,
+            4.0461e-02,
+            1.2467e-01,
+            -9.4575e-02,
+            6.2182e-04,
+            1.3108e-01,
+            6.5035e-02,
+            -5.7784e-02,
+            4.2513e-02,
+            -7.8118e-02,
+            1.1274e-01,
+            8.9333e-02,
+            -4.9991e-02,
+            7.4639e-02,
+            3.0876e-01,
+            7.3974e-02,
+            -1.4039e-02,
+            -7.4526e-02,
+            6.5200e-02,
+            -3.6421e-02,
+            -1.6730e-01,
+            4.1713e-02,
+            -1.5316e-02,
+            1.0056e-01,
+            -5.0370e-02,
+            -4.4531e-02,
+            -5.0233e-02,
+            3.2636e-02,
+            -1.7609e-02,
+            -1.1354e-01,
+            8.1922e-02,
+            -2.0907e-02,
+            -1.3264e-02,
+            1.5115e-01,
+            1.2777e-02,
+            3.8261e-02,
+            2.1829e-02,
+            5.5874e-02,
+            6.8772e-03,
+            -2.7772e-02,
+            6.7979e-02,
+            -3.9099e-02,
+            -2.8619e-02,
+            -5.6936e-02,
+            1.3366e-01,
+            9.4269e-02,
+            4.4852e-02,
+            7.6881e-02,
+            9.6787e-02,
+            -1.2885e-01,
+            7.9722e-02,
+            6.4991e-02,
+        ],
+        [
+            1.0575e-01,
+            6.4565e-02,
+            1.2479e-01,
+            5.5616e-02,
+            -5.5869e-02,
+            8.9005e-02,
+            -1.9612e-02,
+            -3.3032e-02,
+            -3.9719e-02,
+            1.3884e-01,
+            1.2010e-01,
+            -6.3224e-02,
+            1.0893e-01,
+            4.4172e-02,
+            -1.4283e-01,
+            -3.8972e-02,
+            1.1470e-01,
+            7.7830e-02,
+            -4.6450e-02,
+            -8.7246e-02,
+            -5.7215e-03,
+            1.7476e-02,
+            -9.0388e-02,
+            7.4607e-02,
+            -1.3105e-01,
+            2.9371e-02,
+            -1.2681e-01,
+            1.6014e-01,
+            -1.3655e-02,
+            1.8492e-01,
+            1.9841e-01,
+            1.0942e-01,
+            -1.1261e-01,
+            6.8326e-02,
+            -5.1130e-02,
+            -7.1296e-02,
+            9.3895e-02,
+            -3.1097e-02,
+            8.7043e-02,
+            -1.2087e-02,
+            7.3054e-02,
+            -1.0602e-02,
+            -3.6347e-02,
+            3.1324e-02,
+            7.5686e-02,
+            5.3221e-02,
+            -3.1793e-02,
+            -1.3818e-01,
+            1.4972e-01,
+            9.0217e-02,
+            2.6365e-02,
+            1.1733e-01,
+            -2.6107e-02,
+            -1.1933e-01,
+            6.5575e-02,
+            -1.6907e-02,
+            -1.0186e-01,
+            3.3719e-02,
+            -3.4021e-02,
+            -1.0597e-02,
+            -1.0478e-01,
+            -3.4765e-02,
+            -3.0666e-03,
+            -1.2035e-01,
+            1.1053e-01,
+            5.2212e-02,
+            -1.8638e-02,
+            1.2226e-01,
+            1.3614e-01,
+            -4.3806e-02,
+            4.7251e-02,
+            9.2144e-02,
+            -1.2505e-01,
+            8.3020e-02,
+            -1.1203e-01,
+            -4.9654e-02,
+            -1.1939e-01,
+            4.4126e-02,
+            1.1188e-01,
+            -7.8910e-02,
+            -4.1917e-03,
+            1.3279e-01,
+            5.0291e-02,
+            -6.1770e-02,
+            5.1914e-02,
+            -6.8796e-02,
+            1.2148e-01,
+            7.6837e-02,
+            -4.0098e-02,
+            8.3547e-02,
+            3.0367e-01,
+            6.9357e-02,
+            -2.5710e-02,
+            -7.2644e-02,
+            5.2612e-02,
+            -5.3110e-02,
+            -1.7360e-01,
+            5.2791e-02,
+            -1.6490e-02,
+            9.9930e-02,
+            -4.8174e-02,
+            -3.0230e-02,
+            -6.0898e-02,
+            3.4064e-02,
+            8.2446e-03,
+            -1.0747e-01,
+            6.9406e-02,
+            -9.2400e-03,
+            2.1804e-02,
+            1.5491e-01,
+            2.0964e-02,
+            3.4691e-02,
+            3.5648e-02,
+            4.4363e-02,
+            -3.2605e-03,
+            -5.9988e-02,
+            9.3354e-02,
+            -7.4183e-02,
+            -3.5824e-02,
+            -7.1163e-02,
+            1.4776e-01,
+            1.0179e-01,
+            4.1685e-02,
+            6.7698e-02,
+            8.4053e-02,
+            -1.4797e-01,
+            9.9115e-02,
+            7.2061e-02,
+        ],
+        [
+            8.3987e-02,
+            6.9951e-02,
+            1.3728e-01,
+            6.4616e-02,
+            -5.5591e-02,
+            8.8007e-02,
+            -1.7926e-02,
+            -4.8462e-02,
+            -4.8522e-02,
+            1.5932e-01,
+            1.0766e-01,
+            -4.9589e-02,
+            1.0511e-01,
+            4.4010e-02,
+            -1.3653e-01,
+            -4.1380e-02,
+            1.3629e-01,
+            9.2007e-02,
+            -5.2961e-02,
+            -7.1865e-02,
+            -2.8329e-03,
+            1.3417e-02,
+            -5.2565e-02,
+            6.2258e-02,
+            -1.2135e-01,
+            2.7531e-02,
+            -9.4902e-02,
+            1.6567e-01,
+            -1.6613e-02,
+            1.8443e-01,
+            1.6290e-01,
+            1.1109e-01,
+            -1.1476e-01,
+            6.5613e-02,
+            -1.7283e-02,
+            -8.7120e-02,
+            1.0862e-01,
+            -2.5325e-02,
+            8.9025e-02,
+            -1.8290e-02,
+            7.2979e-02,
+            1.1878e-04,
+            -4.1007e-02,
+            1.9679e-02,
+            6.1887e-02,
+            4.0949e-02,
+            -3.6016e-02,
+            -1.5569e-01,
+            1.4839e-01,
+            1.1372e-01,
+            1.8360e-02,
+            9.3030e-02,
+            -1.8070e-02,
+            -1.5578e-01,
+            6.1180e-02,
+            -1.9569e-02,
+            -1.1221e-01,
+            4.0037e-02,
+            -2.3712e-02,
+            -1.0311e-02,
+            -1.0530e-01,
+            -1.2916e-02,
+            3.1546e-02,
+            -1.1647e-01,
+            1.2109e-01,
+            5.8714e-02,
+            -1.9454e-02,
+            9.9438e-02,
+            1.1088e-01,
+            -7.4443e-02,
+            1.6456e-02,
+            8.7619e-02,
+            -1.1460e-01,
+            9.7142e-02,
+            -1.1083e-01,
+            -5.2367e-02,
+            -1.2157e-01,
+            4.6784e-02,
+            8.5968e-02,
+            -7.2488e-02,
+            -1.2658e-02,
+            1.4136e-01,
+            1.7934e-02,
+            -6.4075e-02,
+            5.3809e-02,
+            -5.8649e-02,
+            1.2410e-01,
+            6.7307e-02,
+            -5.3157e-02,
+            8.5616e-02,
+            2.8266e-01,
+            6.4775e-02,
+            -2.9330e-02,
+            -8.7733e-02,
+            1.3974e-02,
+            -7.9092e-02,
+            -1.6691e-01,
+            5.9037e-02,
+            4.1784e-03,
+            9.8219e-02,
+            -3.6970e-02,
+            8.6754e-03,
+            -5.5234e-02,
+            3.4627e-02,
+            3.5126e-02,
+            -9.5508e-02,
+            6.4539e-02,
+            -9.1506e-04,
+            2.9144e-02,
+            1.3975e-01,
+            1.5472e-02,
+            3.9015e-02,
+            4.6270e-02,
+            5.4053e-02,
+            -1.2525e-02,
+            -6.6168e-02,
+            1.2463e-01,
+            -1.0793e-01,
+            -4.0659e-02,
+            -8.1224e-02,
+            1.5702e-01,
+            1.2323e-01,
+            3.3473e-02,
+            6.1162e-02,
+            8.3375e-02,
+            -1.7863e-01,
+            1.1536e-01,
+            6.9906e-02,
+        ],
+        [
+            6.2508e-02,
+            8.5713e-02,
+            1.3566e-01,
+            8.5148e-02,
+            -2.2642e-02,
+            5.8889e-02,
+            -3.3654e-02,
+            -6.6948e-02,
+            -5.1541e-02,
+            1.7082e-01,
+            1.1639e-01,
+            -3.9750e-02,
+            8.5338e-02,
+            2.8182e-02,
+            -1.1108e-01,
+            -3.5316e-02,
+            1.4888e-01,
+            1.0210e-01,
+            -3.4537e-02,
+            -1.1950e-02,
+            5.0482e-03,
+            1.8051e-02,
+            -2.2067e-02,
+            1.8682e-02,
+            -9.7924e-02,
+            2.6326e-02,
+            -7.6128e-02,
+            1.7675e-01,
+            -1.6092e-02,
+            1.6355e-01,
+            1.1133e-01,
+            1.1036e-01,
+            -1.4070e-01,
+            6.9723e-02,
+            2.0482e-02,
+            -1.0421e-01,
+            1.2920e-01,
+            -3.1272e-02,
+            7.6070e-02,
+            -3.1384e-02,
+            6.7374e-02,
+            -8.3764e-05,
+            -2.1820e-02,
+            1.4192e-02,
+            2.3975e-02,
+            2.3384e-02,
+            -2.3589e-02,
+            -1.7196e-01,
+            1.4986e-01,
+            1.1465e-01,
+            2.0120e-02,
+            4.8034e-02,
+            -3.9982e-02,
+            -1.8704e-01,
+            6.2222e-02,
+            -1.4799e-02,
+            -1.0783e-01,
+            5.4098e-02,
+            -1.2046e-02,
+            -1.5444e-02,
+            -1.4001e-01,
+            1.5944e-02,
+            7.9476e-02,
+            -9.7107e-02,
+            1.0629e-01,
+            7.1731e-02,
+            -8.1614e-03,
+            6.2800e-02,
+            5.5151e-02,
+            -1.0514e-01,
+            1.7916e-02,
+            8.1229e-02,
+            -7.8761e-02,
+            1.1222e-01,
+            -1.1032e-01,
+            -4.8039e-02,
+            -1.5340e-01,
+            6.0481e-02,
+            3.2276e-02,
+            -6.2797e-02,
+            -3.3934e-02,
+            1.3062e-01,
+            -6.3068e-03,
+            -6.5936e-02,
+            4.2361e-02,
+            -4.6102e-02,
+            1.2669e-01,
+            6.1909e-02,
+            -5.7771e-02,
+            7.8385e-02,
+            2.4836e-01,
+            7.5433e-02,
+            -2.0151e-02,
+            -1.2117e-01,
+            -4.1874e-02,
+            -1.1751e-01,
+            -1.4471e-01,
+            6.3096e-02,
+            2.9419e-02,
+            1.0021e-01,
+            -1.2180e-02,
+            5.0896e-02,
+            -4.0094e-02,
+            4.1059e-02,
+            8.0853e-02,
+            -7.4614e-02,
+            7.5243e-02,
+            1.4928e-02,
+            1.1477e-02,
+            1.2385e-01,
+            -1.4719e-02,
+            5.3751e-02,
+            5.2591e-02,
+            7.2013e-02,
+            -1.6333e-02,
+            -5.0256e-02,
+            1.4199e-01,
+            -1.3433e-01,
+            -7.1574e-02,
+            -7.4449e-02,
+            1.6136e-01,
+            1.5659e-01,
+            2.1466e-02,
+            6.9775e-02,
+            1.0352e-01,
+            -1.9641e-01,
+            1.3156e-01,
+            5.5903e-02,
+        ],
+        [
+            4.4988e-02,
+            1.0390e-01,
+            1.1266e-01,
+            1.1661e-01,
+            2.1212e-02,
+            1.4641e-02,
+            -3.1484e-02,
+            -8.3443e-02,
+            -4.4991e-02,
+            1.5687e-01,
+            1.3158e-01,
+            -3.3434e-02,
+            6.5104e-02,
+            2.0808e-02,
+            -7.0680e-02,
+            -4.2508e-02,
+            1.4304e-01,
+            9.8483e-02,
+            -2.6904e-02,
+            3.2902e-02,
+            1.1152e-02,
+            2.1118e-02,
+            -1.4842e-02,
+            -2.2356e-02,
+            -8.6942e-02,
+            2.6609e-02,
+            -6.2533e-02,
+            1.7495e-01,
+            -1.0246e-02,
+            1.2697e-01,
+            7.9345e-02,
+            1.0189e-01,
+            -1.7178e-01,
+            7.5006e-02,
+            4.5051e-02,
+            -1.0753e-01,
+            1.2829e-01,
+            -3.3422e-02,
+            7.8522e-02,
+            -4.5604e-02,
+            5.2822e-02,
+            8.3934e-03,
+            -1.4676e-02,
+            1.0482e-02,
+            -6.9389e-03,
+            2.2131e-03,
+            -2.0089e-02,
+            -1.7261e-01,
+            1.5468e-01,
+            1.1537e-01,
+            2.7489e-02,
+            2.9948e-02,
+            -7.4234e-02,
+            -1.9946e-01,
+            7.4632e-02,
+            2.2249e-03,
+            -9.7061e-02,
+            6.4988e-02,
+            -1.1239e-02,
+            -2.6152e-02,
+            -1.6496e-01,
+            3.0305e-02,
+            1.0491e-01,
+            -6.7692e-02,
+            7.3323e-02,
+            7.0258e-02,
+            1.2020e-02,
+            4.1457e-02,
+            3.7126e-03,
+            -1.0465e-01,
+            2.7035e-02,
+            8.1376e-02,
+            -4.0890e-02,
+            1.0791e-01,
+            -1.2141e-01,
+            -1.6070e-02,
+            -1.7111e-01,
+            7.4218e-02,
+            -2.4698e-02,
+            -3.7496e-02,
+            -4.5666e-02,
+            1.0478e-01,
+            -1.7289e-03,
+            -7.1823e-02,
+            2.7480e-02,
+            -5.5305e-02,
+            1.2323e-01,
+            5.3008e-02,
+            -5.6903e-02,
+            6.6087e-02,
+            2.1523e-01,
+            8.2657e-02,
+            -9.3275e-03,
+            -1.3791e-01,
+            -6.7856e-02,
+            -1.4959e-01,
+            -1.1916e-01,
+            6.2014e-02,
+            3.8735e-02,
+            9.7171e-02,
+            2.6820e-02,
+            7.4336e-02,
+            -2.9922e-02,
+            5.3645e-02,
+            1.0981e-01,
+            -3.7215e-02,
+            9.4254e-02,
+            4.2902e-02,
+            -2.0969e-03,
+            1.1838e-01,
+            -5.4015e-02,
+            5.3799e-02,
+            4.8568e-02,
+            8.0494e-02,
+            -2.1410e-02,
+            -2.0629e-02,
+            1.3477e-01,
+            -1.4177e-01,
+            -1.1144e-01,
+            -5.1781e-02,
+            1.5780e-01,
+            1.9357e-01,
+            2.5627e-02,
+            7.5402e-02,
+            1.1982e-01,
+            -1.8961e-01,
+            1.4646e-01,
+            4.2665e-02,
+        ],
+        [
+            5.1689e-02,
+            1.0737e-01,
+            1.1035e-01,
+            1.1308e-01,
+            2.0363e-02,
+            1.9703e-02,
+            -3.2107e-02,
+            -9.0838e-02,
+            -3.7066e-02,
+            1.5839e-01,
+            1.3202e-01,
+            -3.9661e-02,
+            6.3263e-02,
+            1.4509e-02,
+            -6.0145e-02,
+            -4.1459e-02,
+            1.4670e-01,
+            9.4701e-02,
+            -2.7710e-02,
+            3.4281e-02,
+            1.2373e-02,
+            2.0087e-02,
+            -2.0821e-02,
+            -1.6274e-02,
+            -7.9908e-02,
+            2.6209e-02,
+            -7.0731e-02,
+            1.7365e-01,
+            -8.5357e-03,
+            1.2784e-01,
+            7.9591e-02,
+            9.7291e-02,
+            -1.6958e-01,
+            7.5774e-02,
+            4.4993e-02,
+            -1.0499e-01,
+            1.2133e-01,
+            -4.2564e-02,
+            7.7296e-02,
+            -3.9049e-02,
+            5.4639e-02,
+            3.7831e-03,
+            -1.2108e-02,
+            1.7782e-02,
+            -1.8517e-03,
+            1.1160e-02,
+            -1.9819e-02,
+            -1.7390e-01,
+            1.5853e-01,
+            1.1225e-01,
+            2.6839e-02,
+            3.1746e-02,
+            -8.1013e-02,
+            -1.9950e-01,
+            7.7933e-02,
+            4.3883e-03,
+            -9.1775e-02,
+            6.3761e-02,
+            -2.0220e-02,
+            -2.0873e-02,
+            -1.6981e-01,
+            3.2587e-02,
+            1.0188e-01,
+            -7.2338e-02,
+            7.9739e-02,
+            7.2336e-02,
+            1.2712e-02,
+            4.6304e-02,
+            1.1977e-02,
+            -9.9069e-02,
+            3.2361e-02,
+            7.4647e-02,
+            -3.2075e-02,
+            1.0834e-01,
+            -1.2421e-01,
+            -1.4408e-02,
+            -1.7591e-01,
+            7.6520e-02,
+            -2.1893e-02,
+            -4.0591e-02,
+            -3.9790e-02,
+            1.0072e-01,
+            2.1907e-03,
+            -7.5756e-02,
+            3.4002e-02,
+            -6.2571e-02,
+            1.2714e-01,
+            5.2011e-02,
+            -5.6154e-02,
+            6.2876e-02,
+            2.0992e-01,
+            8.9547e-02,
+            -1.2875e-02,
+            -1.3528e-01,
+            -6.6275e-02,
+            -1.4322e-01,
+            -1.0560e-01,
+            5.9555e-02,
+            4.4426e-02,
+            1.0047e-01,
+            3.3182e-02,
+            7.2001e-02,
+            -3.6514e-02,
+            5.0708e-02,
+            1.0564e-01,
+            -4.1174e-02,
+            1.0378e-01,
+            4.2651e-02,
+            -2.6353e-03,
+            1.1964e-01,
+            -6.2808e-02,
+            5.9555e-02,
+            5.4234e-02,
+            8.3802e-02,
+            -1.2534e-02,
+            -1.7949e-02,
+            1.3403e-01,
+            -1.3682e-01,
+            -1.1048e-01,
+            -5.1154e-02,
+            1.5967e-01,
+            1.9743e-01,
+            3.2240e-02,
+            7.4049e-02,
+            1.2060e-01,
+            -1.9087e-01,
+            1.3942e-01,
+            4.1168e-02,
+        ],
+        [
+            7.2669e-02,
+            1.0343e-01,
+            1.1966e-01,
+            9.3777e-02,
+            -4.3358e-03,
+            5.0720e-02,
+            -3.8923e-02,
+            -9.1626e-02,
+            -3.0056e-02,
+            1.6863e-01,
+            1.2676e-01,
+            -5.3184e-02,
+            7.4056e-02,
+            1.5517e-02,
+            -7.5657e-02,
+            -3.7712e-02,
+            1.5280e-01,
+            9.0719e-02,
+            -3.1311e-02,
+            1.7038e-02,
+            1.0062e-02,
+            1.4759e-02,
+            -3.4438e-02,
+            1.5206e-02,
+            -7.7897e-02,
+            2.7431e-02,
+            -9.0602e-02,
+            1.7164e-01,
+            -9.9721e-03,
+            1.4776e-01,
+            1.0475e-01,
+            9.0160e-02,
+            -1.5270e-01,
+            7.4838e-02,
+            2.4463e-02,
+            -1.0180e-01,
+            1.1501e-01,
+            -5.0178e-02,
+            7.2532e-02,
+            -1.8656e-02,
+            6.9202e-02,
+            -9.6694e-03,
+            -1.8219e-02,
+            3.5159e-02,
+            2.5262e-02,
+            3.4003e-02,
+            -1.4053e-02,
+            -1.7146e-01,
+            1.6299e-01,
+            1.0628e-01,
+            2.3901e-02,
+            4.2174e-02,
+            -7.0690e-02,
+            -1.8996e-01,
+            7.4249e-02,
+            -2.3634e-03,
+            -8.9371e-02,
+            5.7919e-02,
+            -3.2949e-02,
+            -2.4951e-03,
+            -1.6144e-01,
+            3.1264e-02,
+            8.4885e-02,
+            -9.7903e-02,
+            1.0806e-01,
+            7.3583e-02,
+            1.4389e-03,
+            6.4343e-02,
+            5.4305e-02,
+            -9.2823e-02,
+            3.4699e-02,
+            6.5887e-02,
+            -4.1142e-02,
+            1.1057e-01,
+            -1.2077e-01,
+            -3.4178e-02,
+            -1.7040e-01,
+            7.2546e-02,
+            1.9399e-02,
+            -6.3051e-02,
+            -3.0156e-02,
+            1.1501e-01,
+            3.9456e-03,
+            -7.8346e-02,
+            5.0548e-02,
+            -6.3681e-02,
+            1.3261e-01,
+            6.1531e-02,
+            -5.1743e-02,
+            6.4762e-02,
+            2.3264e-01,
+            9.4896e-02,
+            -2.0536e-02,
+            -1.1807e-01,
+            -4.9545e-02,
+            -1.1460e-01,
+            -1.0079e-01,
+            5.7836e-02,
+            4.4144e-02,
+            1.1039e-01,
+            2.1311e-02,
+            4.8694e-02,
+            -4.7126e-02,
+            3.7200e-02,
+            8.3108e-02,
+            -7.0069e-02,
+            9.7451e-02,
+            2.0983e-02,
+            -4.1628e-04,
+            1.2463e-01,
+            -5.0752e-02,
+            7.0680e-02,
+            6.3588e-02,
+            8.6767e-02,
+            6.3361e-04,
+            -2.5236e-02,
+            1.3658e-01,
+            -1.1925e-01,
+            -8.3861e-02,
+            -6.3728e-02,
+            1.6767e-01,
+            1.7967e-01,
+            3.3894e-02,
+            7.2830e-02,
+            1.1590e-01,
+            -1.9359e-01,
+            1.2254e-01,
+            4.7767e-02,
+        ],
+        [
+            8.4189e-02,
+            9.3299e-02,
+            1.1532e-01,
+            8.4670e-02,
+            -9.7533e-03,
+            6.4423e-02,
+            -4.3742e-02,
+            -9.0334e-02,
+            -1.9307e-02,
+            1.7434e-01,
+            1.2522e-01,
+            -6.4064e-02,
+            7.7841e-02,
+            2.5594e-02,
+            -9.4764e-02,
+            -3.8088e-02,
+            1.5148e-01,
+            8.7002e-02,
+            -3.8634e-02,
+            1.6073e-02,
+            1.0152e-02,
+            8.9035e-03,
+            -3.9583e-02,
+            3.3614e-02,
+            -7.4457e-02,
+            2.6379e-02,
+            -9.8885e-02,
+            1.6510e-01,
+            -1.1651e-02,
+            1.5700e-01,
+            1.3055e-01,
+            7.7680e-02,
+            -1.3690e-01,
+            7.0143e-02,
+            5.4565e-03,
+            -1.0524e-01,
+            1.1096e-01,
+            -4.6091e-02,
+            6.8167e-02,
+            -5.6453e-03,
+            7.2094e-02,
+            -1.8229e-02,
+            -2.5539e-02,
+            3.9328e-02,
+            4.3699e-02,
+            4.4686e-02,
+            -1.2072e-02,
+            -1.6738e-01,
+            1.6301e-01,
+            1.0474e-01,
+            1.4918e-02,
+            5.4164e-02,
+            -7.2489e-02,
+            -1.7993e-01,
+            7.7292e-02,
+            -2.8187e-03,
+            -8.9506e-02,
+            5.6049e-02,
+            -3.1474e-02,
+            4.7455e-03,
+            -1.5338e-01,
+            2.2827e-02,
+            6.7546e-02,
+            -1.1578e-01,
+            1.2026e-01,
+            8.2049e-02,
+            -1.1262e-02,
+            7.4682e-02,
+            7.3023e-02,
+            -8.6490e-02,
+            3.1955e-02,
+            6.6658e-02,
+            -4.6417e-02,
+            1.0408e-01,
+            -1.1210e-01,
+            -3.3793e-02,
+            -1.6235e-01,
+            7.3353e-02,
+            4.5024e-02,
+            -7.6526e-02,
+            -3.0443e-02,
+            1.2222e-01,
+            1.1513e-02,
+            -8.5835e-02,
+            5.6084e-02,
+            -6.5380e-02,
+            1.3535e-01,
+            7.2861e-02,
+            -4.8121e-02,
+            6.6781e-02,
+            2.5041e-01,
+            9.3076e-02,
+            -2.0334e-02,
+            -1.0416e-01,
+            -3.9805e-02,
+            -1.0051e-01,
+            -9.2133e-02,
+            6.3271e-02,
+            3.4768e-02,
+            1.1286e-01,
+            1.3493e-02,
+            3.0643e-02,
+            -5.5867e-02,
+            2.4630e-02,
+            7.8038e-02,
+            -7.5769e-02,
+            9.0636e-02,
+            1.8840e-02,
+            1.3688e-02,
+            1.2780e-01,
+            -4.3306e-02,
+            7.6513e-02,
+            5.6990e-02,
+            8.7927e-02,
+            1.6438e-03,
+            -1.0833e-02,
+            1.3722e-01,
+            -9.4365e-02,
+            -7.4534e-02,
+            -6.2697e-02,
+            1.6868e-01,
+            1.6850e-01,
+            4.3229e-02,
+            7.4762e-02,
+            1.1225e-01,
+            -1.9559e-01,
+            1.1981e-01,
+            5.0025e-02,
+        ],
+        [
+            1.0233e-01,
+            6.7397e-02,
+            1.0132e-01,
+            1.1631e-01,
+            4.3957e-02,
+            5.7017e-02,
+            -4.0769e-02,
+            -1.2329e-01,
+            -6.1235e-03,
+            1.3821e-01,
+            1.3135e-01,
+            -7.2640e-02,
+            6.7380e-02,
+            4.7518e-02,
+            -1.0261e-01,
+            -8.4497e-02,
+            1.1074e-01,
+            7.9965e-02,
+            4.3208e-03,
+            -7.9881e-03,
+            3.0558e-02,
+            3.8130e-02,
+            -7.0863e-02,
+            2.6484e-02,
+            -8.1795e-02,
+            3.9067e-02,
+            -1.5670e-01,
+            1.2963e-01,
+            4.3868e-03,
+            1.3081e-01,
+            1.7308e-01,
+            5.3891e-02,
+            -1.2683e-01,
+            6.0787e-02,
+            -7.8336e-02,
+            -1.0347e-01,
+            5.8472e-02,
+            -2.7212e-02,
+            7.1385e-02,
+            -6.8297e-03,
+            8.7485e-02,
+            -7.1364e-02,
+            -1.9182e-02,
+            8.7217e-02,
+            8.2944e-02,
+            3.4400e-02,
+            -8.5778e-03,
+            -1.1191e-01,
+            1.4871e-01,
+            1.0602e-01,
+            5.5060e-02,
+            4.0554e-02,
+            -1.0835e-01,
+            -1.4749e-01,
+            1.1336e-01,
+            8.1897e-03,
+            -5.1693e-02,
+            4.9600e-02,
+            -3.6377e-02,
+            -1.6864e-03,
+            -1.4149e-01,
+            -3.8906e-02,
+            5.3481e-02,
+            -1.1160e-01,
+            1.3861e-01,
+            6.3040e-02,
+            -2.2296e-02,
+            8.1680e-02,
+            7.9567e-02,
+            -6.1083e-02,
+            6.7573e-02,
+            9.8196e-02,
+            -7.0956e-02,
+            8.7717e-02,
+            -1.3439e-01,
+            -4.4087e-02,
+            -1.7677e-01,
+            5.1971e-02,
+            1.1013e-01,
+            -1.2838e-01,
+            -1.8450e-02,
+            9.7373e-02,
+            1.5415e-02,
+            -6.8535e-02,
+            5.8374e-02,
+            -1.0450e-01,
+            1.3363e-01,
+            1.2089e-01,
+            -4.1817e-02,
+            3.1015e-02,
+            2.7185e-01,
+            6.9890e-02,
+            -3.8259e-03,
+            -7.7094e-02,
+            2.4651e-02,
+            -1.0086e-01,
+            -8.4380e-02,
+            4.9062e-02,
+            4.7470e-02,
+            1.0739e-01,
+            -2.4693e-03,
+            -5.9990e-02,
+            -6.2990e-02,
+            1.2398e-03,
+            4.8100e-03,
+            -6.0338e-02,
+            6.4579e-02,
+            4.2505e-04,
+            -2.9926e-02,
+            1.3627e-01,
+            -3.0724e-02,
+            4.2972e-02,
+            4.6971e-02,
+            1.2441e-01,
+            -2.4336e-02,
+            6.9954e-02,
+            1.0403e-01,
+            -3.2450e-02,
+            -4.3154e-02,
+            -4.9959e-02,
+            1.5666e-01,
+            1.3688e-01,
+            5.3450e-02,
+            7.9968e-02,
+            1.3858e-01,
+            -1.6817e-01,
+            1.2637e-01,
+            7.3937e-02,
+        ],
+    ]
+)

backend/vespa_app.py

@@ -0,0 +1,458 @@
+import os
+import time
+from typing import Any, Dict, Tuple
+import asyncio
+import numpy as np
+import torch
+from dotenv import load_dotenv
+from vespa.application import Vespa
+from vespa.io import VespaQueryResponse
+from .colpali import SimMapGenerator
+import backend.stopwords
+import logging
+
+
+class VespaQueryClient:
+    MAX_QUERY_TERMS = 64
+    VESPA_SCHEMA_NAME = "pdf_page"
+    SELECT_FIELDS = "id,title,url,blur_image,page_number,snippet,text"
+
+    def __init__(self, logger: logging.Logger):
+        """
+        Initialize the VespaQueryClient by loading environment variables and establishing a connection to the Vespa application.
+        """
+        load_dotenv()
+        self.logger = logger
+
+        if os.environ.get("USE_MTLS") == "true":
+            self.logger.info("Connected using mTLS")
+            mtls_key = os.environ.get("VESPA_CLOUD_MTLS_KEY")
+            mtls_cert = os.environ.get("VESPA_CLOUD_MTLS_CERT")
+
+            self.vespa_app_url = os.environ.get("VESPA_APP_MTLS_URL")
+            if not self.vespa_app_url:
+                raise ValueError(
+                    "Please set the VESPA_APP_MTLS_URL environment variable"
+                )
+
+            if not mtls_cert or not mtls_key:
+                raise ValueError(
+                    "USE_MTLS was true, but VESPA_CLOUD_MTLS_KEY and VESPA_CLOUD_MTLS_CERT were not set"
+                )
+
+            # write the key and cert to a file
+            mtls_key_path = "/tmp/vespa-data-plane-private-key.pem"
+            with open(mtls_key_path, "w") as f:
+                f.write(mtls_key)
+
+            mtls_cert_path = "/tmp/vespa-data-plane-public-cert.pem"
+            with open(mtls_cert_path, "w") as f:
+                f.write(mtls_cert)
+
+            # Instantiate Vespa connection
+            self.app = Vespa(
+                url=self.vespa_app_url, key=mtls_key_path, cert=mtls_cert_path
+            )
+        else:
+            self.logger.info("Connected using token")
+            self.vespa_app_url = os.environ.get("VESPA_APP_TOKEN_URL")
+            if not self.vespa_app_url:
+                raise ValueError(
+                    "Please set the VESPA_APP_TOKEN_URL environment variable"
+                )
+
+            self.vespa_cloud_secret_token = os.environ.get("VESPA_CLOUD_SECRET_TOKEN")
+
+            if not self.vespa_cloud_secret_token:
+                raise ValueError(
+                    "Please set the VESPA_CLOUD_SECRET_TOKEN environment variable"
+                )
+
+            # Instantiate Vespa connection
+            self.app = Vespa(
+                url=self.vespa_app_url,
+                vespa_cloud_secret_token=self.vespa_cloud_secret_token,
+            )
+
+        self.app.wait_for_application_up()
+        self.logger.info(f"Connected to Vespa at {self.vespa_app_url}")
+
+    def get_fields(self, sim_map: bool = False):
+        if not sim_map:
+            return self.SELECT_FIELDS
+        else:
+            return "summaryfeatures"
+
+    def format_query_results(
+        self, query: str, response: VespaQueryResponse, hits: int = 5
+    ) -> dict:
+        """
+        Format the Vespa query results.
+
+        Args:
+            query (str): The query text.
+            response (VespaQueryResponse): The response from Vespa.
+            hits (int, optional): Number of hits to display. Defaults to 5.
+
+        Returns:
+            dict: The JSON content of the response.
+        """
+        query_time = response.json.get("timing", {}).get("searchtime", -1)
+        query_time = round(query_time, 2)
+        count = response.json.get("root", {}).get("fields", {}).get("totalCount", 0)
+        result_text = f"Query text: '{query}', query time {query_time}s, count={count}, top results:\n"
+        self.logger.debug(result_text)
+        return response.json
+
+    async def query_vespa_bm25(
+        self,
+        query: str,
+        q_emb: torch.Tensor,
+        hits: int = 3,
+        timeout: str = "10s",
+        sim_map: bool = False,
+        **kwargs,
+    ) -> dict:
+        """
+        Query Vespa using the BM25 ranking profile.
+        This corresponds to the "BM25" radio button in the UI.
+
+        Args:
+            query (str): The query text.
+            q_emb (torch.Tensor): Query embeddings.
+            hits (int, optional): Number of hits to retrieve. Defaults to 3.
+            timeout (str, optional): Query timeout. Defaults to "10s".
+
+        Returns:
+            dict: The formatted query results.
+        """
+        async with self.app.asyncio(connections=1) as session:
+            query_embedding = self.format_q_embs(q_emb)
+
+            start = time.perf_counter()
+            response: VespaQueryResponse = await session.query(
+                body={
+                    "yql": (
+                        f"select {self.get_fields(sim_map=sim_map)} from {self.VESPA_SCHEMA_NAME} where userQuery();"
+                    ),
+                    "ranking": self.get_rank_profile("bm25", sim_map),
+                    "query": query,
+                    "timeout": timeout,
+                    "hits": hits,
+                    "input.query(qt)": query_embedding,
+                    "presentation.timing": True,
+                    **kwargs,
+                },
+            )
+            assert response.is_successful(), response.json
+            stop = time.perf_counter()
+            self.logger.debug(
+                f"Query time + data transfer took: {stop - start} s, Vespa reported searchtime was "
+                f"{response.json.get('timing', {}).get('searchtime', -1)} s"
+            )
+        return self.format_query_results(query, response)
+
+    def float_to_binary_embedding(self, float_query_embedding: dict) -> dict:
+        """
+        Convert float query embeddings to binary embeddings.
+
+        Args:
+            float_query_embedding (dict): Dictionary of float embeddings.
+
+        Returns:
+            dict: Dictionary of binary embeddings.
+        """
+        binary_query_embeddings = {}
+        for key, vector in float_query_embedding.items():
+            binary_vector = (
+                np.packbits(np.where(np.array(vector) > 0, 1, 0))
+                .astype(np.int8)
+                .tolist()
+            )
+            binary_query_embeddings[key] = binary_vector
+            if len(binary_query_embeddings) >= self.MAX_QUERY_TERMS:
+                self.logger.warning(
+                    f"Warning: Query has more than {self.MAX_QUERY_TERMS} terms. Truncating."
+                )
+                break
+        return binary_query_embeddings
+
+    def create_nn_query_strings(
+        self, binary_query_embeddings: dict, target_hits_per_query_tensor: int = 20
+    ) -> Tuple[str, dict]:
+        """
+        Create nearest neighbor query strings for Vespa.
+
+        Args:
+            binary_query_embeddings (dict): Binary query embeddings.
+            target_hits_per_query_tensor (int, optional): Target hits per query tensor. Defaults to 20.
+
+        Returns:
+            Tuple[str, dict]: Nearest neighbor query string and query tensor dictionary.
+        """
+        nn_query_dict = {}
+        for i in range(len(binary_query_embeddings)):
+            nn_query_dict[f"input.query(rq{i})"] = binary_query_embeddings[i]
+        nn = " OR ".join(
+            [
+                f"({{targetHits:{target_hits_per_query_tensor}}}nearestNeighbor(embedding,rq{i}))"
+                for i in range(len(binary_query_embeddings))
+            ]
+        )
+        return nn, nn_query_dict
+
+    def format_q_embs(self, q_embs: torch.Tensor) -> dict:
+        """
+        Convert query embeddings to a dictionary of lists.
+
+        Args:
+            q_embs (torch.Tensor): Query embeddings tensor.
+
+        Returns:
+            dict: Dictionary where each key is an index and value is the embedding list.
+        """
+        return {idx: emb.tolist() for idx, emb in enumerate(q_embs)}
+
+    async def get_result_from_query(
+        self,
+        query: str,
+        q_embs: torch.Tensor,
+        ranking: str,
+        idx_to_token: dict,
+    ) -> Dict[str, Any]:
+        """
+        Get query results from Vespa based on the ranking method.
+
+        Args:
+            query (str): The query text.
+            q_embs (torch.Tensor): Query embeddings.
+            ranking (str): The ranking method to use.
+            idx_to_token (dict): Index to token mapping.
+
+        Returns:
+            Dict[str, Any]: The query results.
+        """
+
+        # Remove stopwords from the query to avoid visual emphasis on irrelevant words (e.g., "the", "and", "of")
+        query = backend.stopwords.filter(query)
+
+        rank_method = ranking.split("_")[0]
+        sim_map: bool = len(ranking.split("_")) > 1 and ranking.split("_")[1] == "sim"
+        if rank_method == "colpali":  # ColPali
+            result = await self.query_vespa_colpali(
+                query=query, ranking=rank_method, q_emb=q_embs, sim_map=sim_map
+            )
+        elif rank_method == "hybrid":  # Hybrid ColPali+BM25
+            result = await self.query_vespa_colpali(
+                query=query, ranking=rank_method, q_emb=q_embs, sim_map=sim_map
+            )
+        elif rank_method == "bm25":
+            result = await self.query_vespa_bm25(query, q_embs, sim_map=sim_map)
+        else:
+            raise ValueError(f"Unsupported ranking: {rank_method}")
+        if "root" not in result or "children" not in result["root"]:
+            result["root"] = {"children": []}
+            return result
+        for single_result in result["root"]["children"]:
+            self.logger.debug(single_result["fields"].keys())
+        return result
+
+    def get_sim_maps_from_query(
+        self, query: str, q_embs: torch.Tensor, ranking: str, idx_to_token: dict
+    ):
+        """
+        Get similarity maps from Vespa based on the ranking method.
+
+        Args:
+            query (str): The query text.
+            q_embs (torch.Tensor): Query embeddings.
+            ranking (str): The ranking method to use.
+            idx_to_token (dict): Index to token mapping.
+
+        Returns:
+            Dict[str, Any]: The query results.
+        """
+        # Get the result by calling asyncio.run
+        result = asyncio.run(
+            self.get_result_from_query(query, q_embs, ranking, idx_to_token)
+        )
+        vespa_sim_maps = []
+        for single_result in result["root"]["children"]:
+            vespa_sim_map = single_result["fields"].get("summaryfeatures", None)
+            if vespa_sim_map is not None:
+                vespa_sim_maps.append(vespa_sim_map)
+            else:
+                raise ValueError("No sim_map found in Vespa response")
+        return vespa_sim_maps
+
+    async def get_full_image_from_vespa(self, doc_id: str) -> str:
+        """
+        Retrieve the full image from Vespa for a given document ID.
+
+        Args:
+            doc_id (str): The document ID.
+
+        Returns:
+            str: The full image data.
+        """
+        async with self.app.asyncio(connections=1) as session:
+            start = time.perf_counter()
+            response: VespaQueryResponse = await session.query(
+                body={
+                    "yql": f'select full_image from {self.VESPA_SCHEMA_NAME} where id contains "{doc_id}"',
+                    "ranking": "unranked",
+                    "presentation.timing": True,
+                    "ranking.matching.numThreadsPerSearch": 1,
+                },
+            )
+            assert response.is_successful(), response.json
+            stop = time.perf_counter()
+            self.logger.debug(
+                f"Getting image from Vespa took: {stop - start} s, Vespa reported searchtime was "
+                f"{response.json.get('timing', {}).get('searchtime', -1)} s"
+            )
+        return response.json["root"]["children"][0]["fields"]["full_image"]
+
+    def get_results_children(self, result: VespaQueryResponse) -> list:
+        return result["root"]["children"]
+
+    def results_to_search_results(
+        self, result: VespaQueryResponse, idx_to_token: dict
+    ) -> list:
+        # Initialize sim_map_ fields in the result
+        fields_to_add = [
+            f"sim_map_{token}_{idx}"
+            for idx, token in idx_to_token.items()
+            if not SimMapGenerator.should_filter_token(token)
+        ]
+        for child in result["root"]["children"]:
+            for sim_map_key in fields_to_add:
+                child["fields"][sim_map_key] = None
+        return self.get_results_children(result)
+
+    async def get_suggestions(self, query: str) -> list:
+        async with self.app.asyncio(connections=1) as session:
+            start = time.perf_counter()
+            yql = f'select questions from {self.VESPA_SCHEMA_NAME} where questions matches (".*{query}.*")'
+            response: VespaQueryResponse = await session.query(
+                body={
+                    "yql": yql,
+                    "query": query,
+                    "ranking": "unranked",
+                    "presentation.timing": True,
+                    "presentation.summary": "suggestions",
+                    "ranking.matching.numThreadsPerSearch": 1,
+                },
+            )
+            assert response.is_successful(), response.json
+            stop = time.perf_counter()
+            self.logger.debug(
+                f"Getting suggestions from Vespa took: {stop - start} s, Vespa reported searchtime was "
+                f"{response.json.get('timing', {}).get('searchtime', -1)} s"
+            )
+            search_results = (
+                response.json["root"]["children"]
+                if "root" in response.json and "children" in response.json["root"]
+                else []
+            )
+            questions = [
+                result["fields"]["questions"]
+                for result in search_results
+                if "questions" in result["fields"]
+            ]
+
+            unique_questions = set([item for sublist in questions for item in sublist])
+
+            # remove an artifact from our data generation
+            if "string" in unique_questions:
+                unique_questions.remove("string")
+
+            return list(unique_questions)
+
+    def get_rank_profile(self, ranking: str, sim_map: bool) -> str:
+        if sim_map:
+            return f"{ranking}_sim"
+        else:
+            return ranking
+
+    async def query_vespa_colpali(
+        self,
+        query: str,
+        ranking: str,
+        q_emb: torch.Tensor,
+        target_hits_per_query_tensor: int = 100,
+        hnsw_explore_additional_hits: int = 300,
+        hits: int = 3,
+        timeout: str = "10s",
+        sim_map: bool = False,
+        **kwargs,
+    ) -> dict:
+        """
+        Query Vespa using nearest neighbor search with mixed tensors for MaxSim calculations.
+        This corresponds to the "ColPali" radio button in the UI.
+
+        Args:
+            query (str): The query text.
+            q_emb (torch.Tensor): Query embeddings.
+            target_hits_per_query_tensor (int, optional): Target hits per query tensor. Defaults to 20.
+            hits (int, optional): Number of hits to retrieve. Defaults to 3.
+            timeout (str, optional): Query timeout. Defaults to "10s".
+
+        Returns:
+            dict: The formatted query results.
+        """
+        async with self.app.asyncio(connections=1) as session:
+            float_query_embedding = self.format_q_embs(q_emb)
+            binary_query_embeddings = self.float_to_binary_embedding(
+                float_query_embedding
+            )
+
+            # Mixed tensors for MaxSim calculations
+            query_tensors = {
+                "input.query(qtb)": binary_query_embeddings,
+                "input.query(qt)": float_query_embedding,
+            }
+            nn_string, nn_query_dict = self.create_nn_query_strings(
+                binary_query_embeddings, target_hits_per_query_tensor
+            )
+            query_tensors.update(nn_query_dict)
+            response: VespaQueryResponse = await session.query(
+                body={
+                    **query_tensors,
+                    "presentation.timing": True,
+                    "yql": (
+                        f"select {self.get_fields(sim_map=sim_map)} from {self.VESPA_SCHEMA_NAME} where {nn_string} or userQuery()"
+                    ),
+                    "ranking.profile": self.get_rank_profile(
+                        ranking=ranking, sim_map=sim_map
+                    ),
+                    "timeout": timeout,
+                    "hits": hits,
+                    "query": query,
+                    "hnsw.exploreAdditionalHits": hnsw_explore_additional_hits,
+                    "ranking.rerankCount": 100,
+                    **kwargs,
+                },
+            )
+            assert response.is_successful(), response.json
+        return self.format_query_results(query, response)
+
+    async def keepalive(self) -> bool:
+        """
+        Query Vespa to keep the connection alive.
+
+        Returns:
+            bool: True if the connection is alive.
+        """
+        async with self.app.asyncio(connections=1) as session:
+            response: VespaQueryResponse = await session.query(
+                body={
+                    "yql": f"select title from {self.VESPA_SCHEMA_NAME} where true limit 1;",
+                    "ranking": "unranked",
+                    "query": "keepalive",
+                    "timeout": "3s",
+                    "hits": 1,
+                },
+            )
+            assert response.is_successful(), response.json
+        return True

colpali.py

@@ -0,0 +1,521 @@
+#!/usr/bin/env python3
+
+import torch
+from PIL import Image
+import numpy as np
+from typing import cast
+import pprint
+from pathlib import Path
+import base64
+from io import BytesIO
+from typing import Union, Tuple
+import matplotlib
+import re
+
+from colpali_engine.models import ColPali, ColPaliProcessor
+from colpali_engine.utils.torch_utils import get_torch_device
+from einops import rearrange
+from vidore_benchmark.interpretability.plot_utils import plot_similarity_heatmap
+from vidore_benchmark.interpretability.torch_utils import (
+    normalize_similarity_map_per_query_token,
+)
+from vidore_benchmark.interpretability.vit_configs import VIT_CONFIG
+from vidore_benchmark.utils.image_utils import scale_image
+from vespa.application import Vespa
+from vespa.io import VespaQueryResponse
+
+matplotlib.use("Agg")
+
+MAX_QUERY_TERMS = 64
+# OUTPUT_DIR = Path(__file__).parent.parent / "output" / "sim_maps"
+# OUTPUT_DIR.mkdir(exist_ok=True)
+
+COLPALI_GEMMA_MODEL_ID = "vidore--colpaligemma-3b-pt-448-base"
+COLPALI_GEMMA_MODEL_SNAPSHOT = "12c59eb7e23bc4c26876f7be7c17760d5d3a1ffa"
+COLPALI_GEMMA_MODEL_PATH = (
+    Path().home()
+    / f".cache/huggingface/hub/models--{COLPALI_GEMMA_MODEL_ID}/snapshots/{COLPALI_GEMMA_MODEL_SNAPSHOT}"
+)
+COLPALI_MODEL_ID = "vidore--colpali-v1.2"
+COLPALI_MODEL_SNAPSHOT = "9912ce6f8a462d8cf2269f5606eabbd2784e764f"
+COLPALI_MODEL_PATH = (
+    Path().home()
+    / f".cache/huggingface/hub/models--{COLPALI_MODEL_ID}/snapshots/{COLPALI_MODEL_SNAPSHOT}"
+)
+COLPALI_GEMMA_MODEL_NAME = COLPALI_GEMMA_MODEL_ID.replace("--", "/")
+
+
+def load_model() -> Tuple[ColPali, ColPaliProcessor]:
+    model_name = "vidore/colpali-v1.2"
+
+    device = get_torch_device("auto")
+    print(f"Using device: {device}")
+
+    # Load the model
+    model = cast(
+        ColPali,
+        ColPali.from_pretrained(
+            model_name,
+            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+            device_map=device,
+        ),
+    ).eval()
+
+    # Load the processor
+    processor = cast(ColPaliProcessor, ColPaliProcessor.from_pretrained(model_name))
+    return model, processor
+
+
+def load_vit_config(model):
+    # Load the ViT config
+    print(f"VIT config: {VIT_CONFIG}")
+    vit_config = VIT_CONFIG[COLPALI_GEMMA_MODEL_NAME]
+    return vit_config
+
+
+# Create dummy image
+dummy_image = Image.new("RGB", (448, 448), (255, 255, 255))
+
+
+def gen_similarity_map(
+    model, processor, device, vit_config, query, image: Union[Path, str]
+):
+    # Should take in the b64 image from Vespa query result
+    # And possibly the tensor representing the output_image
+    if isinstance(image, Path):
+        # image is a file path
+        try:
+            image = Image.open(image)
+        except Exception as e:
+            raise ValueError(f"Failed to open image from path: {e}")
+    elif isinstance(image, str):
+        # image is b64 string
+        try:
+            image = Image.open(BytesIO(base64.b64decode(image)))
+        except Exception as e:
+            raise ValueError(f"Failed to open image from b64: {e}")
+
+    # Preview the image
+    scale_image(image, 512)
+    # Preprocess inputs
+    input_text_processed = processor.process_queries([query]).to(device)
+    input_image_processed = processor.process_images([image]).to(device)
+    # Forward passes
+    with torch.no_grad():
+        output_text = model.forward(**input_text_processed)
+        output_image = model.forward(**input_image_processed)
+    # output_image is the tensor that we could get from the Vespa query
+    # Print shape of output_text and output_image
+    # Output image shape: torch.Size([1, 1030, 128])
+    # Remove the special tokens from the output
+    output_image = output_image[
+        :, : processor.image_seq_length, :
+    ]  # (1, n_patches_x * n_patches_y, dim)
+
+    # Rearrange the output image tensor to explicitly represent the 2D grid of patches
+    output_image = rearrange(
+        output_image,
+        "b (h w) c -> b h w c",
+        h=vit_config.n_patch_per_dim,
+        w=vit_config.n_patch_per_dim,
+    )  # (1, n_patches_x, n_patches_y, dim)
+    # Get the similarity map
+    similarity_map = torch.einsum(
+        "bnk,bijk->bnij", output_text, output_image
+    )  # (1, query_tokens, n_patches_x, n_patches_y)
+
+    # Normalize the similarity map
+    similarity_map_normalized = normalize_similarity_map_per_query_token(
+        similarity_map
+    )  # (1, query_tokens, n_patches_x, n_patches_y)
+    # Use this cell output to choose a token using its index
+    query_tokens = processor.tokenizer.tokenize(
+        processor.decode(input_text_processed.input_ids[0])
+    )
+    # Choose a token
+    token_idx = (
+        10  # e.g. if "12: '▁Kazakhstan',", set 12 to choose the token 'Kazakhstan'
+    )
+    selected_token = processor.decode(input_text_processed.input_ids[0, token_idx])
+    # strip whitespace
+    selected_token = selected_token.strip()
+    print(f"Selected token: `{selected_token}`")
+    # Retrieve the similarity map for the chosen token
+    pprint.pprint({idx: val for idx, val in enumerate(query_tokens)})
+    # Resize the image to square
+    input_image_square = image.resize((vit_config.resolution, vit_config.resolution))
+
+    # Plot the similarity map
+    fig, ax = plot_similarity_heatmap(
+        input_image_square,
+        patch_size=vit_config.patch_size,
+        image_resolution=vit_config.resolution,
+        similarity_map=similarity_map_normalized[0, token_idx, :, :],
+    )
+    ax = annotate_plot(ax, selected_token)
+    return fig, ax
+
+
+# def save_figure(fig, filename: str = "similarity_map.png"):
+#     fig.savefig(
+#         OUTPUT_DIR / filename,
+#         bbox_inches="tight",
+#         pad_inches=0,
+#     )
+
+
+def annotate_plot(ax, query, selected_token):
+    # Add the query text
+    ax.set_title(query, fontsize=18)
+    # Add annotation with selected token
+    ax.annotate(
+        f"Selected token:`{selected_token}`",
+        xy=(0.5, 0.95),
+        xycoords="axes fraction",
+        ha="center",
+        va="center",
+        fontsize=18,
+        color="black",
+        bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="black", lw=1),
+    )
+    return ax
+
+
+def gen_similarity_map_new(
+    processor: ColPaliProcessor,
+    model: ColPali,
+    device,
+    vit_config,
+    query: str,
+    query_embs: torch.Tensor,
+    token_idx_map: dict,
+    token_to_show: str,
+    image: Union[Path, str],
+):
+    if isinstance(image, Path):
+        # image is a file path
+        try:
+            image = Image.open(image)
+        except Exception as e:
+            raise ValueError(f"Failed to open image from path: {e}")
+    elif isinstance(image, str):
+        # image is b64 string
+        try:
+            image = Image.open(BytesIO(base64.b64decode(image)))
+        except Exception as e:
+            raise ValueError(f"Failed to open image from b64: {e}")
+    token_idx = token_idx_map[token_to_show]
+    print(f"Selected token: `{token_to_show}`")
+    # strip whitespace
+    # Preview the image
+    # scale_image(image, 512)
+    # Preprocess inputs
+    input_image_processed = processor.process_images([image]).to(device)
+    # Forward passes
+    with torch.no_grad():
+        output_image = model.forward(**input_image_processed)
+    # output_image is the tensor that we could get from the Vespa query
+    # Print shape of output_text and output_image
+    # Output image shape: torch.Size([1, 1030, 128])
+    # Remove the special tokens from the output
+    print(f"Output image shape before dim: {output_image.shape}")
+    output_image = output_image[
+        :, : processor.image_seq_length, :
+    ]  # (1, n_patches_x * n_patches_y, dim)
+    print(f"Output image shape after dim: {output_image.shape}")
+    # Rearrange the output image tensor to explicitly represent the 2D grid of patches
+    output_image = rearrange(
+        output_image,
+        "b (h w) c -> b h w c",
+        h=vit_config.n_patch_per_dim,
+        w=vit_config.n_patch_per_dim,
+    )  # (1, n_patches_x, n_patches_y, dim)
+    # Get the similarity map
+    print(f"Query embs shape: {query_embs.shape}")
+    # Add 1 extra dim to start of query_embs
+    query_embs = query_embs.unsqueeze(0).to(device)
+    print(f"Output image shape: {output_image.shape}")
+    similarity_map = torch.einsum(
+        "bnk,bijk->bnij", query_embs, output_image
+    )  # (1, query_tokens, n_patches_x, n_patches_y)
+    print(f"Similarity map shape: {similarity_map.shape}")
+    # Normalize the similarity map
+    similarity_map_normalized = normalize_similarity_map_per_query_token(
+        similarity_map
+    )  # (1, query_tokens, n_patches_x, n_patches_y)
+    print(f"Similarity map normalized shape: {similarity_map_normalized.shape}")
+    # Use this cell output to choose a token using its index
+    input_image_square = image.resize((vit_config.resolution, vit_config.resolution))
+
+    # Plot the similarity map
+    fig, ax = plot_similarity_heatmap(
+        input_image_square,
+        patch_size=vit_config.patch_size,
+        image_resolution=vit_config.resolution,
+        similarity_map=similarity_map_normalized[0, token_idx, :, :],
+    )
+    ax = annotate_plot(ax, query, token_to_show)
+    # save the figure
+    # save_figure(fig, f"similarity_map_{token_to_show}.png")
+    return fig, ax
+
+
+def get_query_embeddings_and_token_map(
+    processor, model, query, image
+) -> Tuple[torch.Tensor, dict]:
+    inputs = processor.process_queries([query]).to(model.device)
+    with torch.no_grad():
+        embeddings_query = model(**inputs)
+        q_emb = embeddings_query.to("cpu")[0]  # Extract the single embedding
+    # Use this cell output to choose a token using its index
+    query_tokens = processor.tokenizer.tokenize(processor.decode(inputs.input_ids[0]))
+    # reverse key, values in dictionary
+    print(query_tokens)
+    token_to_idx = {val: idx for idx, val in enumerate(query_tokens)}
+    return q_emb, token_to_idx
+
+
+def format_query_results(query, response, hits=5) -> dict:
+    query_time = response.json.get("timing", {}).get("searchtime", -1)
+    query_time = round(query_time, 2)
+    count = response.json.get("root", {}).get("fields", {}).get("totalCount", 0)
+    result_text = f"Query text: '{query}', query time {query_time}s, count={count}, top results:\n"
+    print(result_text)
+    return response.json
+
+
+async def query_vespa_default(
+    app: Vespa,
+    query: str,
+    q_emb: torch.Tensor,
+    hits: int = 3,
+    timeout: str = "10s",
+    **kwargs,
+) -> dict:
+    async with app.asyncio(connections=1, total_timeout=120) as session:
+        query_embedding = format_q_embs(q_emb)
+        response: VespaQueryResponse = await session.query(
+            body={
+                "yql": "select id,title,url,image,page_number,text from pdf_page where userQuery();",
+                "ranking": "default",
+                "query": query,
+                "timeout": timeout,
+                "hits": hits,
+                "input.query(qt)": query_embedding,
+                "presentation.timing": True,
+                **kwargs,
+            },
+        )
+        assert response.is_successful(), response.json
+    return format_query_results(query, response)
+
+
+def float_to_binary_embedding(float_query_embedding: dict) -> dict:
+    binary_query_embeddings = {}
+    for k, v in float_query_embedding.items():
+        binary_vector = (
+            np.packbits(np.where(np.array(v) > 0, 1, 0)).astype(np.int8).tolist()
+        )
+        binary_query_embeddings[k] = binary_vector
+        if len(binary_query_embeddings) >= MAX_QUERY_TERMS:
+            print(f"Warning: Query has more than {MAX_QUERY_TERMS} terms. Truncating.")
+            break
+    return binary_query_embeddings
+
+
+def create_nn_query_strings(
+    binary_query_embeddings: dict, target_hits_per_query_tensor: int = 20
+) -> Tuple[str, dict]:
+    # Query tensors for nearest neighbor calculations
+    nn_query_dict = {}
+    for i in range(len(binary_query_embeddings)):
+        nn_query_dict[f"input.query(rq{i})"] = binary_query_embeddings[i]
+    nn = " OR ".join(
+        [
+            f"({{targetHits:{target_hits_per_query_tensor}}}nearestNeighbor(embedding,rq{i}))"
+            for i in range(len(binary_query_embeddings))
+        ]
+    )
+    return nn, nn_query_dict
+
+
+def format_q_embs(q_embs: torch.Tensor) -> dict:
+    float_query_embedding = {k: v.tolist() for k, v in enumerate(q_embs)}
+    return float_query_embedding
+
+
+async def query_vespa_nearest_neighbor(
+    app: Vespa,
+    query: str,
+    q_emb: torch.Tensor,
+    target_hits_per_query_tensor: int = 20,
+    hits: int = 3,
+    timeout: str = "10s",
+    **kwargs,
+) -> dict:
+    # Hyperparameter for speed vs. accuracy
+    async with app.asyncio(connections=1, total_timeout=180) as session:
+        float_query_embedding = format_q_embs(q_emb)
+        binary_query_embeddings = float_to_binary_embedding(float_query_embedding)
+
+        # Mixed tensors for MaxSim calculations
+        query_tensors = {
+            "input.query(qtb)": binary_query_embeddings,
+            "input.query(qt)": float_query_embedding,
+        }
+        nn_string, nn_query_dict = create_nn_query_strings(
+            binary_query_embeddings, target_hits_per_query_tensor
+        )
+        query_tensors.update(nn_query_dict)
+        response: VespaQueryResponse = await session.query(
+            body={
+                **query_tensors,
+                "presentation.timing": True,
+                "yql": f"select id,title,text,url,image,page_number from pdf_page where {nn_string}",
+                "ranking.profile": "retrieval-and-rerank",
+                "timeout": timeout,
+                "hits": hits,
+                **kwargs,
+            },
+        )
+        assert response.is_successful(), response.json
+    return format_query_results(query, response)
+
+
+def is_special_token(token: str) -> bool:
+    # Pattern for tokens that start with '<', numbers, whitespace, or single characters
+    pattern = re.compile(r"^<.*$|^\d+$|^\s+$|^.$")
+    if pattern.match(token):
+        return True
+    return False
+
+
+async def get_result_from_query(
+    app: Vespa,
+    processor: ColPaliProcessor,
+    model: ColPali,
+    query: str,
+    nn=False,
+    gen_sim_map=False,
+):
+    # Get the query embeddings and token map
+    print(query)
+    q_embs, token_to_idx = get_query_embeddings_and_token_map(
+        processor, model, query, dummy_image
+    )
+    print(token_to_idx)
+    # Use the token map to choose a token randomly for now
+    # Dynamically select a token containing 'water'
+
+    if nn:
+        result = await query_vespa_nearest_neighbor(app, query, q_embs)
+    else:
+        result = await query_vespa_default(app, query, q_embs)
+    # Print score, title id and text of the results
+    for idx, child in enumerate(result["root"]["children"]):
+        print(
+            f"Result {idx+1}: {child['relevance']}, {child['fields']['title']}, {child['fields']['id']}"
+        )
+
+    if gen_sim_map:
+        for single_result in result["root"]["children"]:
+            img = single_result["fields"]["image"]
+            for token in token_to_idx:
+                if is_special_token(token):
+                    print(f"Skipping special token: {token}")
+                    continue
+                fig, ax = gen_similarity_map_new(
+                    processor,
+                    model,
+                    model.device,
+                    load_vit_config(model),
+                    query,
+                    q_embs,
+                    token_to_idx,
+                    token,
+                    img,
+                )
+                sim_map = base64.b64encode(fig.canvas.tostring_rgb()).decode("utf-8")
+                single_result["fields"][f"sim_map_{token}"] = sim_map
+    return result
+
+
+def get_result_dummy(query: str, nn: bool = False):
+    result = {}
+    result["timing"] = {}
+    result["timing"]["querytime"] = 0.23700000000000002
+    result["timing"]["summaryfetchtime"] = 0.001
+    result["timing"]["searchtime"] = 0.23900000000000002
+    result["root"] = {}
+    result["root"]["id"] = "toplevel"
+    result["root"]["relevance"] = 1
+    result["root"]["fields"] = {}
+    result["root"]["fields"]["totalCount"] = 59
+    result["root"]["coverage"] = {}
+    result["root"]["coverage"]["coverage"] = 100
+    result["root"]["coverage"]["documents"] = 155
+    result["root"]["coverage"]["full"] = True
+    result["root"]["coverage"]["nodes"] = 1
+    result["root"]["coverage"]["results"] = 1
+    result["root"]["coverage"]["resultsFull"] = 1
+    result["root"]["children"] = []
+    elt0 = {}
+    elt0["id"] = "index:colpalidemo_content/0/424c85e7dece761d226f060f"
+    elt0["relevance"] = 2354.050122871995
+    elt0["source"] = "colpalidemo_content"
+    elt0["fields"] = {}
+    elt0["fields"]["id"] = "a767cb1868be9a776cd56b768347b089"
+    elt0["fields"]["url"] = (
+        "https://static.conocophillips.com/files/resources/conocophillips-2023-sustainability-report.pdf"
+    )
+    elt0["fields"]["title"] = "ConocoPhillips 2023 Sustainability Report"
+    elt0["fields"]["page_number"] = 50
+    elt0["fields"]["image"] = "empty for now - is base64 encoded image"
+    result["root"]["children"].append(elt0)
+    elt1 = {}
+    elt1["id"] = "index:colpalidemo_content/0/b927c4979f0beaf0d7fab8e9"
+    elt1["relevance"] = 2313.7529950886965
+    elt1["source"] = "colpalidemo_content"
+    elt1["fields"] = {}
+    elt1["fields"]["id"] = "9f2fc0aa02c9561adfaa1451c875658f"
+    elt1["fields"]["url"] = (
+        "https://static.conocophillips.com/files/resources/conocophillips-2023-managing-climate-related-risks.pdf"
+    )
+    elt1["fields"]["title"] = "ConocoPhillips Managing Climate Related Risks"
+    elt1["fields"]["page_number"] = 44
+    elt1["fields"]["image"] = "empty for now - is base64 encoded image"
+    result["root"]["children"].append(elt1)
+    elt2 = {}
+    elt2["id"] = "index:colpalidemo_content/0/9632d72238829d6afefba6c9"
+    elt2["relevance"] = 2312.230182081461
+    elt2["source"] = "colpalidemo_content"
+    elt2["fields"] = {}
+    elt2["fields"]["id"] = "d638ded1ddcb446268b289b3f65430fd"
+    elt2["fields"]["url"] = (
+        "https://static.conocophillips.com/files/resources/24-0976-sustainability-highlights_nature.pdf"
+    )
+    elt2["fields"]["title"] = (
+        "ConocoPhillips Sustainability Highlights - Nature (24-0976)"
+    )
+    elt2["fields"]["page_number"] = 0
+    elt2["fields"]["image"] = "empty for now - is base64 encoded image"
+    result["root"]["children"].append(elt2)
+    return result
+
+
+if __name__ == "__main__":
+    model, processor = load_model()
+    vit_config = load_vit_config(model)
+    query = "How many percent of source water is fresh water?"
+    image_filepath = (
+        Path(__file__).parent.parent
+        / "static"
+        / "assets"
+        / "ConocoPhillips Sustainability Highlights - Nature (24-0976).png"
+    )
+    gen_similarity_map(
+        model, processor, model.device, vit_config, query=query, image=image_filepath
+    )
+    result = get_result_dummy("dummy query")
+    print(result)
+    print("Done")

deploy_vespa_app.py

@@ -0,0 +1,208 @@
+#!/usr/bin/env python3
+
+import argparse
+from vespa.package import (
+    ApplicationPackage,
+    Field,
+    Schema,
+    Document,
+    HNSW,
+    RankProfile,
+    Function,
+    AuthClient,
+    Parameter,
+    FieldSet,
+    SecondPhaseRanking,
+)
+from vespa.deployment import VespaCloud
+import os
+from pathlib import Path
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Deploy Vespa application")
+    parser.add_argument("--tenant_name", required=True, help="Vespa Cloud tenant name")
+    parser.add_argument(
+        "--vespa_application_name", required=True, help="Vespa application name"
+    )
+    parser.add_argument(
+        "--token_id_write", required=True, help="Vespa Cloud token ID for write access"
+    )
+    parser.add_argument(
+        "--token_id_read", required=True, help="Vespa Cloud token ID for read access"
+    )
+
+    args = parser.parse_args()
+    tenant_name = args.tenant_name
+    vespa_app_name = args.vespa_application_name
+    token_id_write = args.token_id_write
+    token_id_read = args.token_id_read
+
+    # Define the Vespa schema
+    colpali_schema = Schema(
+        name="pdf_page",
+        document=Document(
+            fields=[
+                Field(
+                    name="id",
+                    type="string",
+                    indexing=["summary", "index"],
+                    match=["word"],
+                ),
+                Field(name="url", type="string", indexing=["summary", "index"]),
+                Field(
+                    name="title",
+                    type="string",
+                    indexing=["summary", "index"],
+                    match=["text"],
+                    index="enable-bm25",
+                ),
+                Field(
+                    name="page_number", type="int", indexing=["summary", "attribute"]
+                ),
+                Field(name="image", type="raw", indexing=["summary"]),
+                Field(name="full_image", type="raw", indexing=["summary"]),
+                Field(
+                    name="text",
+                    type="string",
+                    indexing=["summary", "index"],
+                    match=["text"],
+                    index="enable-bm25",
+                ),
+                Field(
+                    name="embedding",
+                    type="tensor<int8>(patch{}, v[16])",
+                    indexing=[
+                        "attribute",
+                        "index",
+                    ],  # adds HNSW index for candidate retrieval.
+                    ann=HNSW(
+                        distance_metric="hamming",
+                        max_links_per_node=32,
+                        neighbors_to_explore_at_insert=400,
+                    ),
+                ),
+            ]
+        ),
+        fieldsets=[
+            FieldSet(name="default", fields=["title", "url", "page_number", "text"]),
+            FieldSet(name="image", fields=["image"]),
+        ],
+    )
+
+    # Define rank profiles
+    colpali_profile = RankProfile(
+        name="default",
+        inputs=[("query(qt)", "tensor<float>(querytoken{}, v[128])")],
+        functions=[
+            Function(
+                name="max_sim",
+                expression="""
+                    sum(
+                        reduce(
+                            sum(
+                                query(qt) * unpack_bits(attribute(embedding)) , v
+                            ),
+                            max, patch
+                        ),
+                        querytoken
+                    )
+                """,
+            ),
+            Function(name="bm25_score", expression="bm25(title) + bm25(text)"),
+        ],
+        first_phase="bm25_score",
+        second_phase=SecondPhaseRanking(expression="max_sim", rerank_count=10),
+    )
+    colpali_schema.add_rank_profile(colpali_profile)
+
+    # Add retrieval-and-rerank rank profile
+    input_query_tensors = []
+    MAX_QUERY_TERMS = 64
+    for i in range(MAX_QUERY_TERMS):
+        input_query_tensors.append((f"query(rq{i})", "tensor<int8>(v[16])"))
+
+    input_query_tensors.append(("query(qt)", "tensor<float>(querytoken{}, v[128])"))
+    input_query_tensors.append(("query(qtb)", "tensor<int8>(querytoken{}, v[16])"))
+
+    colpali_retrieval_profile = RankProfile(
+        name="retrieval-and-rerank",
+        inputs=input_query_tensors,
+        functions=[
+            Function(
+                name="max_sim",
+                expression="""
+                    sum(
+                        reduce(
+                            sum(
+                                query(qt) * unpack_bits(attribute(embedding)) , v
+                            ),
+                            max, patch
+                        ),
+                        querytoken
+                    )
+                """,
+            ),
+            Function(
+                name="max_sim_binary",
+                expression="""
+                    sum(
+                      reduce(
+                        1/(1 + sum(
+                            hamming(query(qtb), attribute(embedding)) ,v)
+                        ),
+                        max,
+                        patch
+                      ),
+                      querytoken
+                    )
+                """,
+            ),
+        ],
+        first_phase="max_sim_binary",
+        second_phase=SecondPhaseRanking(expression="max_sim", rerank_count=10),
+    )
+    colpali_schema.add_rank_profile(colpali_retrieval_profile)
+
+    # Create the Vespa application package
+    vespa_application_package = ApplicationPackage(
+        name=vespa_app_name,
+        schema=[colpali_schema],
+        auth_clients=[
+            AuthClient(
+                id="mtls",  # Note that you still need to include the mtls client.
+                permissions=["read", "write"],
+                parameters=[Parameter("certificate", {"file": "security/clients.pem"})],
+            ),
+            AuthClient(
+                id="token_write",
+                permissions=["read", "write"],
+                parameters=[Parameter("token", {"id": token_id_write})],
+            ),
+            AuthClient(
+                id="token_read",
+                permissions=["read"],
+                parameters=[Parameter("token", {"id": token_id_read})],
+            ),
+        ],
+    )
+    vespa_team_api_key = os.getenv("VESPA_TEAM_API_KEY")
+    # Deploy the application to Vespa Cloud
+    vespa_cloud = VespaCloud(
+        tenant=tenant_name,
+        application=vespa_app_name,
+        key_content=vespa_team_api_key,
+        application_root="colpali-with-snippets",
+        #application_package=vespa_application_package,
+    )
+
+    #app = vespa_cloud.deploy()
+    vespa_cloud.deploy_from_disk("default", "colpali-with-snippets")
+
+    # Output the endpoint URL
+    endpoint_url = vespa_cloud.get_token_endpoint()
+    print(f"Application deployed. Token endpoint URL: {endpoint_url}")
+
+
+if __name__ == "__main__":
+    main()

feed_vespa.py

@@ -0,0 +1,209 @@
+#!/usr/bin/env python3
+
+import argparse
+import torch
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from io import BytesIO
+from typing import cast
+import os
+import json
+import hashlib
+
+from colpali_engine.models import ColPali, ColPaliProcessor
+from colpali_engine.utils.torch_utils import get_torch_device
+from vidore_benchmark.utils.image_utils import scale_image, get_base64_image
+import requests
+from pdf2image import convert_from_path
+from pypdf import PdfReader
+import numpy as np
+from vespa.application import Vespa
+from vespa.io import VespaResponse
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Feed data into Vespa application")
+    parser.add_argument(
+        "--application_name",
+        required=True,
+        default="colpalidemo",
+        help="Vespa application name",
+    )
+    parser.add_argument(
+        "--vespa_schema_name",
+        required=True,
+        default="pdf_page",
+        help="Vespa schema name",
+    )
+    args = parser.parse_args()
+
+    vespa_app_url = os.getenv("VESPA_APP_URL")
+    vespa_cloud_secret_token = os.getenv("VESPA_CLOUD_SECRET_TOKEN")
+    # Set application and schema names
+    application_name = args.application_name
+    schema_name = args.vespa_schema_name
+    # Instantiate Vespa connection using token
+    app = Vespa(url=vespa_app_url, vespa_cloud_secret_token=vespa_cloud_secret_token)
+    app.get_application_status()
+    model_name = "vidore/colpali-v1.2"
+
+    device = get_torch_device("auto")
+    print(f"Using device: {device}")
+
+    # Load the model
+    model = cast(
+        ColPali,
+        ColPali.from_pretrained(
+            model_name,
+            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+            device_map=device,
+        ),
+    ).eval()
+
+    # Load the processor
+    processor = cast(ColPaliProcessor, ColPaliProcessor.from_pretrained(model_name))
+
+    # Define functions to work with PDFs
+    def download_pdf(url):
+        response = requests.get(url)
+        if response.status_code == 200:
+            return BytesIO(response.content)
+        else:
+            raise Exception(
+                f"Failed to download PDF: Status code {response.status_code}"
+            )
+
+    def get_pdf_images(pdf_url):
+        # Download the PDF
+        pdf_file = download_pdf(pdf_url)
+        # Save the PDF temporarily to disk (pdf2image requires a file path)
+        temp_file = "temp.pdf"
+        with open(temp_file, "wb") as f:
+            f.write(pdf_file.read())
+        reader = PdfReader(temp_file)
+        page_texts = []
+        for page_number in range(len(reader.pages)):
+            page = reader.pages[page_number]
+            text = page.extract_text()
+            page_texts.append(text)
+        images = convert_from_path(temp_file)
+        assert len(images) == len(page_texts)
+        return (images, page_texts)
+
+    # Define sample PDFs
+    sample_pdfs = [
+        {
+            "title": "ConocoPhillips Sustainability Highlights - Nature (24-0976)",
+            "url": "https://static.conocophillips.com/files/resources/24-0976-sustainability-highlights_nature.pdf",
+        },
+        {
+            "title": "ConocoPhillips Managing Climate Related Risks",
+            "url": "https://static.conocophillips.com/files/resources/conocophillips-2023-managing-climate-related-risks.pdf",
+        },
+        {
+            "title": "ConocoPhillips 2023 Sustainability Report",
+            "url": "https://static.conocophillips.com/files/resources/conocophillips-2023-sustainability-report.pdf",
+        },
+    ]
+
+    # Check if vespa_feed.json exists
+    if os.path.exists("vespa_feed.json"):
+        print("Loading vespa_feed from vespa_feed.json")
+        with open("vespa_feed.json", "r") as f:
+            vespa_feed_saved = json.load(f)
+        vespa_feed = []
+        for doc in vespa_feed_saved:
+            put_id = doc["put"]
+            fields = doc["fields"]
+            # Extract document_id from put_id
+            # Format: 'id:application_name:schema_name::document_id'
+            parts = put_id.split("::")
+            document_id = parts[1] if len(parts) > 1 else ""
+            page = {"id": document_id, "fields": fields}
+            vespa_feed.append(page)
+    else:
+        print("Generating vespa_feed")
+        # Process PDFs
+        for pdf in sample_pdfs:
+            page_images, page_texts = get_pdf_images(pdf["url"])
+            pdf["images"] = page_images
+            pdf["texts"] = page_texts
+
+        # Generate embeddings
+        for pdf in sample_pdfs:
+            page_embeddings = []
+            dataloader = DataLoader(
+                pdf["images"],
+                batch_size=2,
+                shuffle=False,
+                collate_fn=lambda x: processor.process_images(x),
+            )
+            for batch_doc in tqdm(dataloader):
+                with torch.no_grad():
+                    batch_doc = {k: v.to(model.device) for k, v in batch_doc.items()}
+                    embeddings_doc = model(**batch_doc)
+                    page_embeddings.extend(list(torch.unbind(embeddings_doc.to("cpu"))))
+            pdf["embeddings"] = page_embeddings
+
+        # Prepare Vespa feed
+        vespa_feed = []
+        for pdf in sample_pdfs:
+            url = pdf["url"]
+            title = pdf["title"]
+            for page_number, (page_text, embedding, image) in enumerate(
+                zip(pdf["texts"], pdf["embeddings"], pdf["images"])
+            ):
+                base_64_image = get_base64_image(
+                    scale_image(image, 640), add_url_prefix=False
+                )
+                base_64_full_image = get_base64_image(image, add_url_prefix=False)
+                embedding_dict = dict()
+                for idx, patch_embedding in enumerate(embedding):
+                    binary_vector = (
+                        np.packbits(np.where(patch_embedding > 0, 1, 0))
+                        .astype(np.int8)
+                        .tobytes()
+                        .hex()
+                    )
+                    embedding_dict[idx] = binary_vector
+                # id_hash should be md5 hash of url and page_number
+                id_hash = hashlib.md5(f"{url}_{page_number}".encode()).hexdigest()
+                page = {
+                    "id": id_hash,
+                    "fields": {
+                        "id": id_hash,
+                        "url": url,
+                        "title": title,
+                        "page_number": page_number,
+                        "image": base_64_image,
+                        "full_image": base_64_full_image,
+                        "text": page_text,
+                        "embedding": embedding_dict,
+                    },
+                }
+                vespa_feed.append(page)
+
+        # Save vespa_feed to vespa_feed.json in the specified format
+        vespa_feed_to_save = []
+        for page in vespa_feed:
+            document_id = page["id"]
+            put_id = f"id:{application_name}:{schema_name}::{document_id}"
+            vespa_feed_to_save.append({"put": put_id, "fields": page["fields"]})
+        with open("vespa_feed.json", "w") as f:
+            json.dump(vespa_feed_to_save, f)
+
+    def callback(response: VespaResponse, id: str):
+        if not response.is_successful():
+            print(
+                f"Failed to feed document {id} with status code {response.status_code}: Reason {response.get_json()}"
+            )
+
+    # Feed data into Vespa
+    app.feed_iterable(vespa_feed, schema=schema_name, callback=callback)
+
+
+if __name__ == "__main__":
+    main()

frontend/app.py

@@ -0,0 +1,768 @@
+from typing import Optional
+from urllib.parse import quote_plus
+
+from fasthtml.components import (
+    H1,
+    H2,
+    H3,
+    Br,
+    Div,
+    Form,
+    Img,
+    NotStr,
+    P,
+    Hr,
+    Span,
+    A,
+    Script,
+    Button,
+    Label,
+    RadioGroup,
+    RadioGroupItem,
+    Separator,
+    Ul,
+    Li,
+    Strong,
+    Iframe,
+)
+from fasthtml.xtend import A, Script
+from lucide_fasthtml import Lucide
+from shad4fast import Badge, Button, Input, Label, RadioGroup, RadioGroupItem, Separator
+
+# JavaScript to check the input value and enable/disable the search button and radio buttons
+check_input_script = Script(
+    """
+        window.onload = function() {
+            const input = document.getElementById('search-input');
+            const button = document.querySelector('[data-button="search-button"]');
+            const radioGroupItems = document.querySelectorAll('button[data-ref="radio-item"]');  // Get all radio buttons
+            
+            function checkInputValue() {
+                const isInputEmpty = input.value.trim() === "";
+                button.disabled = isInputEmpty;  // Disable the submit button
+                radioGroupItems.forEach(item => {
+                    item.disabled = isInputEmpty;  // Disable/enable the radio buttons
+                });
+            }
+
+            input.addEventListener('input', checkInputValue);  // Listen for input changes
+            checkInputValue();  // Initial check when the page loads
+        };
+    """
+)
+
+# JavaScript to handle the image swapping, reset button, and active class toggling
+image_swapping = Script(
+    """
+    document.addEventListener('click', function (e) {
+        if (e.target.classList.contains('sim-map-button') || e.target.classList.contains('reset-button')) {
+            const imgContainer = e.target.closest('.relative'); 
+            const overlayContainer = imgContainer.querySelector('.overlay-container');
+            const newSrc = e.target.getAttribute('data-image-src');
+    
+            // If it's a reset button, remove the overlay image
+            if (e.target.classList.contains('reset-button')) {
+                overlayContainer.innerHTML = '';  // Clear the overlay container, showing only the full image
+            } else {
+                // Create a new overlay image
+                const img = document.createElement('img');
+                img.src = newSrc;
+                img.classList.add('overlay-image', 'absolute', 'top-0', 'left-0', 'w-full', 'h-full');
+                overlayContainer.innerHTML = '';  // Clear any previous overlay
+                overlayContainer.appendChild(img);  // Add the new overlay image
+            }
+    
+            // Toggle active class on buttons
+            const activeButton = document.querySelector('.sim-map-button.active');
+            if (activeButton) {
+                activeButton.classList.remove('active');
+            }
+            if (e.target.classList.contains('sim-map-button')) {
+                e.target.classList.add('active');
+            }
+        }
+    });
+    """
+)
+
+toggle_text_content = Script(
+    """
+    function toggleTextContent(idx) {
+        const textColumn = document.getElementById(`text-column-${idx}`);
+        const imageTextColumns = document.getElementById(`image-text-columns-${idx}`);
+        const toggleButton = document.getElementById(`toggle-button-${idx}`);
+    
+        if (textColumn.classList.contains('md-grid-text-column')) {
+          // Hide the text column
+          textColumn.classList.remove('md-grid-text-column');
+          imageTextColumns.classList.remove('grid-image-text-columns');
+          toggleButton.innerText = `Show Text`;
+        } else {
+          // Show the text column
+          textColumn.classList.add('md-grid-text-column');
+          imageTextColumns.classList.add('grid-image-text-columns');
+          toggleButton.innerText = `Hide Text`;
+        }
+    }
+    """
+)
+
+autocomplete_script = Script(
+    """
+    document.addEventListener('DOMContentLoaded', function() {
+        const input = document.querySelector('#search-input');
+        const awesomplete = new Awesomplete(input, { minChars: 1, maxItems: 5 });
+
+        input.addEventListener('input', function() {
+            if (this.value.length >= 1) {
+                // Use template literals to insert the input value dynamically in the query parameter
+                fetch(`/suggestions?query=${encodeURIComponent(this.value)}`)
+                    .then(response => response.json())
+                    .then(data => {
+                        // Update the Awesomplete list dynamically with fetched suggestions
+                        awesomplete.list = data.suggestions;
+                    })
+                    .catch(err => console.error('Error fetching suggestions:', err));
+            }
+        });
+    });
+    """
+)
+
+dynamic_elements_scrollbars = Script(
+    """
+    (function () {
+        const { applyOverlayScrollbars, getScrollbarTheme } = OverlayScrollbarsManager;
+
+        function applyScrollbarsToDynamicElements() {
+            const scrollbarTheme = getScrollbarTheme();
+
+            // Apply scrollbars to dynamically loaded result-text-full and result-text-snippet elements
+            const resultTextFullElements = document.querySelectorAll('[id^="result-text-full"]');
+            const resultTextSnippetElements = document.querySelectorAll('[id^="result-text-snippet"]');
+
+            resultTextFullElements.forEach(element => {
+                applyOverlayScrollbars(element, scrollbarTheme);
+            });
+
+            resultTextSnippetElements.forEach(element => {
+                applyOverlayScrollbars(element, scrollbarTheme);
+            });
+        }
+
+        // Apply scrollbars after dynamic content is loaded (e.g., after search results)
+        applyScrollbarsToDynamicElements();
+
+        // Observe changes in the 'dark' class to adjust the theme dynamically if needed
+        const observer = new MutationObserver(applyScrollbarsToDynamicElements);
+        observer.observe(document.documentElement, { attributes: true, attributeFilter: ['class'] });
+    })();
+    """
+)
+
+submit_form_on_radio_change = Script(
+    """
+    document.addEventListener('click', function (e) {
+        // if target has data-ref="radio-item" and type is button
+        if (e.target.getAttribute('data-ref') === 'radio-item' && e.target.type === 'button') {
+            console.log('Radio button clicked');
+            const form = e.target.closest('form');
+            form.submit();
+        }
+    });
+    """
+)
+
+
+def ShareButtons():
+    title = "Visual RAG over PDFs with Vespa and ColPali"
+    url = "https://huggingface.co/spaces/vespa-engine/colpali-vespa-visual-retrieval"
+    return Div(
+        A(
+            Img(src="/static/img/linkedin.svg", aria_hidden="true", cls="h-[21px]"),
+            "Share on LinkedIn",
+            href=f"https://www.linkedin.com/sharing/share-offsite/?url={quote_plus(url)}",
+            rel="noopener noreferrer",
+            target="_blank",
+            cls="bg-[#0A66C2] text-white inline-flex items-center gap-x-1.5 px-2.5 py-1.5 border rounded-md text-sm font-semibold",
+        ),
+        A(
+            Img(src="/static/img/x.svg", aria_hidden="true", cls="h-[21px]"),
+            "Share on X",
+            href=f"https://twitter.com/intent/tweet?text={quote_plus(title)}&url={quote_plus(url)}",
+            rel="noopener noreferrer",
+            target="_blank",
+            cls="bg-black text-white inline-flex items-center gap-x-1.5 px-2.5 py-1.5 border rounded-md text-sm font-semibold",
+        ),
+        cls="flex items-center justify-center space-x-8 mt-5",
+    )
+
+
+def SearchBox(with_border=False, query_value="", ranking_value="hybrid"):
+    grid_cls = "grid gap-2 items-center p-3 bg-muted w-full"
+
+    if with_border:
+        grid_cls = "grid gap-2 p-3 rounded-md border border-input bg-muted w-full ring-offset-background focus-within:outline-none focus-within:ring-2 focus-within:ring-ring focus-within:ring-offset-2 focus-within:border-input"
+
+    return Form(
+        Div(
+            Lucide(
+                icon="search", cls="absolute left-2 top-2 text-muted-foreground z-10"
+            ),
+            Input(
+                placeholder="Enter your search query...",
+                name="query",
+                value=query_value,
+                id="search-input",
+                cls="text-base pl-10 border-transparent ring-offset-transparent ring-0 focus-visible:ring-transparent bg-white dark:bg-background awesomplete",
+                data_list="#suggestions",
+                style="font-size: 1rem",
+                autofocus=True,
+            ),
+            cls="relative",
+        ),
+        Div(
+            Div(
+                Span("Ranking by:", cls="text-muted-foreground text-xs font-semibold"),
+                RadioGroup(
+                    Div(
+                        RadioGroupItem(value="colpali", id="colpali"),
+                        Label("ColPali", htmlFor="ColPali"),
+                        cls="flex items-center space-x-2",
+                    ),
+                    Div(
+                        RadioGroupItem(value="bm25", id="bm25"),
+                        Label("BM25", htmlFor="BM25"),
+                        cls="flex items-center space-x-2",
+                    ),
+                    Div(
+                        RadioGroupItem(value="hybrid", id="hybrid"),
+                        Label("Hybrid ColPali + BM25", htmlFor="Hybrid ColPali + BM25"),
+                        cls="flex items-center space-x-2",
+                    ),
+                    name="ranking",
+                    default_value=ranking_value,
+                    cls="grid-flow-col gap-x-5 text-muted-foreground",
+                    # Submit form when radio button is clicked
+                ),
+                cls="grid grid-flow-col items-center gap-x-3 border border-input px-3 rounded-sm",
+            ),
+            Button(
+                Lucide(icon="arrow-right", size="21"),
+                size="sm",
+                type="submit",
+                data_button="search-button",
+                disabled=True,
+            ),
+            cls="flex justify-between",
+        ),
+        check_input_script,
+        autocomplete_script,
+        submit_form_on_radio_change,
+        action=f"/search?query={quote_plus(query_value)}&ranking={quote_plus(ranking_value)}",
+        method="GET",
+        hx_get="/fetch_results",  # As the component is a form, input components query and ranking are sent as query parameters automatically, see https://htmx.org/docs/#parameters
+        hx_trigger="load",
+        hx_target="#search-results",
+        hx_swap="outerHTML",
+        hx_indicator="#loading-indicator",
+        cls=grid_cls,
+    )
+
+
+def SampleQueries():
+    sample_queries = [
+        "What percentage of the funds unlisted real estate investments were in Switzerland 2023?",
+        "Gender balance at level 4 or above in NY office 2023?",
+        "Number of graduate applications trend 2021-2023",
+        "Total amount of fixed salaries paid in 2023?",
+        "Proportion of female new hires 2021-2023?",
+        "child jumping over puddle",
+        "hula hoop kid",
+    ]
+
+    query_badges = []
+    for query in sample_queries:
+        query_badges.append(
+            A(
+                Badge(
+                    Div(
+                        Lucide(
+                            icon="text-search", size="18", cls="text-muted-foreground"
+                        ),
+                        Span(query, cls="text-base font-normal"),
+                        cls="flex gap-2 items-center",
+                    ),
+                    variant="outline",
+                    cls="text-base font-normal text-muted-foreground hover:border-black dark:hover:border-white",
+                ),
+                href=f"/search?query={quote_plus(query)}",
+                cls="no-underline",
+            )
+        )
+
+    return Div(*query_badges, cls="grid gap-2 justify-items-center")
+
+
+def Hero():
+    return Div(
+        H1(
+            "Visual RAG over PDFs",
+            cls="text-5xl md:text-6xl font-bold tracking-wide md:tracking-wider bg-clip-text text-transparent bg-gradient-to-r from-black to-slate-700 dark:from-white dark:to-slate-300 animate-fade-in",
+        ),
+        P(
+            "See how Vespa and ColPali can be used for Visual RAG in this demo",
+            cls="text-base md:text-2xl text-muted-foreground md:tracking-wide",
+        ),
+        cls="grid gap-5 text-center",
+    )
+
+
+def Home():
+    return Div(
+        Div(
+            Hero(),
+            SearchBox(with_border=True),
+            SampleQueries(),
+            ShareButtons(),
+            cls="grid gap-8 content-start mt-[13vh]",
+        ),
+        cls="grid w-full h-full max-w-screen-md gap-4 mx-auto",
+    )
+
+
+def LinkResource(text, href):
+    return Li(
+        A(
+            Lucide(icon="external-link", size="18"),
+            text,
+            href=href,
+            target="_blank",
+            cls="flex items-center gap-1.5 hover:underline bold text-md",
+        ),
+    )
+
+
+def AboutThisDemo():
+    resources = [
+        {
+            "text": "Vespa Blog: How we built this demo",
+            "href": "https://blog.vespa.ai/visual-rag-in-practice",
+        },
+        {
+            "text": "Notebook to set up Vespa application and feed dataset",
+            "href": "https://pyvespa.readthedocs.io/en/latest/examples/visual_pdf_rag_with_vespa_colpali_cloud.html",
+        },
+        {
+            "text": "Web App (FastHTML) Code",
+            "href": "https://github.com/vespa-engine/sample-apps/tree/master/visual-retrieval-colpali",
+        },
+        {
+            "text": "Vespa Blog: Scaling ColPali to Billions",
+            "href": "https://blog.vespa.ai/scaling-colpali-to-billions/",
+        },
+        {
+            "text": "Vespa Blog: Retrieval with Vision Language Models",
+            "href": "https://blog.vespa.ai/retrieval-with-vision-language-models-colpali/",
+        },
+    ]
+    return Div(
+        H1(
+            "About This Demo",
+            cls="text-3xl md:text-5xl font-bold tracking-wide md:tracking-wider",
+        ),
+        P(
+            "This demo showcases a Visual Retrieval-Augmented Generation (RAG) application over PDFs using ColPali embeddings in Vespa, built entirely in Python, using FastHTML. The code is fully open source.",
+            cls="text-base",
+        ),
+        Img(
+            src="/static/img/colpali_child.png",
+            alt="Example of token level similarity map",
+            cls="w-full",
+        ),
+        H2("Resources", cls="text-2xl font-semibold"),
+        Ul(
+            *[
+                LinkResource(resource["text"], resource["href"])
+                for resource in resources
+            ],
+            cls="space-y-2 list-disc pl-5",
+        ),
+        H2("Architecture Overview", cls="text-2xl font-semibold"),
+        Img(
+            src="/static/img/visual-retrieval-demoapp-arch.png",
+            alt="Architecture Overview",
+            cls="w-full",
+        ),
+        Ul(
+            Li(
+                Strong("Vespa Application: "),
+                "Vespa Application that handles indexing, search, ranking and queries, leveraging features like phased ranking and multivector MaxSim calculations.",
+            ),
+            Li(
+                Strong("Frontend: "),
+                "Built with FastHTML, offering a professional and responsive user interface without the complexity of separate frontend frameworks.",
+            ),
+            Li(
+                Strong("Backend: "),
+                "Also built with FastHTML. Handles query embedding inference using ColPali, serves static files, and is responsible for orchestrating interactions between Vespa and the frontend.",
+            ),
+            Li(
+                Strong("Gemini API: "),
+                "VLM for the AI response, providing responses based on the top results from Vespa.",
+                cls="list-disc list-inside",
+            ),
+            H2("User Experience Highlights", cls="text-2xl font-semibold"),
+            Ul(
+                Li(
+                    Strong("Fast and Responsive: "),
+                    "Optimized for quick loading times, with phased content delivery to display essential information immediately while loading detailed data in the background.",
+                ),
+                Li(
+                    Strong("Similarity Maps: "),
+                    "Provides visual highlights of the most relevant parts of a page in response to a query, enhancing interpretability.",
+                ),
+                Li(
+                    Strong("Type-Ahead Suggestions: "),
+                    "Offers query suggestions to assist users in formulating effective searches.",
+                ),
+                cls="list-disc list-inside",
+            ),
+            cls="grid gap-5",
+        ),
+        H2("Dataset", cls="text-2xl font-semibold"),
+        P(
+            "The dataset used in this demo is retrieved from reports published by the Norwegian Government Pension Fund Global. It contains 6,992 pages from 116 PDF reports (2000–2024). The information is often presented in visual formats, making it an ideal dataset for visual retrieval applications.",
+            cls="text-base",
+        ),
+        Iframe(
+            src="https://huggingface.co/datasets/vespa-engine/gpfg-QA/embed/viewer",
+            frameborder="0",
+            width="100%",
+            height="500",
+        ),
+        Hr(),  # To add some margin to bottom. Probably a much better way to do this, but the mb-[16vh] class doesn't seem to be applied
+        cls="w-full h-full max-w-screen-md gap-4 mx-auto mt-[8vh] mb-[16vh] grid gap-8 content-start",
+    )
+
+
+def Search(request, search_results=[]):
+    query_value = request.query_params.get("query", "").strip()
+    ranking_value = request.query_params.get("ranking", "hybrid")
+    return Div(
+        Div(
+            Div(
+                SearchBox(query_value=query_value, ranking_value=ranking_value),
+                Div(
+                    LoadingMessage(),
+                    id="search-results",  # This will be replaced by the search results
+                ),
+                cls="grid",
+            ),
+            cls="grid",
+        ),
+    )
+
+
+def LoadingMessage(display_text="Retrieving search results"):
+    return Div(
+        Lucide(icon="loader-circle", cls="size-5 mr-1.5 animate-spin"),
+        Span(display_text, cls="text-base text-center"),
+        cls="p-10 text-muted-foreground flex items-center justify-center",
+        id="loading-indicator",
+    )
+
+
+def LoadingSkeleton():
+    return Div(
+        Div(cls="h-5 bg-muted"),
+        Div(cls="h-5 bg-muted"),
+        Div(cls="h-5 bg-muted"),
+        cls="grid gap-2 animate-pulse",
+    )
+
+
+def SimMapButtonReady(query_id, idx, token, token_idx, img_src):
+    return Button(
+        token.replace("\u2581", ""),
+        size="sm",
+        data_image_src=img_src,
+        id=f"sim-map-button-{query_id}-{idx}-{token_idx}-{token}",
+        cls="sim-map-button pointer-events-auto font-mono text-xs h-5 rounded-none px-2",
+    )
+
+
+def SimMapButtonPoll(query_id, idx, token, token_idx):
+    return Button(
+        Lucide(icon="loader-circle", size="15", cls="animate-spin"),
+        size="sm",
+        disabled=True,
+        hx_get=f"/get_sim_map?query_id={query_id}&idx={idx}&token={token}&token_idx={token_idx}",
+        hx_trigger="every 0.5s",
+        hx_swap="outerHTML",
+        cls="pointer-events-auto text-xs h-5 rounded-none px-2",
+    )
+
+
+def SearchInfo(search_time, total_count):
+    return (
+        Div(
+            Span(
+                "Retrieved ",
+                Strong(total_count),
+                Span(" results"),
+                Span(" in "),
+                Strong(f"{search_time:.3f}"),  # 3 significant digits
+                Span(" seconds."),
+            ),
+            cls="grid bg-background border-t text-sm text-center p-3",
+        ),
+    )
+
+
+def SearchResult(
+    results: list,
+    query: str,
+    query_id: Optional[str] = None,
+    search_time: float = 0,
+    total_count: int = 0,
+):
+    if not results:
+        return Div(
+            P(
+                "No results found for your query.",
+                cls="text-muted-foreground text-base text-center",
+            ),
+            cls="grid p-10",
+        )
+
+    doc_ids = []
+    # Otherwise, display the search results
+    result_items = []
+    for idx, result in enumerate(results):
+        fields = result["fields"]  # Extract the 'fields' part of each result
+        doc_id = fields["id"]
+        doc_ids.append(doc_id)
+        blur_image_base64 = f"data:image/jpeg;base64,{fields['blur_image']}"
+
+        sim_map_fields = {
+            key: value
+            for key, value in fields.items()
+            if key.startswith(
+                "sim_map_"
+            )  # filtering is done before creating with 'should_filter_token'-function
+        }
+
+        # Generate buttons for the sim_map fields
+        sim_map_buttons = []
+        for key, value in sim_map_fields.items():
+            token = key.split("_")[-2]
+            token_idx = int(key.split("_")[-1])
+            if value is not None:
+                sim_map_base64 = f"data:image/jpeg;base64,{value}"
+                sim_map_buttons.append(
+                    SimMapButtonReady(
+                        query_id=query_id,
+                        idx=idx,
+                        token=token,
+                        token_idx=token_idx,
+                        img_src=sim_map_base64,
+                    )
+                )
+            else:
+                sim_map_buttons.append(
+                    SimMapButtonPoll(
+                        query_id=query_id,
+                        idx=idx,
+                        token=token,
+                        token_idx=token_idx,
+                    )
+                )
+
+        # Add "Reset Image" button to restore the full image
+        reset_button = Button(
+            "Reset",
+            variant="outline",
+            size="sm",
+            data_image_src=blur_image_base64,
+            cls="reset-button pointer-events-auto font-mono text-xs h-5 rounded-none px-2",
+        )
+
+        tokens_icon = Lucide(icon="images", size="15")
+
+        # Add "Tokens" button - this has no action, just a placeholder
+        tokens_button = Button(
+            tokens_icon,
+            "Tokens",
+            size="sm",
+            cls="tokens-button flex gap-[3px] font-bold pointer-events-none font-mono text-xs h-5 rounded-none px-2",
+        )
+
+        result_items.append(
+            Div(
+                Div(
+                    Div(
+                        Lucide(icon="file-text"),
+                        H2(fields["title"], cls="text-xl md:text-2xl font-semibold"),
+                        Separator(orientation="vertical"),
+                        Badge(
+                            f"Relevance score: {result['relevance']:.4f}",
+                            cls="flex gap-1.5 items-center justify-center",
+                        ),
+                        cls="flex items-center gap-2",
+                    ),
+                    Div(
+                        Button(
+                            "Hide Text",
+                            size="sm",
+                            id=f"toggle-button-{idx}",
+                            onclick=f"toggleTextContent({idx})",
+                            cls="hidden md:block",
+                        ),
+                    ),
+                    cls="flex flex-wrap items-center justify-between bg-background px-3 py-4",
+                ),
+                Div(
+                    Div(
+                        Div(
+                            tokens_button,
+                            *sim_map_buttons,
+                            reset_button,
+                            cls="flex flex-wrap gap-px w-full pointer-events-none",
+                        ),
+                        Div(
+                            Div(
+                                Div(
+                                    Img(
+                                        src=blur_image_base64,
+                                        hx_get=f"/full_image?doc_id={doc_id}",
+                                        style="backdrop-filter: blur(5px);",
+                                        hx_trigger="load",
+                                        hx_swap="outerHTML",
+                                        alt=fields["title"],
+                                        cls="result-image w-full h-full object-contain",
+                                    ),
+                                    Div(
+                                        cls="overlay-container absolute top-0 left-0 w-full h-full pointer-events-none"
+                                    ),
+                                    cls="relative w-full h-full",
+                                ),
+                                cls="grid bg-muted p-2",
+                            ),
+                            cls="block",
+                        ),
+                        id=f"image-column-{idx}",
+                        cls="image-column relative bg-background px-3 py-5 grid-image-column",
+                    ),
+                    Div(
+                        Div(
+                            A(
+                                Lucide(icon="external-link", size="18"),
+                                f"PDF Source (Page {fields['page_number'] + 1})",
+                                href=f"{fields['url']}#page={fields['page_number'] + 1}",
+                                target="_blank",
+                                cls="flex items-center gap-1.5 font-mono bold text-sm",
+                            ),
+                            cls="flex items-center justify-end",
+                        ),
+                        Div(
+                            Div(
+                                Div(
+                                    Div(
+                                        Div(
+                                            H3(
+                                                "Dynamic summary",
+                                                cls="text-base font-semibold",
+                                            ),
+                                            P(
+                                                NotStr(fields.get("snippet", "")),
+                                                cls="text-highlight text-muted-foreground",
+                                            ),
+                                            cls="grid grid-rows-[auto_0px] content-start gap-y-3",
+                                        ),
+                                        id=f"result-text-snippet-{idx}",
+                                        cls="grid gap-y-3 p-8 border border-dashed",
+                                    ),
+                                    Div(
+                                        Div(
+                                            Div(
+                                                H3(
+                                                    "Full text",
+                                                    cls="text-base font-semibold",
+                                                ),
+                                                Div(
+                                                    P(
+                                                        NotStr(fields.get("text", "")),
+                                                        cls="text-highlight text-muted-foreground",
+                                                    ),
+                                                    Br(),
+                                                ),
+                                                cls="grid grid-rows-[auto_0px] content-start gap-y-3",
+                                            ),
+                                            id=f"result-text-full-{idx}",
+                                            cls="grid gap-y-3 p-8 border border-dashed",
+                                        ),
+                                        Div(
+                                            cls="absolute inset-x-0 bottom-0 bg-gradient-to-t from-[#fcfcfd] dark:from-[#1c2024] pt-[7%]"
+                                        ),
+                                        cls="relative grid",
+                                    ),
+                                    cls="grid grid-rows-[1fr_1fr] xl:grid-rows-[1fr_2fr] gap-y-8 p-8 text-sm",
+                                ),
+                                cls="grid bg-background",
+                            ),
+                            cls="grid bg-muted p-2",
+                        ),
+                        id=f"text-column-{idx}",
+                        cls="text-column relative bg-background px-3 py-5 hidden md-grid-text-column",
+                    ),
+                    id=f"image-text-columns-{idx}",
+                    cls="relative grid grid-cols-1 border-t grid-image-text-columns",
+                ),
+                cls="grid grid-cols-1 grid-rows-[auto_auto_1fr]",
+            ),
+        )
+
+    return [
+        Div(
+            SearchInfo(search_time, total_count),
+            *result_items,
+            image_swapping,
+            toggle_text_content,
+            dynamic_elements_scrollbars,
+            id="search-results",
+            cls="grid grid-cols-1 gap-px bg-border min-h-0",
+        ),
+        Div(
+            ChatResult(query_id=query_id, query=query, doc_ids=doc_ids),
+            hx_swap_oob="true",
+            id="chat_messages",
+        ),
+    ]
+
+
+def ChatResult(query_id: str, query: str, doc_ids: Optional[list] = None):
+    messages = Div(LoadingSkeleton())
+
+    if doc_ids:
+        messages = Div(
+            LoadingSkeleton(),
+            hx_ext="sse",
+            sse_connect=f"/get-message?query_id={query_id}&doc_ids={','.join(doc_ids)}&query={quote_plus(query)}",
+            sse_swap="message",
+            sse_close="close",
+            hx_swap="innerHTML",
+        )
+
+    return Div(
+        Div("AI-response (Gemini-2.0)", cls="text-xl font-semibold p-5"),
+        Div(
+            Div(
+                messages,
+            ),
+            id="chat-messages",
+            cls="overflow-auto min-h-0 grid items-end px-5",
+        ),
+        id="chat_messages",
+        cls="h-full grid grid-rows-[auto_1fr_auto] min-h-0 gap-3",
+    )

frontend/layout.py

@@ -0,0 +1,171 @@
+from fasthtml.components import Body, Div, Header, Img, Nav, Title
+from fasthtml.xtend import A, Script
+from lucide_fasthtml import Lucide
+from shad4fast import Button, Separator
+
+layout_script = Script(
+    """
+    document.addEventListener("DOMContentLoaded", function () {
+          const main = document.querySelector('main');
+          const aside = document.querySelector('aside');
+          const body = document.body;
+        
+          if (main && aside && main.nextElementSibling === aside) {
+            // If we have both main and aside, adjust the layout for larger screens
+            body.classList.remove('grid-cols-1'); // Remove single-column layout
+            body.classList.add('md:grid-cols-[minmax(0,_45fr)_minmax(0,_15fr)]'); // Two-column layout on larger screens
+          } else if (main) {
+            // If only main, keep it full width
+            body.classList.add('grid-cols-1');
+          }
+    });
+    """
+)
+
+overlay_scrollbars_manager = Script(
+    """
+    (function () {
+        const { OverlayScrollbars } = OverlayScrollbarsGlobal;
+
+        function getPreferredTheme() {
+            return localStorage.theme === 'dark' || (!('theme' in localStorage) && window.matchMedia('(prefers-color-scheme: dark)').matches)
+                ? 'dark'
+                : 'light';
+        }
+
+        function applyOverlayScrollbars(element, scrollbarTheme) {
+            // Destroy existing OverlayScrollbars instance if it exists
+            const instance = OverlayScrollbars(element);
+            if (instance) {
+                instance.destroy();
+            }
+
+            // Reinitialize OverlayScrollbars with the correct theme and settings
+            OverlayScrollbars(element, {
+                overflow: {
+                    x: 'hidden',
+                    y: 'scroll'
+                },
+                scrollbars: {
+                    theme: scrollbarTheme,
+                    visibility: 'auto',
+                    autoHide: 'leave',
+                    autoHideDelay: 800
+                }
+            });
+        }
+
+        // Function to get the current scrollbar theme (light or dark)
+        function getScrollbarTheme() {
+            const isDarkMode = getPreferredTheme() === 'dark';
+            return isDarkMode ? 'os-theme-light' : 'os-theme-dark';  // Light theme in dark mode, dark theme in light mode
+        }
+
+        // Expose the common functions globally for reuse
+        window.OverlayScrollbarsManager = {
+            applyOverlayScrollbars: applyOverlayScrollbars,
+            getScrollbarTheme: getScrollbarTheme
+        };
+    })();
+    """
+)
+
+static_elements_scrollbars = Script(
+    """
+    (function () {
+        const { applyOverlayScrollbars, getScrollbarTheme } = OverlayScrollbarsManager;
+
+        function applyScrollbarsToStaticElements() {
+            const mainElement = document.querySelector('main');
+            const chatMessagesElement = document.querySelector('#chat-messages');
+
+            const scrollbarTheme = getScrollbarTheme();
+
+            if (mainElement) {
+                applyOverlayScrollbars(mainElement, scrollbarTheme);
+            }
+
+            if (chatMessagesElement) {
+                applyOverlayScrollbars(chatMessagesElement, scrollbarTheme);
+            }
+        }
+
+        // Apply the scrollbars on page load
+        applyScrollbarsToStaticElements();
+
+        // Observe changes in the 'dark' class on the <html> element to adjust the theme dynamically
+        const observer = new MutationObserver(applyScrollbarsToStaticElements);
+        observer.observe(document.documentElement, { attributes: true, attributeFilter: ['class'] });
+    })();
+    """
+)
+
+
+def Logo():
+    return Div(
+        Img(
+            src="https://assets.vespa.ai/logos/vespa-logo-black.svg",
+            alt="Vespa Logo",
+            cls="h-full dark:hidden",
+        ),
+        Img(
+            src="https://assets.vespa.ai/logos/vespa-logo-white.svg",
+            alt="Vespa Logo Dark Mode",
+            cls="h-full hidden dark:block",
+        ),
+        cls="h-[27px]",
+    )
+
+
+def ThemeToggle(variant="ghost", cls=None, **kwargs):
+    return Button(
+        Lucide("sun", cls="dark:flex hidden"),
+        Lucide("moon", cls="dark:hidden"),
+        variant=variant,
+        size="icon",
+        cls=f"theme-toggle {cls}",
+        **kwargs,
+    )
+
+
+def Links():
+    return Nav(
+        A(
+            Button("About this demo?", variant="link"),
+            href="/about-this-demo",
+        ),
+        Separator(orientation="vertical"),
+        A(
+            Button(Lucide(icon="github"), size="icon", variant="ghost"),
+            href="https://github.com/vespa-engine/vespa",
+            target="_blank",
+        ),
+        A(
+            Button(Lucide(icon="slack"), size="icon", variant="ghost"),
+            href="https://slack.vespa.ai",
+            target="_blank",
+        ),
+        Separator(orientation="vertical"),
+        ThemeToggle(),
+        cls="flex items-center space-x-2",
+    )
+
+
+def Layout(*c, is_home=False, **kwargs):
+    return (
+        Title("Visual Retrieval ColPali"),
+        Body(
+            Header(
+                A(Logo(), href="/"),
+                Links(),
+                cls="min-h-[55px] h-[55px] w-full flex items-center justify-between px-4",
+            ),
+            *c,
+            **kwargs,
+            data_is_home=str(is_home).lower(),
+            cls="grid grid-rows-[minmax(0,55px)_minmax(0,1fr)] min-h-0",
+        ),
+        layout_script,
+        overlay_scrollbars_manager,
+        static_elements_scrollbars,
+    )

hello.py

@@ -0,0 +1,17 @@
+from fasthtml.common import *
+from importlib.util import find_spec
+
+# Run find_spec for all the modules (imports will be removed by ruff if not used. This is just to check if the modules are available, and should be removed)ß
+for module in ["torch", "einops", "PIL", "vidore_benchmark", "colpali_engine"]:
+    spec = find_spec(module)
+    assert spec is not None, f"Module {module} not found"
+
+app, rt = fast_app()
+
+
+@rt("/")
+def get():
+    return Div(P("Hello World!"), hx_get="/change")
+
+
+serve()

icons.py

@@ -0,0 +1 @@
+ICONS = {"chevrons-right": "<path d=\"m6 17 5-5-5-5\"></path><path d=\"m13 17 5-5-5-5\"></path>", "moon": "<path d=\"M12 3a6 6 0 0 0 9 9 9 9 0 1 1-9-9Z\"></path>", "sun": "<circle cx=\"12\" cy=\"12\" r=\"4\"></circle><path d=\"M12 2v2\"></path><path d=\"M12 20v2\"></path><path d=\"m4.93 4.93 1.41 1.41\"></path><path d=\"m17.66 17.66 1.41 1.41\"></path><path d=\"M2 12h2\"></path><path d=\"M20 12h2\"></path><path d=\"m6.34 17.66-1.41 1.41\"></path><path d=\"m19.07 4.93-1.41 1.41\"></path>", "github": "<path d=\"M15 22v-4a4.8 4.8 0 0 0-1-3.5c3 0 6-2 6-5.5.08-1.25-.27-2.48-1-3.5.28-1.15.28-2.35 0-3.5 0 0-1 0-3 1.5-2.64-.5-5.36-.5-8 0C6 2 5 2 5 2c-.3 1.15-.3 2.35 0 3.5A5.403 5.403 0 0 0 4 9c0 3.5 3 5.5 6 5.5-.39.49-.68 1.05-.85 1.65-.17.6-.22 1.23-.15 1.85v4\"></path><path d=\"M9 18c-4.51 2-5-2-7-2\"></path>", "slack": "<rect height=\"8\" rx=\"1.5\" width=\"3\" x=\"13\" y=\"2\"></rect><path d=\"M19 8.5V10h1.5A1.5 1.5 0 1 0 19 8.5\"></path><rect height=\"8\" rx=\"1.5\" width=\"3\" x=\"8\" y=\"14\"></rect><path d=\"M5 15.5V14H3.5A1.5 1.5 0 1 0 5 15.5\"></path><rect height=\"3\" rx=\"1.5\" width=\"8\" x=\"14\" y=\"13\"></rect><path d=\"M15.5 19H14v1.5a1.5 1.5 0 1 0 1.5-1.5\"></path><rect height=\"3\" rx=\"1.5\" width=\"8\" x=\"2\" y=\"8\"></rect><path d=\"M8.5 5H10V3.5A1.5 1.5 0 1 0 8.5 5\"></path>", "settings": "<path d=\"M12.22 2h-.44a2 2 0 0 0-2 2v.18a2 2 0 0 1-1 1.73l-.43.25a2 2 0 0 1-2 0l-.15-.08a2 2 0 0 0-2.73.73l-.22.38a2 2 0 0 0 .73 2.73l.15.1a2 2 0 0 1 1 1.72v.51a2 2 0 0 1-1 1.74l-.15.09a2 2 0 0 0-.73 2.73l.22.38a2 2 0 0 0 2.73.73l.15-.08a2 2 0 0 1 2 0l.43.25a2 2 0 0 1 1 1.73V20a2 2 0 0 0 2 2h.44a2 2 0 0 0 2-2v-.18a2 2 0 0 1 1-1.73l.43-.25a2 2 0 0 1 2 0l.15.08a2 2 0 0 0 2.73-.73l.22-.39a2 2 0 0 0-.73-2.73l-.15-.08a2 2 0 0 1-1-1.74v-.5a2 2 0 0 1 1-1.74l.15-.09a2 2 0 0 0 .73-2.73l-.22-.38a2 2 0 0 0-2.73-.73l-.15.08a2 2 0 0 1-2 0l-.43-.25a2 2 0 0 1-1-1.73V4a2 2 0 0 0-2-2z\"></path><circle cx=\"12\" cy=\"12\" r=\"3\"></circle>", "arrow-right": "<path d=\"M5 12h14\"></path><path d=\"m12 5 7 7-7 7\"></path>", "search": "<circle cx=\"11\" cy=\"11\" r=\"8\"></circle><path d=\"m21 21-4.3-4.3\"></path>", "file-search": "<path d=\"M14 2v4a2 2 0 0 0 2 2h4\"></path><path d=\"M4.268 21a2 2 0 0 0 1.727 1H18a2 2 0 0 0 2-2V7l-5-5H6a2 2 0 0 0-2 2v3\"></path><path d=\"m9 18-1.5-1.5\"></path><circle cx=\"5\" cy=\"14\" r=\"3\"></circle>", "message-circle-question": "<path d=\"M7.9 20A9 9 0 1 0 4 16.1L2 22Z\"></path><path d=\"M9.09 9a3 3 0 0 1 5.83 1c0 2-3 3-3 3\"></path><path d=\"M12 17h.01\"></path>", "text-search": "<path d=\"M21 6H3\"></path><path d=\"M10 12H3\"></path><path d=\"M10 18H3\"></path><circle cx=\"17\" cy=\"15\" r=\"3\"></circle><path d=\"m21 19-1.9-1.9\"></path>", "maximize": "<path d=\"M8 3H5a2 2 0 0 0-2 2v3\"></path><path d=\"M21 8V5a2 2 0 0 0-2-2h-3\"></path><path d=\"M3 16v3a2 2 0 0 0 2 2h3\"></path><path d=\"M16 21h3a2 2 0 0 0 2-2v-3\"></path>", "expand": "<path d=\"m21 21-6-6m6 6v-4.8m0 4.8h-4.8\"></path><path d=\"M3 16.2V21m0 0h4.8M3 21l6-6\"></path><path d=\"M21 7.8V3m0 0h-4.8M21 3l-6 6\"></path><path d=\"M3 7.8V3m0 0h4.8M3 3l6 6\"></path>", "fullscreen": "<path d=\"M3 7V5a2 2 0 0 1 2-2h2\"></path><path d=\"M17 3h2a2 2 0 0 1 2 2v2\"></path><path d=\"M21 17v2a2 2 0 0 1-2 2h-2\"></path><path d=\"M7 21H5a2 2 0 0 1-2-2v-2\"></path><rect height=\"8\" rx=\"1\" width=\"10\" x=\"7\" y=\"8\"></rect>", "images": "<path d=\"M18 22H4a2 2 0 0 1-2-2V6\"></path><path d=\"m22 13-1.296-1.296a2.41 2.41 0 0 0-3.408 0L11 18\"></path><circle cx=\"12\" cy=\"8\" r=\"2\"></circle><rect height=\"16\" rx=\"2\" width=\"16\" x=\"6\" y=\"2\"></rect>", "circle": "<circle cx=\"12\" cy=\"12\" r=\"10\"></circle>", "loader-circle": "<path d=\"M21 12a9 9 0 1 1-6.219-8.56\"></path>", "file-text": "<path d=\"M15 2H6a2 2 0 0 0-2 2v16a2 2 0 0 0 2 2h12a2 2 0 0 0 2-2V7Z\"></path><path d=\"M14 2v4a2 2 0 0 0 2 2h4\"></path><path d=\"M10 9H8\"></path><path d=\"M16 13H8\"></path><path d=\"M16 17H8\"></path>", "file-question": "<path d=\"M12 17h.01\"></path><path d=\"M15 2H6a2 2 0 0 0-2 2v16a2 2 0 0 0 2 2h12a2 2 0 0 0 2-2V7z\"></path><path d=\"M9.1 9a3 3 0 0 1 5.82 1c0 2-3 3-3 3\"></path>", "external-link": "<path d=\"M15 3h6v6\"></path><path d=\"M10 14 21 3\"></path><path d=\"M18 13v6a2 2 0 0 1-2 2H5a2 2 0 0 1-2-2V8a2 2 0 0 1 2-2h6\"></path>", "linkedin": "<path d=\"M16 8a6 6 0 0 1 6 6v7h-4v-7a2 2 0 0 0-2-2 2 2 0 0 0-2 2v7h-4v-7a6 6 0 0 1 6-6z\"></path><rect height=\"12\" width=\"4\" x=\"2\" y=\"9\"></rect><circle cx=\"4\" cy=\"4\" r=\"2\"></circle>"}

main.py

@@ -0,0 +1,420 @@
+import asyncio
+import base64
+import os
+import time
+import uuid
+import logging
+import sys
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+
+import google.generativeai as genai
+from fastcore.parallel import threaded
+from fasthtml.common import (
+    Aside,
+    Div,
+    FileResponse,
+    HighlightJS,
+    Img,
+    JSONResponse,
+    Link,
+    Main,
+    P,
+    RedirectResponse,
+    Script,
+    StreamingResponse,
+    fast_app,
+    serve,
+)
+from PIL import Image
+from shad4fast import ShadHead
+from vespa.application import Vespa
+
+from backend.colpali import SimMapGenerator
+from backend.vespa_app import VespaQueryClient
+from frontend.app import (
+    AboutThisDemo,
+    ChatResult,
+    Home,
+    Search,
+    SearchBox,
+    SearchResult,
+    SimMapButtonPoll,
+    SimMapButtonReady,
+)
+from frontend.layout import Layout
+
+highlight_js_theme_link = Link(id="highlight-theme", rel="stylesheet", href="")
+highlight_js_theme = Script(src="/static/js/highlightjs-theme.js")
+highlight_js = HighlightJS(
+    langs=["python", "javascript", "java", "json", "xml"],
+    dark="github-dark",
+    light="github",
+)
+
+overlayscrollbars_link = Link(
+    rel="stylesheet",
+    href="https://cdnjs.cloudflare.com/ajax/libs/overlayscrollbars/2.10.0/styles/overlayscrollbars.min.css",
+    type="text/css",
+)
+overlayscrollbars_js = Script(
+    src="https://cdnjs.cloudflare.com/ajax/libs/overlayscrollbars/2.10.0/browser/overlayscrollbars.browser.es5.min.js"
+)
+awesomplete_link = Link(
+    rel="stylesheet",
+    href="https://cdnjs.cloudflare.com/ajax/libs/awesomplete/1.1.7/awesomplete.min.css",
+    type="text/css",
+)
+awesomplete_js = Script(
+    src="https://cdnjs.cloudflare.com/ajax/libs/awesomplete/1.1.7/awesomplete.min.js"
+)
+sselink = Script(src="https://unpkg.com/[email protected]/sse.js")
+
+# Get log level from environment variable, default to INFO
+LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO").upper()
+# Configure logger
+logger = logging.getLogger("vespa_app")
+handler = logging.StreamHandler(sys.stdout)
+handler.setFormatter(
+    logging.Formatter(
+        "%(levelname)s: \t %(asctime)s \t %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+)
+logger.addHandler(handler)
+logger.setLevel(getattr(logging, LOG_LEVEL))
+
+app, rt = fast_app(
+    htmlkw={"cls": "grid h-full"},
+    pico=False,
+    hdrs=(
+        highlight_js,
+        highlight_js_theme_link,
+        highlight_js_theme,
+        overlayscrollbars_link,
+        overlayscrollbars_js,
+        awesomplete_link,
+        awesomplete_js,
+        sselink,
+        ShadHead(tw_cdn=False, theme_handle=True),
+    ),
+)
+vespa_app: Vespa = VespaQueryClient(logger=logger)
+thread_pool = ThreadPoolExecutor()
+# Gemini config
+
+genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
+GEMINI_SYSTEM_PROMPT = """If the user query is a question, try your best to answer it based on the provided images. 
+If the user query can not be interpreted as a question, or if the answer to the query can not be inferred from the images,
+answer with the exact phrase "I am sorry, I can't find enough relevant information on these pages to answer your question.".
+Your response should be HTML formatted, but only simple tags, such as <b>. <p>, <i>, <br> <ul> and <li> are allowed. No HTML tables.
+This means that newlines will be replaced with <br> tags, bold text will be enclosed in <b> tags, and so on.
+Do NOT include backticks (`) in your response. Only simple HTML tags and text.
+"""
+gemini_model = genai.GenerativeModel(
+    "gemini-2.0-flash", system_instruction=GEMINI_SYSTEM_PROMPT
+)
+STATIC_DIR = Path("static")
+IMG_DIR = STATIC_DIR / "full_images"
+SIM_MAP_DIR = STATIC_DIR / "sim_maps"
+os.makedirs(IMG_DIR, exist_ok=True)
+os.makedirs(SIM_MAP_DIR, exist_ok=True)
+
+
+@app.on_event("startup")
+def load_model_on_startup():
+    app.sim_map_generator = SimMapGenerator(logger=logger)
+    return
+
+
+@app.on_event("startup")
+async def keepalive():
+    asyncio.create_task(poll_vespa_keepalive())
+    return
+
+
+def generate_query_id(query, ranking_value):
+    hash_input = (query + ranking_value).encode("utf-8")
+    return hash(hash_input)
+
+
+@rt("/static/{filepath:path}")
+def serve_static(filepath: str):
+    return FileResponse(STATIC_DIR / filepath)
+
+
+@rt("/")
+def get(session):
+    if "session_id" not in session:
+        session["session_id"] = str(uuid.uuid4())
+    return Layout(Main(Home()), is_home=True)
+
+
+@rt("/about-this-demo")
+def get():
+    return Layout(Main(AboutThisDemo()))
+
+
+@rt("/search")
+def get(request, query: str = "", ranking: str = "hybrid"):
+    logger.info(f"/search: Fetching results for query: {query}, ranking: {ranking}")
+
+    # Always render the SearchBox first
+    if not query:
+        # Show SearchBox and a message for missing query
+        return Layout(
+            Main(
+                Div(
+                    SearchBox(query_value=query, ranking_value=ranking),
+                    Div(
+                        P(
+                            "No query provided. Please enter a query.",
+                            cls="text-center text-muted-foreground",
+                        ),
+                        cls="p-10",
+                    ),
+                    cls="grid",
+                )
+            )
+        )
+    # Generate a unique query_id based on the query and ranking value
+    query_id = generate_query_id(query, ranking)
+    # Show the loading message if a query is provided
+    return Layout(
+        Main(Search(request), data_overlayscrollbars_initialize=True, cls="border-t"),
+        Aside(
+            ChatResult(query_id=query_id, query=query),
+            cls="border-t border-l hidden md:block",
+        ),
+    )  # Show SearchBox and Loading message initially
+
+
+@rt("/fetch_results")
+async def get(session, request, query: str, ranking: str):
+    if "hx-request" not in request.headers:
+        return RedirectResponse("/search")
+
+    # Get the hash of the query and ranking value
+    query_id = generate_query_id(query, ranking)
+    logger.info(f"Query id in /fetch_results: {query_id}")
+    # Run the embedding and query against Vespa app
+    start_inference = time.perf_counter()
+    q_embs, idx_to_token = app.sim_map_generator.get_query_embeddings_and_token_map(
+        query
+    )
+    end_inference = time.perf_counter()
+    logger.info(
+        f"Inference time for query_id: {query_id} \t {end_inference - start_inference:.2f} seconds"
+    )
+
+    start = time.perf_counter()
+    # Fetch real search results from Vespa
+    result = await vespa_app.get_result_from_query(
+        query=query,
+        q_embs=q_embs,
+        ranking=ranking,
+        idx_to_token=idx_to_token,
+    )
+    end = time.perf_counter()
+    logger.info(
+        f"Search results fetched in {end - start:.2f} seconds. Vespa search time: {result['timing']['searchtime']}"
+    )
+    search_time = result["timing"]["searchtime"]
+    # Safely get total_count with a default of 0
+    total_count = result.get("root", {}).get("fields", {}).get("totalCount", 0)
+
+    search_results = vespa_app.results_to_search_results(result, idx_to_token)
+
+    get_and_store_sim_maps(
+        query_id=query_id,
+        query=query,
+        q_embs=q_embs,
+        ranking=ranking,
+        idx_to_token=idx_to_token,
+        doc_ids=[result["fields"]["id"] for result in search_results],
+    )
+    return SearchResult(search_results, query, query_id, search_time, total_count)
+
+
+def get_results_children(result):
+    search_results = (
+        result["root"]["children"]
+        if "root" in result and "children" in result["root"]
+        else []
+    )
+    return search_results
+
+
+async def poll_vespa_keepalive():
+    while True:
+        await asyncio.sleep(5)
+        await vespa_app.keepalive()
+        logger.debug(f"Vespa keepalive: {time.time()}")
+
+
+@threaded
+def get_and_store_sim_maps(
+    query_id, query: str, q_embs, ranking, idx_to_token, doc_ids
+):
+    ranking_sim = ranking + "_sim"
+    vespa_sim_maps = vespa_app.get_sim_maps_from_query(
+        query=query,
+        q_embs=q_embs,
+        ranking=ranking_sim,
+        idx_to_token=idx_to_token,
+    )
+    img_paths = [IMG_DIR / f"{doc_id}.jpg" for doc_id in doc_ids]
+    # All images should be downloaded, but best to wait 5 secs
+    max_wait = 5
+    start_time = time.time()
+    while (
+        not all([os.path.exists(img_path) for img_path in img_paths])
+        and time.time() - start_time < max_wait
+    ):
+        time.sleep(0.2)
+    if not all([os.path.exists(img_path) for img_path in img_paths]):
+        logger.warning(f"Images not ready in 5 seconds for query_id: {query_id}")
+        return False
+    sim_map_generator = app.sim_map_generator.gen_similarity_maps(
+        query=query,
+        query_embs=q_embs,
+        token_idx_map=idx_to_token,
+        images=img_paths,
+        vespa_sim_maps=vespa_sim_maps,
+    )
+    for idx, token, token_idx, blended_img_base64 in sim_map_generator:
+        with open(SIM_MAP_DIR / f"{query_id}_{idx}_{token_idx}.png", "wb") as f:
+            f.write(base64.b64decode(blended_img_base64))
+        logger.debug(
+            f"Sim map saved to disk for query_id: {query_id}, idx: {idx}, token: {token}"
+        )
+    return True
+
+
+@app.get("/get_sim_map")
+async def get_sim_map(query_id: str, idx: int, token: str, token_idx: int):
+    """
+    Endpoint that each of the sim map button polls to get the sim map image
+    when it is ready. If it is not ready, returns a SimMapButtonPoll, that
+    continues to poll every 1 second.
+    """
+    sim_map_path = SIM_MAP_DIR / f"{query_id}_{idx}_{token_idx}.png"
+    if not os.path.exists(sim_map_path):
+        logger.debug(
+            f"Sim map not ready for query_id: {query_id}, idx: {idx}, token: {token}"
+        )
+        return SimMapButtonPoll(
+            query_id=query_id, idx=idx, token=token, token_idx=token_idx
+        )
+    else:
+        return SimMapButtonReady(
+            query_id=query_id,
+            idx=idx,
+            token=token,
+            token_idx=token_idx,
+            img_src=sim_map_path,
+        )
+
+
+@app.get("/full_image")
+async def full_image(doc_id: str):
+    """
+    Endpoint to get the full quality image for a given result id.
+    """
+    img_path = IMG_DIR / f"{doc_id}.jpg"
+    if not os.path.exists(img_path):
+        image_data = await vespa_app.get_full_image_from_vespa(doc_id)
+        # image data is base 64 encoded string. Save it to disk as jpg.
+        with open(img_path, "wb") as f:
+            f.write(base64.b64decode(image_data))
+        logger.debug(f"Full image saved to disk for doc_id: {doc_id}")
+    else:
+        with open(img_path, "rb") as f:
+            image_data = base64.b64encode(f.read()).decode("utf-8")
+    return Img(
+        src=f"data:image/jpeg;base64,{image_data}",
+        alt="something",
+        cls="result-image w-full h-full object-contain",
+    )
+
+
+@rt("/suggestions")
+async def get_suggestions(query: str = ""):
+    """Endpoint to get suggestions as user types in the search box"""
+    query = query.lower().strip()
+
+    if query:
+        suggestions = await vespa_app.get_suggestions(query)
+        if len(suggestions) > 0:
+            return JSONResponse({"suggestions": suggestions})
+
+    return JSONResponse({"suggestions": []})
+
+
+async def message_generator(query_id: str, query: str, doc_ids: list):
+    """Generator function to yield SSE messages for chat response"""
+    images = []
+    num_images = 3  # Number of images before firing chat request
+    max_wait = 10  # seconds
+    start_time = time.time()
+    # Check if full images are ready on disk
+    while (
+        len(images) < min(num_images, len(doc_ids))
+        and time.time() - start_time < max_wait
+    ):
+        images = []
+        for idx in range(num_images):
+            image_filename = IMG_DIR / f"{doc_ids[idx]}.jpg"
+            if not os.path.exists(image_filename):
+                logger.debug(
+                    f"Message generator: Full image not ready for query_id: {query_id}, idx: {idx}"
+                )
+                continue
+            else:
+                logger.debug(
+                    f"Message generator: image ready for query_id: {query_id}, idx: {idx}"
+                )
+                images.append(Image.open(image_filename))
+        if len(images) < num_images:
+            await asyncio.sleep(0.2)
+
+    # yield message with number of images ready
+    yield f"event: message\ndata: Generating response based on {len(images)} images...\n\n"
+    if not images:
+        yield "event: message\ndata: Failed to send images to Gemini 2.0!\n\n"
+        yield "event: close\ndata: \n\n"
+        return
+
+    # If newlines are present in the response, the connection will be closed.
+    def replace_newline_with_br(text):
+        return text.replace("\n", "<br>")
+
+    response_text = ""
+    async for chunk in await gemini_model.generate_content_async(
+        images + ["\n\n Query: ", query], stream=True
+    ):
+        if chunk.text:
+            response_text += chunk.text
+            response_text = replace_newline_with_br(response_text)
+            yield f"event: message\ndata: {response_text}\n\n"
+            await asyncio.sleep(0.1)
+    yield "event: close\ndata: \n\n"
+
+
+@app.get("/get-message")
+async def get_message(query_id: str, query: str, doc_ids: str):
+    return StreamingResponse(
+        message_generator(query_id=query_id, query=query, doc_ids=doc_ids.split(",")),
+        media_type="text/event-stream",
+    )
+
+
+@rt("/app")
+def get():
+    return Layout(Main(Div(P(f"Connected to Vespa at {vespa_app.url}"), cls="p-4")))
+
+
+if __name__ == "__main__":
+    HOT_RELOAD = os.getenv("HOT_RELOAD", "False").lower() == "true"
+    logger.info(f"Starting app with hot reload: {HOT_RELOAD}")
+    serve(port=7860, reload=HOT_RELOAD)

prepare_feed_deploy.py

@@ -0,0 +1,956 @@
+# # Visual PDF Retrieval - demo application
+#
+# In this notebook, we will prepare the Vespa backend application for our visual retrieval demo.
+# We will use ColPali as the model to extract patch vectors from images of pdf pages.
+# At query time, we use MaxSim to retrieve and/or (based on the configuration) rank the page results.
+#
+# To see the application in action, visit TODO:
+#
+# The web application is written in FastHTML, meaning the complete application is written in python.
+#
+# The steps we will take in this notebook are:
+#
+# 0. Setup and configuration
+# 1. Download the data
+# 2. Prepare the data
+# 3. Generate queries for evaluation and typeahead search suggestions
+# 4. Deploy the Vespa application
+# 5. Create the Vespa application
+# 6. Feed the data to the Vespa application
+#
+# All the steps that are needed to provision the Vespa application, including feeding the data, can be done from this notebook.
+# We have tried to make it easy for others to run this notebook, to create your own PDF Enterprise Search application using Vespa.
+#
+
+# ## 0. Setup and Configuration
+#
+
+# +
+import os
+import asyncio
+import json
+from typing import Tuple
+import hashlib
+import numpy as np
+
+# Vespa
+from vespa.package import (
+    ApplicationPackage,
+    Field,
+    Schema,
+    Document,
+    HNSW,
+    RankProfile,
+    Function,
+    FieldSet,
+    SecondPhaseRanking,
+    Summary,
+    DocumentSummary,
+)
+from vespa.deployment import VespaCloud
+from vespa.application import Vespa
+from vespa.io import VespaResponse
+
+# Google Generative AI
+import google.generativeai as genai
+
+# Torch and other ML libraries
+import torch
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from pdf2image import convert_from_path
+from pypdf import PdfReader
+
+# ColPali model and processor
+from colpali_engine.models import ColPali, ColPaliProcessor
+from colpali_engine.utils.torch_utils import get_torch_device
+from vidore_benchmark.utils.image_utils import scale_image, get_base64_image
+
+# Other utilities
+from bs4 import BeautifulSoup
+import httpx
+from urllib.parse import urljoin, urlparse
+
+# Load environment variables
+from dotenv import load_dotenv
+
+load_dotenv()
+
+# Avoid warning from huggingface tokenizers
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+# -
+
+# ### Create a free trial in Vespa Cloud
+#
+# Create a tenant from [here](https://vespa.ai/free-trial/).
+# The trial includes $300 credit.
+# Take note of your tenant name.
+#
+
+VESPA_TENANT_NAME = "vespa-team"
+
+# Here, set your desired application name. (Will be created in later steps)
+# Note that you can not have hyphen `-` or underscore `_` in the application name.
+#
+
+VESPA_APPLICATION_NAME = "colpalidemo"
+VESPA_SCHEMA_NAME = "pdf_page"
+
+# Next, you need to create some tokens for feeding data, and querying the application.
+# We recommend separate tokens for feeding and querying, (the former with write permission, and the latter with read permission).
+# The tokens can be created from the [Vespa Cloud console](https://console.vespa-cloud.com/) in the 'Account' -> 'Tokens' section.
+#
+
+VESPA_TOKEN_ID_WRITE = "colpalidemo_write"
+
+# We also need to set the value of the write token to be able to feed data to the Vespa application.
+#
+
+VESPA_CLOUD_SECRET_TOKEN = os.getenv("VESPA_CLOUD_SECRET_TOKEN") or input(
+    "Enter Vespa cloud secret token: "
+)
+
+# We will also use the Gemini API to create sample queries for our images.
+# You can also use other VLM's to create these queries.
+# Create a Gemini API key from [here](https://aistudio.google.com/app/apikey).
+#
+
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY") or input(
+    "Enter Google Generative AI API key: "
+)
+
+# +
+MODEL_NAME = "vidore/colpali-v1.2"
+
+# Configure Google Generative AI
+genai.configure(api_key=GEMINI_API_KEY)
+
+# Set device for Torch
+device = get_torch_device("auto")
+print(f"Using device: {device}")
+
+# Load the ColPali model and processor
+model = ColPali.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+    device_map=device,
+).eval()
+
+processor = ColPaliProcessor.from_pretrained(MODEL_NAME)
+# -
+
+# ## 1. Download PDFs
+#
+# We are going to use public reports from the Norwegian Government Pension Fund Global (also known as the Oil Fund).
+# The fund puts transparency at the forefront and publishes reports on its investments, holdings, and returns, as well as its strategy and governance.
+#
+# These reports are the ones we are going to use for this showcase.
+# Here are some sample images:
+#
+# ![Sample1](./static/img/gfpg-sample-1.png)
+# ![Sample2](./static/img/gfpg-sample-2.png)
+#
+
+# As we can see, a lot of the information is in the form of tables, charts and numbers.
+# These are not easily extractable using pdf-readers or OCR tools.
+#
+
+# +
+import requests
+
+url = "https://www.nbim.no/en/publications/reports/"
+response = requests.get(url)
+response.raise_for_status()
+html_content = response.text
+
+# Parse with BeautifulSoup
+soup = BeautifulSoup(html_content, "html.parser")
+
+links = []
+url_to_year = {}
+
+# Find all 'div's with id starting with 'year-'
+for year_div in soup.find_all("div", id=lambda x: x and x.startswith("year-")):
+    year_id = year_div.get("id", "")
+    year = year_id.replace("year-", "")
+
+    # Within this div, find all 'a' elements with the specific classes
+    for a_tag in year_div.select("a.button.button--download-secondary[href]"):
+        href = a_tag["href"]
+        full_url = urljoin(url, href)
+        links.append(full_url)
+        url_to_year[full_url] = year
+links, url_to_year
+# -
+
+# Limit the number of PDFs to download
+NUM_PDFS = 2  # Set to None to download all PDFs
+links = links[:NUM_PDFS] if NUM_PDFS else links
+links
+
+# +
+from nest_asyncio import apply
+from typing import List
+
+apply()
+
+max_attempts = 3
+
+
+async def download_pdf(session, url, filename):
+    attempt = 0
+    while attempt < max_attempts:
+        try:
+            response = await session.get(url)
+            response.raise_for_status()
+
+            # Use Content-Disposition header to get the filename if available
+            content_disposition = response.headers.get("Content-Disposition")
+            if content_disposition:
+                import re
+
+                fname = re.findall('filename="(.+)"', content_disposition)
+                if fname:
+                    filename = fname[0]
+
+            # Ensure the filename is safe to use on the filesystem
+            safe_filename = filename.replace("/", "_").replace("\\", "_")
+            if not safe_filename or safe_filename == "_":
+                print(f"Invalid filename: {filename}")
+                continue
+
+            filepath = os.path.join("pdfs", safe_filename)
+            with open(filepath, "wb") as f:
+                f.write(response.content)
+            print(f"Downloaded {safe_filename}")
+            return filepath
+        except Exception as e:
+            print(f"Error downloading {filename}: {e}")
+            print(f"Retrying ({attempt})...")
+            await asyncio.sleep(1)  # Wait a bit before retrying
+            attempt += 1
+    return None
+
+
+async def download_pdfs(links: List[str]) -> List[dict]:
+    """Download PDFs from a list of URLs. Add the filename to the dictionary."""
+    async with httpx.AsyncClient() as client:
+        tasks = []
+
+        for idx, link in enumerate(links):
+            # Try to get the filename from the URL
+            path = urlparse(link).path
+            filename = os.path.basename(path)
+
+            # If filename is empty,skip
+            if not filename:
+                continue
+            tasks.append(download_pdf(client, link, filename))
+
+        # Run the tasks concurrently
+        paths = await asyncio.gather(*tasks)
+        pdf_files = [
+            {"url": link, "path": path} for link, path in zip(links, paths) if path
+        ]
+        return pdf_files
+
+
+# Create the pdfs directory if it doesn't exist
+os.makedirs("pdfs", exist_ok=True)
+# Now run the download_pdfs function with the URL
+pdfs = asyncio.run(download_pdfs(links))
+# -
+
+pdfs
+
+# ## 2. Convert PDFs to Images
+#
+
+
+# +
+def get_pdf_images(pdf_path):
+    reader = PdfReader(pdf_path)
+    page_texts = []
+    for page_number in range(len(reader.pages)):
+        page = reader.pages[page_number]
+        text = page.extract_text()
+        page_texts.append(text)
+    images = convert_from_path(pdf_path)
+    # Convert to PIL images
+    assert len(images) == len(page_texts)
+    return images, page_texts
+
+
+pdf_folder = "pdfs"
+pdf_pages = []
+for pdf in tqdm(pdfs):
+    pdf_file = pdf["path"]
+    title = os.path.splitext(os.path.basename(pdf_file))[0]
+    images, texts = get_pdf_images(pdf_file)
+    for page_no, (image, text) in enumerate(zip(images, texts)):
+        pdf_pages.append(
+            {
+                "title": title,
+                "year": int(url_to_year[pdf["url"]]),
+                "url": pdf["url"],
+                "path": pdf_file,
+                "image": image,
+                "text": text,
+                "page_no": page_no,
+            }
+        )
+# -
+
+len(pdf_pages)
+
+# +
+from collections import Counter
+
+# Print the length of the text fields - mean, max and min
+text_lengths = [len(page["text"]) for page in pdf_pages]
+print(f"Mean text length: {np.mean(text_lengths)}")
+print(f"Max text length: {np.max(text_lengths)}")
+print(f"Min text length: {np.min(text_lengths)}")
+print(f"Median text length: {np.median(text_lengths)}")
+print(f"Number of text with length == 0: {Counter(text_lengths)[0]}")
+# -
+
+# ## 3. Generate Queries
+#
+# In this step, we want to generate queries for each page image.
+# These will be useful for 2 reasons:
+#
+# 1. We can use these queries as typeahead suggestions in the search bar.
+# 2. We can use the queries to generate an evaluation dataset. See [Improving Retrieval with LLM-as-a-judge](https://blog.vespa.ai/improving-retrieval-with-llm-as-a-judge/) for a deeper dive into this topic.
+#
+# The prompt for generating queries is taken from [this](https://danielvanstrien.xyz/posts/post-with-code/colpali/2024-09-23-generate_colpali_dataset.html#an-update-retrieval-focused-prompt) wonderful blog post by Daniel van Strien.
+#
+# We will use the Gemini API to generate these queries, with `gemini-1.5-flash-8b` as the model.
+#
+
+# +
+from pydantic import BaseModel
+
+
+class GeneratedQueries(BaseModel):
+    broad_topical_question: str
+    broad_topical_query: str
+    specific_detail_question: str
+    specific_detail_query: str
+    visual_element_question: str
+    visual_element_query: str
+
+
+def get_retrieval_prompt() -> Tuple[str, GeneratedQueries]:
+    prompt = (
+        prompt
+    ) = """You are an investor, stock analyst and financial expert. You will be presented an image of a document page from a report published by the Norwegian Government Pension Fund Global (GPFG). The report may be annual or quarterly reports, or policy reports, on topics such as responsible investment, risk etc.
+Your task is to generate retrieval queries and questions that you would use to retrieve this document (or ask based on this document) in a large corpus.
+Please generate 3 different types of retrieval queries and questions.
+A retrieval query is a keyword based query, made up of 2-5 words, that you would type into a search engine to find this document.
+A question is a natural language question that you would ask, for which the document contains the answer.
+The queries should be of the following types:
+1. A broad topical query: This should cover the main subject of the document.
+2. A specific detail query: This should cover a specific detail or aspect of the document.
+3. A visual element query: This should cover a visual element of the document, such as a chart, graph, or image.
+
+Important guidelines:
+- Ensure the queries are relevant for retrieval tasks, not just describing the page content.
+- Use a fact-based natural language style for the questions.
+- Frame the queries as if someone is searching for this document in a large corpus.
+- Make the queries diverse and representative of different search strategies.
+
+Format your response as a JSON object with the structure of the following example:
+{
+    "broad_topical_question": "What was the Responsible Investment Policy in 2019?",
+    "broad_topical_query": "responsible investment policy 2019",
+    "specific_detail_question": "What is the percentage of investments in renewable energy?",
+    "specific_detail_query": "renewable energy investments percentage",
+    "visual_element_question": "What is the trend of total holding value over time?",
+    "visual_element_query": "total holding value trend"
+}
+
+If there are no relevant visual elements, provide an empty string for the visual element question and query.
+Here is the document image to analyze:
+Generate the queries based on this image and provide the response in the specified JSON format.
+Only return JSON. Don't return any extra explanation text. """
+
+    return prompt, GeneratedQueries
+
+
+prompt_text, pydantic_model = get_retrieval_prompt()
+
+# +
+gemini_model = genai.GenerativeModel("gemini-1.5-flash-8b")
+
+
+def generate_queries(image, prompt_text, pydantic_model):
+    try:
+        response = gemini_model.generate_content(
+            [image, "\n\n", prompt_text],
+            generation_config=genai.GenerationConfig(
+                response_mime_type="application/json",
+                response_schema=pydantic_model,
+            ),
+        )
+        queries = json.loads(response.text)
+    except Exception as _e:
+        queries = {
+            "broad_topical_question": "",
+            "broad_topical_query": "",
+            "specific_detail_question": "",
+            "specific_detail_query": "",
+            "visual_element_question": "",
+            "visual_element_query": "",
+        }
+    return queries
+
+
+# -
+
+for pdf in tqdm(pdf_pages):
+    image = pdf.get("image")
+    pdf["queries"] = generate_queries(image, prompt_text, pydantic_model)
+
+pdf_pages[46]["image"]
+
+pdf_pages[46]["queries"]
+
+# +
+# Generate queries async - keeping for now as we probably need when applying to the full dataset
+# import asyncio
+# from tenacity import retry, stop_after_attempt, wait_exponential
+# import google.generativeai as genai
+# from tqdm.asyncio import tqdm_asyncio
+
+# max_in_flight = 200  # Maximum number of concurrent requests
+
+
+# async def generate_queries_for_image_async(model, image, semaphore):
+#     @retry(stop=stop_after_attempt(3), wait=wait_exponential(), reraise=True)
+#     async def _generate():
+#         async with semaphore:
+#             result = await model.generate_content_async(
+#                 [image, "\n\n", prompt_text],
+#                 generation_config=genai.GenerationConfig(
+#                     response_mime_type="application/json",
+#                     response_schema=pydantic_model,
+#                 ),
+#             )
+#             return json.loads(result.text)
+
+#     try:
+#         return await _generate()
+#     except Exception as e:
+#         print(f"Error generating queries for image: {e}")
+#         return None  # Return None or handle as needed
+
+
+# async def enrich_pdfs():
+#     gemini_model = genai.GenerativeModel("gemini-1.5-flash-8b")
+#     semaphore = asyncio.Semaphore(max_in_flight)
+#     tasks = []
+#     for pdf in pdf_pages:
+#         pdf["queries"] = []
+#         image = pdf.get("image")
+#         if image:
+#             task = generate_queries_for_image_async(gemini_model, image, semaphore)
+#             tasks.append((pdf, task))
+
+#     # Run the tasks concurrently using asyncio.gather()
+#     for pdf, task in tqdm_asyncio(tasks):
+#         result = await task
+#         if result:
+#             pdf["queries"] = result
+#     return pdf_pages
+
+
+# pdf_pages = asyncio.run(enrich_pdfs())
+
+# +
+# write title, url, page_no, text, queries, not image to JSON
+with open("output/pdf_pages.json", "w") as f:
+    to_write = [{k: v for k, v in pdf.items() if k != "image"} for pdf in pdf_pages]
+    json.dump(to_write, f, indent=2)
+
+# with open("pdfs/pdf_pages.json", "r") as f:
+#     saved_pdf_pages = json.load(f)
+# for pdf, saved_pdf in zip(pdf_pages, saved_pdf_pages):
+#     pdf.update(saved_pdf)
+# -
+
+# ## 4. Generate embeddings
+#
+# Now that we have the queries, we can use the ColPali model to generate embeddings for each page image.
+#
+
+
+def generate_embeddings(images, model, processor, batch_size=2) -> np.ndarray:
+    """
+    Generate embeddings for a list of images.
+    Move to CPU only once per batch.
+
+    Args:
+        images (List[PIL.Image]): List of PIL images.
+        model (nn.Module): The model to generate embeddings.
+        processor: The processor to preprocess images.
+        batch_size (int, optional): Batch size for processing. Defaults to 64.
+
+    Returns:
+        np.ndarray: Embeddings for the images, shape
+                    (len(images), processor.max_patch_length (1030 for ColPali), model.config.hidden_size (Patch embedding dimension - 128 for ColPali)).
+    """
+    embeddings_list = []
+
+    def collate_fn(batch):
+        # Batch is a list of images
+        return processor.process_images(batch)  # Should return a dict of tensors
+
+    dataloader = DataLoader(
+        images,
+        shuffle=False,
+        collate_fn=collate_fn,
+    )
+
+    for batch_doc in tqdm(dataloader, desc="Generating embeddings"):
+        with torch.no_grad():
+            # Move batch to the device
+            batch_doc = {k: v.to(model.device) for k, v in batch_doc.items()}
+            embeddings_batch = model(**batch_doc)
+            embeddings_list.append(torch.unbind(embeddings_batch.to("cpu"), dim=0))
+    # Concatenate all embeddings and create a numpy array
+    all_embeddings = np.concatenate(embeddings_list, axis=0)
+    return all_embeddings
+
+
+# Generate embeddings for all images
+images = [pdf["image"] for pdf in pdf_pages]
+embeddings = generate_embeddings(images, model, processor)
+
+embeddings.shape
+
+# ## 5. Prepare Data on Vespa Format
+#
+# Now, that we have all the data we need, all that remains is to make sure it is in the right format for Vespa.
+#
+
+
+def float_to_binary_embedding(float_query_embedding: dict) -> dict:
+    """Utility function to convert float query embeddings to binary query embeddings."""
+    binary_query_embeddings = {}
+    for k, v in float_query_embedding.items():
+        binary_vector = (
+            np.packbits(np.where(np.array(v) > 0, 1, 0)).astype(np.int8).tolist()
+        )
+        binary_query_embeddings[k] = binary_vector
+    return binary_query_embeddings
+
+
+vespa_feed = []
+for pdf, embedding in zip(pdf_pages, embeddings):
+    url = pdf["url"]
+    year = pdf["year"]
+    title = pdf["title"]
+    image = pdf["image"]
+    text = pdf.get("text", "")
+    page_no = pdf["page_no"]
+    query_dict = pdf["queries"]
+    questions = [v for k, v in query_dict.items() if "question" in k and v]
+    queries = [v for k, v in query_dict.items() if "query" in k and v]
+    base_64_image = get_base64_image(
+        scale_image(image, 32), add_url_prefix=False
+    )  # Scaled down image to return fast on search (~1kb)
+    base_64_full_image = get_base64_image(image, add_url_prefix=False)
+    embedding_dict = {k: v for k, v in enumerate(embedding)}
+    binary_embedding = float_to_binary_embedding(embedding_dict)
+    # id_hash should be md5 hash of url and page_number
+    id_hash = hashlib.md5(f"{url}_{page_no}".encode()).hexdigest()
+    page = {
+        "id": id_hash,
+        "fields": {
+            "id": id_hash,
+            "url": url,
+            "title": title,
+            "year": year,
+            "page_number": page_no,
+            "blur_image": base_64_image,
+            "full_image": base_64_full_image,
+            "text": text,
+            "embedding": binary_embedding,
+            "queries": queries,
+            "questions": questions,
+        },
+    }
+    vespa_feed.append(page)
+
+# +
+# We will prepare the Vespa feed data, including the embeddings and the generated queries
+
+
+# Save vespa_feed to vespa_feed.json
+os.makedirs("output", exist_ok=True)
+with open("output/vespa_feed.json", "w") as f:
+    vespa_feed_to_save = []
+    for page in vespa_feed:
+        document_id = page["id"]
+        put_id = f"id:{VESPA_APPLICATION_NAME}:{VESPA_SCHEMA_NAME}::{document_id}"
+        vespa_feed_to_save.append({"put": put_id, "fields": page["fields"]})
+    json.dump(vespa_feed_to_save, f)
+
+# +
+# import json
+
+# with open("output/vespa_feed.json", "r") as f:
+#     vespa_feed = json.load(f)
+# -
+
+len(vespa_feed)
+
+# ## 5. Prepare Vespa Application
+#
+
+# +
+# Define the Vespa schema
+colpali_schema = Schema(
+    name=VESPA_SCHEMA_NAME,
+    document=Document(
+        fields=[
+            Field(
+                name="id",
+                type="string",
+                indexing=["summary", "index"],
+                match=["word"],
+            ),
+            Field(name="url", type="string", indexing=["summary", "index"]),
+            Field(name="year", type="int", indexing=["summary", "attribute"]),
+            Field(
+                name="title",
+                type="string",
+                indexing=["summary", "index"],
+                match=["text"],
+                index="enable-bm25",
+            ),
+            Field(name="page_number", type="int", indexing=["summary", "attribute"]),
+            Field(name="blur_image", type="raw", indexing=["summary"]),
+            Field(name="full_image", type="raw", indexing=["summary"]),
+            Field(
+                name="text",
+                type="string",
+                indexing=["summary", "index"],
+                match=["text"],
+                index="enable-bm25",
+            ),
+            Field(
+                name="embedding",
+                type="tensor<int8>(patch{}, v[16])",
+                indexing=[
+                    "attribute",
+                    "index",
+                ],
+                ann=HNSW(
+                    distance_metric="hamming",
+                    max_links_per_node=32,
+                    neighbors_to_explore_at_insert=400,
+                ),
+            ),
+            Field(
+                name="questions",
+                type="array<string>",
+                indexing=["summary", "attribute"],
+                summary=Summary(fields=["matched-elements-only"]),
+            ),
+            Field(
+                name="queries",
+                type="array<string>",
+                indexing=["summary", "attribute"],
+                summary=Summary(fields=["matched-elements-only"]),
+            ),
+        ]
+    ),
+    fieldsets=[
+        FieldSet(
+            name="default",
+            fields=["title", "url", "blur_image", "page_number", "text"],
+        ),
+        FieldSet(
+            name="image",
+            fields=["full_image"],
+        ),
+    ],
+    document_summaries=[
+        DocumentSummary(
+            name="default",
+            summary_fields=[
+                Summary(
+                    name="text",
+                    fields=[("bolding", "on")],
+                ),
+                Summary(
+                    name="snippet",
+                    fields=[("source", "text"), "dynamic"],
+                ),
+            ],
+            from_disk=True,
+        ),
+        DocumentSummary(
+            name="suggestions",
+            summary_fields=[
+                Summary(name="questions"),
+            ],
+            from_disk=True,
+        ),
+    ],
+)
+
+# Define similarity functions used in all rank profiles
+mapfunctions = [
+    Function(
+        name="similarities",  # computes similarity scores between each query token and image patch
+        expression="""
+                sum(
+                    query(qt) * unpack_bits(attribute(embedding)), v
+                )
+            """,
+    ),
+    Function(
+        name="normalized",  # normalizes the similarity scores to [-1, 1]
+        expression="""
+                (similarities - reduce(similarities, min)) / (reduce((similarities - reduce(similarities, min)), max)) * 2 - 1
+            """,
+    ),
+    Function(
+        name="quantized",  # quantizes the normalized similarity scores to signed 8-bit integers [-128, 127]
+        expression="""
+                cell_cast(normalized * 127.999, int8)
+            """,
+    ),
+]
+
+# Define the 'bm25' rank profile
+colpali_bm25_profile = RankProfile(
+    name="bm25",
+    inputs=[("query(qt)", "tensor<float>(querytoken{}, v[128])")],
+    first_phase="bm25(title) + bm25(text)",
+    functions=mapfunctions,
+)
+
+
+# A function to create an inherited rank profile which also returns quantized similarity scores
+def with_quantized_similarity(rank_profile: RankProfile) -> RankProfile:
+    return RankProfile(
+        name=f"{rank_profile.name}_sim",
+        first_phase=rank_profile.first_phase,
+        inherits=rank_profile.name,
+        summary_features=["quantized"],
+    )
+
+
+colpali_schema.add_rank_profile(colpali_bm25_profile)
+colpali_schema.add_rank_profile(with_quantized_similarity(colpali_bm25_profile))
+
+# Update the 'default' rank profile
+colpali_profile = RankProfile(
+    name="default",
+    inputs=[("query(qt)", "tensor<float>(querytoken{}, v[128])")],
+    first_phase="bm25_score",
+    second_phase=SecondPhaseRanking(expression="max_sim", rerank_count=10),
+    functions=mapfunctions
+    + [
+        Function(
+            name="max_sim",
+            expression="""
+                sum(
+                    reduce(
+                        sum(
+                            query(qt) * unpack_bits(attribute(embedding)), v
+                        ),
+                        max, patch
+                    ),
+                    querytoken
+                )
+            """,
+        ),
+        Function(name="bm25_score", expression="bm25(title) + bm25(text)"),
+    ],
+)
+colpali_schema.add_rank_profile(colpali_profile)
+colpali_schema.add_rank_profile(with_quantized_similarity(colpali_profile))
+
+# Update the 'retrieval-and-rerank' rank profile
+input_query_tensors = []
+MAX_QUERY_TERMS = 64
+for i in range(MAX_QUERY_TERMS):
+    input_query_tensors.append((f"query(rq{i})", "tensor<int8>(v[16])"))
+
+input_query_tensors.extend(
+    [
+        ("query(qt)", "tensor<float>(querytoken{}, v[128])"),
+        ("query(qtb)", "tensor<int8>(querytoken{}, v[16])"),
+    ]
+)
+
+colpali_retrieval_profile = RankProfile(
+    name="retrieval-and-rerank",
+    inputs=input_query_tensors,
+    first_phase="max_sim_binary",
+    second_phase=SecondPhaseRanking(expression="max_sim", rerank_count=10),
+    functions=mapfunctions
+    + [
+        Function(
+            name="max_sim",
+            expression="""
+                sum(
+                    reduce(
+                        sum(
+                            query(qt) * unpack_bits(attribute(embedding)), v
+                        ),
+                        max, patch
+                    ),
+                    querytoken
+                )
+            """,
+        ),
+        Function(
+            name="max_sim_binary",
+            expression="""
+                sum(
+                    reduce(
+                        1 / (1 + sum(
+                            hamming(query(qtb), attribute(embedding)), v)
+                        ),
+                        max, patch
+                    ),
+                    querytoken
+                )
+            """,
+        ),
+    ],
+)
+colpali_schema.add_rank_profile(colpali_retrieval_profile)
+colpali_schema.add_rank_profile(with_quantized_similarity(colpali_retrieval_profile))
+
+# +
+from vespa.configuration.services import (
+    services,
+    container,
+    search,
+    document_api,
+    document_processing,
+    clients,
+    client,
+    config,
+    content,
+    redundancy,
+    documents,
+    node,
+    certificate,
+    token,
+    document,
+    nodes,
+)
+from vespa.configuration.vt import vt
+from vespa.package import ServicesConfiguration
+
+service_config = ServicesConfiguration(
+    application_name=VESPA_APPLICATION_NAME,
+    services_config=services(
+        container(
+            search(),
+            document_api(),
+            document_processing(),
+            clients(
+                client(
+                    certificate(file="security/clients.pem"),
+                    id="mtls",
+                    permissions="read,write",
+                ),
+                client(
+                    token(id=f"{VESPA_TOKEN_ID_WRITE}"),
+                    id="token_write",
+                    permissions="read,write",
+                ),
+            ),
+            config(
+                vt("tag")(
+                    vt("bold")(
+                        vt("open", "<strong>"),
+                        vt("close", "</strong>"),
+                    ),
+                    vt("separator", "..."),
+                ),
+                name="container.qr-searchers",
+            ),
+            id=f"{VESPA_APPLICATION_NAME}_container",
+            version="1.0",
+        ),
+        content(
+            redundancy("1"),
+            documents(document(type="pdf_page", mode="index")),
+            nodes(node(distribution_key="0", hostalias="node1")),
+            config(
+                vt("max_matches", "2", replace_underscores=False),
+                vt("length", "1000"),
+                vt("surround_max", "500", replace_underscores=False),
+                vt("min_length", "300", replace_underscores=False),
+                name="vespa.config.search.summary.juniperrc",
+            ),
+            id=f"{VESPA_APPLICATION_NAME}_content",
+            version="1.0",
+        ),
+        version="1.0",
+    ),
+)
+# -
+
+# Create the Vespa application package
+vespa_application_package = ApplicationPackage(
+    name=VESPA_APPLICATION_NAME,
+    schema=[colpali_schema],
+    services_config=service_config,
+)
+
+# ## 6. Deploy Vespa Application
+#
+
+VESPA_TEAM_API_KEY = os.getenv("VESPA_TEAM_API_KEY") or input(
+    "Enter Vespa team API key: "
+)
+
+# +
+vespa_cloud = VespaCloud(
+    tenant=VESPA_TENANT_NAME,
+    application=VESPA_APPLICATION_NAME,
+    key_content=VESPA_TEAM_API_KEY,
+    application_package=vespa_application_package,
+)
+
+# Deploy the application
+vespa_cloud.deploy()
+
+# Output the endpoint URL
+endpoint_url = vespa_cloud.get_token_endpoint()
+print(f"Application deployed. Token endpoint URL: {endpoint_url}")
+# -
+
+# Make sure to take note of the token endpoint_url.
+# You need to put this in your `.env` file - `VESPA_APP_URL=https://abcd.vespa-app.cloud` - to access the Vespa application from your web application.
+#
+
+# ## 8. Feed Data to Vespa
+#
+
+# Instantiate Vespa connection using token
+app = Vespa(url=endpoint_url, vespa_cloud_secret_token=VESPA_CLOUD_SECRET_TOKEN)
+app.get_application_status()
+
+
+# +
+def callback(response: VespaResponse, id: str):
+    if not response.is_successful():
+        print(
+            f"Failed to feed document {id} with status code {response.status_code}: Reason {response.get_json()}"
+        )
+
+
+# Feed data into Vespa asynchronously
+app.feed_async_iterable(vespa_feed, schema=VESPA_SCHEMA_NAME, callback=callback)

prepare_feed_deploy_v2.py

@@ -0,0 +1,956 @@
+# # Visual PDF Retrieval - demo application
+#
+# In this notebook, we will prepare the Vespa backend application for our visual retrieval demo.
+# We will use ColPali as the model to extract patch vectors from images of pdf pages.
+# At query time, we use MaxSim to retrieve and/or (based on the configuration) rank the page results.
+#
+# To see the application in action, visit TODO:
+#
+# The web application is written in FastHTML, meaning the complete application is written in python.
+#
+# The steps we will take in this notebook are:
+#
+# 0. Setup and configuration
+# 1. Download the data
+# 2. Prepare the data
+# 3. Generate queries for evaluation and typeahead search suggestions
+# 4. Deploy the Vespa application
+# 5. Create the Vespa application
+# 6. Feed the data to the Vespa application
+#
+# All the steps that are needed to provision the Vespa application, including feeding the data, can be done from this notebook.
+# We have tried to make it easy for others to run this notebook, to create your own PDF Enterprise Search application using Vespa.
+#
+
+# ## 0. Setup and Configuration
+#
+
+# +
+import os
+import asyncio
+import json
+from typing import Tuple
+import hashlib
+import numpy as np
+
+# Vespa
+from vespa.package import (
+    ApplicationPackage,
+    Field,
+    Schema,
+    Document,
+    HNSW,
+    RankProfile,
+    Function,
+    FieldSet,
+    SecondPhaseRanking,
+    Summary,
+    DocumentSummary,
+)
+from vespa.deployment import VespaCloud
+from vespa.application import Vespa
+from vespa.io import VespaResponse
+
+# Google Generative AI
+import google.generativeai as genai
+
+# Torch and other ML libraries
+import torch
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from pdf2image import convert_from_path
+from pypdf import PdfReader
+
+# ColPali model and processor
+from colpali_engine.models import ColPali, ColPaliProcessor
+from colpali_engine.utils.torch_utils import get_torch_device
+from vidore_benchmark.utils.image_utils import scale_image, get_base64_image
+
+# Other utilities
+from bs4 import BeautifulSoup
+import httpx
+from urllib.parse import urljoin, urlparse
+
+# Load environment variables
+from dotenv import load_dotenv
+
+load_dotenv()
+
+# Avoid warning from huggingface tokenizers
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+# -
+
+# ### Create a free trial in Vespa Cloud
+#
+# Create a tenant from [here](https://vespa.ai/free-trial/).
+# The trial includes $300 credit.
+# Take note of your tenant name.
+#
+
+VESPA_TENANT_NAME = "vespa-team"
+
+# Here, set your desired application name. (Will be created in later steps)
+# Note that you can not have hyphen `-` or underscore `_` in the application name.
+#
+
+VESPA_APPLICATION_NAME = "colpalidemo"
+VESPA_SCHEMA_NAME = "pdf_page"
+
+# Next, you need to create some tokens for feeding data, and querying the application.
+# We recommend separate tokens for feeding and querying, (the former with write permission, and the latter with read permission).
+# The tokens can be created from the [Vespa Cloud console](https://console.vespa-cloud.com/) in the 'Account' -> 'Tokens' section.
+#
+
+VESPA_TOKEN_ID_WRITE = "colpalidemo_write"
+
+# We also need to set the value of the write token to be able to feed data to the Vespa application.
+#
+
+VESPA_CLOUD_SECRET_TOKEN = os.getenv("VESPA_CLOUD_SECRET_TOKEN") or input(
+    "Enter Vespa cloud secret token: "
+)
+
+# We will also use the Gemini API to create sample queries for our images.
+# You can also use other VLM's to create these queries.
+# Create a Gemini API key from [here](https://aistudio.google.com/app/apikey).
+#
+
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY") or input(
+    "Enter Google Generative AI API key: "
+)
+
+# +
+MODEL_NAME = "vidore/colpali-v1.2"
+
+# Configure Google Generative AI
+genai.configure(api_key=GEMINI_API_KEY)
+
+# Set device for Torch
+device = get_torch_device("auto")
+print(f"Using device: {device}")
+
+# Load the ColPali model and processor
+model = ColPali.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+    device_map=device,
+).eval()
+
+processor = ColPaliProcessor.from_pretrained(MODEL_NAME)
+# -
+
+# ## 1. Download PDFs
+#
+# We are going to use public reports from the Norwegian Government Pension Fund Global (also known as the Oil Fund).
+# The fund puts transparency at the forefront and publishes reports on its investments, holdings, and returns, as well as its strategy and governance.
+#
+# These reports are the ones we are going to use for this showcase.
+# Here are some sample images:
+#
+# ![Sample1](./static/img/gfpg-sample-1.png)
+# ![Sample2](./static/img/gfpg-sample-2.png)
+#
+
+# As we can see, a lot of the information is in the form of tables, charts and numbers.
+# These are not easily extractable using pdf-readers or OCR tools.
+#
+
+# +
+import requests
+
+url = "https://www.nbim.no/en/publications/reports/"
+response = requests.get(url)
+response.raise_for_status()
+html_content = response.text
+
+# Parse with BeautifulSoup
+soup = BeautifulSoup(html_content, "html.parser")
+
+links = []
+url_to_year = {}
+
+# Find all 'div's with id starting with 'year-'
+for year_div in soup.find_all("div", id=lambda x: x and x.startswith("year-")):
+    year_id = year_div.get("id", "")
+    year = year_id.replace("year-", "")
+
+    # Within this div, find all 'a' elements with the specific classes
+    for a_tag in year_div.select("a.button.button--download-secondary[href]"):
+        href = a_tag["href"]
+        full_url = urljoin(url, href)
+        links.append(full_url)
+        url_to_year[full_url] = year
+links, url_to_year
+# -
+
+# Limit the number of PDFs to download
+NUM_PDFS = 2  # Set to None to download all PDFs
+links = links[:NUM_PDFS] if NUM_PDFS else links
+links
+
+# +
+from nest_asyncio import apply
+from typing import List
+
+apply()
+
+max_attempts = 3
+
+
+async def download_pdf(session, url, filename):
+    attempt = 0
+    while attempt < max_attempts:
+        try:
+            response = await session.get(url)
+            response.raise_for_status()
+
+            # Use Content-Disposition header to get the filename if available
+            content_disposition = response.headers.get("Content-Disposition")
+            if content_disposition:
+                import re
+
+                fname = re.findall('filename="(.+)"', content_disposition)
+                if fname:
+                    filename = fname[0]
+
+            # Ensure the filename is safe to use on the filesystem
+            safe_filename = filename.replace("/", "_").replace("\\", "_")
+            if not safe_filename or safe_filename == "_":
+                print(f"Invalid filename: {filename}")
+                continue
+
+            filepath = os.path.join("pdfs", safe_filename)
+            with open(filepath, "wb") as f:
+                f.write(response.content)
+            print(f"Downloaded {safe_filename}")
+            return filepath
+        except Exception as e:
+            print(f"Error downloading {filename}: {e}")
+            print(f"Retrying ({attempt})...")
+            await asyncio.sleep(1)  # Wait a bit before retrying
+            attempt += 1
+    return None
+
+
+async def download_pdfs(links: List[str]) -> List[dict]:
+    """Download PDFs from a list of URLs. Add the filename to the dictionary."""
+    async with httpx.AsyncClient() as client:
+        tasks = []
+
+        for idx, link in enumerate(links):
+            # Try to get the filename from the URL
+            path = urlparse(link).path
+            filename = os.path.basename(path)
+
+            # If filename is empty,skip
+            if not filename:
+                continue
+            tasks.append(download_pdf(client, link, filename))
+
+        # Run the tasks concurrently
+        paths = await asyncio.gather(*tasks)
+        pdf_files = [
+            {"url": link, "path": path} for link, path in zip(links, paths) if path
+        ]
+        return pdf_files
+
+
+# Create the pdfs directory if it doesn't exist
+os.makedirs("pdfs", exist_ok=True)
+# Now run the download_pdfs function with the URL
+pdfs = asyncio.run(download_pdfs(links))
+# -
+
+pdfs
+
+# ## 2. Convert PDFs to Images
+#
+
+
+# +
+def get_pdf_images(pdf_path):
+    reader = PdfReader(pdf_path)
+    page_texts = []
+    for page_number in range(len(reader.pages)):
+        page = reader.pages[page_number]
+        text = page.extract_text()
+        page_texts.append(text)
+    images = convert_from_path(pdf_path)
+    # Convert to PIL images
+    assert len(images) == len(page_texts)
+    return images, page_texts
+
+
+pdf_folder = "pdfs"
+pdf_pages = []
+for pdf in tqdm(pdfs):
+    pdf_file = pdf["path"]
+    title = os.path.splitext(os.path.basename(pdf_file))[0]
+    images, texts = get_pdf_images(pdf_file)
+    for page_no, (image, text) in enumerate(zip(images, texts)):
+        pdf_pages.append(
+            {
+                "title": title,
+                "year": int(url_to_year[pdf["url"]]),
+                "url": pdf["url"],
+                "path": pdf_file,
+                "image": image,
+                "text": text,
+                "page_no": page_no,
+            }
+        )
+# -
+
+len(pdf_pages)
+
+# +
+from collections import Counter
+
+# Print the length of the text fields - mean, max and min
+text_lengths = [len(page["text"]) for page in pdf_pages]
+print(f"Mean text length: {np.mean(text_lengths)}")
+print(f"Max text length: {np.max(text_lengths)}")
+print(f"Min text length: {np.min(text_lengths)}")
+print(f"Median text length: {np.median(text_lengths)}")
+print(f"Number of text with length == 0: {Counter(text_lengths)[0]}")
+# -
+
+# ## 3. Generate Queries
+#
+# In this step, we want to generate queries for each page image.
+# These will be useful for 2 reasons:
+#
+# 1. We can use these queries as typeahead suggestions in the search bar.
+# 2. We can use the queries to generate an evaluation dataset. See [Improving Retrieval with LLM-as-a-judge](https://blog.vespa.ai/improving-retrieval-with-llm-as-a-judge/) for a deeper dive into this topic.
+#
+# The prompt for generating queries is taken from [this](https://danielvanstrien.xyz/posts/post-with-code/colpali/2024-09-23-generate_colpali_dataset.html#an-update-retrieval-focused-prompt) wonderful blog post by Daniel van Strien.
+#
+# We will use the Gemini API to generate these queries, with `gemini-1.5-flash-8b` as the model.
+#
+
+# +
+from pydantic import BaseModel
+
+
+class GeneratedQueries(BaseModel):
+    broad_topical_question: str
+    broad_topical_query: str
+    specific_detail_question: str
+    specific_detail_query: str
+    visual_element_question: str
+    visual_element_query: str
+
+
+def get_retrieval_prompt() -> Tuple[str, GeneratedQueries]:
+    prompt = (
+        prompt
+    ) = """You are an investor, stock analyst and financial expert. You will be presented an image of a document page from a report published by the Norwegian Government Pension Fund Global (GPFG). The report may be annual or quarterly reports, or policy reports, on topics such as responsible investment, risk etc.
+Your task is to generate retrieval queries and questions that you would use to retrieve this document (or ask based on this document) in a large corpus.
+Please generate 3 different types of retrieval queries and questions.
+A retrieval query is a keyword based query, made up of 2-5 words, that you would type into a search engine to find this document.
+A question is a natural language question that you would ask, for which the document contains the answer.
+The queries should be of the following types:
+1. A broad topical query: This should cover the main subject of the document.
+2. A specific detail query: This should cover a specific detail or aspect of the document.
+3. A visual element query: This should cover a visual element of the document, such as a chart, graph, or image.
+
+Important guidelines:
+- Ensure the queries are relevant for retrieval tasks, not just describing the page content.
+- Use a fact-based natural language style for the questions.
+- Frame the queries as if someone is searching for this document in a large corpus.
+- Make the queries diverse and representative of different search strategies.
+
+Format your response as a JSON object with the structure of the following example:
+{
+    "broad_topical_question": "What was the Responsible Investment Policy in 2019?",
+    "broad_topical_query": "responsible investment policy 2019",
+    "specific_detail_question": "What is the percentage of investments in renewable energy?",
+    "specific_detail_query": "renewable energy investments percentage",
+    "visual_element_question": "What is the trend of total holding value over time?",
+    "visual_element_query": "total holding value trend"
+}
+
+If there are no relevant visual elements, provide an empty string for the visual element question and query.
+Here is the document image to analyze:
+Generate the queries based on this image and provide the response in the specified JSON format.
+Only return JSON. Don't return any extra explanation text. """
+
+    return prompt, GeneratedQueries
+
+
+prompt_text, pydantic_model = get_retrieval_prompt()
+
+# +
+gemini_model = genai.GenerativeModel("gemini-1.5-flash-8b")
+
+
+def generate_queries(image, prompt_text, pydantic_model):
+    try:
+        response = gemini_model.generate_content(
+            [image, "\n\n", prompt_text],
+            generation_config=genai.GenerationConfig(
+                response_mime_type="application/json",
+                response_schema=pydantic_model,
+            ),
+        )
+        queries = json.loads(response.text)
+    except Exception as _e:
+        queries = {
+            "broad_topical_question": "",
+            "broad_topical_query": "",
+            "specific_detail_question": "",
+            "specific_detail_query": "",
+            "visual_element_question": "",
+            "visual_element_query": "",
+        }
+    return queries
+
+
+# -
+
+for pdf in tqdm(pdf_pages):
+    image = pdf.get("image")
+    pdf["queries"] = generate_queries(image, prompt_text, pydantic_model)
+
+pdf_pages[46]["image"]
+
+pdf_pages[46]["queries"]
+
+# +
+# Generate queries async - keeping for now as we probably need when applying to the full dataset
+# import asyncio
+# from tenacity import retry, stop_after_attempt, wait_exponential
+# import google.generativeai as genai
+# from tqdm.asyncio import tqdm_asyncio
+
+# max_in_flight = 200  # Maximum number of concurrent requests
+
+
+# async def generate_queries_for_image_async(model, image, semaphore):
+#     @retry(stop=stop_after_attempt(3), wait=wait_exponential(), reraise=True)
+#     async def _generate():
+#         async with semaphore:
+#             result = await model.generate_content_async(
+#                 [image, "\n\n", prompt_text],
+#                 generation_config=genai.GenerationConfig(
+#                     response_mime_type="application/json",
+#                     response_schema=pydantic_model,
+#                 ),
+#             )
+#             return json.loads(result.text)
+
+#     try:
+#         return await _generate()
+#     except Exception as e:
+#         print(f"Error generating queries for image: {e}")
+#         return None  # Return None or handle as needed
+
+
+# async def enrich_pdfs():
+#     gemini_model = genai.GenerativeModel("gemini-1.5-flash-8b")
+#     semaphore = asyncio.Semaphore(max_in_flight)
+#     tasks = []
+#     for pdf in pdf_pages:
+#         pdf["queries"] = []
+#         image = pdf.get("image")
+#         if image:
+#             task = generate_queries_for_image_async(gemini_model, image, semaphore)
+#             tasks.append((pdf, task))
+
+#     # Run the tasks concurrently using asyncio.gather()
+#     for pdf, task in tqdm_asyncio(tasks):
+#         result = await task
+#         if result:
+#             pdf["queries"] = result
+#     return pdf_pages
+
+
+# pdf_pages = asyncio.run(enrich_pdfs())
+
+# +
+# write title, url, page_no, text, queries, not image to JSON
+with open("output/pdf_pages.json", "w") as f:
+    to_write = [{k: v for k, v in pdf.items() if k != "image"} for pdf in pdf_pages]
+    json.dump(to_write, f, indent=2)
+
+# with open("pdfs/pdf_pages.json", "r") as f:
+#     saved_pdf_pages = json.load(f)
+# for pdf, saved_pdf in zip(pdf_pages, saved_pdf_pages):
+#     pdf.update(saved_pdf)
+# -
+
+# ## 4. Generate embeddings
+#
+# Now that we have the queries, we can use the ColPali model to generate embeddings for each page image.
+#
+
+
+def generate_embeddings(images, model, processor, batch_size=2) -> np.ndarray:
+    """
+    Generate embeddings for a list of images.
+    Move to CPU only once per batch.
+
+    Args:
+        images (List[PIL.Image]): List of PIL images.
+        model (nn.Module): The model to generate embeddings.
+        processor: The processor to preprocess images.
+        batch_size (int, optional): Batch size for processing. Defaults to 64.
+
+    Returns:
+        np.ndarray: Embeddings for the images, shape
+                    (len(images), processor.max_patch_length (1030 for ColPali), model.config.hidden_size (Patch embedding dimension - 128 for ColPali)).
+    """
+    embeddings_list = []
+
+    def collate_fn(batch):
+        # Batch is a list of images
+        return processor.process_images(batch)  # Should return a dict of tensors
+
+    dataloader = DataLoader(
+        images,
+        shuffle=False,
+        collate_fn=collate_fn,
+    )
+
+    for batch_doc in tqdm(dataloader, desc="Generating embeddings"):
+        with torch.no_grad():
+            # Move batch to the device
+            batch_doc = {k: v.to(model.device) for k, v in batch_doc.items()}
+            embeddings_batch = model(**batch_doc)
+            embeddings_list.append(torch.unbind(embeddings_batch.to("cpu"), dim=0))
+    # Concatenate all embeddings and create a numpy array
+    all_embeddings = np.concatenate(embeddings_list, axis=0)
+    return all_embeddings
+
+
+# Generate embeddings for all images
+images = [pdf["image"] for pdf in pdf_pages]
+embeddings = generate_embeddings(images, model, processor)
+
+embeddings.shape
+
+# ## 5. Prepare Data on Vespa Format
+#
+# Now, that we have all the data we need, all that remains is to make sure it is in the right format for Vespa.
+#
+
+
+def float_to_binary_embedding(float_query_embedding: dict) -> dict:
+    """Utility function to convert float query embeddings to binary query embeddings."""
+    binary_query_embeddings = {}
+    for k, v in float_query_embedding.items():
+        binary_vector = (
+            np.packbits(np.where(np.array(v) > 0, 1, 0)).astype(np.int8).tolist()
+        )
+        binary_query_embeddings[k] = binary_vector
+    return binary_query_embeddings
+
+
+vespa_feed = []
+for pdf, embedding in zip(pdf_pages, embeddings):
+    url = pdf["url"]
+    year = pdf["year"]
+    title = pdf["title"]
+    image = pdf["image"]
+    text = pdf.get("text", "")
+    page_no = pdf["page_no"]
+    query_dict = pdf["queries"]
+    questions = [v for k, v in query_dict.items() if "question" in k and v]
+    queries = [v for k, v in query_dict.items() if "query" in k and v]
+    base_64_image = get_base64_image(
+        scale_image(image, 32), add_url_prefix=False
+    )  # Scaled down image to return fast on search (~1kb)
+    base_64_full_image = get_base64_image(image, add_url_prefix=False)
+    embedding_dict = {k: v for k, v in enumerate(embedding)}
+    binary_embedding = float_to_binary_embedding(embedding_dict)
+    # id_hash should be md5 hash of url and page_number
+    id_hash = hashlib.md5(f"{url}_{page_no}".encode()).hexdigest()
+    page = {
+        "id": id_hash,
+        "fields": {
+            "id": id_hash,
+            "url": url,
+            "title": title,
+            "year": year,
+            "page_number": page_no,
+            "blur_image": base_64_image,
+            "full_image": base_64_full_image,
+            "text": text,
+            "embedding": binary_embedding,
+            "queries": queries,
+            "questions": questions,
+        },
+    }
+    vespa_feed.append(page)
+
+# +
+# We will prepare the Vespa feed data, including the embeddings and the generated queries
+
+
+# Save vespa_feed to vespa_feed.json
+os.makedirs("output", exist_ok=True)
+with open("output/vespa_feed.json", "w") as f:
+    vespa_feed_to_save = []
+    for page in vespa_feed:
+        document_id = page["id"]
+        put_id = f"id:{VESPA_APPLICATION_NAME}:{VESPA_SCHEMA_NAME}::{document_id}"
+        vespa_feed_to_save.append({"put": put_id, "fields": page["fields"]})
+    json.dump(vespa_feed_to_save, f)
+
+# +
+# import json
+
+# with open("output/vespa_feed.json", "r") as f:
+#     vespa_feed = json.load(f)
+# -
+
+len(vespa_feed)
+
+# ## 5. Prepare Vespa Application
+#
+
+# +
+# Define the Vespa schema
+colpali_schema = Schema(
+    name=VESPA_SCHEMA_NAME,
+    document=Document(
+        fields=[
+            Field(
+                name="id",
+                type="string",
+                indexing=["summary", "index"],
+                match=["word"],
+            ),
+            Field(name="url", type="string", indexing=["summary", "index"]),
+            Field(name="year", type="int", indexing=["summary", "attribute"]),
+            Field(
+                name="title",
+                type="string",
+                indexing=["summary", "index"],
+                match=["text"],
+                index="enable-bm25",
+            ),
+            Field(name="page_number", type="int", indexing=["summary", "attribute"]),
+            Field(name="blur_image", type="raw", indexing=["summary"]),
+            Field(name="full_image", type="raw", indexing=["summary"]),
+            Field(
+                name="text",
+                type="string",
+                indexing=["summary", "index"],
+                match=["text"],
+                index="enable-bm25",
+            ),
+            Field(
+                name="embedding",
+                type="tensor<int8>(patch{}, v[16])",
+                indexing=[
+                    "attribute",
+                    "index",
+                ],
+                ann=HNSW(
+                    distance_metric="hamming",
+                    max_links_per_node=32,
+                    neighbors_to_explore_at_insert=400,
+                ),
+            ),
+            Field(
+                name="questions",
+                type="array<string>",
+                indexing=["summary", "attribute"],
+                summary=Summary(fields=["matched-elements-only"]),
+            ),
+            Field(
+                name="queries",
+                type="array<string>",
+                indexing=["summary", "attribute"],
+                summary=Summary(fields=["matched-elements-only"]),
+            ),
+        ]
+    ),
+    fieldsets=[
+        FieldSet(
+            name="default",
+            fields=["title", "url", "blur_image", "page_number", "text"],
+        ),
+        FieldSet(
+            name="image",
+            fields=["full_image"],
+        ),
+    ],
+    document_summaries=[
+        DocumentSummary(
+            name="default",
+            summary_fields=[
+                Summary(
+                    name="text",
+                    fields=[("bolding", "on")],
+                ),
+                Summary(
+                    name="snippet",
+                    fields=[("source", "text"), "dynamic"],
+                ),
+            ],
+            from_disk=True,
+        ),
+        DocumentSummary(
+            name="suggestions",
+            summary_fields=[
+                Summary(name="questions"),
+            ],
+            from_disk=True,
+        ),
+    ],
+)
+
+# Define similarity functions used in all rank profiles
+mapfunctions = [
+    Function(
+        name="similarities",  # computes similarity scores between each query token and image patch
+        expression="""
+                sum(
+                    query(qt) * unpack_bits(attribute(embedding)), v
+                )
+            """,
+    ),
+    Function(
+        name="normalized",  # normalizes the similarity scores to [-1, 1]
+        expression="""
+                (similarities - reduce(similarities, min)) / (reduce((similarities - reduce(similarities, min)), max)) * 2 - 1
+            """,
+    ),
+    Function(
+        name="quantized",  # quantizes the normalized similarity scores to signed 8-bit integers [-128, 127]
+        expression="""
+                cell_cast(normalized * 127.999, int8)
+            """,
+    ),
+]
+
+# Define the 'bm25' rank profile
+colpali_bm25_profile = RankProfile(
+    name="bm25",
+    inputs=[("query(qt)", "tensor<float>(querytoken{}, v[128])")],
+    first_phase="bm25(title) + bm25(text)",
+    functions=mapfunctions,
+)
+
+
+# A function to create an inherited rank profile which also returns quantized similarity scores
+def with_quantized_similarity(rank_profile: RankProfile) -> RankProfile:
+    return RankProfile(
+        name=f"{rank_profile.name}_sim",
+        first_phase=rank_profile.first_phase,
+        inherits=rank_profile.name,
+        summary_features=["quantized"],
+    )
+
+
+colpali_schema.add_rank_profile(colpali_bm25_profile)
+colpali_schema.add_rank_profile(with_quantized_similarity(colpali_bm25_profile))
+
+# Update the 'default' rank profile
+colpali_profile = RankProfile(
+    name="default",
+    inputs=[("query(qt)", "tensor<float>(querytoken{}, v[128])")],
+    first_phase="bm25_score",
+    second_phase=SecondPhaseRanking(expression="max_sim", rerank_count=10),
+    functions=mapfunctions
+    + [
+        Function(
+            name="max_sim",
+            expression="""
+                sum(
+                    reduce(
+                        sum(
+                            query(qt) * unpack_bits(attribute(embedding)), v
+                        ),
+                        max, patch
+                    ),
+                    querytoken
+                )
+            """,
+        ),
+        Function(name="bm25_score", expression="bm25(title) + bm25(text)"),
+    ],
+)
+colpali_schema.add_rank_profile(colpali_profile)
+colpali_schema.add_rank_profile(with_quantized_similarity(colpali_profile))
+
+# Update the 'retrieval-and-rerank' rank profile
+input_query_tensors = []
+MAX_QUERY_TERMS = 64
+for i in range(MAX_QUERY_TERMS):
+    input_query_tensors.append((f"query(rq{i})", "tensor<int8>(v[16])"))
+
+input_query_tensors.extend(
+    [
+        ("query(qt)", "tensor<float>(querytoken{}, v[128])"),
+        ("query(qtb)", "tensor<int8>(querytoken{}, v[16])"),
+    ]
+)
+
+colpali_retrieval_profile = RankProfile(
+    name="retrieval-and-rerank",
+    inputs=input_query_tensors,
+    first_phase="max_sim_binary",
+    second_phase=SecondPhaseRanking(expression="max_sim", rerank_count=10),
+    functions=mapfunctions
+    + [
+        Function(
+            name="max_sim",
+            expression="""
+                sum(
+                    reduce(
+                        sum(
+                            query(qt) * unpack_bits(attribute(embedding)), v
+                        ),
+                        max, patch
+                    ),
+                    querytoken
+                )
+            """,
+        ),
+        Function(
+            name="max_sim_binary",
+            expression="""
+                sum(
+                    reduce(
+                        1 / (1 + sum(
+                            hamming(query(qtb), attribute(embedding)), v)
+                        ),
+                        max, patch
+                    ),
+                    querytoken
+                )
+            """,
+        ),
+    ],
+)
+colpali_schema.add_rank_profile(colpali_retrieval_profile)
+colpali_schema.add_rank_profile(with_quantized_similarity(colpali_retrieval_profile))
+
+# +
+from vespa.configuration.services import (
+    services,
+    container,
+    search,
+    document_api,
+    document_processing,
+    clients,
+    client,
+    config,
+    content,
+    redundancy,
+    documents,
+    node,
+    certificate,
+    token,
+    document,
+    nodes,
+)
+from vespa.configuration.vt import vt
+from vespa.package import ServicesConfiguration
+
+service_config = ServicesConfiguration(
+    application_name=VESPA_APPLICATION_NAME,
+    services_config=services(
+        container(
+            search(),
+            document_api(),
+            document_processing(),
+            clients(
+                client(
+                    certificate(file="security/clients.pem"),
+                    id="mtls",
+                    permissions="read,write",
+                ),
+                client(
+                    token(id=f"{VESPA_TOKEN_ID_WRITE}"),
+                    id="token_write",
+                    permissions="read,write",
+                ),
+            ),
+            config(
+                vt("tag")(
+                    vt("bold")(
+                        vt("open", "<strong>"),
+                        vt("close", "</strong>"),
+                    ),
+                    vt("separator", "..."),
+                ),
+                name="container.qr-searchers",
+            ),
+            id=f"{VESPA_APPLICATION_NAME}_container",
+            version="1.0",
+        ),
+        content(
+            redundancy("1"),
+            documents(document(type="pdf_page", mode="index")),
+            nodes(node(distribution_key="0", hostalias="node1")),
+            config(
+                vt("max_matches", "2", replace_underscores=False),
+                vt("length", "1000"),
+                vt("surround_max", "500", replace_underscores=False),
+                vt("min_length", "300", replace_underscores=False),
+                name="vespa.config.search.summary.juniperrc",
+            ),
+            id=f"{VESPA_APPLICATION_NAME}_content",
+            version="1.0",
+        ),
+        version="1.0",
+    ),
+)
+# -
+
+# Create the Vespa application package
+vespa_application_package = ApplicationPackage(
+    name=VESPA_APPLICATION_NAME,
+    schema=[colpali_schema],
+    services_config=service_config,
+)
+
+# ## 6. Deploy Vespa Application
+#
+
+VESPA_TEAM_API_KEY = os.getenv("VESPA_TEAM_API_KEY") or input(
+    "Enter Vespa team API key: "
+)
+
+# +
+vespa_cloud = VespaCloud(
+    tenant=VESPA_TENANT_NAME,
+    application=VESPA_APPLICATION_NAME,
+    key_content=VESPA_TEAM_API_KEY,
+    application_package=vespa_application_package,
+)
+
+# Deploy the application
+vespa_cloud.deploy()
+
+# Output the endpoint URL
+endpoint_url = vespa_cloud.get_token_endpoint()
+print(f"Application deployed. Token endpoint URL: {endpoint_url}")
+# -
+
+# Make sure to take note of the token endpoint_url.
+# You need to put this in your `.env` file - `VESPA_APP_URL=https://abcd.vespa-app.cloud` - to access the Vespa application from your web application.
+#
+
+# ## 8. Feed Data to Vespa
+#
+
+# Instantiate Vespa connection using token
+app = Vespa(url=endpoint_url, vespa_cloud_secret_token=VESPA_CLOUD_SECRET_TOKEN)
+app.get_application_status()
+
+
+# +
+def callback(response: VespaResponse, id: str):
+    if not response.is_successful():
+        print(
+            f"Failed to feed document {id} with status code {response.status_code}: Reason {response.get_json()}"
+        )
+
+
+# Feed data into Vespa asynchronously
+app.feed_async_iterable(vespa_feed, schema=VESPA_SCHEMA_NAME, callback=callback)

pyproject.toml

@@ -0,0 +1,119 @@
+[project]
+name = "visual-retrieval-colpali"
+version = "0.1.0"
+description = "Visual retrieval with ColPali"
+readme = "README.md"
+requires-python = ">=3.10, <3.13" 
+license = { text = "Apache-2.0" }
+dependencies = [
+    "python-fasthtml",
+    "huggingface-hub",
+    "pyvespa>=0.50.0",
+    "vespacli",
+    "torch",
+    "vidore-benchmark[interpretability]>=4.0.0,<5.0.0",
+    "colpali-engine",
+    "einops",
+    "pypdf",
+    "setuptools",
+    "python-dotenv",
+    "shad4fast>=1.2.1",
+    "google-generativeai>=0.7.2",
+    "spacy",
+    "pip",
+    "matplotlib"
+]
+
+# dev-dependencies
+[project.optional-dependencies]
+dev = [
+    "ruff",
+    "python-dotenv",
+    "huggingface_hub[cli]"
+]
+feed = [
+    "ipykernel",
+    "jupytext",
+    "pydantic",
+    "beautifulsoup4",
+    "pdf2image",
+    "google-generativeai"
+]
+[tool.ruff]
+# Exclude a variety of commonly ignored directories.
+exclude = [
+    ".bzr",
+    ".direnv",
+    ".eggs",
+    ".git",
+    ".git-rewrite",
+    ".hg",
+    ".ipynb_checkpoints",
+    ".mypy_cache",
+    ".nox",
+    ".pants.d",
+    ".pyenv",
+    ".pytest_cache",
+    ".pytype",
+    ".ruff_cache",
+    ".svn",
+    ".tox",
+    ".venv",
+    ".vscode",
+    "__pypackages__",
+    "_build",
+    "buck-out",
+    "build",
+    "dist",
+    "node_modules",
+    "site-packages",
+    "venv",
+]
+
+# Same as Black.
+line-length = 88
+indent-width = 4
+
+# Assume Python 3.8
+target-version = "py38"
+
+[tool.ruff.lint]
+# Enable Pyflakes (`F`) and a subset of the pycodestyle (`E`)  codes by default.
+# Unlike Flake8, Ruff doesn't enable pycodestyle warnings (`W`) or
+# McCabe complexity (`C901`) by default.
+select = ["E4", "E7", "E9", "F"]
+ignore = []
+
+# Allow fix for all enabled rules (when `--fix`) is provided.
+fixable = ["ALL"]
+unfixable = []
+
+# Allow unused variables when underscore-prefixed.
+dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
+
+[tool.ruff.format]
+# Like Black, use double quotes for strings.
+quote-style = "double"
+
+# Like Black, indent with spaces, rather than tabs.
+indent-style = "space"
+
+# Like Black, respect magic trailing commas.
+skip-magic-trailing-comma = false
+
+# Like Black, automatically detect the appropriate line ending.
+line-ending = "auto"
+
+# Enable auto-formatting of code examples in docstrings. Markdown,
+# reStructuredText code/literal blocks and doctests are all supported.
+#
+# This is currently disabled by default, but it is planned for this
+# to be opt-out in the future.
+docstring-code-format = false
+
+# Set the line length limit used when formatting code snippets in
+# docstrings.
+#
+# This only has an effect when the `docstring-code-format` setting is
+# enabled.
+docstring-code-line-length = "dynamic"

query_vespa.py

@@ -0,0 +1,193 @@
+#!/usr/bin/env python3
+
+import os
+import torch
+from torch.utils.data import DataLoader
+from PIL import Image
+import numpy as np
+from typing import cast
+import asyncio
+
+from colpali_engine.models import ColPali, ColPaliProcessor
+from colpali_engine.utils.torch_utils import get_torch_device
+from vespa.application import Vespa
+from vespa.io import VespaQueryResponse
+from dotenv import load_dotenv
+from pathlib import Path
+
+MAX_QUERY_TERMS = 64
+SAVEDIR = Path(__file__) / "output" / "images"
+load_dotenv()
+
+
+def process_queries(processor, queries, image):
+    inputs = processor(
+        images=[image] * len(queries), text=queries, return_tensors="pt", padding=True
+    )
+    return inputs
+
+
+def display_query_results(query, response, hits=5):
+    query_time = response.json.get("timing", {}).get("searchtime", -1)
+    query_time = round(query_time, 2)
+    count = response.json.get("root", {}).get("fields", {}).get("totalCount", 0)
+    result_text = f"Query text: '{query}', query time {query_time}s, count={count}, top results:\n"
+
+    for i, hit in enumerate(response.hits[:hits]):
+        title = hit["fields"]["title"]
+        url = hit["fields"]["url"]
+        page = hit["fields"]["page_number"]
+        image = hit["fields"]["image"]
+        _id = hit["id"]
+        score = hit["relevance"]
+
+        result_text += f"\nPDF Result {i + 1}\n"
+        result_text += f"Title: {title}, page {page+1} with score {score:.2f}\n"
+        result_text += f"URL: {url}\n"
+        result_text += f"ID: {_id}\n"
+        # Optionally, save or display the image
+        # img_data = base64.b64decode(image)
+        # img_path = SAVEDIR / f"{title}.png"
+        # with open(f"{img_path}", "wb") as f:
+        #     f.write(img_data)
+    print(result_text)
+
+
+async def query_vespa_default(app, queries, qs):
+    async with app.asyncio(connections=1, total_timeout=120) as session:
+        for idx, query in enumerate(queries):
+            query_embedding = {k: v.tolist() for k, v in enumerate(qs[idx])}
+            response: VespaQueryResponse = await session.query(
+                yql="select documentid,title,url,image,page_number from pdf_page where userInput(@userQuery)",
+                ranking="default",
+                userQuery=query,
+                timeout=120,
+                hits=3,
+                body={"input.query(qt)": query_embedding, "presentation.timing": True},
+            )
+            assert response.is_successful()
+            display_query_results(query, response)
+
+
+async def query_vespa_nearest_neighbor(app, queries, qs):
+    # Using nearestNeighbor for retrieval
+    target_hits_per_query_tensor = (
+        20  # this is a hyper parameter that can be tuned for speed versus accuracy
+    )
+    async with app.asyncio(connections=1, total_timeout=180) as session:
+        for idx, query in enumerate(queries):
+            float_query_embedding = {k: v.tolist() for k, v in enumerate(qs[idx])}
+            binary_query_embeddings = dict()
+            for k, v in float_query_embedding.items():
+                binary_vector = (
+                    np.packbits(np.where(np.array(v) > 0, 1, 0))
+                    .astype(np.int8)
+                    .tolist()
+                )
+                binary_query_embeddings[k] = binary_vector
+                if len(binary_query_embeddings) >= MAX_QUERY_TERMS:
+                    print(
+                        f"Warning: Query has more than {MAX_QUERY_TERMS} terms. Truncating."
+                    )
+                    break
+
+            # The mixed tensors used in MaxSim calculations
+            # We use both binary and float representations
+            query_tensors = {
+                "input.query(qtb)": binary_query_embeddings,
+                "input.query(qt)": float_query_embedding,
+            }
+            # The query tensors used in the nearest neighbor calculations
+            for i in range(0, len(binary_query_embeddings)):
+                query_tensors[f"input.query(rq{i})"] = binary_query_embeddings[i]
+            nn = []
+            for i in range(0, len(binary_query_embeddings)):
+                nn.append(
+                    f"({{targetHits:{target_hits_per_query_tensor}}}nearestNeighbor(embedding,rq{i}))"
+                )
+            # We use an OR operator to combine the nearest neighbor operator
+            nn = " OR ".join(nn)
+            response: VespaQueryResponse = await session.query(
+                body={
+                    **query_tensors,
+                    "presentation.timing": True,
+                    "yql": f"select documentid, title, url, image, page_number from pdf_page where {nn}",
+                    "ranking.profile": "retrieval-and-rerank",
+                    "timeout": 120,
+                    "hits": 3,
+                },
+            )
+            assert response.is_successful(), response.json
+            display_query_results(query, response)
+
+
+def main():
+    vespa_app_url = os.environ.get(
+        "VESPA_APP_URL"
+    )  # Ensure this is set to your Vespa app URL
+    vespa_cloud_secret_token = os.environ.get("VESPA_CLOUD_SECRET_TOKEN")
+    if not vespa_app_url or not vespa_cloud_secret_token:
+        raise ValueError(
+            "Please set the VESPA_APP_URL and VESPA_CLOUD_SECRET_TOKEN environment variables"
+        )
+    # Instantiate Vespa connection
+    app = Vespa(url=vespa_app_url, vespa_cloud_secret_token=vespa_cloud_secret_token)
+    status_resp = app.get_application_status()
+    if status_resp.status_code != 200:
+        print(f"Failed to connect to Vespa at {vespa_app_url}")
+        return
+    else:
+        print(f"Connected to Vespa at {vespa_app_url}")
+    # Load the model
+    device = get_torch_device("auto")
+    print(f"Using device: {device}")
+
+    model_name = "vidore/colpali-v1.2"
+    processor_name = "google/paligemma-3b-mix-448"
+
+    model = cast(
+        ColPali,
+        ColPali.from_pretrained(
+            model_name,
+            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+            device_map=device,
+        ),
+    ).eval()
+
+    processor = cast(ColPaliProcessor, ColPaliProcessor.from_pretrained(processor_name))
+
+    # Create dummy image
+    dummy_image = Image.new("RGB", (448, 448), (255, 255, 255))
+
+    # Define queries
+    queries = [
+        "Percentage of non-fresh water as source?",
+        "Policies related to nature risk?",
+        "How much of produced water is recycled?",
+    ]
+
+    # Obtain query embeddings
+    dataloader = DataLoader(
+        queries,
+        batch_size=1,
+        shuffle=False,
+        collate_fn=lambda x: process_queries(processor, x, dummy_image),
+    )
+    qs = []
+    for batch_query in dataloader:
+        with torch.no_grad():
+            batch_query = {k: v.to(model.device) for k, v in batch_query.items()}
+            embeddings_query = model(**batch_query)
+            qs.extend(list(torch.unbind(embeddings_query.to("cpu"))))
+
+    # Perform queries using default rank profile
+    print("Performing queries using default rank profile:")
+    asyncio.run(query_vespa_default(app, queries, qs))
+
+    # Perform queries using nearestNeighbor
+    print("Performing queries using nearestNeighbor:")
+    asyncio.run(query_vespa_nearest_neighbor(app, queries, qs))
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,540 @@
+# This file was autogenerated by uv via the following command:
+#    uv pip compile pyproject.toml -o src/requirements.txt
+accelerate==0.34.2
+    # via peft
+aiohappyeyeballs==2.4.3
+    # via aiohttp
+aiohttp==3.10.11
+    # via
+    #   datasets
+    #   fsspec
+    #   pyvespa
+aiosignal==1.3.1
+    # via aiohttp
+annotated-types==0.7.0
+    # via pydantic
+anyio==4.6.0
+    # via
+    #   httpx
+    #   starlette
+    #   watchfiles
+async-timeout==4.0.3
+    # via aiohttp
+attrs==24.2.0
+    # via aiohttp
+beautifulsoup4==4.12.3
+    # via python-fasthtml
+blis==0.7.11
+    # via thinc
+cachetools==5.5.0
+    # via google-auth
+catalogue==2.0.10
+    # via
+    #   spacy
+    #   srsly
+    #   thinc
+certifi==2024.8.30
+    # via
+    #   httpcore
+    #   httpx
+    #   requests
+cffi==1.17.1
+    # via cryptography
+charset-normalizer==3.3.2
+    # via requests
+click==8.1.7
+    # via
+    #   typer
+    #   uvicorn
+cloudpathlib==0.20.0
+    # via weasel
+colpali-engine==0.3.1
+    # via
+    #   visual-retrieval-colpali (pyproject.toml)
+    #   vidore-benchmark
+confection==0.1.5
+    # via
+    #   thinc
+    #   weasel
+contourpy==1.3.0
+    # via matplotlib
+cryptography==43.0.1
+    # via pyvespa
+cycler==0.12.1
+    # via matplotlib
+cymem==2.0.8
+    # via
+    #   preshed
+    #   spacy
+    #   thinc
+datasets==2.21.0
+    # via
+    #   mteb
+    #   vidore-benchmark
+dill==0.3.8
+    # via
+    #   datasets
+    #   multiprocess
+docker==7.1.0
+    # via pyvespa
+einops==0.8.0
+    # via
+    #   visual-retrieval-colpali (pyproject.toml)
+    #   vidore-benchmark
+eval-type-backport==0.2.0
+    # via mteb
+exceptiongroup==1.2.2
+    # via anyio
+fastcore==1.7.11
+    # via
+    #   fastlite
+    #   python-fasthtml
+    #   pyvespa
+    #   sqlite-minutils
+fastlite==0.0.11
+    # via python-fasthtml
+filelock==3.16.1
+    # via
+    #   datasets
+    #   huggingface-hub
+    #   torch
+    #   transformers
+fonttools==4.54.1
+    # via matplotlib
+frozenlist==1.4.1
+    # via
+    #   aiohttp
+    #   aiosignal
+fsspec==2024.6.1
+    # via
+    #   datasets
+    #   huggingface-hub
+    #   torch
+google-ai-generativelanguage==0.6.10
+    # via google-generativeai
+google-api-core==2.21.0
+    # via
+    #   google-ai-generativelanguage
+    #   google-api-python-client
+    #   google-generativeai
+google-api-python-client==2.149.0
+    # via google-generativeai
+google-auth==2.35.0
+    # via
+    #   google-ai-generativelanguage
+    #   google-api-core
+    #   google-api-python-client
+    #   google-auth-httplib2
+    #   google-generativeai
+google-auth-httplib2==0.2.0
+    # via google-api-python-client
+google-generativeai==0.8.3
+    # via visual-retrieval-colpali (pyproject.toml)
+googleapis-common-protos==1.65.0
+    # via
+    #   google-api-core
+    #   grpcio-status
+gputil==1.4.0
+    # via
+    #   colpali-engine
+    #   vidore-benchmark
+grpcio==1.67.0
+    # via
+    #   google-api-core
+    #   grpcio-status
+grpcio-status==1.67.0
+    # via google-api-core
+h11==0.14.0
+    # via
+    #   httpcore
+    #   uvicorn
+h2==4.1.0
+    # via httpx
+hpack==4.0.0
+    # via h2
+httpcore==1.0.6
+    # via httpx
+httplib2==0.22.0
+    # via
+    #   google-api-python-client
+    #   google-auth-httplib2
+httptools==0.6.1
+    # via uvicorn
+httpx==0.27.2
+    # via
+    #   python-fasthtml
+    #   pyvespa
+huggingface-hub==0.25.1
+    # via
+    #   visual-retrieval-colpali (pyproject.toml)
+    #   accelerate
+    #   datasets
+    #   peft
+    #   sentence-transformers
+    #   tokenizers
+    #   transformers
+hyperframe==6.0.1
+    # via h2
+idna==3.10
+    # via
+    #   anyio
+    #   httpx
+    #   requests
+    #   yarl
+itsdangerous==2.2.0
+    # via python-fasthtml
+jinja2==3.1.5
+    # via
+    #   pyvespa
+    #   spacy
+    #   torch
+joblib==1.4.2
+    # via scikit-learn
+kiwisolver==1.4.7
+    # via matplotlib
+langcodes==3.4.1
+    # via spacy
+language-data==1.2.0
+    # via langcodes
+loguru==0.7.2
+    # via vidore-benchmark
+lucide-fasthtml==0.0.9
+    # via shad4fast
+lxml==5.3.0
+    # via
+    #   lucide-fasthtml
+    #   pyvespa
+marisa-trie==1.2.1
+    # via language-data
+markdown-it-py==3.0.0
+    # via rich
+markupsafe==2.1.5
+    # via jinja2
+matplotlib==3.9.2
+    # via
+    #   seaborn
+    #   vidore-benchmark
+mdurl==0.1.2
+    # via markdown-it-py
+mpmath==1.3.0
+    # via sympy
+mteb==1.15.3
+    # via vidore-benchmark
+multidict==6.1.0
+    # via
+    #   aiohttp
+    #   yarl
+multiprocess==0.70.16
+    # via datasets
+murmurhash==1.0.10
+    # via
+    #   preshed
+    #   spacy
+    #   thinc
+networkx==3.3
+    # via torch
+numpy==1.26.4
+    # via
+    #   accelerate
+    #   blis
+    #   colpali-engine
+    #   contourpy
+    #   datasets
+    #   matplotlib
+    #   mteb
+    #   pandas
+    #   peft
+    #   pyarrow
+    #   scikit-learn
+    #   scipy
+    #   seaborn
+    #   spacy
+    #   thinc
+    #   transformers
+    #   vidore-benchmark
+oauthlib==3.2.2
+    # via python-fasthtml
+packaging==24.1
+    # via
+    #   accelerate
+    #   datasets
+    #   fastcore
+    #   huggingface-hub
+    #   matplotlib
+    #   peft
+    #   spacy
+    #   thinc
+    #   transformers
+    #   weasel
+pandas==2.2.3
+    # via
+    #   datasets
+    #   seaborn
+pdf2image==1.17.0
+    # via vidore-benchmark
+peft==0.11.1
+    # via
+    #   colpali-engine
+    #   vidore-benchmark
+pillow==10.4.0
+    # via
+    #   colpali-engine
+    #   matplotlib
+    #   pdf2image
+    #   sentence-transformers
+    #   vidore-benchmark
+pip==24.3.1
+    # via visual-retrieval-colpali (pyproject.toml)
+polars==1.9.0
+    # via mteb
+preshed==3.0.9
+    # via
+    #   spacy
+    #   thinc
+proto-plus==1.24.0
+    # via
+    #   google-ai-generativelanguage
+    #   google-api-core
+protobuf==5.28.3
+    # via
+    #   google-ai-generativelanguage
+    #   google-api-core
+    #   google-generativeai
+    #   googleapis-common-protos
+    #   grpcio-status
+    #   proto-plus
+psutil==6.0.0
+    # via
+    #   accelerate
+    #   peft
+pyarrow==17.0.0
+    # via datasets
+pyasn1==0.6.1
+    # via
+    #   pyasn1-modules
+    #   rsa
+pyasn1-modules==0.4.1
+    # via google-auth
+pycparser==2.22
+    # via cffi
+pydantic==2.9.2
+    # via
+    #   confection
+    #   google-generativeai
+    #   mteb
+    #   spacy
+    #   thinc
+    #   weasel
+pydantic-core==2.23.4
+    # via pydantic
+pygments==2.18.0
+    # via rich
+pyparsing==3.1.4
+    # via
+    #   httplib2
+    #   matplotlib
+pypdf==5.0.1
+    # via visual-retrieval-colpali (pyproject.toml)
+python-dateutil==2.9.0.post0
+    # via
+    #   matplotlib
+    #   pandas
+    #   python-fasthtml
+    #   pyvespa
+python-dotenv==1.0.1
+    # via
+    #   visual-retrieval-colpali (pyproject.toml)
+    #   uvicorn
+    #   vidore-benchmark
+python-fasthtml==0.6.9
+    # via
+    #   visual-retrieval-colpali (pyproject.toml)
+    #   lucide-fasthtml
+    #   shad4fast
+python-multipart==0.0.18
+    # via python-fasthtml
+pytrec-eval-terrier==0.5.6
+    # via mteb
+pytz==2024.2
+    # via pandas
+pyvespa==0.50.0
+    # via visual-retrieval-colpali (pyproject.toml)
+pyyaml==6.0.2
+    # via
+    #   accelerate
+    #   datasets
+    #   huggingface-hub
+    #   peft
+    #   transformers
+    #   uvicorn
+regex==2024.9.11
+    # via transformers
+requests==2.32.3
+    # via
+    #   colpali-engine
+    #   datasets
+    #   docker
+    #   google-api-core
+    #   huggingface-hub
+    #   lucide-fasthtml
+    #   mteb
+    #   pyvespa
+    #   requests-toolbelt
+    #   spacy
+    #   transformers
+    #   weasel
+requests-toolbelt==1.0.0
+    # via pyvespa
+rich==13.9.2
+    # via
+    #   mteb
+    #   typer
+rsa==4.9
+    # via google-auth
+safetensors==0.4.5
+    # via
+    #   accelerate
+    #   peft
+    #   transformers
+scikit-learn==1.5.2
+    # via
+    #   mteb
+    #   sentence-transformers
+scipy==1.14.1
+    # via
+    #   mteb
+    #   scikit-learn
+    #   sentence-transformers
+seaborn==0.13.2
+    # via vidore-benchmark
+sentence-transformers==3.1.1
+    # via
+    #   mteb
+    #   vidore-benchmark
+sentencepiece==0.2.0
+    # via vidore-benchmark
+setuptools==75.1.0
+    # via
+    #   visual-retrieval-colpali (pyproject.toml)
+    #   marisa-trie
+    #   spacy
+    #   thinc
+shad4fast==1.2.1
+    # via visual-retrieval-colpali (pyproject.toml)
+shellingham==1.5.4
+    # via typer
+six==1.16.0
+    # via python-dateutil
+smart-open==7.0.5
+    # via weasel
+sniffio==1.3.1
+    # via
+    #   anyio
+    #   httpx
+soupsieve==2.6
+    # via beautifulsoup4
+spacy==3.7.5
+    # via visual-retrieval-colpali (pyproject.toml)
+spacy-legacy==3.0.12
+    # via spacy
+spacy-loggers==1.0.5
+    # via spacy
+sqlite-minutils==3.37.0.post3
+    # via fastlite
+srsly==2.4.8
+    # via
+    #   confection
+    #   spacy
+    #   thinc
+    #   weasel
+starlette==0.39.2
+    # via python-fasthtml
+sympy==1.13.3
+    # via torch
+tenacity==9.0.0
+    # via pyvespa
+thinc==8.2.5
+    # via spacy
+threadpoolctl==3.5.0
+    # via scikit-learn
+tokenizers==0.20.0
+    # via transformers
+torch==2.4.1
+    # via
+    #   visual-retrieval-colpali (pyproject.toml)
+    #   accelerate
+    #   colpali-engine
+    #   mteb
+    #   peft
+    #   sentence-transformers
+    #   vidore-benchmark
+tqdm==4.66.5
+    # via
+    #   datasets
+    #   google-generativeai
+    #   huggingface-hub
+    #   mteb
+    #   peft
+    #   sentence-transformers
+    #   spacy
+    #   transformers
+transformers==4.45.1
+    # via
+    #   colpali-engine
+    #   peft
+    #   sentence-transformers
+    #   vidore-benchmark
+typer==0.12.5
+    # via
+    #   spacy
+    #   vidore-benchmark
+    #   weasel
+typing-extensions==4.12.2
+    # via
+    #   anyio
+    #   cloudpathlib
+    #   google-generativeai
+    #   huggingface-hub
+    #   mteb
+    #   multidict
+    #   pydantic
+    #   pydantic-core
+    #   pypdf
+    #   pyvespa
+    #   rich
+    #   torch
+    #   typer
+    #   uvicorn
+tzdata==2024.2
+    # via pandas
+uritemplate==4.1.1
+    # via google-api-python-client
+urllib3==2.2.3
+    # via
+    #   docker
+    #   requests
+uvicorn==0.31.0
+    # via python-fasthtml
+uvloop==0.20.0
+    # via uvicorn
+vespacli==8.391.23
+    # via visual-retrieval-colpali (pyproject.toml)
+vidore-benchmark==4.0.0
+    # via visual-retrieval-colpali (pyproject.toml)
+wasabi==1.1.3
+    # via
+    #   spacy
+    #   thinc
+    #   weasel
+watchfiles==0.24.0
+    # via uvicorn
+weasel==0.4.1
+    # via spacy
+websockets==13.1
+    # via uvicorn
+wrapt==1.16.0
+    # via smart-open
+xxhash==3.5.0
+    # via datasets
+yarl==1.13.1
+    # via aiohttp

ruff.toml

@@ -0,0 +1,77 @@
+# Exclude a variety of commonly ignored directories.
+exclude = [
+    ".bzr",
+    ".direnv",
+    ".eggs",
+    ".git",
+    ".git-rewrite",
+    ".hg",
+    ".ipynb_checkpoints",
+    ".mypy_cache",
+    ".nox",
+    ".pants.d",
+    ".pyenv",
+    ".pytest_cache",
+    ".pytype",
+    ".ruff_cache",
+    ".svn",
+    ".tox",
+    ".venv",
+    ".vscode",
+    "__pypackages__",
+    "_build",
+    "buck-out",
+    "build",
+    "dist",
+    "node_modules",
+    "site-packages",
+    "venv",
+]
+
+# Same as Black.
+line-length = 88
+indent-width = 4
+
+# Assume Python 3.8
+target-version = "py38"
+
+[lint]
+# Enable Pyflakes (`F`) and a subset of the pycodestyle (`E`)  codes by default.
+# Unlike Flake8, Ruff doesn't enable pycodestyle warnings (`W`) or
+# McCabe complexity (`C901`) by default.
+select = ["E4", "E7", "E9", "F"]
+ignore = []
+
+# Allow fix for all enabled rules (when `--fix`) is provided.
+fixable = ["ALL"]
+unfixable = []
+
+# Allow unused variables when underscore-prefixed.
+dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
+
+[format]
+# Like Black, use double quotes for strings.
+quote-style = "double"
+
+# Like Black, indent with spaces, rather than tabs.
+indent-style = "space"
+
+# Like Black, respect magic trailing commas.
+skip-magic-trailing-comma = false
+
+# Like Black, automatically detect the appropriate line ending.
+line-ending = "auto"
+
+# Enable auto-formatting of code examples in docstrings. Markdown,
+# reStructuredText code/literal blocks and doctests are all supported.
+#
+# This is currently disabled by default, but it is planned for this
+# to be opt-out in the future.
+docstring-code-format = false
+
+# Set the line length limit used when formatting code snippets in
+# docstrings.
+#
+# This only has an effect when the `docstring-code-format` setting is
+# enabled.
+docstring-code-line-length = "dynamic"

setup.py

@@ -0,0 +1,104 @@
+#!/usr/bin/env python3
+"""
+Quick setup script for ColPali-Vespa Visual Retrieval System
+"""
+
+import os
+import sys
+from pathlib import Path
+
+
+def create_env_file():
+    """Create a sample .env file if it doesn't exist"""
+    env_path = Path(".env")
+    if env_path.exists():
+        print("✅ .env file already exists")
+        return
+
+    env_content = """# Vespa Configuration
+# Choose one authentication method:
+
+# Option 1: Token Authentication (Recommended)
+VESPA_APP_TOKEN_URL=https://your-app.your-tenant.vespa-cloud.com
+VESPA_CLOUD_SECRET_TOKEN=your_vespa_secret_token_here
+
+# Option 2: mTLS Authentication
+# USE_MTLS=true
+# VESPA_APP_MTLS_URL=https://your-app.your-tenant.vespa-cloud.com
+# VESPA_CLOUD_MTLS_KEY="-----BEGIN PRIVATE KEY-----
+# Your private key content here
+# -----END PRIVATE KEY-----"
+# VESPA_CLOUD_MTLS_CERT="-----BEGIN CERTIFICATE-----
+# Your certificate content here
+# -----END CERTIFICATE-----"
+
+# Google Gemini Configuration (Optional - for AI chat features)
+GEMINI_API_KEY=your_gemini_api_key_here
+
+# Application Configuration
+LOG_LEVEL=INFO
+HOT_RELOAD=false
+
+# Development Configuration
+# Uncomment for development mode
+# HOT_RELOAD=true
+# LOG_LEVEL=DEBUG
+"""
+
+    with open(env_path, "w") as f:
+        f.write(env_content)
+
+    print("✅ Created .env file with sample configuration")
+    print("   Please edit .env with your actual credentials")
+
+
+def create_directories():
+    """Create necessary directories"""
+    directories = ["static", "static/full_images", "static/sim_maps"]
+
+    for directory in directories:
+        Path(directory).mkdir(parents=True, exist_ok=True)
+
+    print("✅ Created necessary directories")
+
+
+def check_python_version():
+    """Check if Python version is compatible"""
+    version = sys.version_info
+    if version.major != 3 or version.minor < 10 or version.minor >= 13:
+        print("❌ Python 3.10, 3.11, or 3.12 is required")
+        print(f"   Current version: {version.major}.{version.minor}.{version.micro}")
+        return False
+
+    print(
+        f"✅ Python version {version.major}.{version.minor}.{version.micro} is compatible"
+    )
+    return True
+
+
+def main():
+    """Main setup function"""
+    print("🚀 ColPali-Vespa Visual Retrieval Setup")
+    print("=" * 40)
+
+    # Check Python version
+    if not check_python_version():
+        sys.exit(1)
+
+    # Create directories
+    create_directories()
+
+    # Create .env file
+    create_env_file()
+
+    print("\n📋 Next steps:")
+    print("1. Edit .env file with your Vespa and Gemini credentials")
+    print("2. Install dependencies: pip install -e .")
+    print("3. Deploy Vespa application: python deploy_vespa_app.py ...")
+    print("4. Upload documents: python feed_vespa.py ...")
+    print("5. Run the application: python main.py")
+    print("\n📖 See README.md for detailed instructions")
+
+
+if __name__ == "__main__":
+    main()