mirror of
https://github.com/qdrant/qdrant.git
synced 2026-09-25 15:37:42 -05:00
96 lines
2.7 KiB
Python
96 lines
2.7 KiB
Python
#!/usr/bin/env python3
|
|
"""ColBERT-style late-interaction search: each document is stored as a matrix of per-token embeddings in a named multi-vector field"""
|
|
|
|
import os
|
|
import shutil
|
|
from pathlib import Path
|
|
|
|
from qdrant_edge import (
|
|
Distance,
|
|
EdgeConfig,
|
|
EdgeShard,
|
|
EdgeVectorParams,
|
|
MultiVectorComparator,
|
|
MultiVectorConfig,
|
|
Point,
|
|
Query,
|
|
QueryRequest,
|
|
UpdateOperation,
|
|
)
|
|
|
|
|
|
DATA_DIR = Path(__file__).parent.parent.parent / "data"
|
|
TMP_DIR = DATA_DIR / "tmp"
|
|
path = TMP_DIR / "qdrant_edge_multivector_colbert"
|
|
shutil.rmtree(path, ignore_errors=True)
|
|
os.makedirs(path)
|
|
|
|
# The "colbert" field stores a matrix of vectors per point instead of one,
|
|
# with MaxSim comparing every query row against every stored row.
|
|
config = EdgeConfig(
|
|
vectors={
|
|
"colbert": EdgeVectorParams(
|
|
size=4,
|
|
distance=Distance.Cosine,
|
|
multivector_config=MultiVectorConfig(MultiVectorComparator.MaxSim),
|
|
),
|
|
},
|
|
)
|
|
shard = EdgeShard.create(path, config)
|
|
|
|
documents = [
|
|
(1, "mountain hiking trail guide", [
|
|
[0.85, 0.10, 0.30, 0.05],
|
|
[0.80, 0.15, 0.35, 0.08],
|
|
[0.82, 0.12, 0.32, 0.10],
|
|
]),
|
|
(2, "sourdough bread baking recipe", [
|
|
[0.10, 0.85, 0.15, 0.30],
|
|
[0.15, 0.80, 0.10, 0.35],
|
|
[0.12, 0.88, 0.18, 0.28],
|
|
]),
|
|
(3, "vintage synthesizer sound design tutorial", [
|
|
[0.20, 0.10, 0.85, 0.15],
|
|
[0.25, 0.08, 0.80, 0.20],
|
|
[0.18, 0.12, 0.82, 0.18],
|
|
[0.22, 0.10, 0.78, 0.22],
|
|
]),
|
|
(4, "alpine trekking gear checklist", [
|
|
[0.88, 0.08, 0.28, 0.10],
|
|
[0.83, 0.12, 0.30, 0.12],
|
|
[0.86, 0.10, 0.25, 0.08],
|
|
]),
|
|
]
|
|
|
|
shard.update(UpdateOperation.upsert_points([
|
|
Point(point_id, {"colbert": token_vectors}, {"text": text})
|
|
for point_id, text, token_vectors in documents
|
|
]))
|
|
|
|
# The query is also a matrix - one row per query token.
|
|
query_tokens = [
|
|
[0.85, 0.10, 0.30, 0.05],
|
|
[0.84, 0.11, 0.27, 0.09],
|
|
]
|
|
|
|
results = shard.query(QueryRequest(
|
|
query=Query.Nearest(query_tokens, using="colbert"),
|
|
limit=4,
|
|
with_payload=True,
|
|
))
|
|
|
|
print("=== ColBERT-style MaxSim search ===")
|
|
for point in results:
|
|
print(f"id={point.id} score={point.score:.4f} text={point.payload['text']}")
|
|
|
|
assert results[0].id == 1, "the hiking document should score highest for a hiking-like query"
|
|
|
|
# Round-trip check: the stored matrix keeps one row per token (rows are
|
|
# normalized on write since the field uses cosine distance).
|
|
records = shard.retrieve([1], with_payload=False, with_vector=True)
|
|
stored = records[0].vector["colbert"]
|
|
print(f"\nStored matrix for point 1 has {len(stored)} rows")
|
|
assert len(stored) == len(documents[0][2])
|
|
|
|
print("\nDone!")
|