Added utilities (#3)

Author: spietras

kilroy_module_pytorch_py_sdk/poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "kilroy-module-pytorch-py-sdk"
-version = "0.1.0"
+version = "0.2.0"
 description = "SDK for kilroy modules using PyTorch 🧰"
 readme = "README.md"
 authors = ["kilroy <[email protected]>"]
@@ -11,6 +11,8 @@ documentation = "https://kilroybot.github.io/kilroy-module-pytorch-py-sdk"
 
 [tool.poetry.dependencies]
 python = "^3.9"
+torch = "~1"
+diskcache = "~5"
 
 # dev
 

kilroy_module_pytorch_py_sdk/pyproject.toml

@@ -1 +1,40 @@
-from kilroy_module_pytorch_py_sdk.resources import *
+from kilroy_module_pytorch_py_sdk.resources import (
+    resource,
+    resource_bytes,
+    resource_text,
+)
+from kilroy_module_pytorch_py_sdk.buffer import TensorBuffer
+from kilroy_module_pytorch_py_sdk.cache import (
+    Cache,
+    Disk,
+    CacheLike,
+    TensorDisk,
+    TensorCache,
+    ProgressiveCache,
+    ProgressiveTensorCache,
+    SelfCleaningCache,
+)
+from kilroy_module_pytorch_py_sdk.generator import (
+    GenerationResult,
+    SequenceGenerator,
+)
+from kilroy_module_pytorch_py_sdk.samplers import (
+    SampleResult,
+    Sampler,
+    ProportionalCategoricalSampler,
+    TopKCategoricalSampler,
+    NucleusCategoricalSampler,
+    WithEpsilon,
+    EpsilonProportionalCategoricalSampler,
+    EpsilonTopKCategoricalSampler,
+    EpsilonNucleusCategoricalSampler,
+)
+from kilroy_module_pytorch_py_sdk.utils import (
+    pad,
+    unpad,
+    pack_padded,
+    pack_list,
+    unpack_to_padded,
+    unpack_to_list,
+    squash_packed,
+)

kilroy_module_pytorch_py_sdk/src/kilroy_module_pytorch_py_sdk/__init__.py

@@ -0,0 +1,41 @@
+from typing import Generic, MutableMapping, TypeVar
+from uuid import uuid4
+
+from torch import Tensor, nn
+from torch.autograd.graph import saved_tensors_hooks
+
+from kilroy_module_pytorch_py_sdk.utils import SelfCleaningKey
+
+StoreType = TypeVar("StoreType", bound=MutableMapping)
+ModelType = TypeVar("ModelType", bound=nn.Module)
+
+
+class TensorBuffer(Generic[StoreType]):
+    _store: StoreType
+
+    def __init__(self, store: StoreType) -> None:
+        self._store = store
+        self._hooks_ctx = saved_tensors_hooks(self.pack_hook, self.unpack_hook)
+
+    @property
+    def store(self) -> StoreType:
+        return self._store
+
+    def pack_hook(self, x: Tensor) -> SelfCleaningKey:
+        # TODO: remove SelfCleaningKey and just return UUID int
+        # on unpack_hook add this int to some set
+        # and make flush() method to delete all keys from this set
+        # this method should be called after backwards() is called
+        key = SelfCleaningKey(uuid4().hex, self._store)
+        self._store[key.key] = x.clone()
+        return key
+
+    def unpack_hook(self, x: SelfCleaningKey) -> Tensor:
+        return self._store[x.key]
+
+    def __enter__(self) -> "TensorBuffer":
+        self._hooks_ctx.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
+        return self._hooks_ctx.__exit__(exc_type, exc_value, traceback)