Simplify changes further, revert transformers changes

But the transformers path now works with eager/sdpa too

README.md

@@ -2,6 +2,7 @@
 license: cc-by-nc-sa-4.0
 tags:
   - sentence-transformers
+  - transformers
   - splade
   - sparse-encoder
   - code
@@ -51,3 +52,49 @@ print(decoded)
 #     ("ĠGroup", 2.1875),
 # ]]
 ```
+
+### Using Transformers
+
+```bash
+pip install transformers
+```
+
+```python
+from transformers import AutoModelForCausalLM, AutoModel
+import os
+import torch
+
+splade = AutoModelForCausalLM.from_pretrained("naver/splade-code-06B", trust_remote_code=True)
+device = (torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu"))
+splade.to(device)
+splade.eval()
+queries = ["SELECT *\nFROM Student\nWHERE Age = (\nSELECT MAX(Age)\nFROM Student\nWHERE Group = 'specific_group'\n)\nAND Group = 'specific_group';"]
+bow_dict = splade.encode(queries, prompt_type="query", top_k_q=10, return_dict=True, print_dict=True)
+```
+
+```
++--------------------------------------------------------------------+
+|                        TOP ACTIVATED WORDS                         |
++--------------------------------------------------------------------+
+
+
+* INPUT: SELECT *
+FROM Student
+WHERE Age = (
+SELECT MAX(Age)
+FROM Student
+WHERE Group = 'specific_group'
+)
+AND Group = 'specific_group';
+
+Ġgroup                    | ████████████████████ 2.34
+Ġage                      | ███████████████████ 2.33
+ĠAge                      | ███████████████████ 2.33
+_group                    | ███████████████████ 2.30
+ĠStudent                  | ███████████████████ 2.30
+Ġspecific                 | ███████████████████ 2.28
+Ġmax                      | ██████████████████ 2.22
+ĠMax                      | ██████████████████ 2.22
+Ġstudent                  | ██████████████████ 2.20
+ĠGroup                    | ██████████████████ 2.19
+```

config.json

@@ -29,6 +29,7 @@
   "use_sliding_window": false,
   "vocab_size": 151936,
   "auto_map": {
-    "AutoModelForMaskedLM": "modeling_splade.Qwen3ForCausalLM"
+    "AutoModelForMaskedLM": "splade.Qwen3ForCausalLM",
+    "AutoModelForCausalLM": "splade.Splade"
   }
 }

modeling_splade.py

@@ -1,28 +0,0 @@
-"""
-This file exists solely to allow loading the Qwen3ForCausalLM via the AutoModelForMaskedLM class.
-Compared to standard Qwen3, we're using bidirectional attention and not causal attention, but it's specified
-with `is_causal=False` in the config.
-"""
-
-from transformers import Qwen3ForCausalLM as _Qwen3ForCausalLM
-
-
-class Qwen3ForCausalLM(_Qwen3ForCausalLM):
-    def tie_weights(self, *args, **kwargs):
-        """Explicitly re-tie lm_head to embed_tokens to hopefully avoid meta tensor errors."""
-        super().tie_weights(*args, **kwargs)
-        if (
-            self.config.tie_word_embeddings
-            and hasattr(self, "lm_head")
-            and hasattr(self, "model")
-        ):
-            self.lm_head.weight = self.model.embed_tokens.weight
-
-    def _init_weights(self, module):
-        """Skip lm_head init when it will be tied to embed_tokens later."""
-        if module is getattr(self, "lm_head", None) and self.config.tie_word_embeddings:
-            return
-        super()._init_weights(module)
-
-
-__all__ = ["Qwen3ForCausalLM"]

splade.py

@@ -0,0 +1,130 @@
+"""
+Compared to standard Qwen3, we're using bidirectional attention and not causal attention, but it's specified
+with `is_causal=False` in the config.
+
+This file supports two loading paths:
+1. Sentence Transformers: `SparseEncoder("naver/splade-code-06B", trust_remote_code=True)` via AutoModelForMaskedLM -> Qwen3ForCausalLM
+2. Transformers: `AutoModelForCausalLM.from_pretrained("naver/splade-code-06B", trust_remote_code=True)` -> Splade
+"""
+
+import torch
+import os
+from transformers import Qwen3ForCausalLM as TransformersQwen3ForCausalLM
+from transformers import PretrainedConfig, PreTrainedModel, AutoConfig
+from transformers.utils import is_flash_attn_2_available
+from .utils import prepare_tokenizer, splade_max, similarity, encode
+
+
+class Qwen3ForCausalLM(TransformersQwen3ForCausalLM):
+    def tie_weights(self, *args, **kwargs):
+        """Explicitly re-tie lm_head to embed_tokens to hopefully avoid meta tensor errors."""
+        if (
+            self.config.tie_word_embeddings
+            and hasattr(self, "lm_head")
+            and hasattr(self, "model")
+        ):
+            self.lm_head.weight = self.model.embed_tokens.weight
+            missing_keys = kwargs.get("missing_keys")
+            if missing_keys is not None:
+                missing_keys.discard("lm_head.weight")
+        else:
+            super().tie_weights(*args, **kwargs)
+
+    def _init_weights(self, module):
+        """Skip lm_head init when it will be tied to embed_tokens later."""
+        if module is getattr(self, "lm_head", None) and self.config.tie_word_embeddings:
+            return
+        super()._init_weights(module)
+
+
+class SpladeConfig(PretrainedConfig):
+    model_type = "qwen3"
+
+    def __init__(
+        self,
+        model_name_or_path: str = "Qwen/Qwen3-0.6B",
+        attn_implementation: str = "flash_attention_2",
+        bidirectional: bool = True,  # only for decoder models
+        padding_side: str = "left",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.model_name_or_path = model_name_or_path
+        self.attn_implementation = attn_implementation
+        self.bidirectional = bidirectional
+        self.padding_side = padding_side
+
+
+class Splade(PreTrainedModel):
+    config_class = SpladeConfig
+
+    # methods for MTEB's interface
+    similarity = similarity
+    encode = encode
+
+    def __init__(self, config, weights_path=None, token=None):
+        super().__init__(config)
+        self.name = "splade"
+
+        base_cfg = AutoConfig.from_pretrained(
+            weights_path,
+            attn_implementation=config.attn_implementation,
+            torch_dtype="auto",
+        )
+
+        self.tokenizer = prepare_tokenizer(
+            weights_path, padding_side=config.padding_side
+        )
+
+        if is_flash_attn_2_available():
+            config.attn_implementation = "flash_attention_2"
+        else:
+            config.attn_implementation = "sdpa"
+
+        self.model = Qwen3ForCausalLM.from_pretrained(
+            weights_path,
+            config=base_cfg,
+            torch_dtype=torch.bfloat16,
+            attn_implementation=config.attn_implementation,
+            token=token,
+        )
+
+    def save_pretrained(self, save_directory, *args, **kwargs):
+        self.model.save_pretrained(os.path.join(save_directory, "lora"))
+        self.config.save_pretrained(save_directory)
+
+    @classmethod
+    def from_pretrained(cls, model_name_or_path, *args, **kwargs):
+        token = kwargs.get("token", None)
+
+        config = SpladeConfig.from_pretrained(
+            model_name_or_path,
+            token=token,
+        )
+
+        model = cls(config, weights_path=model_name_or_path, token=token)
+
+        model.reverse_voc = {v: k for k, v in model.tokenizer.vocab.items()}
+        return model
+
+    def forward(self, **tokens):
+        output = self.model(**tokens)
+        splade_reps, _ = splade_max(output.logits, tokens["attention_mask"])
+        return (splade_reps,)
+
+    def get_width(self):
+        return self.model.config.vocab_size
+
+    def create_batch_dict(self, input_texts, max_length):
+        return self.tokenizer(
+            input_texts,
+            add_special_tokens=True,
+            padding="longest",
+            truncation=True,
+            max_length=max_length,
+            return_attention_mask=True,
+            return_tensors="pt",
+        )
+
+
+__all__ = ["Qwen3ForCausalLM", "Splade"]

utils.py

@@ -0,0 +1,128 @@
+import numpy as np
+import torch
+
+from typing import Any
+from transformers import AutoTokenizer
+
+
+def splade_max(features, attention_mask):
+    """
+    SPLADE pooling operation
+    """
+    relu = torch.nn.ReLU(inplace=False)
+    values, ids_ = torch.max(
+        torch.log(1 + relu(features)) * attention_mask.unsqueeze(-1), dim=1
+    )
+    return values, ids_
+
+
+def encode(
+    self,
+    sentences: list[str],
+    max_length: int = 1024,
+    prompt_type: str = "document",
+    return_dict: bool = False,
+    print_dict: bool = False,
+    batch_size: int = 8,
+    top_k_q: int = -1,
+    top_k_d: int = -1,
+    **kwargs: Any,
+) -> np.ndarray:
+    all_embeddings = []
+    for i in range(0, len(sentences), batch_size):
+        batch_texts = sentences[i : i + batch_size]
+        batch_dict = self.create_batch_dict(batch_texts, max_length)
+        batch_dict = {
+            key: value.to(self.model.device) for key, value in batch_dict.items()
+        }
+        with torch.no_grad():
+            splare_reps = self(**batch_dict)[0]
+            if prompt_type == "query" and top_k_q > 0:
+                splare_reps = top_k(splare_reps, top_k_q)
+            if prompt_type == "document" and top_k_d > 0:
+                splare_reps = top_k(splare_reps, top_k_d)
+            all_embeddings.append(splare_reps.cpu().float().numpy())
+    if return_dict:
+        d = bow_dict(self, np.concatenate(all_embeddings, axis=0))
+        if print_dict:
+            print_bow_bars(sentences, d)
+        return d
+    else:
+        return np.concatenate(all_embeddings, axis=0)
+
+
+def bow_dict(self, embeddings):
+    out = []
+    for vector in embeddings:
+        idx = np.nonzero(vector)[0]
+        weights = vector[idx]
+        d = {k: v for k, v in zip(idx.tolist(), weights.tolist())}
+        sorted_d = {
+            self.reverse_voc[k]: float(v)
+            for k, v in sorted(d.items(), key=lambda item: item[1], reverse=True)
+        }
+        out.append(sorted_d)
+    return out
+
+
+def print_bow_bars(sentences, bow_list, width=20):
+    ascii_header("TOP ACTIVATED WORDS")
+    for sent, bow in zip(sentences, bow_list):
+        print(f"* INPUT: {sent}\n")
+        max_w = max(bow.values())
+        for k, v in sorted(bow.items(), key=lambda x: x[1], reverse=True):
+            bar = "█" * int(v / max_w * width)
+            print(f"{k[:25]:25} | {bar} {v:.2f}")
+        print("\n")
+
+
+def ascii_header(title, width=70):
+    title = f" {title} "
+    print("+" + "-" * (width - 2) + "+")
+    print("|" + title.center(width - 2) + "|")
+    print("+" + "-" * (width - 2) + "+")
+    print("\n")
+
+
+def similarity(self, a, b) -> torch.Tensor:
+    """
+    MTEB eval requires this
+    """
+    if not isinstance(a, torch.Tensor):
+        a = torch.tensor(a)
+    if not isinstance(b, torch.Tensor):
+        b = torch.tensor(b)
+
+    def _dot_score_core(a_tensor, b_tensor):
+        if len(a_tensor.shape) == 1:
+            a_tensor = a_tensor.unsqueeze(0)
+        if len(b_tensor.shape) == 1:
+            b_tensor = b_tensor.unsqueeze(0)
+        return a_tensor @ b_tensor.transpose(0, 1)
+
+    return _dot_score_core(a, b)
+
+
+def prepare_tokenizer(tokenizer_name: str, padding_side="right"):
+    """
+    loads and prepares tokenizer
+    """
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+    tokenizer.pad_token = (
+        tokenizer.bos_token or tokenizer.pad_token or tokenizer.eos_token
+    )
+    tokenizer.padding_side = padding_side
+    return tokenizer
+
+
+def top_k(x: torch.Tensor, k: int) -> torch.Tensor:
+    """
+    zeroes out all but the top-k values in the last dimension of x
+    """
+    _, topk_indices = x.topk(k, dim=-1)
+    # create a zero tensor of the same shape as x
+    mask = torch.zeros_like(x, dtype=torch.bool)
+    # use scatter along the last dimension
+    mask.scatter_(-1, topk_indices, True)
+    # zero out all but the top-k
+    return x * mask