Update so huge that I won't even bother and will commit all at once (#30)

Author: spietras

kilroy_module_pytorch_py_sdk/poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "kilroy-module-pytorch-py-sdk"
-version = "0.7.2"
+version = "0.8.0"
 description = "SDK for kilroy modules using PyTorch 🧰"
 readme = "README.md"
 authors = ["kilroy <[email protected]>"]
@@ -13,7 +13,7 @@ documentation = "https://kilroybot.github.io/kilroy-module-pytorch-py-sdk"
 python = "^3.10"
 torch = "~1"
 numpy = "~1"
-kilroy-module-server-py-sdk = "~0.9"
+kilroy-module-server-py-sdk = "~0.10"
 
 [tool.poetry.group.poe.dependencies]
 poethepoet = "^0.16"

kilroy_module_pytorch_py_sdk/pyproject.toml

@@ -1,21 +1,14 @@
 from kilroy_module_server_py_sdk import *
-from kilroy_module_pytorch_py_sdk.codec import Codec
-from kilroy_module_pytorch_py_sdk.generator import GenerationResult, Generator
-from kilroy_module_pytorch_py_sdk.models import LanguageModel, RewardModel
-from kilroy_module_pytorch_py_sdk.modules.basic import (
-    BasicModule,
-    MetricsState as BasicModuleMetricsState,
-    ReportsState as BasicModuleReportsState,
-    State as BasicModuleState,
-)
-from kilroy_module_pytorch_py_sdk.modules.reward import (
-    LanguageModelState as RewardModelModuleLanguageModelState,
-    MetricsState as RewardModelModuleMetricsState,
-    ReportsState as RewardModelModuleReportsState,
-    RewardModelModule,
-    RewardModelState as RewardModelModuleRewardModelState,
-    State as RewardModelModuleState,
+from kilroy_module_pytorch_py_sdk.generator import Generator
+from kilroy_module_pytorch_py_sdk.models.abc import SequentialModel
+from kilroy_module_pytorch_py_sdk.models.loader import ModelLoader
+from kilroy_module_pytorch_py_sdk.models.registry import ModelsRegistry
+from kilroy_module_pytorch_py_sdk.module.module import PytorchModule
+from kilroy_module_pytorch_py_sdk.trainers import (
+    Trainer,
+    VanillaTrainer,
 )
+from kilroy_module_pytorch_py_sdk.metrics import LineMetric
 from kilroy_module_pytorch_py_sdk.optimizers import (
     AdamOptimizer,
     Optimizer,
@@ -27,15 +20,6 @@
     resource_bytes,
     resource_text,
 )
-from kilroy_module_pytorch_py_sdk.samplers import (
-    EpsilonNucleusSampler,
-    EpsilonProportionalSampler,
-    EpsilonTopKSampler,
-    NucleusSampler,
-    ProportionalSampler,
-    Sampler,
-    TopKSampler,
-)
 from kilroy_module_pytorch_py_sdk.schedulers import (
     ConstantScheduler,
     CosineAnnealingScheduler,
@@ -51,7 +35,6 @@
 )
 from kilroy_module_pytorch_py_sdk.tokenizer import Tokenizer
 from kilroy_module_pytorch_py_sdk.utils import (
-    freeze,
     pack_list,
     pack_padded,
     pad,
@@ -62,4 +45,5 @@
     unpack_to_list,
     unpack_to_padded,
     unpad,
+    freeze,
 )

kilroy_module_pytorch_py_sdk/src/kilroy_module_pytorch_py_sdk/__init__.py

@@ -0,0 +1,88 @@
+from typing import Dict, Any, Callable, Awaitable
+
+import torch
+from pydantic import Field
+from torch.nn.utils.rnn import PackedSequence
+
+from kilroy_module_pytorch_py_sdk.utils import unpack_to_list, pack_list
+from kilroy_module_server_py_sdk import SerializableState
+from kilroy_server_py_utils import Configurable, Parameter, classproperty
+
+
+class State(SerializableState):
+    gamma: float = 0.99
+    lambda_: float = Field(0.95, alias="lambda")
+
+
+class GeneralizedAdvantageEstimator(Configurable[State]):
+    class GammaParameter(Parameter[State, float]):
+        # noinspection PyMethodParameters
+        @classproperty
+        def schema(cls) -> Dict[str, Any]:
+            return {
+                "type": "number",
+                "minimum": 0,
+                "maximum": 1,
+                "default": 0.99,
+                "title": cls.pretty_name,
+            }
+
+    class LambdaParameter(Parameter[State, float]):
+        # noinspection PyMethodParameters
+        @classproperty
+        def schema(cls) -> Dict[str, Any]:
+            return {
+                "type": "number",
+                "minimum": 0,
+                "maximum": 1,
+                "default": 0.95,
+                "title": cls.pretty_name,
+            }
+
+        @classmethod
+        async def _get(cls, state: State) -> float:
+            return state.lambda_
+
+        @classmethod
+        async def _set(
+            cls, state: State, value: float
+        ) -> Callable[[], Awaitable]:
+            original_value = state.lambda_
+
+            async def undo():
+                state.lambda_ = original_value
+
+            state.lambda_ = value
+            return undo
+
+    async def calculate(
+        self, rewards: PackedSequence, values: PackedSequence
+    ) -> PackedSequence:
+        async with self.state.read_lock() as state:
+            gamma = state.gamma
+            lambda_ = state.lambda_
+
+        batch_rewards = unpack_to_list(rewards)
+        batch_values = unpack_to_list(values)
+
+        batch_advantages = []
+
+        for rewards, values in zip(batch_rewards, batch_values):
+            advantages = []
+            advantage = 0
+
+            for i in reversed(range(len(rewards))):
+                reward = rewards[i]
+                value = values[i]
+                next_value = (
+                    torch.zeros(1) if i == len(rewards) - 1 else values[i + 1]
+                )
+
+                delta = reward + gamma * next_value - value
+                advantage = delta + gamma * lambda_ * advantage
+                advantages.append(advantage)
+
+            advantages = torch.stack(advantages[::-1])
+            batch_advantages.append(advantages)
+
+        return pack_list(batch_advantages)

kilroy_module_pytorch_py_sdk/src/kilroy_module_pytorch_py_sdk/codec.py

@@ -1,2 +1 @@
 from kilroy_module_pytorch_py_sdk.generator.generator import Generator
-from kilroy_module_pytorch_py_sdk.generator.utils import GenerationResult