Skip to content

Add support for WARP#211

Open
pau-mensa wants to merge 6 commits intolightonai:mainfrom
pau-mensa:main
Open

Add support for WARP#211
pau-mensa wants to merge 6 commits intolightonai:mainfrom
pau-mensa:main

Conversation

@pau-mensa
Copy link
Copy Markdown

Summary

  • Add a new WARP index backend wrapping xtr-warp-rs and export it from pylate.indexes
  • Implement add_documents, remove_documents (with auto-compaction), and search with optional document ID filtering (shared or per-query)
  • Add update_documents and compact as new public methods not present in the base Index class. update_documents replaces embeddings in-place without delete+add, compact physically removes tombstoned passages, empty centroids and reclaims disk space
  • Expose all xtr-warp-rs hyperparameters via two config dataclasses (WARPSearchConfig, WARPIndexingConfig) to keep the constructor clean
  • Extract shared embedding conversion utilities (convert_embeddings_to_torch) into pylate/indexes/utils.py for reuse across WARP and FastPLAID
  • Add xtr-warp-rs as an optional dependency
  • Add 12 focused WARP unit tests

Tests

  • python -m pytest tests/test_warp.py

@pau-mensa
Copy link
Copy Markdown
Author

Notes

  • WARP supports metadata filtering now, but I did not add it because I saw that the other indices don't have it
  • WARP also supports streamed indexing, which can cut VRAM usage by 66%, but that would require changing things that seem a bit out of scope. Still, it is a powerful feature, let me know if you want me to add it

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant