多模态 RAG：ColPali 让 VLM 直接读 PDF

# Late Interaction: 不在索引阶段做 dot product
# 而在查询阶段做 max-sim
def score(query_embeddings, doc_patch_embeddings):
    """
    query_embeddings: [Q, d]  Q 个 query token embedding
    doc_patch_embeddings: [P, d]  P 个 patch embedding
    """
    # 每个 query token 与最相似的 patch 取 max
    sim_matrix = query_embeddings @ doc_patch_embeddings.T  # [Q, P]
    max_sim_per_query = sim_matrix.max(dim=-1).values  # [Q]
    return max_sim_per_query.sum()  # 总分

from colpali_engine.models import ColQwen2_5, ColQwen2_5_Processor
from PIL import Image
import torch

# 1. 加载模型
model = ColQwen2_5.from_pretrained(
    "vidore/colqwen2.5-v0.2",
    torch_dtype=torch.bfloat16,
    device_map="cuda",
)
processor = ColQwen2_5_Processor.from_pretrained("vidore/colqwen2.5-v0.2")

# 2. PDF 转图片
def pdf_to_images(pdf_path, dpi=150):
    images = []
    # 用 pdf2image 或 PyMuPDF
    for page_num in range(num_pages):
        img = render_page(pdf_path, page_num, dpi=dpi)
        images.append(img)
    return images

# 3. 生成 patch embeddings
def index_documents(pdf_paths):
    all_embeddings = []
    for pdf_path in pdf_paths:
        images = pdf_to_images(pdf_path)
        # 每个 page 独立编码
        batch = processor.process_images(images).to(model.device)
        with torch.no_grad():
            embeddings = model(**batch)  # [num_pages, num_patches, dim]
        all_embeddings.append({
            "pdf": pdf_path,
            "embeddings": embeddings.cpu().numpy(),
            "page_count": len(images),
        })
    return all_embeddings

import numpy as np
from qdrant_client import QdrantClient

# 1. 用 Qdrant 存储 patch embeddings（按 page 分组）
client = QdrantClient("localhost", port=6333)
client.create_collection(
    collection_name="multimodal_docs",
    vectors_config=models.VectorParams(
        size=128,  # ColQwen2.5 patch dim
        distance=models.Distance.COSINE,
    ),
)

# 索引每个 patch
def index_patches(pdf_path, embeddings, page_count):
    points = []
    for page_idx in range(page_count):
        page_emb = embeddings[page_idx]  # [num_patches, dim]
        for patch_idx, emb in enumerate(page_emb):
            points.append(models.PointStruct(
                id=hash(f"{pdf_path}_{page_idx}_{patch_idx}"),
                vector=emb.tolist(),
                payload={
                    "pdf": pdf_path,
                    "page": page_idx,
                    "patch_idx": patch_idx,
                },
            ))
    client.upsert(collection_name="multimodal_docs", points=points)

# 2. 检索：Late Interaction
def search(query, top_k=10):
    # 编码 query
    query_emb = processor.process_queries([query]).to(model.device)
    with torch.no_grad():
        query_features = model(**query_emb)  # [1, num_query_tokens, dim]
    
    # 从向量库取候选 page 的所有 patches
    candidates = client.search(
        collection_name="multimodal_docs",
        query_vector=query_features[0].mean(dim=0).tolist(),  # 初筛用均值
        limit=top_k * 5,
    )
    
    # Late Interaction 精排
    scores = []
    for candidate in candidates:
        doc_patches = np.array(candidate.vector)  # 简化：实际是多向量
        # 完整 late interaction 需重算
        sim = late_interaction_score(query_features[0], doc_patches)
        scores.append((candidate.payload, sim))
    
    # 按 score 排序返回
    scores.sort(key=lambda x: x[1], reverse=True)
    return scores[:top_k]

from openai import OpenAI

def generate_answer(query, retrieved_pages):
    """把检索到的 PDF 页面截图直接喂给 VLM"""
    client = OpenAI()
    
    content = [{"type": "text", "text": f"参考以下文档页面回答问题：\n\n问题：{query}\n\n"}]
    
    for page_info in retrieved_pages:
        # 把检索到的 page 图片加入
        page_image = load_page_image(page_info["pdf"], page_info["page"])
        content.append({
            "type": "image_url",
            "image_url": {"url": f"data:image/png;base64,{encode_image(page_image)}"},
        })
    
    response = client.chat.completions.create(
        model="gpt-4o",  # 或 Claude 3.5 Sonnet
        messages=[{"role": "user", "content": content}],
        max_tokens=1000,
    )
    return response.choices[0].message.content

def hybrid_retrieve(query, k=10):
    # 1. 传统管线（快、便宜）
    text_results = bge_retrieve(query, k=k)
    
    # 2. ColPali（慢、准）
    visual_results = colpali_retrieve(query, k=k)
    
    # 3. Reciprocal Rank Fusion
    fused_scores = {}
    for rank, doc in enumerate(text_results):
        fused_scores[doc.id] = fused_scores.get(doc.id, 0) + 1 / (rank + 60)
    for rank, doc in enumerate(visual_results):
        fused_scores[doc.id] = fused_scores.get(doc.id, 0) + 1 / (rank + 60)
    
    # 4. 排序返回
    sorted_docs = sorted(fused_scores.items(), key=lambda x: x[1], reverse=True)
    return [doc for doc_id, score in sorted_docs[:k]]

def fast_retrieval(query, top_k=10):
    # 1. 粗筛：用平均 embedding 在 IVF 索引中找 top-100 候选页
    query_mean = query_embeddings.mean(dim=0)  # [dim]
    candidates = ivf_index.search(query_mean, k=100)  # 5ms
    
    # 2. 精排：只在 100 个候选页上跑完整 late interaction
    scores = []
    for page_id in candidates:
        page_patches = load_page_patches(page_id)  # 从 SSD 读
        sim = late_interaction(query_embeddings, page_patches)  # 1ms × 100 = 100ms
        scores.append((page_id, sim))
    
    # 3. 排序返回
    scores.sort(key=lambda x: x[1], reverse=True)
    return scores[:top_k]

class IncrementalColPaliIndex:
    def __init__(self, base_index_path):
        self.base = load_index(base_index_path)  # 已有索引
        self.pending_updates = []
    
    def add_pdf(self, pdf_path):
        # 1. 只对新 PDF 生成 patch embedding
        embeddings = colpali.encode(pdf_to_images(pdf_path))
        
        # 2. 暂存不写入
        self.pending_updates.append({
            "pdf": pdf_path,
            "embeddings": embeddings,
            "timestamp": now(),
        })
    
    def flush(self, threshold=100):
        if len(self.pending_updates) >= threshold:
            self.base.merge(self.pending_updates)
            self.base.save()
            self.pending_updates = []