Implement auto-register seed-scheme factory

seed_scheme_factory.py

@@ -0,0 +1,44 @@
+from typing import Union, Callable
+
+import torch
+
+
+class SeedSchemeFactory:
+    registry = {}
+
+    @classmethod
+    def register(cls, name: str):
+        """
+        Register the hash scheme by name. Hash scheme must be callable.
+
+        Args:
+            name: name of seed scheme.
+        """
+
+        def wrapper(wrapped_class):
+            if name in cls.registry:
+                print(f"Override {name} in SeedSchemeFactory")
+            cls.registry[name] = wrapped_class
+            return wrapped_class
+
+        return wrapper
+
+    @classmethod
+    def get_instance(cls, name: str, *args, **kwargs):
+        """
+        Get the hash scheme by name.
+
+        Args:
+            name: name of seed scheme.
+        """
+        if name in cls.registry:
+            return cls.registry[name](*args, **kwargs)
+        else:
+            return None
+
+
+class SeedScheme():
+    def __call__(self, input_ids: torch.Tensor) -> int:
+        return 0
+
+from seed_schemes import *

seed_schemes/__init__.py

@@ -0,0 +1,13 @@
+from inspect import isclass
+from pkgutil import iter_modules
+from importlib import import_module
+import os
+
+pkg_dir = os.path.dirname(__file__)
+for _, module_name, _ in iter_modules([pkg_dir]):
+    module = import_module(f"{__name__}.{module_name}")
+
+    for attr_name in dir(module):
+        attr = getattr(module, attr_name)
+        if isclass(attr):
+            globals()[attr_name] = attr

seed_schemes/dummy_hash.py

@@ -0,0 +1,12 @@
+import torch
+
+from seed_scheme_factory import SeedSchemeFactory, SeedScheme
+
+@SeedSchemeFactory.register("dummy_hash")
+class DummyHash(SeedScheme):
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def __call__(self, input_ids: torch.Tensor):
+        return int(input_ids[-1].item())
+

seed_schemes/sha_hash.py

@@ -0,0 +1,33 @@
+from typing import Union
+import torch
+import io
+
+from cryptography.hazmat.primitives import hashes
+
+from seed_scheme_factory import SeedSchemeFactory, SeedScheme
+
+
+@SeedSchemeFactory.register("sha_left_hash")
+class SHALeftHash(SeedScheme):
+    def __init__(self, private_key: Union[int, None] = None, *args, **kwargs):
+        self.private_key = (
+            private_key.to_bytes(8, "big") if private_key is not None else None
+        )
+
+    def __call__(self, input_ids: torch.Tensor):
+        buff = io.BytesIO()
+        if self.private_key is not None:
+            buff.write(self.private_key)
+        for input_id in input_ids:
+            buff.write(int(input_id.item()).to_bytes(8, "big"))
+        buff.seek(0)
+        input_ids_bytes = buff.read()
+
+        digest = hashes.Hash(hashes.SHA224())
+        digest.update(input_ids_bytes)
+        hashed_value = digest.finalize()
+
+        # Only take the first 8 bytes because seed in torch rng only accept int64 seed
+        seed = int.from_bytes(hashed_value[:8], byteorder="big")
+
+        return seed