Fix short-conv padding masking on transformers >=5

README.md

@@ -19,7 +19,6 @@ tags:
 - ColBERT
 - PyLate
 - sentence-transformers
-- multi-vector
 - sentence-similarity
 - feature-extraction
 pipeline_tag: sentence-similarity
@@ -27,7 +26,6 @@ library_name: PyLate
 license: other
 license_name: lfm1.0
 license_link: LICENSE
-base_model: LiquidAI/LFM2.5-350M-Base
 ---
 
 <center>
@@ -56,10 +54,7 @@ Both models are 350M params and the first bidirectional members of the LFM famil
 
 Find more details about the bidirectional architecture and training recipe in our [blog post](https://www.liquid.ai/blog/lfm2-5-retrievers).
 
-> [!NOTE]
-> 💻 **Demo**: https://huggingface.co/spaces/LiquidAI/colbert-tool-selection
-
-![colb](https://cdn-uploads.huggingface.co/production/uploads/63f389fda096536aeaae0a66/snq-O6VmWRJNMYCNHGtDD.png)
+![colbert_retrieval](https://cdn-uploads.huggingface.co/production/uploads/63f389fda096536aeaae0a66/8J_3tsS-wDq77qeDpDgOd.png)
 
 ## 📄 Model details
 
@@ -69,11 +64,13 @@ Find more details about the bidirectional architecture and training recipe in ou
 | **Total parameters**  | ~353M                                  | ~354M                               |
 | **Backbone**          | [LFM2.5-350M-Base](https://huggingface.co/LiquidAI/LFM2.5-350M-Base) + bi-directional patches  |  [LFM2.5-350M-Base](https://huggingface.co/LiquidAI/LFM2.5-350M-Base) + bi-directional patches |
 | **Layers**            | 17 (10 conv + 6 attn + 1 dense)        | 17 (10 conv + 6 attn + 1 pool)      |
+| **Context length**    | 32,768 tokens                          | 32,768 tokens                       |
 | **Vocabulary size**   | 64,402                                 | 65,536                              |
 | **Output**            | 128-dim per token                      | 1024-dim CLS vector                 |
 | **Similarity**        | MaxSim                                 | Cosine                              |
 | **Training precision**| BF16                                   | BF16                                |
 | **License**           | LFM Open License v1.0                  | LFM Open License v1.0               |
+| **Architecture**      | <img src="https://cdn-uploads.huggingface.co/production/uploads/63f389fda096536aeaae0a66/neXdblhfapp1Oln8M9J4V.png" width="250">| <img src="https://cdn-uploads.huggingface.co/production/uploads/63f389fda096536aeaae0a66/y4Lx5eS3hZCRv1g3r1e9i.png" width="250"> |
 
 **Document length:** 512 tokens &nbsp;&nbsp; 
 
@@ -101,45 +98,15 @@ 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>
 
-### 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`.
+First, install the PyLate and transformers libraries:
 
 ```bash
-pip install "sentence-transformers>=6.0.0"
-```
-
-```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)
-
-# 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]])
+pip install -U pylate
 ```
 
-### Using PyLate
+### Retrieval
 
-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
-```
+Use this model with PyLate to index and retrieve documents. The index uses [FastPLAID](https://github.com/lightonai/fast-plaid) for efficient similarity search.
 
 #### Indexing documents
 

modeling_lfm2_bidirectional.py

@@ -71,11 +71,7 @@ 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