rxnn 0.2.49__py3-none-any.whl → 0.2.51__py3-none-any.whl

rxnn/training/models.py

@@ -208,17 +208,18 @@ class MrlActorModel(nn.Module):
 
 
 class MrlCriticModel(nn.Module, PyTorchModelHubMixin, license="apache-2.0", pipeline_tag="text-classification"):
-    def __init__(self, encoder: nn.Module, embed_dim: int,
-                 out_activation: Literal['sigmoid', 'tanh', 'linear'] = 'sigmoid', output_scale: float = 1.0, **kwargs):
+    def __init__(self, encoder: nn.Module, embed_dim: int, **kwargs):
         super(MrlCriticModel, self).__init__(**kwargs)
         self.encoder = encoder
         self.value_head = nn.Sequential(
             GatedLinearUnit(embed_dim, embed_dim, nn.SiLU()),
             nn.LayerNorm(embed_dim),
-            nn.Linear(embed_dim, 1),
-            get_activation_layer(out_activation)
+            nn.Linear(embed_dim, 1)
         )
-        self.output_scale = output_scale
+        # Learnable scaling parameters
+        self.scale = nn.Parameter(torch.tensor(1.0))
+        self.shift = nn.Parameter(torch.tensor(0.0))
+
 
     def head_parameters(self) -> Iterator[nn.Parameter]:
         return self.value_head.parameters()
@@ -235,4 +236,4 @@ class MrlCriticModel(nn.Module, PyTorchModelHubMixin, license="apache-2.0", pipe
         else:
             x = x.mean(dim=1)
 
-        return self.value_head(x) * self.output_scale
+        return self.value_head(x) * self.scale + self.shift

rxnn/training/reward.py

@@ -242,8 +242,8 @@ class MrlRewardModel:
         return self.neg_cos_saved_factor * (1 - gen_and_saved) + self.neg_cos_ref_factor * gen_and_ref
 
     def len_reward(self, generated: TokenizedDict, reference: TokenizedDict) -> torch.Tensor:
-        target_lens = reference['attention_mask'].sum(dim=1) if self.target_len_as_ref else self.max_rewarded_len
-        lens = generated['attention_mask'].sum(dim=1)
+        target_lens = reference['attention_mask'].to(self.device).sum(dim=1) if self.target_len_as_ref else self.max_rewarded_len
+        lens = generated['attention_mask'].to(self.device).sum(dim=1)
         neg_lens = target_lens / lens if self.neg_reward_len else 1.0
         len_reward = torch.where(lens >= target_lens, neg_lens, lens / target_lens)
         return len_reward

rxnn/training/rl.py

@@ -55,6 +55,7 @@ class PPOAlgorithm(RlAlgorithm):
         # Critic loss with clipped values
         if self.clip_critic_values:
             values = torch.clamp(values, -self.critic_value_clip, self.critic_value_clip)
+            ref_values = torch.clamp(ref_values, -self.critic_value_clip, self.critic_value_clip)
         return self.critic_loss_fn(values, ref_values)
 
     def policy_loss(self, query: TokenizedDict, answer: TokenizedDict, logits: torch.Tensor,

{rxnn-0.2.49.dist-info → rxnn-0.2.51.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: rxnn
-Version: 0.2.49
+Version: 0.2.51
 Summary: RxNN: Reactive Neural Networks Platform
 License: Apache-2.0
 Keywords: deep-learning,ai,machine-learning

{rxnn-0.2.49.dist-info → rxnn-0.2.51.dist-info}/RECORD

@@ -16,10 +16,10 @@ rxnn/training/bml.py,sha256=hw6gLpLkGvqLzxIvBg4MvCc5r8cHpEm2RDyh7nH6CtE,16914
 rxnn/training/callbacks.py,sha256=rS8leuVFPVVfE5Zc8DMkUZhRIPN-vpPbUjowXE5TSBw,36779
 rxnn/training/dataset.py,sha256=7hTilFWPpqUEc6zNcMqBPjxFKxCfvTKKF3E8tVlwccQ,51250
 rxnn/training/ddp.py,sha256=VsNBjn3cY-uUj8hbsW7oKvb0_ZKnXnJ2KgObm-Mr9i4,836
-rxnn/training/models.py,sha256=L2emJM06u7B9f9T1dFsGXzXX-rsV77ND7L1pAM9Z_Ow,9051
+rxnn/training/models.py,sha256=CS6mjD338knXmCbMZ3bCpOlA-DR3kmQUOSj5u5F6jII,9002
 rxnn/training/mrl.py,sha256=IOi_xbQ47RPgv_2ucT9EkPeWLGBRlgPxKHFeQsYc3Pw,61074
-rxnn/training/reward.py,sha256=dq3b5DRhBLHOvtlHX3eSSuxYBGYCyV5gVqbzCam4uP8,16112
-rxnn/training/rl.py,sha256=q4NzIZAmXRHVToT13IHrPTtEikWQUvT0NO0IjApjAO8,6171
+rxnn/training/reward.py,sha256=uiSsBXmjMw2yv-1Bssy3RTlpU6zP8ape3490Sl-aT0M,16144
+rxnn/training/rl.py,sha256=FiOag3kaI4I40ylXE9Yx5iHWmprINBSMBbarKudABEE,6269
 rxnn/training/scheduler.py,sha256=LcjU35mEwz2U5x3U6tLfeeYlBqMxbFSxYzJYuXkWbSY,1408
 rxnn/training/tokenizer.py,sha256=umaLByMBx_NMrQElA45HLm9gkuzyKWDTFaKVd-CjXl0,8344
 rxnn/training/utils.py,sha256=Bw8nZLKIt7NQpUVCYkb_79kWKChVFOYgYXwODo4SvNc,5718
@@ -33,7 +33,7 @@ rxnn/transformers/moe.py,sha256=j6jEx6Ip0zttlUZKKn82azxo95lkLZs-H2GLSMD88hY,5859
 rxnn/transformers/positional.py,sha256=1PjcJybUzeQlIKJI4tahAGZcYgCRCL0otxs7mpsNuzM,4410
 rxnn/transformers/sampler.py,sha256=t6iiQTdLQ0TakUWnnhKkb5DKF2F_9-thXHBydDF3fxg,17389
 rxnn/utils.py,sha256=ihb6OTyDtPiocB_lOvnq7eOkjjpCkgs8wxvXUBNQ7mM,996
-rxnn-0.2.49.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.49.dist-info/METADATA,sha256=Yd5xCJVA_rFdzYkTkHZ8tyronArNMOgUQ6VqNF9-vqs,25997
-rxnn-0.2.49.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.49.dist-info/RECORD,,
+rxnn-0.2.51.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.51.dist-info/METADATA,sha256=vDeCYIrxa3o0Pe09n_nppoMvyAIHnSyRJw4Q74ofBIQ,25997
+rxnn-0.2.51.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.51.dist-info/RECORD,,