mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-18 12:08:35 +00:00
fix: harden deterministic IVF-PQ training
This commit is contained in:
@@ -772,7 +772,9 @@ class IvfPq:
|
||||
seed: int, optional
|
||||
Seed used for deterministic sampling and training. Given identical data in
|
||||
the same row order and identical index parameters, the same seed produces
|
||||
the same IVF centroids and PQ codebook. If omitted, training remains random.
|
||||
the same IVF centroids and PQ codebook. This option is supported for local
|
||||
CPU index builds; remote and accelerator builds reject it explicitly. If
|
||||
omitted, training remains random.
|
||||
|
||||
target_partition_size: int, default is 8192
|
||||
|
||||
@@ -794,6 +796,12 @@ class IvfPq:
|
||||
# create_index() dispatches to pylance to build the index on the accelerator.
|
||||
accelerator: Optional[str] = None
|
||||
|
||||
def __post_init__(self):
|
||||
if self.seed is not None and self.accelerator is not None:
|
||||
raise ValueError(
|
||||
"IvfPq seed is not supported with accelerator-based index training"
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class IvfRq:
|
||||
|
||||
@@ -395,6 +395,11 @@ async def test_create_4bit_ivfpq_index(some_table: AsyncTable):
|
||||
assert stats.num_indices == 1
|
||||
|
||||
|
||||
def test_seeded_ivfpq_rejects_accelerator():
|
||||
with pytest.raises(ValueError, match="seed is not supported with accelerator"):
|
||||
IvfPq(seed=42, accelerator="cuda")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_ivfrq_index(some_table: AsyncTable):
|
||||
await some_table.create_index("vector", config=IvfRq(num_bits=1))
|
||||
|
||||
Reference in New Issue
Block a user