Added ColPali / ColQwen2 example [skip ci]

Author: ankane

README.md

@@ -33,6 +33,7 @@ Or check out some examples:
 - [Hybrid search](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search/cross_encoder.py) with SentenceTransformers (cross-encoder)
 - [Sparse search](https://github.com/pgvector/pgvector-python/blob/master/examples/sparse_search/example.py) with Transformers
 - [Late interaction search](https://github.com/pgvector/pgvector-python/blob/master/examples/colbert/exact.py) with ColBERT
+- [Document retrieval](https://github.com/pgvector/pgvector-python/blob/master/examples/colpali/exact.py) with ColPali
 - [Image search](https://github.com/pgvector/pgvector-python/blob/master/examples/image_search/example.py) with PyTorch
 - [Image search](https://github.com/pgvector/pgvector-python/blob/master/examples/imagehash/example.py) with perceptual hashing
 - [Morgan fingerprints](https://github.com/pgvector/pgvector-python/blob/master/examples/rdkit/example.py) with RDKit

examples/colpali/exact.py

@@ -0,0 +1,52 @@
+from colpali_engine.models import ColQwen2, ColQwen2Processor
+from datasets import load_dataset
+from pgvector.psycopg import register_vector
+import psycopg
+import torch
+
+conn = psycopg.connect(dbname='pgvector_example', autocommit=True)
+
+conn.execute('CREATE EXTENSION IF NOT EXISTS vector')
+register_vector(conn)
+
+conn.execute('DROP TABLE IF EXISTS documents')
+conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, embeddings vector(128)[])')
+conn.execute("""
+CREATE OR REPLACE FUNCTION max_sim(document vector[], query vector[]) RETURNS double precision AS $$
+    WITH queries AS (
+        SELECT row_number() OVER () AS query_number, * FROM (SELECT unnest(query) AS query)
+    ),
+    documents AS (
+        SELECT unnest(document) AS document
+    ),
+    similarities AS (
+        SELECT query_number, 1 - (document <=> query) AS similarity FROM queries CROSS JOIN documents
+    ),
+    max_similarities AS (
+        SELECT MAX(similarity) AS max_similarity FROM similarities GROUP BY query_number
+    )
+    SELECT SUM(max_similarity) FROM max_similarities
+$$ LANGUAGE SQL
+""")
+
+
+device = 'mps' if torch.backends.mps.is_available() else 'cpu'
+model = ColQwen2.from_pretrained('vidore/colqwen2-v1.0', torch_dtype=torch.bfloat16, device_map=device).eval()
+processor = ColQwen2Processor.from_pretrained('vidore/colqwen2-v1.0')
+
+
+def generate_embeddings(processed):
+    with torch.no_grad():
+        return model(**processed.to(model.device)).to(device='cpu', dtype=torch.float32)
+
+
+input = load_dataset('vidore/docvqa_test_subsampled', split='test[:3]')['image']
+for content in input:
+    embeddings = [e.numpy() for e in generate_embeddings(processor.process_images([content]))[0]]
+    conn.execute('INSERT INTO documents (embeddings) VALUES (%s)', (embeddings,))
+
+query = 'dividend'
+query_embeddings = [e.numpy() for e in generate_embeddings(processor.process_queries([query]))[0]]
+result = conn.execute('SELECT id, max_sim(embeddings, %s) AS max_sim FROM documents ORDER BY max_sim DESC LIMIT 5', (query_embeddings,)).fetchall()
+for row in result:
+    print(row)

examples/colpali/requirements.txt

@@ -0,0 +1,4 @@
+colpali-engine
+datasets
+pgvector
+psycopg[binary]