Updated ColPali example to use binary quantization [skip ci]

Author: ankane

examples/colbert/exact_binary.py

@@ -1,6 +1,6 @@
 from colpali_engine.models import ColQwen2, ColQwen2Processor
 from datasets import load_dataset
-from pgvector.psycopg import register_vector
+from pgvector.psycopg import register_vector, Bit
 import psycopg
 import torch
 
@@ -10,17 +10,17 @@
 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 TABLE documents (id bigserial PRIMARY KEY, embeddings bit(128)[])')
 conn.execute("""
-CREATE OR REPLACE FUNCTION max_sim(document vector[], query vector[]) RETURNS double precision AS $$
+CREATE OR REPLACE FUNCTION max_sim(document bit[], query bit[]) 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
+        SELECT query_number, 1 - ((document <~> query) / bit_length(query)) AS similarity FROM queries CROSS JOIN documents
     ),
     max_similarities AS (
         SELECT MAX(similarity) AS max_similarity FROM similarities GROUP BY query_number
@@ -40,13 +40,17 @@ def generate_embeddings(processed):
         return model(**processed.to(model.device)).to(device='cpu', dtype=torch.float32)
 
 
+def binary_quantize(embedding):
+    return Bit(embedding > 0)
+
+
 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]]
+    embeddings = [binary_quantize(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]]
+query_embeddings = [binary_quantize(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)