Upload SSE_quantize.py

SSE_quantize.py

@@ -0,0 +1,228 @@
+"""
+coding = utf-8
+Copyright 2026 Rikka Botan. All rights reserved
+Licensed under "MIT License"
+Stable Static Embedding official PyTorch implementation
+"""
+
+from __future__ import annotations
+import os
+from pathlib import Path
+from safetensors.torch import save_file as save_safetensors_file
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from typing import Dict
+from dataclasses import dataclass
+from tokenizers import Tokenizer
+from transformers import PreTrainedTokenizerFast
+from sentence_transformers.models.InputModule import InputModule
+from safetensors.torch import load_file
+
+
+def quantize_q4_k_m(weight: torch.Tensor):
+    """
+    weight: (vocab, dim)
+    returns: packed uint8 + scale + zero
+    """
+    w = weight.detach().cpu().numpy().astype(np.float32)
+
+    scales = np.max(np.abs(w), axis=1, keepdims=True) + 1e-8
+    w_norm = w / scales
+
+    q = np.clip(np.round((w_norm + 1) * 7.5), 0, 15).astype(np.uint8)
+
+    # pack 2x4bit -> uint8
+    packed = (q[:, 0::2] << 4) | q[:, 1::2]
+
+    return {
+        "packed": packed,
+        "scales": scales.astype(np.float32),
+    }
+
+
+def dequantize_q4_k_m(packed: np.ndarray, scales: np.ndarray):
+    hi = (packed >> 4) & 0xF
+    lo = packed & 0xF
+
+    q = np.empty((packed.shape[0], packed.shape[1]*2), dtype=np.uint8)
+    q[:, 0::2] = hi
+    q[:, 1::2] = lo
+
+    w = (q.astype(np.float32) / 7.5) - 1.0
+    w = w * scales
+    return torch.from_numpy(w)
+
+
+class SeparableDyT(nn.Module):
+    def __init__(
+        self,
+        hidden_dim: int,
+        alpha_init: float = 0.5
+    ):
+        super().__init__()
+        self.alpha = nn.Parameter(alpha_init*torch.ones(hidden_dim))
+        self.beta = nn.Parameter(torch.ones(hidden_dim))
+        self.bias = nn.Parameter(torch.zeros(hidden_dim))
+    
+    def forward(
+        self,
+        x: torch.Tensor
+    ) -> torch.Tensor:
+        x = self.beta * F.tanh(self.alpha * x + self.bias)
+        return x
+
+
+class SSEQ(InputModule):
+    """
+    Stable Static Embedding (SSE)
+    StaticEmbedding-compatible Sentence-Transformers module
+    """
+
+    def __init__(
+        self,
+        tokenizer: Tokenizer | PreTrainedTokenizerFast,
+        vocab_size: int,
+        hidden_dim: int = 1024,
+        **kwargs,
+    ):
+        super().__init__()
+
+        if isinstance(tokenizer, PreTrainedTokenizerFast):
+            tokenizer = tokenizer._tokenizer
+        elif not isinstance(tokenizer, Tokenizer):
+            raise ValueError("Tokenizer must be a fast (Rust) tokenizer")
+
+        self.tokenizer: Tokenizer = tokenizer
+        self.tokenizer.no_padding()
+
+        self.embedding = nn.EmbeddingBag(vocab_size, hidden_dim)
+        self.dyt = SeparableDyT(hidden_dim)
+
+        self.embedding_dim = hidden_dim
+
+        # For model card compatibility
+        self.base_model = kwargs.get("base_model", None)
+
+    # Tokenization (StaticEmbedding-compatible)
+    def tokenize(
+        self,
+        texts: list[str],
+        **kwargs
+    ) -> dict[str, torch.Tensor]:
+        encodings = self.tokenizer.encode_batch(texts, add_special_tokens=False)
+        encodings_ids = [encoding.ids for encoding in encodings]
+
+        offsets = torch.from_numpy(
+            np.cumsum(
+                [0] + [len(token_ids) for token_ids in encodings_ids[:-1]]
+            )
+        )
+        input_ids = torch.tensor(
+            [token_id for token_ids in encodings_ids for token_id in token_ids],
+            dtype=torch.long
+        )
+        return {
+            "input_ids": input_ids,
+            "offsets": offsets
+        }
+
+    # Forward
+    def forward(
+        self,
+        features: Dict[str, torch.Tensor],
+        **kwargs,
+    ) -> Dict[str, torch.Tensor]:
+        x = self.embedding(features["input_ids"], features["offsets"])
+        x = self.dyt(x)
+        features["sentence_embedding"] = x
+        return features
+
+    # Required APIs
+    def get_sentence_embedding_dimension(self) -> int:
+        return self.embedding_dim
+
+    @property
+    def max_seq_length(self) -> int:
+        return torch.inf
+    
+    def save(self, output_path: str):
+        os.makedirs(output_path, exist_ok=True)
+
+        state = self.state_dict()
+
+        emb = state["embedding.weight"]
+        q = quantize_q4_k_m(emb)
+
+        del state["embedding.weight"]
+
+        save_safetensors_file(
+            state,
+            os.path.join(output_path, "model_rest.safetensors"),
+        )
+
+        with open(os.path.join(output_path, "embedding.q4_k_m.bin"), "wb") as f:
+            f.write(q["packed"].tobytes())
+            f.write(q["scales"].tobytes())
+
+        self.tokenizer.save(
+            str(Path(output_path) / "tokenizer.json")
+        )
+    
+    @classmethod
+    def load(cls, model_path: str):
+
+        tokenizer = Tokenizer.from_file(
+            os.path.join(model_path, "tokenizer.json")
+        )
+
+        state = load_file(
+            os.path.join(model_path, "model_rest.safetensors"),
+            device="cpu"
+        )
+
+        # read q4 binary
+        bin_path = os.path.join(model_path, "embedding.q4_k_m.bin")
+        with open(bin_path, "rb") as f:
+            raw = f.read()
+
+        hidden = state["dyt.alpha"].shape[0]
+        total_uint8 = len(raw)
+
+        bytes_per_row = hidden // 2 + 4
+        vocab = total_uint8 // bytes_per_row
+
+        packed_size = vocab * hidden // 2
+
+        packed = np.frombuffer(raw[:packed_size], dtype=np.uint8)
+        scales = np.frombuffer(raw[packed_size:], dtype=np.float32)
+
+        packed = packed.reshape(vocab, hidden // 2)
+        scales = scales.reshape(vocab, 1)
+
+        emb = dequantize_q4_k_m(packed, scales)
+
+        # rebuild model
+        model = cls(
+            tokenizer=tokenizer,
+            vocab_size=emb.shape[0],
+            hidden_dim=emb.shape[1]
+        )
+
+        state["embedding.weight"] = emb
+        model.load_state_dict(state)
+
+        return model
+
+
+@dataclass
+class SSESforzandoConfig:
+    hidden_dim: int = 512
+    vocab_size: int = 30522
+
+
+@dataclass
+class SSEForzandoConfig:
+    hidden_dim: int = 384
+    vocab_size: int = 30522