Initial release

.gitattributes

@@ -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

@@ -0,0 +1,99 @@
+---
+library_name: sentence-transformers
+pipeline_tag: feature-extraction
+base_model: lightonai/DenseOn
+language:
+- en
+datasets:
+- lightonai/embeddings-fine-tuning-filtered-en
+tags:
+- sentence-transformers
+- custom-code
+- sparse-autoencoder
+- sparse-retrieval
+---
+<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)
+
+[SlideOn](https://huggingface.co/lightonai/SlideOn) | [mSlideOn](https://huggingface.co/lightonai/mSlideOn) | 📚 [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-SAE</h1>
+<h3 align="center">English Sparse Autoencoder by LightOn</h3>
+
+> 🎯 **TL;DR**: The pretrained sparse autoencoder used to initialize **[SlideOn](https://huggingface.co/lightonai/SlideOn)**. It learns **32,768 latent features** from frozen DenseOn token representations and provides a starting point for supervised sparse retrieval fine-tuning. Use the final SlideOn checkpoint for the retrieval results reported in our blog.
+
+## Model Details
+
+SlideOn-SAE maps L2-normalized DenseOn token representations to their top-32 latent activations, followed by ReLU. Its decoder reconstructs the original token representations during SAE pretraining. For retrieval, the included adapter max-pools token activations into one sparse vector per query or document.
+
+| Property | Value |
+| --- | --- |
+| Backbone | [DenseOn](https://huggingface.co/lightonai/DenseOn) / ModernBERT-base |
+| Training stage | SAE pretraining, before supervised retrieval fine-tuning |
+| Token representation | 768 dimensions, L2-normalized before sparse encoding |
+| Latent dimensionality | 32,768 learned features |
+| Token sparsity | Top-32 latent selection followed by ReLU |
+| Sequence pooling | Maximum over token activations |
+| Similarity function | Dot product |
+| Language | English |
+
+The checkpoint includes the DenseOn backbone, SAE encoder, and original SAE decoder in a single `model.safetensors` file. The model and tokenizer load directly from this repository.
+
+## 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). Training code and the recorded training configuration are available in the [code repository](https://github.com/lightonai/slideon-mslideon/).
+
+## Usage
+
+Install the supported library versions:
+
+```bash
+pip install -U "sentence-transformers>=6.1,<6.2" "transformers>=5.3,<6"
+```
+
+Load the SAE initialization and encode text into sparse vectors:
+
+```python
+from sentence_transformers import SparseEncoder
+
+model = SparseEncoder("lightonai/SlideOn-SAE", trust_remote_code=True)
+
+query_embeddings = model.encode_query(["What is sparse retrieval?"])
+document_embeddings = model.encode_document([
+    "Sparse retrieval uses inverted indexes to search sparse vectors."
+])
+
+print(query_embeddings.shape, document_embeddings.shape)
+# torch.Size([1, 32768]) torch.Size([1, 32768])
+```
+
+The model loaded above can also be fine-tuned with Sentence Transformers' `SparseEncoderTrainer`. Our [code repository](https://github.com/lightonai/slideon-mslideon/) provides the custom retrieval loss and training scripts used for our fine-tuning recipe.
+
+## Related Checkpoints
+
+| Model | Description |
+| --- | --- |
+| **[SlideOn-SAE](https://huggingface.co/lightonai/SlideOn-SAE)** *(this model)* | DenseOn-based SAE initialization |
+| [SlideOn](https://huggingface.co/lightonai/SlideOn) | English sparse retriever after supervised fine-tuning |
+| [mSlideOn-SAE](https://huggingface.co/lightonai/mSlideOn-SAE) | mLateOn-based multilingual SAE initialization |
+| [mSlideOn](https://huggingface.co/lightonai/mSlideOn) | Multilingual sparse retriever after supervised fine-tuning |
+
+## Citation
+
+```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

@@ -0,0 +1,120 @@
+{
+  "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": "sae",
+  "head_config": {
+    "d_lift": 0,
+    "d_model": 768,
+    "has_decoder": true,
+    "n_latents": 32768
+  },
+  "inference": {
+    "document_length": 8192,
+    "document_top_k": 0,
+    "head_chunk_size": 256,
+    "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_revision": "669feba609eeb5d4ee59ef791679c7e393092677",
+  "sae_source": "lightonai/SlideOn-SAE",
+  "skip_token_ids": [],
+  "tokenizer_settings": {
+    "document_prefix": "",
+    "query_prefix": "",
+    "skiplist_words": []
+  },
+  "transformers_version": "5.3.0"
+}

config_sentence_transformers.json

@@ -0,0 +1,14 @@
+{
+  "__version__": {
+    "pytorch": "2.14.1+cu130",
+    "sentence_transformers": "6.1.0",
+    "transformers": "5.19.0"
+  },
+  "default_prompt_name": null,
+  "model_type": "SparseEncoder",
+  "prompts": {
+    "document": "",
+    "query": ""
+  },
+  "similarity_fn_name": "dot"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:81dde50d52aeadad52af09af3bb28905ff3c6c5c7b8a26421ca0e7b4a0ec5d6d
+size 797925808

modeling_slideon.py

@@ -0,0 +1,219 @@
+"""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)
+        # Loading tensors can reset the decoder's requires_grad flag.
+        if hasattr(model.head, "W_dec"):
+            model.head.W_dec.requires_grad_(False)
+        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

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

slideon.py

@@ -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_config.json

@@ -0,0 +1,24 @@
+{
+  "backend": "tokenizers",
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "[CLS]",
+  "is_local": true,
+  "local_files_only": false,
+  "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]"
+}