vortex_embed_v4_5.py

"""Vortex-Embed v4.5 — Standalone Single-File Native 4-Bit Embedding Engine.

Lightweight, 4.72 MB RAM native 4-bit embedding engine with Matryoshka support.
Repo: https://huggingface.co/VTXAI/Vortex-Embed-v4-5-sentence
"""
from __future__ import annotations

import json
from dataclasses import dataclass, field
from pathlib import Path
from typing import List, Optional, Sequence, Tuple, Union

import numpy as np
from safetensors.numpy import load_file, save_file

try:
    from tokenizers import Tokenizer
except ImportError:
    Tokenizer = None


class VortexEmbedConfig:
    def __init__(
        self,
        vocab_size: int = 29528,
        embedding_dim: int = 256,
        block_size: int = 32,
        num_blocks: int = 8,
        model_type: str = "vortex-embed",
        architectures: Optional[List[str]] = None,
        quantization: str = "lf4",
        bits: int = 4,
        sif_a: float = 0.05,
        sif_pc: float = 1.0,
        pc_k: int = 1,
        matryoshka_dim: Optional[int] = None,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.embedding_dim = embedding_dim
        self.block_size = block_size
        self.num_blocks = num_blocks
        self.model_type = model_type
        self.architectures = architectures or ["VortexEmbedV4_5"]
        self.quantization = quantization
        self.bits = bits
        self.sif_a = sif_a
        self.sif_pc = sif_pc
        self.pc_k = pc_k
        self.matryoshka_dim = matryoshka_dim

    @classmethod
    def from_dict(cls, d: dict) -> "VortexEmbedConfig":
        return cls(**d)

    def to_dict(self) -> dict:
        return {
            "vocab_size": self.vocab_size,
            "embedding_dim": self.embedding_dim,
            "block_size": self.block_size,
            "num_blocks": self.num_blocks,
            "model_type": self.model_type,
            "architectures": self.architectures,
            "quantization": self.quantization,
            "bits": self.bits,
            "sif_a": self.sif_a,
            "sif_pc": self.sif_pc,
            "pc_k": self.pc_k,
            "matryoshka_dim": self.matryoshka_dim,
        }



class VortexEmbedV4_5:
    """Vortex-Embed v4.5 — Native 4-Bit Sentence Embedding Model.

    Features:
      - 4.72 MB RAM Footprint (Zero FP32 matrix in RAM).
      - On-the-fly dequantization per batch.
      - Matryoshka Representation Learning (truncation to 256, 128, 64 dims).
    """

    def __init__(
        self,
        packed: np.ndarray,
        scales: np.ndarray,
        zeros: np.ndarray,
        tokenizer_data: Union[str, Path],
        config: Union[dict, VortexEmbedConfig],
        *,
        matryoshka_dim: Optional[int] = None,
    ) -> None:
        self.packed = np.asarray(packed, dtype=np.uint8)
        self.scales = np.asarray(scales, dtype=np.float16)
        self.zeros = np.asarray(zeros, dtype=np.float16)
        self.tokenizer_data = str(tokenizer_data)
        self.config = config if isinstance(config, VortexEmbedConfig) else VortexEmbedConfig.from_dict(config)
        self.vocab_size = int(self.config.vocab_size)
        self.dim = int(self.config.embedding_dim)
        self.block_size = int(self.config.block_size)
        self.num_blocks = int(self.config.num_blocks)
        self.sif_a = float(self.config.sif_a)
        self.sif_pc = float(self.config.sif_pc)
        self.pc_k = int(self.config.pc_k)
        self.matryoshka_dim = matryoshka_dim or self.config.matryoshka_dim

        self._tokenizer: Optional[Tokenizer] = None
        self._sif_weights: Optional[np.ndarray] = None
        self._pc_directions: Optional[np.ndarray] = None

    @property
    def tokenizer(self) -> Tokenizer:
        if self._tokenizer is None:
            if Tokenizer is None:
                raise RuntimeError("tokenizers required: pip install tokenizers")
            self._tokenizer = Tokenizer.from_file(self.tokenizer_data)
        return self._tokenizer

    @property
    def model_size_mb(self) -> float:
        """Returns actual in-RAM size of stored parameters (4.72 MB)."""
        return (self.packed.nbytes + self.scales.nbytes + self.zeros.nbytes) / 1e6

    @property
    def on_disk_size_mb(self) -> float:
        return (self.packed.nbytes + self.scales.nbytes + self.zeros.nbytes) / 1e6

    @classmethod
    def from_pretrained(
        cls,
        path_or_id: Union[str, Path],
        matryoshka_dim: Optional[int] = None,
        **overrides,
    ) -> "VortexEmbedV4_5":
        path = Path(path_or_id)
        if not path.is_dir():
            from huggingface_hub import snapshot_download
            path = Path(snapshot_download(str(path_or_id)))
        tensors = load_file(str(path / "model.safetensors"))
        config = json.loads((path / "config.json").read_text())
        for k, v in overrides.items():
            if k in VortexEmbedConfig.__dataclass_fields__:
                config[k] = v
        return cls(
            packed=tensors["embedding_packed"],
            scales=tensors["embedding_scales"],
            zeros=tensors["embedding_zeros"],
            tokenizer_data=str(path / "tokenizer.json"),
            config=config,
            matryoshka_dim=matryoshka_dim,
        )

    def save_pretrained(self, path: Union[str, Path]) -> None:
        out = Path(path)
        out.mkdir(parents=True, exist_ok=True)
        save_file(
            {
                "embedding_packed": self.packed,
                "embedding_scales": self.scales,
                "embedding_zeros": self.zeros,
            },
            str(out / "model.safetensors"),
        )
        (out / "config.json").write_text(json.dumps(self.config.to_dict(), indent=2))
        if not (out / "tokenizer.json").exists():
            (out / "tokenizer.json").write_text(Path(self.tokenizer_data).read_text())

    def _dequantize_ids_on_the_fly(self, token_ids: np.ndarray) -> np.ndarray:
        if token_ids.size == 0:
            return np.empty((0, self.dim), dtype=np.float32)

        p = self.packed[token_ids]
        s = self.scales[token_ids].astype(np.float32)[:, :, None]
        z = self.zeros[token_ids].astype(np.float32)[:, :, None]

        low = (p & 0x0F).astype(np.float32)
        high = ((p >> 4) & 0x0F).astype(np.float32)

        n = len(token_ids)
        padded = p.shape[1] * 2
        unpacked = np.empty((n, padded), dtype=np.float32)
        unpacked[:, 0::2] = low
        unpacked[:, 1::2] = high

        blocked = unpacked.reshape(n, self.num_blocks, self.block_size)
        out = (blocked * s + z).reshape(n, padded)
        return out[:, : self.dim]

    def fit_idf(self, corpus_token_lists: Sequence[Sequence[int]]) -> "VortexEmbedV4_5":
        flat = (np.concatenate(corpus_token_lists) if corpus_token_lists else np.empty(0, dtype=np.int64))
        total = max(int(flat.size), 1)
        counts = np.bincount(flat, minlength=self.vocab_size).astype(np.float64)
        p = counts / total
        denom = self.sif_a + p
        with np.errstate(divide="ignore", invalid="ignore"):
            weights = np.where(p > 0, self.sif_a / denom, 1.0)
        self._sif_weights = weights.astype(np.float32)
        return self

    def fit_pc(self, corpus_embeddings: np.ndarray, k: Optional[int] = None) -> "VortexEmbedV4_5":
        if k is None:
            k = self.pc_k
        if corpus_embeddings.size == 0 or k <= 0:
            return self
        x = corpus_embeddings.astype(np.float32)
        x = x - x.mean(axis=0, keepdims=True)
        try:
            _, _, vt = np.linalg.svd(x, full_matrices=False)
            pcs = vt[:k].astype(np.float32)
            pcs = pcs / (np.linalg.norm(pcs, axis=1, keepdims=True) + 1e-12)
            self._pc_directions = pcs
        except np.linalg.LinAlgError:
            self._pc_directions = None
        return self

    def _apply_pc(self, x: np.ndarray) -> np.ndarray:
        if self.sif_pc <= 0 or self._pc_directions is None:
            return x
        out = x
        for pc in self._pc_directions:
            proj = (out @ pc)[:, None] * pc[None, :]
            out = out - self.sif_pc * proj
        return out

    def _tokenize_batch(self, texts: Sequence[str]) -> List[List[int]]:
        encoded = self.tokenizer.encode_batch(list(texts))
        return [[tid for tid in item.ids if 0 <= int(tid) < self.vocab_size] for item in encoded]

    def encode_batch(
        self,
        texts: Sequence[str],
        *,
        normalize: bool = True,
        truncate_dim: Optional[int] = None,
    ) -> np.ndarray:
        if not texts:
            return np.zeros((0, self.dim), dtype=np.float32)

        token_lists = self._tokenize_batch(list(texts))
        n = len(token_lists)
        flat = (np.concatenate(token_lists) if token_lists else np.empty(0, dtype=np.int64))

        if flat.size == 0:
            return np.zeros((n, self.dim), dtype=np.float32)

        unique_ids, inverse_indices = np.unique(flat, return_inverse=True)
        unique_embs = self._dequantize_ids_on_the_fly(unique_ids)
        token_embs = unique_embs[inverse_indices]

        if self._sif_weights is not None:
            w = self._sif_weights[flat].astype(np.float32)[:, None]
            token_embs = token_embs * w

        chunk_lens = np.array([len(ids) for ids in token_lists], dtype=np.int64)
        chunk_ends = np.cumsum(chunk_lens)
        boundaries = np.empty(n + 1, dtype=np.int64)
        boundaries[0] = 0
        boundaries[1:] = chunk_ends

        sums = np.add.reduceat(token_embs, boundaries[:-1], axis=0)

        if self._sif_weights is not None:
            w_full = self._sif_weights[flat].astype(np.float32)
            w_per_row = np.add.reduceat(w_full, boundaries[:-1])
            w_per_row = np.maximum(w_per_row, 1e-12)
        else:
            w_per_row = np.maximum(chunk_lens.astype(np.float32), 1.0)

        embeddings = sums / w_per_row[:, None]
        embeddings = self._apply_pc(embeddings)

        dim = truncate_dim if truncate_dim is not None else self.matryoshka_dim
        if dim is not None and 0 < dim < self.dim:
            embeddings = embeddings[:, :dim]

        if normalize and embeddings.shape[0] > 0:
            norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
            np.divide(embeddings, np.maximum(norms, 1e-12), out=embeddings)

        return embeddings

    def encode(
        self,
        texts: Union[str, Sequence[str]],
        *,
        normalize: bool = True,
        truncate_dim: Optional[int] = None,
    ) -> np.ndarray:
        if isinstance(texts, str):
            return self.encode_batch([texts], normalize=normalize, truncate_dim=truncate_dim)[0]
        return self.encode_batch(list(texts), normalize=normalize, truncate_dim=truncate_dim)