feat: create standalone IVF training API in python #2553
Conversation
…ng with IVF training
| f"Vector column {c} must be FixedSizeListArray " | ||
| f"1-dimensional FixedShapeTensorArray, got {field.type}" | ||
| ) | ||
| if not pa.types.is_floating(field.type.value_type): |
There was a problem hiding this comment.
We can support uint8 as well
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## main #2553 +/- ##
==========================================
+ Coverage 79.95% 80.01% +0.05%
==========================================
Files 207 209 +2
Lines 59500 59821 +321
Branches 59500 59821 +321
==========================================
+ Hits 47576 47865 +289
- Misses 9122 9131 +9
- Partials 2802 2825 +23
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. |
| self.distance_type = distance_type | ||
| """The distance type used to train the IVF model""" | ||
|
|
||
| def num_partitions(self) -> int: |
There was a problem hiding this comment.
Seems like this should be a property:
| def num_partitions(self) -> int: | |
| @property | |
| def num_partitions(self) -> int: |
There was a problem hiding this comment.
Changed to a property
|
|
||
| K-means is an iterative algorithm that can be computationally expensive. The | ||
| accelerator argument can be used to offload the computation to a hardware | ||
| accelerator such as a GPU or TPU. |
There was a problem hiding this comment.
Is there a progress bar for this? Should we mention the LANCE_LOG variable here?
There was a problem hiding this comment.
Yes, there is a tqdm progress bar if an accelerator is used. I started to add one if an accelerator is not used but it ended up being a fair amount of code. Will open a different PR.
| from .dependencies import torch | ||
|
|
||
|
|
||
| class IvfModel: |
There was a problem hiding this comment.
Not necessary for now, but it would be very cool to be able to get this from an existing index.
|
I'm going to merge this on green. I have a few PRs built on top of this one already. |
This is the start of breaking IVF / PQ index training up into four standalone methods with check pointing in between each method.