Instructions to use KaLM-Embedding/KaLM-Reranker-V1-Small with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Transformers
How to use KaLM-Embedding/KaLM-Reranker-V1-Small with Transformers:
# Load model directly from transformers import AutoProcessor, AutoModelForMultimodalLM processor = AutoProcessor.from_pretrained("KaLM-Embedding/KaLM-Reranker-V1-Small") model = AutoModelForMultimodalLM.from_pretrained("KaLM-Embedding/KaLM-Reranker-V1-Small", device_map="auto") - Notebooks
- Google Colab
- Kaggle
fix(reranker): avoid re-computing the first batch in predict()'s batch-size probe
Existing Issue:
KaLMReranker.predict() runs a full forward pass on the first batch twice: once as a throwaway OOM-size probe, and once again for the real computation. The probe's result is computed but never used. This roughly doubles reranking latency for any call where the total number of documents fits within a single batch (the common case, since batch_size defaults to 32).
Proposed Fix:
Keep the probe's result and reuse it as the first batch's scores, starting the real loop after the first batch instead of from the beginning.
This PR introduces a bug whereby, if CUDA OOM errors occur for every batch size up to tested_batch_size=1, the first batch will never be computed. While this is likely to be a rare edge case, it is a newly introduced bug. I will therefore close my PR, thoroughly test my adapted solution, and submit a new PR.