Upload scalar tokenizer to HF

Author: LTMeyer

aion/codecs/tokenizers/scalar.py

@@ -1,14 +1,19 @@
 from typing import Dict, Optional
 
 import torch
+from huggingface_hub import PyTorchModelHubMixin
 from jaxtyping import Float
 from torch import Tensor
 
 from aion.codecs.quantizers import Quantizer
+from aion.codecs.quantizers.scalar import (
+    ScalarLogReservoirQuantizer,
+    ScalarReservoirQuantizer,
+)
 from aion.codecs.tokenizers.base import QuantizedCodec
 
 
-class ScalarIdentityCodec(QuantizedCodec):
+class BaseScalarIdentityCodec(QuantizedCodec, PyTorchModelHubMixin):
     """Codec for scalar quantities.
 
     A codec that embeds scalar quantities through an identity mapping. A
@@ -34,3 +39,21 @@ def _encode(self, x: Dict[str, Dict[str, Float[Tensor, "b t"]]]) -> Tensor:
 
     def _decode(self, z: torch.FloatTensor) -> Dict[str, torch.FloatTensor]:
         return {self.modality: z}
+
+
+class ScalarReservoirCodec(BaseScalarIdentityCodec):
+    def __init__(self, modality: str, codebook_size: int, reservoir_size: int):
+        quantizer = ScalarReservoirQuantizer(
+            codebook_size=codebook_size,
+            reservoir_size=reservoir_size,
+        )
+        super().__init__(modality, quantizer)
+
+
+class ScalarLogReservoirCodec(BaseScalarIdentityCodec):
+    def __init__(self, modality: str, codebook_size: int, reservoir_size: int):
+        quantizer = ScalarLogReservoirQuantizer(
+            codebook_size=codebook_size,
+            reservoir_size=reservoir_size,
+        )
+        super().__init__(modality, quantizer)

tests/tokenizers/test_scalar_tokenizer.py

@@ -1,11 +1,10 @@
 import pytest
 import torch
 
-from aion.codecs.quantizers.scalar import (
-    ScalarLogReservoirQuantizer,
-    ScalarReservoirQuantizer,
+from aion.codecs.tokenizers.scalar import (
+    ScalarLogReservoirCodec,
+    ScalarReservoirCodec,
 )
-from aion.codecs.tokenizers.scalar import ScalarIdentityCodec
 
 
 @pytest.mark.parametrize(
@@ -23,14 +22,10 @@
     ],
 )
 def test_log_reservoir_tokenizer(data_dir, modality):
-    codec = ScalarIdentityCodec(
-        modality=modality,
-        quantizer=ScalarLogReservoirQuantizer(
-            codebook_size=1024, reservoir_size=100000
-        ),
+    codec = ScalarLogReservoirCodec.from_pretrained(
+        f"polymathic-ai/aion-scalar-{modality.lower().replace('_', '-')}-codec"
     )
     codec.eval()
-    codec.load_state_dict(torch.load(data_dir / f"{modality}_codec.pt"))
 
     input_batch = torch.load(data_dir / f"{modality}_input.pt")
     output_batch = torch.load(data_dir / f"{modality}_output.pt")
@@ -42,9 +37,8 @@ def test_log_reservoir_tokenizer(data_dir, modality):
 
 @pytest.mark.parametrize("modality", ["SHAPE_E1", "SHAPE_E2", "EBV"])
 def test_reservoir_tokenizer(data_dir, modality):
-    codec = ScalarIdentityCodec(
-        modality=modality,
-        quantizer=ScalarReservoirQuantizer(codebook_size=1024, reservoir_size=100000),
+    codec = ScalarReservoirCodec.from_pretrained(
+        f"polymathic-ai/aion-scalar-{modality.lower().replace('_', '-')}-codec"
     )
     codec.eval()
     codec.load_state_dict(torch.load(data_dir / f"{modality}_codec.pt"))