File size: 5,525 Bytes
bff06c9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | """GreenLeaf Law Embed — quantization modules.
Provides native int8 and binary quantization for embedding vectors,
enabling efficient storage and fast similarity search without
requiring post-hoc compression.
Quantizers use straight-through estimation: the forward pass applies
hard quantization while gradients flow through the soft (continuous)
approximation during training.
"""
import torch
import numpy as np
from typing import Literal
from sentence_transformers.models import Module
class BaseQuantizer(torch.nn.Module):
"""Base quantizer with straight-through gradient estimation.
During inference (hard=True), applies discrete quantization.
During training, gradients bypass the discretization step so
the model can learn through the quantization bottleneck.
"""
def __init__(self, hard: bool = True):
"""
Args:
hard: If True, output is discretized. If False, output is
the soft continuous approximation.
"""
super().__init__()
self._hard = hard
def _hard_quantize(self, x, *args, **kwargs) -> torch.Tensor:
raise NotImplementedError
def _soft_quantize(self, x, *args, **kwargs) -> torch.Tensor:
raise NotImplementedError
def forward(self, x, *args, **kwargs) -> torch.Tensor:
soft = self._soft_quantize(x, *args, **kwargs)
if not self._hard:
return soft
# Straight-through estimator: forward uses hard quantization,
# backward passes gradients through the soft path unchanged.
return (
self._hard_quantize(x, *args, **kwargs).detach()
+ soft - soft.detach()
)
class Int8EmbeddingQuantizer(BaseQuantizer):
"""Quantizes embeddings to signed 8-bit integers via tanh scaling.
The soft path applies tanh to squash values into [-1, 1], then
the hard path scales to [-128, 127] and rounds to integers.
"""
def __init__(self, hard: bool = True):
super().__init__(hard=hard)
self.qmin = -128
self.qmax = 127
def _soft_quantize(self, x, *args, **kwargs):
return torch.tanh(x)
def _hard_quantize(self, x, *args, **kwargs):
soft = self._soft_quantize(x)
int_x = torch.round(soft * self.qmax)
return torch.clamp(int_x, self.qmin, self.qmax)
class BinaryEmbeddingQuantizer(BaseQuantizer):
"""Quantizes embeddings to {-1, +1} via sign function.
The soft path uses scaled tanh as a differentiable approximation
to the sign function. The hard path applies the actual sign.
"""
def __init__(self, hard: bool = True, scale: float = 1.0):
super().__init__(hard)
self._scale = scale
def _soft_quantize(self, x, *args, **kwargs):
return torch.tanh(self._scale * x)
def _hard_quantize(self, x, *args, **kwargs):
return torch.where(x >= 0, 1.0, -1.0)
class PackedBinaryEncoder:
"""Packs binary {-1, +1} embeddings into uint8 bit-arrays.
Each embedding dimension maps to 1 bit, reducing storage by 32x
compared to float32. Useful for large-scale retrieval with
Hamming distance.
"""
def __call__(self, x: torch.Tensor) -> torch.Tensor:
bits = np.where(x.cpu().numpy() >= 0, True, False)
packed = np.packbits(bits, axis=-1)
return torch.from_numpy(packed).to(x.device)
class FlexibleQuantizer(Module):
"""Sentence-transformers module that applies optional quantization.
Default behavior: pass through raw float embeddings (bf16/fp32).
Quantization is applied only when explicitly requested.
Supported modes:
- None: raw float embeddings (default, no quantization)
- "int8": signed 8-bit integer embeddings
- "binary": {-1, +1} binary embeddings
- "ubinary": packed binary as uint8 bit-arrays
"""
def __init__(self):
super().__init__()
self._int8_quantizer = Int8EmbeddingQuantizer()
self._binary_quantizer = BinaryEmbeddingQuantizer()
self._packed_binary_encoder = PackedBinaryEncoder()
def forward(
self,
features: dict[str, torch.Tensor],
quantization: Literal["int8", "binary", "ubinary"] | None = None,
**kwargs,
) -> dict[str, torch.Tensor]:
if quantization is None:
# Default: return raw float embeddings, no quantization
return features
elif quantization == "int8":
features["sentence_embedding"] = self._int8_quantizer(
features["sentence_embedding"]
)
elif quantization == "binary":
features["sentence_embedding"] = self._binary_quantizer(
features["sentence_embedding"]
)
elif quantization == "ubinary":
features["sentence_embedding"] = self._packed_binary_encoder(
features["sentence_embedding"]
)
else:
raise ValueError(
f"Unknown quantization mode: '{quantization}'. "
f"Supported modes: 'int8', 'binary', 'ubinary'."
)
return features
@classmethod
def load(
cls,
model_name_or_path: str,
subfolder: str = "",
token: bool | str | None = None,
cache_folder: str | None = None,
revision: str | None = None,
local_files_only: bool = False,
**kwargs,
):
return cls()
def save(self, output_path: str, *args, **kwargs) -> None:
return
|