Update st_quantize.py

st_quantize.py

@@ -90,4 +90,32 @@ class NormalizedInt8TanhQuantizer(Int8TanhQuantizer):
     @classmethod
     def load(cls, input_path: str) -> "PoolAndQuantize":
         return cls()
+
+
+class Binarizer(Quantizer):
+    def __init__(self, tanh_scale: float = 1.0, **kwargs):
+        super().__init__(**kwargs)
+        self._tanh_scale = tanh_scale
+
+    def _hard_quantize(self, x, *args, **kwargs) -> torch.Tensor:
+        return torch.where(x > 0, 1.0, -1.0)
+
+    def _soft_quantize(self, x, *args, **kwargs) -> torch.Tensor:
+        return torch.tanh(x * self._tanh_scale)
+
+
+class UnnormalizedBinarizer(nn.Module):
+    def __init__(self, tanh_scale: float = 1.0, hard: bool = True):
+        super().__init__()
+        self.quantizer = Binarizer(tanh_scale=tanh_scale, hard=hard)
+
+    def forward(self, features: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        features["sentence_embedding"] = self.quantizer(
+            features["sentence_embedding"]
+        )
+        return features
+
+    @classmethod
+    def load(cls, input_path: str) -> "UnnormalizedBinarizer":
+        return cls()