Colbert

[michaelfeil]

Description

Please provide a clear and concise description of the changes in this PR.

Related Issue

If applicable, link the issue this PR addresses.

Types of Change

• Bug fix
• New feature
• Documentation update

Checklist

• I have read the CONTRIBUTING guidelines.
• My code follows the code style of this project.
• I have added tests to cover my changes.
• All new and existing tests passed.
• My changes generate no new warnings.
• I have updated the documentation accordingly.

Additional Notes

Add any other context about the PR here.

License

By submitting this PR, I confirm that my contribution is made under the terms of the MIT license.

[greptile-apps]

PR Summary

Here's a concise summary of the key changes in this PR that adds ColBERT model support:

Added support for multiple model replicas and device placement across GPUs:

• Introduced new --device-id CLI argument to specify GPU/MPS device placement (e.g. "0,1") for model distribution
• Added LoadingStrategy class to manage device mapping, dtype selection, and quantization across replicas
• Modified BatchHandler to support multiple ModelWorker instances running parallel model replicas
• Updated AsyncEmbeddingEngine to handle multiple model replicas instead of single instance

Added ColBERT-specific functionality:

• Introduced mode_colbert flag in SentenceTransformerPatched for ColBERT model detection
• Modified embedding processing to handle ColBERT's token-level embeddings and attention masks
• Added normalize_embeddings flag to control embedding normalization behavior
• Added test coverage for ColBERT models with token-level embedding verification

The changes enable efficient parallel processing across devices while maintaining backward compatibility with existing models.

_{25 file(s) reviewed, 9 comment(s)
Edit PR Review Bot Settings | Greptile}

[greptile-apps]

logic: device_id is not included in from_env() method's zip_longest parameters, causing potential initialization issues

[greptile-apps]

logic: redundant check for not self.mode_colbert since it's already in an if not self.mode_colbert block

[codecov-commenter]

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 91.66667% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.19%. Comparing base (317e809) to head (e6fe213).
⚠️ Report is 110 commits behind head on main.

Files with missing lines	Patch %	Lines
libs/infinity_emb/infinity_emb/infinity_server.py	75.00%	2 Missing ⚠️
❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #456      +/-   ##
==========================================
+ Coverage   79.12%   79.19%   +0.07%     
==========================================
  Files          42       42              
  Lines        3367     3379      +12     
==========================================
+ Hits         2664     2676      +12     
  Misses        703      703              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.