Fix transformers 5.x loading and add Sentence Transformers usage

README.md

@@ -19,6 +19,7 @@ tags:
 - ColBERT
 - PyLate
 - sentence-transformers
+- multi-vector
 - sentence-similarity
 - feature-extraction
 pipeline_tag: sentence-similarity
@@ -100,15 +101,45 @@ We recommend LFM2.5-Embedding-350M and LFM2.5-ColBERT-350M for short-context ret
 
 <a href="https://colab.research.google.com/drive/1uLswYrRTNw4P2P2qZ8JG-b8j-KDkQqwL?usp=sharing"><img src="https://cdn-uploads.huggingface.co/production/uploads/61b8e2ba285851687028d395/vlOyMEjwHa_b_LXysEu2E.png" width=120 alt="Colab link"></a>
 
-First, install the PyLate and transformers libraries:
+### Using Sentence Transformers
+
+This model can be used as a multi-vector (ColBERT-style late interaction) retriever directly with [Sentence Transformers](https://www.sbert.net/) via the `MultiVectorEncoder`.
 
 ```bash
-pip install -U pylate
+pip install "sentence-transformers>=6.0.0"
 ```
 
-### Retrieval
+```python
+from sentence_transformers import MultiVectorEncoder
+
+model = MultiVectorEncoder("LiquidAI/LFM2.5-ColBERT-350M", trust_remote_code=True)
+
+query = "Which planet is known as the Red Planet?"
+documents = [
+    "Venus is often called Earth's twin because of its similar size and proximity.",
+    "Mars, known for its reddish appearance, is often referred to as the Red Planet.",
+    "Jupiter, the largest planet in our solar system, has a prominent red spot.",
+    "Saturn, famous for its rings, is sometimes mistaken for the Red Planet.",
+]
+
+query_embeddings = model.encode_query([query])
+document_embeddings = model.encode_document(documents)
+print(query_embeddings[0].shape, document_embeddings[0].shape)
+# (32, 128) (17, 128)
 
-Use this model with PyLate to index and retrieve documents. The index uses [FastPLAID](https://github.com/lightonai/fast-plaid) for efficient similarity search.
+# MaxSim late-interaction scoring (the Mars document ranks highest)
+scores = model.similarity(query_embeddings, document_embeddings)
+print(scores)
+# tensor([[27.1621, 28.2578, 27.7266, 28.1992]])
+```
+
+### Using PyLate
+
+Use this model with PyLate to index and retrieve documents. The index uses [FastPLAID](https://github.com/lightonai/fast-plaid) for efficient similarity search. First, install PyLate and transformers:
+
+```bash
+pip install -U pylate
+```
 
 #### Indexing documents
 

modeling_lfm2_bidirectional.py

@@ -71,7 +71,11 @@ def _noncausal_shortconv_forward(
     past_key_values=None,
     cache_position=None,
     attention_mask: Optional[torch.Tensor] = None,
+    **kwargs,
 ) -> torch.Tensor:
+    # transformers >=5.x passes seq_idx (packed-sample conv-state reset) to the conv. This full
+    # sequence non-causal conv has no cache and no packing, so it is ignored, like the cache args
+    # above. **kwargs absorbs it and any future additions rather than breaking on each new one.
     # Only the flash_attention_2 path expects padding states zeroed before the
     # conv. On eager/sdpa the checkpoints were trained WITHOUT zeroing: under
     # transformers 4.56 the conv received the 4D additive mask, on which