Initial release

Browse files

Files changed (10) hide show

.gitattributes +35 -0
README.md +248 -0
config.json +185 -0
config_sentence_transformers.json +12 -0
model.safetensors +3 -0
modeling_slideon.py +216 -0
modules.json +8 -0
slideon.py +96 -0
tokenizer.json +0 -0
tokenizer_config.json +23 -0

.gitattributes ADDED Viewed

	@@ -0,0 +1,35 @@

+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md ADDED Viewed

	@@ -0,0 +1,248 @@

+---
+library_name: sentence-transformers
+pipeline_tag: feature-extraction
+base_model: lightonai/DenseOn
+base_model_relation: finetune
+language:
+- en
+datasets:
+- lightonai/embeddings-fine-tuning-filtered-en
+tags:
+- sentence-transformers
+- sentence-similarity
+- custom-code
+- sparse-retrieval
+- matryoshka
+---
+<div align="center">
+<img src="https://cdn-uploads.huggingface.co/production/uploads/67e184d667350657a7d68aac/5PySy2cTawrkFWwaqNtDA.png" alt="SlideOn: a puppy sliding through a library with a bag of glowing stars" width="512" style="max-width: 100%; height: auto;">
+[![Website](https://img.shields.io/badge/LightOn-Website-blue?logo=google-chrome)](https://lighton.ai)
+[![LinkedIn](https://img.shields.io/badge/LightOn-LinkedIn-0A66C2?logo=linkedin)](https://www.linkedin.com/company/lighton/)
+[![X](https://img.shields.io/badge/@LightOnIO-X-black?logo=x)](https://x.com/LightOnIO)
+📚 [Collection](https://huggingface.co/collections/lightonai/slideon-and-mslideon) | 📝 [Blogpost](https://huggingface.co/blog/lightonai/slideon-mslideon) | 🛠️ [Training and Evaluation](https://github.com/lightonai/slideon-mslideon/)
+</div>
+<h1 align="center">SlideOn</h1>
+<h3 align="center">English Sparse Retrieval Model by LightOn</h3>
+<p align="center">
+<a href="https://huggingface.co/lightonai/SlideOn">SlideOn</a> |
+<a href="https://huggingface.co/lightonai/mSlideOn">mSlideOn</a> |
+<a href="https://huggingface.co/lightonai/SlideOn-SAE">SlideOn-SAE</a> |
+<a href="https://huggingface.co/lightonai/mSlideOn-SAE">mSlideOn-SAE</a>
+</p>
+> 🎯 **TL;DR**: A **174M-parameter English sparse retrieval model** built from DenseOn through sparse autoencoder (SAE) pretraining and supervised retrieval fine-tuning. SlideOn reaches **56.68 nDCG@10 on BEIR-15**, surpassing DenseOn's 56.20, while producing sparse representations compatible with inverted-index search. Joint query/document **Matryoshka training** preserves retrieval quality under pruning, allowing search and storage costs to be adjusted at inference without retraining.
+## About the SlideOn / mSlideOn Family
+With [DenseOn and LateOn](https://huggingface.co/blog/lightonai/denseon-lateon), then [mDenseOn and mLateOn](https://huggingface.co/blog/lightonai/mdenseon-mlateon), we explored dense and late-interaction retrieval using our open data recipe. SlideOn and mSlideOn reuse these models' retrieval knowledge to learn sparse representations through SAE pretraining followed by supervised fine-tuning.
+**SlideOn** starts from [DenseOn](https://huggingface.co/lightonai/DenseOn) for English retrieval. **[mSlideOn](https://huggingface.co/lightonai/mSlideOn)** starts from [mLateOn](https://huggingface.co/lightonai/mLateOn) for multilingual, long-context, and code retrieval. We release both the pretrained SAE initializations and the final retrieval models.
+See our [blog post](https://huggingface.co/blog/lightonai/slideon-mslideon) for the full method, results, and ablations.
+## SlideOn
+SlideOn encodes queries and documents into **32,768-dimensional nonnegative sparse vectors**. A learned linear encoder maps each normalized DenseOn token representation to its top-32 latent activations, followed by ReLU. Max pooling across tokens produces one vector per query or document, and retrieval uses the **dot product** between these vectors.
+The dimensions are learned features, independent of the tokenizer vocabulary. Active latents can be stored in an inverted index, where query latents select posting lists and their weights contribute to document scores.
+SlideOn notably:
+- Reaches **56.68 on BEIR-15** with unpruned vectors, improving on DenseOn by **0.48 points**.
+- Reaches **56.98** with unpruned queries and documents capped at **512 nonzero latents**.
+- Retains **56.66** with queries capped at **64** and documents at **512**, reducing the measured retrieval FLOPS proxy to **8.72**.
+- Supports inputs of up to **8,192 tokens** and independently adjustable query and document sparsity.
+## Results
+### Headline Results (BEIR-15, nDCG@10)
+BEIR-15 scores are average nDCG@10 across 15 datasets, reported on a 0–100 scale. Sparse models use dot-product scoring and exact search over each full evaluation corpus. All rows use **unpruned query and document vectors**. Dense and late-interaction results are taken from our [DenseOn and LateOn blog post](https://huggingface.co/blog/lightonai/denseon-lateon).
+<div style="width: fit-content; max-width: 100%; margin: 0 auto;">
+  <div style="overflow-x: auto;">
+    <table style="margin: 0 auto; border-collapse: collapse; color: inherit; text-align: right; font-size: 0.85em; line-height: 1.6;">
+      <thead>
+        <tr style="border-bottom: 1px solid rgba(128, 128, 128, 0.35);">
+          <th scope="col" style="padding: 8px 20px 8px 12px; text-align: left;">Model</th>
+          <th scope="col" style="padding: 8px 16px;">Parameters</th>
+          <th scope="col" style="padding: 8px 16px; white-space: nowrap;">BEIR-15 ↑</th>
+          <th scope="col" style="padding: 8px 12px 8px 16px; white-space: nowrap;">Retrieval FLOPS ↓</th>
+        </tr>
+      </thead>
+      <tbody>
+        <tr>
+          <th colspan="4" scope="rowgroup" style="padding: 16px 12px 6px; text-align: left; font-size: 0.8em; letter-spacing: 0.08em; text-transform: uppercase;">Sparse</th>
+        </tr>
+        <tr><td style="padding: 3px 20px 3px 12px; text-align: left; color: #3b82f6;"><strong style="color: #3b82f6;">SlideOn</strong></td><td style="padding: 3px 16px; color: #3b82f6;">174M</td><td style="padding: 3px 16px; color: #3b82f6;">56.68</td><td style="padding: 3px 12px 3px 16px; color: #3b82f6;">30.87</td></tr>
+        <tr><td style="padding: 3px 20px 3px 12px; text-align: left;">LACONIC-8B</td><td style="padding: 3px 16px;">8.03B</td><td style="padding: 3px 16px;"><strong>60.06</strong></td><td style="padding: 3px 12px 3px 16px;">20.07</td></tr>
+        <tr><td style="padding: 3px 20px 3px 12px; text-align: left;">LACONIC-3B</td><td style="padding: 3px 16px;">3.21B</td><td style="padding: 3px 16px;">58.30</td><td style="padding: 3px 12px 3px 16px;">39.53</td></tr>
+        <tr><td style="padding: 3px 20px 3px 12px; text-align: left;">LACONIC-1B</td><td style="padding: 3px 16px;">1.24B</td><td style="padding: 3px 16px;">56.13</td><td style="padding: 3px 12px 3px 16px;">47.95</td></tr>
+        <tr><td style="padding: 3px 20px 3px 12px; text-align: left;">SPARSEUP</td><td style="padding: 3px 16px;">188M</td><td style="padding: 3px 16px;">54.46</td><td style="padding: 3px 12px 3px 16px;">6.20</td></tr>
+        <tr><td style="padding: 3px 20px 3px 12px; text-align: left;">Lion-SP-8B</td><td style="padding: 3px 16px;">8.0B</td><td style="padding: 3px 16px;">53.62</td><td style="padding: 3px 12px 3px 16px;">57.03</td></tr>
+        <tr><td style="padding: 3px 20px 3px 12px; text-align: left;">OpenSearch doc-v3-GTE</td><td style="padding: 3px 16px;">137M</td><td style="padding: 3px 16px;">52.68</td><td style="padding: 3px 12px 3px 16px;"><strong>0.72</strong></td></tr>
+        <tr><td style="padding: 3px 20px 3px 12px; text-align: left;">OpenSearch v2-distill</td><td style="padding: 3px 16px;">67M</td><td style="padding: 3px 16px;">51.20</td><td style="padding: 3px 12px 3px 16px;">5.14</td></tr>
+        <tr><td style="padding: 3px 20px 3px 12px; text-align: left;">SPLADE-v3</td><td style="padding: 3px 16px;">133M</td><td style="padding: 3px 16px;">50.46</td><td style="padding: 3px 12px 3px 16px;">5.55</td></tr>
+        <tr><td style="padding: 3px 20px 3px 12px; text-align: left;">SPLADE-v3-DistilBERT</td><td style="padding: 3px 16px;">90M</td><td style="padding: 3px 16px;">49.04</td><td style="padding: 3px 12px 3px 16px;">5.79</td></tr>
+        <tr><td style="padding: 3px 20px 3px 12px; text-align: left;">Granite-30M</td><td style="padding: 3px 16px;">30M</td><td style="padding: 3px 16px;">48.86</td><td style="padding: 3px 12px 3px 16px;">6.42</td></tr>
+        <tr><td style="padding: 3px 20px 12px 12px; text-align: left;">SAE-SPLADE</td><td style="padding: 3px 16px 12px;">117M</td><td style="padding: 3px 16px 12px;">48.21</td><td style="padding: 3px 12px 12px 16px;">2.67</td></tr>
+      </tbody>
+      <tbody style="border-top: 1px solid rgba(128, 128, 128, 0.25);">
+        <tr>
+          <th colspan="4" scope="rowgroup" style="padding: 16px 12px 6px; text-align: left; font-size: 0.8em; letter-spacing: 0.08em; text-transform: uppercase;">Dense</th>
+        </tr>
+        <tr><td style="padding: 3px 20px 3px 12px; text-align: left; color: #3b82f6;"><strong style="color: #3b82f6;">DenseOn</strong></td><td style="padding: 3px 16px; color: #3b82f6;">149M</td><td style="padding: 3px 16px; color: #3b82f6;"><strong style="color: #3b82f6;">56.20</strong></td><td style="padding: 3px 12px 3px 16px; color: #3b82f6;">—</td></tr>
+        <tr><td style="padding: 3px 20px 3px 12px; text-align: left;">gte-modernbert-base</td><td style="padding: 3px 16px;">149M</td><td style="padding: 3px 16px;">55.19</td><td style="padding: 3px 12px 3px 16px;">—</td></tr>
+        <tr><td style="padding: 3px 20px 3px 12px; text-align: left;">bge-large-en-v1.5</td><td style="padding: 3px 16px;">335M</td><td style="padding: 3px 16px;">54.34</td><td style="padding: 3px 12px 3px 16px;">—</td></tr>
+        <tr><td style="padding: 3px 20px 12px 12px; text-align: left;">modernbert-embed-base</td><td style="padding: 3px 16px 12px;">149M</td><td style="padding: 3px 16px 12px;">52.89</td><td style="padding: 3px 12px 12px 16px;">—</td></tr>
+      </tbody>
+      <tbody style="border-top: 1px solid rgba(128, 128, 128, 0.25); border-bottom: 1px solid rgba(128, 128, 128, 0.35);">
+        <tr>
+          <th colspan="4" scope="rowgroup" style="padding: 16px 12px 6px; text-align: left; font-size: 0.8em; letter-spacing: 0.08em; text-transform: uppercase;">Late interaction</th>
+        </tr>
+        <tr><td style="padding: 3px 20px 3px 12px; text-align: left; color: #3b82f6;"><strong style="color: #3b82f6;">LateOn</strong></td><td style="padding: 3px 16px; color: #3b82f6;">149M</td><td style="padding: 3px 16px; color: #3b82f6;"><strong style="color: #3b82f6;">57.22</strong></td><td style="padding: 3px 12px 3px 16px; color: #3b82f6;">—</td></tr>
+        <tr><td style="padding: 3px 20px 3px 12px; text-align: left;">ColBERT-Zero</td><td style="padding: 3px 16px;">149M</td><td style="padding: 3px 16px;">55.39</td><td style="padding: 3px 12px 3px 16px;">—</td></tr>
+        <tr><td style="padding: 3px 20px 3px 12px; text-align: left;">GTE-ModernColBERT-v1</td><td style="padding: 3px 16px;">149M</td><td style="padding: 3px 16px;">54.75</td><td style="padding: 3px 12px 3px 16px;">—</td></tr>
+        <tr><td style="padding: 3px 20px 3px 12px; text-align: left;">ColBERT-small</td><td style="padding: 3px 16px;">33M</td><td style="padding: 3px 16px;">53.79</td><td style="padding: 3px 12px 3px 16px;">—</td></tr>
+        <tr><td style="padding: 3px 20px 12px 12px; text-align: left;">ColBERTv2</td><td style="padding: 3px 16px 12px;">110M</td><td style="padding: 3px 16px 12px;">48.63</td><td style="padding: 3px 12px 12px 16px;">—</td></tr>
+      </tbody>
+    </table>
+  </div>
+  <p style="max-width: 650px; margin: 14px auto 0; font-size: 0.85em; line-height: 1.5; text-align: center;">
+  <em>Average nDCG@10 on BEIR-15 (↑ higher is better) and retrieval FLOPS (↓ lower is better). <span style="color: #3b82f6;">Blue</span> marks our models. <strong>Bold</strong> marks best per column within each family.</em>
+</p>
+</div>
+SlideOn achieves the highest BEIR-15 score among the sparse models below 3B parameters in this comparison. It improves on SPARSEUP by **2.22 points** and on DenseOn by **0.48 points**, while remaining **0.54 points** behind LateOn. The figures below show how pruning changes the comparison when search and storage costs are taken into account.
+### Retrieval Quality versus Search and Storage Cost
+Query pruning reduces the number of posting lists consulted during search. Document pruning shortens those lists and reduces the number of stored document weights. Matryoshka training concentrates retrieval information in the strongest latents so that both budgets can be adjusted after training.
+<div style="margin: 0 auto; text-align: center;">
+  <img
+    src="https://cdn-uploads.huggingface.co/production/uploads/67e184d667350657a7d68aac/biV_cKSk6mQfoo0x_GSSu.png"
+    alt="BEIR-15 retrieval quality versus retrieval FLOPS for sparse models below 2B parameters, including SlideOn with different query pruning budgets."
+    style="display: block; max-width: 100%; height: auto; margin: 0 auto;"
+  />
+  <p style="max-width: 750px; margin: 12px auto 0; font-size: 0.85em; line-height: 1.5;">
+    <em>Retrieval quality versus search cost for SlideOn and sub-2B baselines. <span style="color: #3b82f6;">Blue</span> points show SlideOn with documents pruned to the top-512 latents and varying query budgets, where q-full denotes unpruned queries. Higher BEIR-15 ↑ scores and lower FLOPS ↓ are better. Retrieval FLOPS are measured on a subset of SciFact, NFCorpus, and FiQA.</em>
+  </p>
+</div>
+<div style="margin: 0 auto; text-align: center;">
+  <img
+    src="https://cdn-uploads.huggingface.co/production/uploads/67e184d667350657a7d68aac/kzPT8yhDHf1lIUh2Z2_7V.png"
+    alt="BEIR-15 retrieval quality versus average document NNZ for sparse models below 2B parameters, including SlideOn with different document pruning budgets."
+    style="display: block; max-width: 100%; height: auto; margin: 0 auto;"
+  />
+  <p style="max-width: 750px; margin: 12px auto 0; font-size: 0.85em; line-height: 1.5;">
+    <em>Retrieval quality versus index storage cost for SlideOn and sub-2B baselines. <span style="color: #3b82f6;">Blue</span> points show SlideOn with unpruned queries and varying document budgets, where d-full denotes unpruned documents. Higher BEIR-15 ↑ scores and lower NNZ ↓ are better. Document NNZ is measured on SciFact, NFCorpus, and FiQA.</em>
+  </p>
+</div>
+**Retrieval FLOPS** estimates the expected number of shared active features per query–document pair. It is a search-cost proxy, not measured latency or encoder FLOPs. **Document NNZ** is the average number of nonzero document features and serves as a storage-cost proxy, rather than a measurement of index size in bytes. Efficiency measurements use the fixed sample of **6,000 documents and 1,271 queries** from SciFact, NFCorpus, and FiQA.
+## Related Checkpoints
+| Model | Description |
+| --- | --- |
+| **[SlideOn](https://huggingface.co/lightonai/SlideOn)** *(this model)* | English sparse retrieval model, after SAE pretraining and retrieval fine-tuning |
+| [SlideOn-SAE](https://huggingface.co/lightonai/SlideOn-SAE) | DenseOn-based SAE initialization before retrieval fine-tuning |
+| [mSlideOn](https://huggingface.co/lightonai/mSlideOn) | Multilingual sparse retrieval model for natural language, long documents, and code |
+| [mSlideOn-SAE](https://huggingface.co/lightonai/mSlideOn-SAE) | mLateOn-based SAE initialization before retrieval fine-tuning |
+| [DenseOn](https://huggingface.co/lightonai/DenseOn) | English dense retriever used to initialize SlideOn's backbone |
+The SAE checkpoints are starting points for training. Use SlideOn or mSlideOn for the retrieval results reported in the blog.
+## Training Data and Recipe
+See our blog post for [training data](https://huggingface.co/blog/lightonai/slideon-mslideon#training-data), [SAE pretraining](https://huggingface.co/blog/lightonai/slideon-mslideon#sparse-autoencoder-pretraining), and [supervised retrieval fine-tuning](https://huggingface.co/blog/lightonai/slideon-mslideon#supervised-fine-tuning-for-sparse-retrieval).
+## Model Details
+### Model Description
+| Property | Value |
+| --- | --- |
+| Model type | Learned sparse retriever; Sentence Transformers `SparseEncoder` |
+| Backbone | DenseOn / ModernBERT-base |
+| Total parameters | 174M: 149M backbone + 25M sparse head |
+| Maximum sequence length | 8,192 tokens, including prompts and special tokens |
+| Output dimensionality | 32,768 learned latent features |
+| Token representation | 768 dimensions, L2-normalized before sparse encoding |
+| Token sparsity | Top-32 latent selection followed by ReLU |
+| Sequence pooling | Maximum over token activations |
+| Similarity function | Dot product |
+| Prompts | `query: ` for queries; `document: ` for documents |
+| Language | English |
+## Usage
+### Direct Usage (Sentence Transformers)
+Install the supported library versions:
+```bash
+pip install -U "sentence-transformers>=6.1,<6.2" "transformers>=5.3,<6"
+```
+Load the model and encode queries and documents. This example uses the **q-64 / d-512** configuration:
+```python
+from sentence_transformers import SparseEncoder
+model = SparseEncoder("lightonai/SlideOn", trust_remote_code=True)
+queries = ["What causes a solar eclipse?"]
+documents = [
+    "A solar eclipse occurs when the Moon passes between the Sun and Earth.",
+    "Wind turbines convert the kinetic energy of wind into electricity.",
+]
+query_embeddings = model.encode_query(queries, top_k=64)
+document_embeddings = model.encode_document(documents, top_k=512)
+print(query_embeddings.shape, document_embeddings.shape)
+# torch.Size([1, 32768]) torch.Size([2, 32768])
+scores = model.similarity(query_embeddings, document_embeddings)
+print(scores)
+```
+`encode_query()` and `encode_document()` apply the appropriate prompts and preprocessing automatically. Pass the raw text without manually adding prefixes. The returned embeddings are PyTorch sparse COO tensors. The model and custom inference code are included in this repository; loading does not require the training package or PyLate.
+### Choosing Query and Document Budgets
+Use `top_k` independently on each encoding call to cap the number of nonzero features **after pooling**. This is separate from the per-token top-32 operation and does not change the vector dimensionality.
+To reproduce the unpruned setting:
+```python
+query_embeddings = model.encode_query(queries, top_k=0)
+document_embeddings = model.encode_document(documents, top_k=0)
+```
+Omitting `top_k` also uses the saved unpruned defaults. To use the **q-full / d-512** setting, keep the query call at `top_k=0` and set the document call to `top_k=512`.
+Dimensions represent **learned latents**, so they do not map directly to tokenizer vocabulary entries. Inspect the sparse indices and weights; `SparseEncoder.decode()` cannot label them as vocabulary tokens. Retrieval uses raw dot products, so leave final vector normalization disabled.
+## Citation
+If you find this work useful for your research, please consider citing us:
+```bibtex
+@misc{moura2026slideon,
+  author       = {Moura Junior, Paulo Roberto and Sourty, Raphaël},
+  title        = {{SlideOn} \& {mSlideOn}: Sparsified late-interaction and dense retrievers},
+  year         = {2026},
+  howpublished = {Hugging Face Blog},
+  url          = {https://huggingface.co/blog/lightonai/slideon-mslideon}
+}
+```

config.json ADDED Viewed

	@@ -0,0 +1,185 @@

+{
+  "architectures": [
+    "SlideOnModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "modeling_slideon.SlideOnConfig",
+    "AutoModel": "modeling_slideon.SlideOnModel"
+  },
+  "backbone_config": {
+    "architectures": [
+      "ModernBertModel"
+    ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "bos_token_id": 50281,
+    "classifier_activation": "gelu",
+    "classifier_bias": false,
+    "classifier_dropout": 0.0,
+    "classifier_pooling": "mean",
+    "cls_token_id": 50281,
+    "decoder_bias": true,
+    "deterministic_flash_attn": false,
+    "dtype": "float32",
+    "embedding_dropout": 0.0,
+    "eos_token_id": 50282,
+    "global_attn_every_n_layers": 3,
+    "gradient_checkpointing": false,
+    "hidden_activation": "gelu",
+    "hidden_size": 768,
+    "initializer_cutoff_factor": 2.0,
+    "initializer_range": 0.02,
+    "intermediate_size": 1152,
+    "layer_norm_eps": 1e-05,
+    "layer_types": [
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention"
+    ],
+    "local_attention": 128,
+    "max_position_embeddings": 8192,
+    "mlp_bias": false,
+    "mlp_dropout": 0.0,
+    "model_type": "modernbert",
+    "norm_bias": false,
+    "norm_eps": 1e-05,
+    "num_attention_heads": 12,
+    "num_hidden_layers": 22,
+    "pad_token_id": 50283,
+    "position_embedding_type": "absolute",
+    "repad_logits_with_grad": false,
+    "rope_parameters": {
+      "full_attention": {
+        "rope_theta": 160000.0,
+        "rope_type": "default"
+      },
+      "sliding_attention": {
+        "rope_theta": 10000.0,
+        "rope_type": "default"
+      }
+    },
+    "sep_token_id": 50282,
+    "sparse_pred_ignore_index": -100,
+    "sparse_prediction": false,
+    "tie_word_embeddings": true,
+    "vocab_size": 50368
+  },
+  "base_model": "lightonai/DenseOn",
+  "base_model_revision": "cb9947ebccb33862d24e3c7ca2edb25e51acd887",
+  "checkpoint_type": "retrieval",
+  "head_config": {
+    "d_lift": 0,
+    "d_model": 768,
+    "has_decoder": false,
+    "n_latents": 32768
+  },
+  "inference": {
+    "document_length": 8192,
+    "document_top_k": 0,
+    "head_chunk_size": 2048,
+    "pooling": "max_relu",
+    "prompts": {
+      "document": "document: ",
+      "query": "query: "
+    },
+    "query_length": 8192,
+    "query_top_k": 0,
+    "score": "dot",
+    "token_k": 32
+  },
+  "model_type": "slideon",
+  "projection_configs": [],
+  "sae_path": "lightonai/SlideOn-SAE",
+  "skip_token_ids": [
+    2,
+    3,
+    4,
+    5,
+    6,
+    7,
+    8,
+    9,
+    10,
+    11,
+    12,
+    13,
+    14,
+    15,
+    16,
+    27,
+    28,
+    29,
+    30,
+    31,
+    32,
+    33,
+    60,
+    61,
+    62,
+    63,
+    64,
+    65,
+    92,
+    93,
+    94,
+    95
+  ],
+  "tokenizer_settings": {
+    "document_prefix": "",
+    "query_prefix": "",
+    "skiplist_words": [
+      "!",
+      "\"",
+      "#",
+      "$",
+      "%",
+      "&",
+      "'",
+      "(",
+      ")",
+      "*",
+      "+",
+      ",",
+      "-",
+      ".",
+      "/",
+      ":",
+      ";",
+      "<",
+      "=",
+      ">",
+      "?",
+      "@",
+      "[",
+      "\\",
+      "]",
+      "^",
+      "_",
+      "`",
+      "{",
+      "|",
+      "}",
+      "~"
+    ]
+  },
+  "transformers_version": "5.3.0"
+}

config_sentence_transformers.json ADDED Viewed

	@@ -0,0 +1,12 @@

+{
+  "model_type": "SparseEncoder",
+  "__version__": {
+    "sentence_transformers": "6.1.0"
+  },
+  "prompts": {
+    "query": "query: ",
+    "document": "document: "
+  },
+  "default_prompt_name": "document",
+  "similarity_fn_name": "dot"
+}

model.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:a06ae32e0a4736077f1997f02eb80b4460626e4ea665d60f1e5335dbb1165f6e
+size 697000192

modeling_slideon.py ADDED Viewed

	@@ -0,0 +1,216 @@

+"""Complete SlideOn model: Transformer, token projections, and sparse latent head."""
+from __future__ import annotations
+import torch
+from torch.utils.checkpoint import checkpoint as activation_checkpoint
+from transformers import AutoConfig, AutoModel, PretrainedConfig, PreTrainedModel
+class SlideOnConfig(PretrainedConfig):
+    model_type = "slideon"
+    sub_configs = {"backbone_config": AutoConfig}
+    def __init__(self, backbone_config=None, projection_configs=None, head_config=None,
+                 inference=None, tokenizer_settings=None, skip_token_ids=None,
+                 checkpoint_type="retrieval", **kwargs):
+        super().__init__(**kwargs)
+        if isinstance(backbone_config, dict):
+            backbone_config = dict(backbone_config)
+            model_type = backbone_config.pop("model_type", "modernbert")
+            backbone_config = AutoConfig.for_model(model_type, **backbone_config)
+        self.backbone_config = backbone_config or AutoConfig.for_model("modernbert")
+        self.projection_configs = projection_configs or []
+        self.head_config = head_config or dict(d_model=768, n_latents=32768, d_lift=0)
+        self.inference = inference or dict(
+            token_k=32, query_top_k=0, document_top_k=0,
+            query_length=8192, document_length=8192, head_chunk_size=2048,
+            prompts={"query": "", "document": ""}, pooling="max_relu", score="dot")
+        self.tokenizer_settings = tokenizer_settings or {}
+        self.skip_token_ids = skip_token_ids or []
+        self.checkpoint_type = checkpoint_type
+        if self.inference.get("pooling") != "max_relu" or self.inference.get("score") != "dot":
+            raise ValueError("SlideOn supports max_relu pooling and dot scoring")
+class TokenProjection(torch.nn.Module):
+    """The token-wise residual Dense layers used by mSlideOn."""
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.linear = torch.nn.Linear(
+            config["in_features"], config["out_features"], bias=config["bias"]
+        )
+        self.use_residual = config.get("use_residual", False)
+        self.residual = None
+        if self.use_residual and config["in_features"] != config["out_features"]:
+            self.residual = torch.nn.Linear(
+                config["in_features"], config["out_features"], bias=False
+            )
+        if config["activation_function"] != "torch.nn.modules.linear.Identity":
+            raise ValueError("Unsupported projection activation")
+    def forward(self, tokens):
+        projected = self.linear(tokens)
+        if self.use_residual:
+            projected = projected + (
+                self.residual(tokens) if self.residual is not None else tokens
+            )
+        return projected
+class LatentHead(torch.nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        d_model, n_latents = config["d_model"], config["n_latents"]
+        d_lift = config.get("d_lift") or 0
+        self.W_enc = torch.nn.Parameter(torch.empty(d_lift or d_model, n_latents))
+        self.b_enc = torch.nn.Parameter(torch.zeros(n_latents))
+        self.b_dec = torch.nn.Parameter(torch.zeros(d_model))
+        self.gate_relu_phi = torch.nn.Parameter(torch.ones(n_latents))
+        self.lift = None
+        if d_lift:
+            hidden = config["hidden_dim"]
+            self.lift = torch.nn.Linear(d_model, d_lift, bias=False)
+            self.mlp = torch.nn.Sequential(
+                torch.nn.Linear(d_lift, hidden), torch.nn.ReLU(),
+                torch.nn.Linear(hidden, d_lift))
+        if config.get("has_decoder", False):
+            # Preserve the pretrained SAE decoder without optimizing it in retrieval SFT.
+            self.W_dec = torch.nn.Parameter(torch.empty(n_latents, d_model), requires_grad=False)
+            self.register_buffer("steps_since_fired", torch.zeros(n_latents, dtype=torch.long))
+    def forward(self, tokens):
+        centered = tokens - self.b_dec
+        if self.lift is not None:
+            lifted = self.lift(centered)
+            centered = lifted + self.mlp(lifted)
+        return centered @ self.W_enc + self.b_enc
+class SlideOnModel(PreTrainedModel):
+    config_class = SlideOnConfig
+    base_model_prefix = "backbone"
+    _supports_sdpa = True
+    _supports_flash_attn = True
+    _no_split_modules = ["ModernBertEncoderLayer", "LatentHead", "TokenProjection"]
+    _keep_in_fp32_modules_strict = ["head"]
+    def __init__(self, config):
+        super().__init__(config)
+        config.backbone_config._attn_implementation = config._attn_implementation
+        self.backbone = AutoModel.from_config(config.backbone_config)
+        self.projections = torch.nn.ModuleList(
+            TokenProjection(item) for item in config.projection_configs)
+        self.head = LatentHead(config.head_config)
+        self.settings = config.inference
+        self.post_init()
+    def _init_weights(self, module):
+        if isinstance(module, LatentHead):
+            torch.nn.init.normal_(module.W_enc, std=0.02)
+            if hasattr(module, "W_dec"):
+                torch.nn.init.normal_(module.W_dec, std=0.02)
+        else:
+            super()._init_weights(module)
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        result = super().from_pretrained(*args, **kwargs)
+        model = result[0] if isinstance(result, tuple) else result
+        # Match the original PyLate encoder's RoPE buffer precision.
+        model.backbone.to(dtype=model.backbone.dtype)
+        return result
+    def get_embedding_dimension(self):
+        return self.config.head_config["n_latents"]
+    @staticmethod
+    def _task(task):
+        task = task or "document"
+        if task not in ("query", "document"):
+            raise ValueError("task must be 'query' or 'document'")
+        return task
+    def forward(self, input_ids, attention_mask, token_type_ids=None, task=None,
+                slideon_task=None, top_k=None, token_k=None, **kwargs):
+        features = {"input_ids": input_ids, "attention_mask": attention_mask}
+        if slideon_task is not None or task is not None:
+            features["slideon_task"] = slideon_task or task
+        if token_type_ids is not None:
+            features["token_type_ids"] = token_type_ids
+        task = self._task(task or features.get("slideon_task"))
+        if features.get("slideon_task", task) != task:
+            raise ValueError("Preprocessing and forward tasks must match")
+        n_latents = self.get_embedding_dimension()
+        cap = self.settings[task + "_top_k"] if top_k is None else top_k
+        k = self.settings["token_k"] if token_k is None else token_k
+        if not isinstance(cap, int) or not 0 <= cap <= n_latents:
+            raise ValueError(f"top_k must be an integer between 0 and {n_latents}")
+        if not isinstance(k, int) or not 1 <= k <= n_latents:
+            raise ValueError(f"token_k must be an integer between 1 and {n_latents}")
+        model_inputs = {
+            key: features[key]
+            for key in ("input_ids", "attention_mask", "token_type_ids")
+            if key in features
+        }
+        tokens = self.backbone(**model_inputs).last_hidden_state
+        for projection in self.projections:
+            tokens = projection(tokens)
+        mask = features["attention_mask"].bool()
+        if task == "document" and self.config.skip_token_ids:
+            skip_ids = torch.tensor(self.config.skip_token_ids, device=input_ids.device)
+            mask = mask & ~torch.isin(input_ids, skip_ids)
+        owners = torch.arange(tokens.shape[0], device=tokens.device)[:, None].expand_as(
+            mask
+        )[mask]
+        tokens = torch.nn.functional.normalize(tokens[mask], p=2, dim=-1)
+        # Match the evaluation loader's FP16 token transfer before head pooling.
+        if not self.training:
+            tokens = tokens.half().float()
+        probability_sum = tokens.new_zeros(n_latents, dtype=torch.float32)
+        usage_sum = tokens.new_zeros(n_latents, dtype=torch.float32)
+        pooled = torch.zeros(
+            mask.shape[0] * n_latents, dtype=torch.float32, device=tokens.device
+        )
+        for start in range(0, tokens.shape[0], self.settings["head_chunk_size"]):
+            end = start + self.settings["head_chunk_size"]
+            # The reference evaluates the head under BF16 autocast on CUDA.
+            def head_chunk(chunk):
+                with torch.autocast(chunk.device.type, dtype=torch.bfloat16, enabled=chunk.is_cuda):
+                    activations = self.head(chunk)
+                    values, indices = activations.topk(k, dim=-1)
+                probabilities = activations.float().softmax(-1).sum(0) if self.training else chunk.new_zeros(n_latents)
+                return values, indices, probabilities
+            if self.training and torch.is_grad_enabled():
+                values, indices, probabilities = activation_checkpoint(head_chunk, tokens[start:end], use_reentrant=False)
+            else:
+                values, indices, probabilities = head_chunk(tokens[start:end])
+            if self.training:
+                probability_sum = probability_sum + probabilities
+                usage_sum = usage_sum + torch.bincount(indices.flatten(), minlength=n_latents)
+            values = values.float().relu()
+            indices = owners[start:end, None] * n_latents + indices
+            pooled = pooled.scatter_reduce(
+                0, indices.flatten(), values.flatten(), reduce="amax", include_self=True
+            )
+        pooled = (
+            pooled.view(mask.shape[0], n_latents)
+            * self.head.gate_relu_phi.float().relu()
+        )
+        if cap:
+            values, indices = pooled.topk(cap, dim=-1)
+            pooled = torch.zeros_like(pooled).scatter(1, indices, values)
+        # The reference stores retrieval weights in FP16; retain that rounding.
+        features["sentence_embedding"] = pooled if self.training else pooled.half().float()
+        if self.training:
+            count = max(tokens.shape[0], 1)
+            features["load_balance"] = n_latents * (
+                probability_sum / count * usage_sum.detach() / (count * k)
+            ).sum()
+        return {key: features[key] for key in ("sentence_embedding", "load_balance") if key in features}
+SlideOnConfig.register_for_auto_class()
+SlideOnModel.register_for_auto_class("AutoModel")

modules.json ADDED Viewed

	@@ -0,0 +1,8 @@

+[
+  {
+    "idx": 0,
+    "name": "0",
+    "path": "",
+    "type": "slideon.SlideOn"
+  }
+]

slideon.py ADDED Viewed

	@@ -0,0 +1,96 @@

+"""Sentence Transformers adapter for the complete, flat SlideOn checkpoint."""
+from pathlib import Path
+import torch
+from sentence_transformers.base.modules import InputModule
+from transformers import AutoConfig, AutoModel, AutoTokenizer
+class SlideOn(InputModule):
+    config_file_name = "config.json"
+    forward_kwargs = {"task", "top_k", "token_k"}
+    save_in_root = True
+    def __init__(self, checkpoint, model_kwargs=None, processor_kwargs=None,
+                 config_kwargs=None, **hub_kwargs):
+        super().__init__()
+        config = AutoConfig.from_pretrained(
+            checkpoint, trust_remote_code=True, **hub_kwargs, **(config_kwargs or {}))
+        if config.model_type != "slideon":
+            raise ValueError("Expected a flat SlideOn checkpoint; use slideon.export for legacy checkpoints")
+        self.model = AutoModel.from_pretrained(
+            checkpoint, config=config, trust_remote_code=True,
+            **hub_kwargs, **(model_kwargs or {}))
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            checkpoint, trust_remote_code=True, **hub_kwargs, **(processor_kwargs or {}))
+        if self.tokenizer.mask_token_id is not None:
+            self.tokenizer.pad_token_id = self.tokenizer.mask_token_id
+        elif self.tokenizer.eos_token_id is not None:
+            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+        self.settings = self.model.config.inference
+        self.encoder_config = self.model.config.tokenizer_settings
+        self.max_seq_length = max(self.settings["query_length"], self.settings["document_length"])
+    def get_embedding_dimension(self):
+        return self.model.get_embedding_dimension()
+    @staticmethod
+    def _task(task):
+        task = task or "document"
+        if task not in ("query", "document"):
+            raise ValueError("task must be 'query' or 'document'")
+        return task
+    def preprocess(
+        self, inputs, prompt=None, task=None, processing_kwargs=None, **kwargs
+    ):
+        task = self._task(task)
+        if not all(isinstance(text, str) for text in inputs):
+            raise TypeError("SlideOn accepts text strings only")
+        if prompt is None:
+            prompt = self.settings["prompts"][task]
+        texts = [(prompt + text).strip() for text in inputs]
+        prefix = self.encoder_config.get(task + "_prefix", "")
+        length = min(self.max_seq_length, self.settings[task + "_length"])
+        options = dict(
+            padding=True,
+            truncation=True,
+            max_length=length - bool(prefix),
+            return_tensors="pt",
+        )
+        options.update(processing_kwargs or {})
+        features = dict(self.tokenizer(texts, **options))
+        if prefix:
+            prefix_id = self.tokenizer.convert_tokens_to_ids(prefix)
+            for name, value in list(features.items()):
+                fill = (
+                    prefix_id
+                    if name == "input_ids"
+                    else (1 if name == "attention_mask" else 0)
+                )
+                marker = torch.full(
+                    (value.shape[0], 1), fill, dtype=value.dtype, device=value.device
+                )
+                features[name] = torch.cat((value[:, :1], marker, value[:, 1:]), dim=1)
+        features["slideon_task"] = task
+        return features
+    def forward(self, features, task=None, top_k=None, token_k=None, **kwargs):
+        inputs = {key: features[key] for key in
+                  ("input_ids", "attention_mask", "token_type_ids", "slideon_task") if key in features}
+        features.update(self.model(**inputs, task=task, top_k=top_k, token_k=token_k))
+        return features
+    def save(self, output_path, safe_serialization=True, **kwargs):
+        self.model.save_pretrained(output_path, safe_serialization=safe_serialization)
+        self.tokenizer.save_pretrained(output_path)
+    @classmethod
+    def load(cls, model_name_or_path, subfolder="", token=None, cache_folder=None,
+             revision=None, local_files_only=False, model_kwargs=None,
+             processor_kwargs=None, config_kwargs=None, backend="torch", **kwargs):
+        if backend != "torch":
+            raise ValueError("SlideOn currently supports the torch backend only")
+        return cls(model_name_or_path, model_kwargs, processor_kwargs, config_kwargs,
+                   subfolder=subfolder, token=token, cache_dir=cache_folder,
+                   revision=revision, local_files_only=local_files_only)

tokenizer.json ADDED Viewed

The diff for this file is too large to render. See raw diff

tokenizer_config.json ADDED Viewed

	@@ -0,0 +1,23 @@

+{
+  "backend": "tokenizers",
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "[CLS]",
+  "is_local": true,
+  "mask_token": "[MASK]",
+  "max_length": 512,
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 512,
+  "pad_to_multiple_of": null,
+  "pad_token": "[MASK]",
+  "pad_token_type_id": 0,
+  "padding_side": "right",
+  "sep_token": "[SEP]",
+  "stride": 0,
+  "tokenizer_class": "TokenizersBackend",
+  "truncation_side": "right",
+  "truncation_strategy": "longest_first",
+  "unk_token": "[UNK]"
+}