rxnn 0.2.31__py3-none-any.whl → 0.2.33__py3-none-any.whl

rxnn/memory/stm.py

@@ -44,7 +44,6 @@ class ShortTermMemory(nn.Module):
 
     def update_all(self, new_stm: torch.Tensor):
         self.memory = new_stm
-        # self.memory.copy_(new_stm)
 
     def make_trainable(self):
         if not self.is_trainable:

rxnn/rxt/models.py

@@ -5,7 +5,7 @@ from huggingface_hub import PyTorchModelHubMixin
 from ..transformers.positional import RotaryPositionalEmbedding
 from ..transformers.attention import init_attention
 from ..transformers.layers import ReactiveTransformerLayer
-from ..transformers.models import ReactiveTransformerBase, ReactiveTransformerEncoder, ReactiveTransformerDecoder
+from ..transformers.models import ReactiveTransformerBase, ReactiveTransformerEncoder, ReactiveTransformerDecoder, ReactiveTransformerEncoderDetachStm
 from ..transformers.ff import get_activation_layer
 from ..memory.stm import ShortTermMemory
 from ..memory.norm import init_memory_norm
@@ -293,3 +293,29 @@ class RxTAlphaMemoryAttention(nn.Module, PyTorchModelHubMixin, license="apache-2
 
     def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:
         return self.model(x, attention_mask=attention_mask)
+
+class RxTAlphaCriticEncoder(RxTAlphaComponentBase, pipeline_tag="text-classification", license="apache-2.0"):
+    """RxT-Alpha (Reactive Transformer) encoder model"""
+
+    def __init__(self, **kwargs: RxTAlphaComponentConfig):
+        super(RxTAlphaCriticEncoder, self).__init__(False, **kwargs)
+
+    def _init_model(
+            self,
+            stm: ShortTermMemory,
+            layers: nn.ModuleList,
+            embedding: nn.Embedding,
+            use_flash_attention: bool,
+            embed_dim: int,
+            vocab_size: int
+    ) -> ReactiveTransformerEncoderDetachStm:
+        return ReactiveTransformerEncoderDetachStm(
+            stm=stm,
+            embedding=embedding,
+            own_layers=layers,
+            use_flash_attention=use_flash_attention,
+        )
+
+    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> tuple[torch.Tensor, torch.Tensor]:
+        return self.model(x, attention_mask=attention_mask)
+

rxnn/training/models.py

@@ -4,6 +4,7 @@ from enum import Enum
 from typing import Literal, Iterator
 from huggingface_hub import PyTorchModelHubMixin
 from ..transformers.models import ReactiveTransformerEncoder, ReactiveTransformerDecoder
+from ..transformers.ff import GatedLinearUnit, get_activation_layer
 
 class MLMHead(nn.Module, PyTorchModelHubMixin, license="apache-2.0"):
     def __init__(self, embed_dim: int, vocab_size: int, *args, **kwargs):
@@ -161,10 +162,16 @@ class MrlActorModel(nn.Module):
             return self.memory_attention(ed, attention_mask=attention_mask)
 
 class MrlCriticModel(nn.Module, PyTorchModelHubMixin, license="apache-2.0", pipeline_tag="text-classification"):
-    def __init__(self, encoder: nn.Module, embed_dim: int, **kwargs):
+    def __init__(self, encoder: nn.Module, embed_dim: int, out_activation: Literal['sigmoid', 'tanh', 'linear'] = 'sigmoid', output_scale: float = 1.0,  **kwargs):
         super(MrlCriticModel, self).__init__(**kwargs)
         self.encoder = encoder
-        self.value_head = nn.Linear(embed_dim, 1)
+        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)
+        )
+        self.output_scale = output_scale
 
     def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:
         x, _ = self.encoder(x, attention_mask=attention_mask)
@@ -175,4 +182,5 @@ class MrlCriticModel(nn.Module, PyTorchModelHubMixin, license="apache-2.0", pipe
         else:
             x = x.mean(dim=1)
 
-        return self.value_head(x)
+        return self.value_head(x) * self.output_scale
+

rxnn/training/mrl.py

@@ -481,7 +481,7 @@ class MRLTrainer:
                 critic_state = smart_concat_critic_states(*state, max_length=self.critic_max_len,
                                                           pad_token_id=self.pad_token_id)
                 values = self.critic(critic_state['input_ids'], attention_mask=critic_state['attention_mask']).squeeze()
-                critic_loss = self.rl_algorithm.critic_loss(values, ref_values)
+                critic_loss = self.rl_algorithm.critic_loss(values, ref_values.detach())
             # 2.2 Run backpropagation with scaler
             self.critic_scaler.scale(critic_loss).backward()
             # 2.3 Unscale and clip gradients
@@ -495,7 +495,7 @@ class MRLTrainer:
             critic_state = smart_concat_critic_states(*state, max_length=self.critic_max_len,
                                                       pad_token_id=self.pad_token_id)
             values = self.critic(critic_state['input_ids'], attention_mask=critic_state['attention_mask']).squeeze()
-            critic_loss = self.rl_algorithm.critic_loss(values, ref_values)
+            critic_loss = self.rl_algorithm.critic_loss(values, ref_values.detach())
             # 2.2 Run backpropagation
             critic_loss.backward()
             # 2.3 Clip gradients

rxnn/training/reward.py

@@ -42,6 +42,7 @@ class MrlRewardModel:
             running_mean_decay: float = 0.2,
             bleu_saved_weights: tuple = (0.5, 0.5),
             bleu_ref_weights: tuple = (0.5, 0.5),
+            tanh_reward_scale: bool = False,
             rewards_scale: float = 1.0,
     ):
         self.shared_embedding = shared_embedding.to(device)
@@ -71,6 +72,7 @@ class MrlRewardModel:
         self.running_mean_decay = running_mean_decay
         self.bleu_ref_weights = bleu_ref_weights
         self.bleu_saved_weights = bleu_saved_weights
+        self.tanh_reward_scale = tanh_reward_scale
         self.rewards_scale = rewards_scale
 
         self.prev_data_running_mean = None
@@ -175,6 +177,12 @@ class MrlRewardModel:
         self.prev_data_running_mean = (1 - self.running_mean_decay) * self._sequence_embedding(
             prev_data) + self.running_mean_decay * self.prev_data_running_mean
 
+    def _pre_scale_rewards(self, rewards: torch.Tensor) -> torch.Tensor:
+        if self.tanh_reward_scale:
+            return (rewards * 2) - 1  # Convert [0,1] to [-1,1]
+        else:
+            return rewards
+
     def __call__(
             self,
             generated: TokenizedDict,
@@ -204,5 +212,5 @@ class MrlRewardModel:
             cosine = self.negative_cosine(generated['input_ids'], reference['input_ids'], saved_data['input_ids'])
             sim_rewards = self.neg_bleu_factor * torch.tensor(bleu, device=self.device) + self.neg_cos_factor * cosine
 
-        rewards = (sim_rewards + self.len_factor * self.len_reward(generated) if self.reward_len else sim_rewards) * self.rewards_scale
+        rewards = self._pre_scale_rewards(sim_rewards + self.len_factor * self.len_reward(generated) if self.reward_len else sim_rewards) * self.rewards_scale
         return rewards.tolist()

rxnn/training/rl.py

@@ -21,8 +21,8 @@ class RlAlgorithm(ABC):
     def calculate_advantages(self, rewards: torch.Tensor, values: torch.Tensor, dones: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         pass
 
-    def critic_loss(self, rewards: torch.Tensor, values: torch.Tensor) -> torch.Tensor:
-        return self.critic_loss(rewards, values)
+    def critic_loss(self, values: torch.Tensor, ref_values: torch.Tensor) -> torch.Tensor:
+        return self.critic_loss(values, ref_values)
 
 
 class PPOConfig(TypedDict):
@@ -31,6 +31,8 @@ class PPOConfig(TypedDict):
     gae_gamma: Optional[float]
     entropy_coef: Optional[float]
     use_distributed_advantage_norm: Optional[bool]
+    clip_critic_values: Optional[bool]
+    critic_value_clip: Optional[float]
 
 
 class PPOAlgorithm(RlAlgorithm):
@@ -46,6 +48,14 @@ class PPOAlgorithm(RlAlgorithm):
         self.gae_gamma = config.get('gae_gamma', 0.99)
         self.entropy_coef = config.get('entropy_coef', 0.01)
         self.use_distributed_advantage_norm = config.get('use_distributed_advantage_norm', False)
+        self.clip_critic_values = config.get('clip_critic_values', True)
+        self.critic_value_clip = config.get('critic_value_clip', 10.0)
+
+    def critic_loss(self, values: torch.Tensor, ref_values: torch.Tensor) -> torch.Tensor:
+        # Critic loss with clipped values
+        if self.clip_critic_values:
+            values = torch.clamp(values, -self.critic_value_clip, self.critic_value_clip)
+        return self.critic_loss(values, ref_values)
 
     def policy_loss(self, query: TokenizedDict, answer: TokenizedDict, logits: torch.Tensor,
                     old_log_probs: torch.Tensor, advantages: torch.Tensor) -> torch.Tensor:
@@ -96,23 +106,24 @@ class PPOAlgorithm(RlAlgorithm):
 
         return policy_loss
 
-    def _compute_gae(self, rewards: torch.Tensor, values: torch.Tensor, next_value: torch.Tensor, dones: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        T, B = rewards.shape
-        advantages = torch.zeros_like(rewards, device=values.device)
+    def _compute_gae(self, rewards: torch.Tensor, values: torch.Tensor,
+                     last_value: torch.Tensor, dones: torch.Tensor):
+        trajectory_len, batch_size = rewards.shape
+        advantages = torch.zeros_like(rewards, device=rewards.device)
         last_advantage = 0
-        last_value = next_value.detach()
-
-        for t in reversed(range(T)):
-            if t == T - 1:
-                next_values = last_value
-            else:
-                next_values = values[t + 1]
-
-            # Mask next values if episode ended
-            next_values = next_values * ~dones[t]
-            delta = rewards[t] + self.gae_gamma * next_values - values[t]
-            advantages[t] = delta + self.gae_gamma * self.gae_lambda * last_advantage
+        next_value = last_value
+        next_done = torch.zeros(batch_size, device=dones.device)  # Last state is terminal
+        dones = dones.float()
+        for t in reversed(range(trajectory_len)):
+            # Check if next state is terminal
+            non_terminal = 1.0 - next_done
+
+            # Delta should not include next_value if next is terminal
+            delta = rewards[t] + self.gae_gamma * next_value * non_terminal - values[t]
+            advantages[t] = delta + self.gae_gamma * self.gae_lambda * non_terminal * last_advantage
             last_advantage = advantages[t]
+            next_value = values[t]
+            next_done = dones[t]
 
         returns = advantages + values
         return advantages, returns

rxnn/transformers/ff.py

@@ -66,6 +66,8 @@ def get_activation_layer(activation: str):
         return nn.SiLU()
     elif activation == 'sigmoid':
         return nn.Sigmoid()
+    elif activation == 'tanh':
+        return nn.Tanh()
     elif activation == 'linear':
         return LinearActivation()
     else:

rxnn/transformers/models.py

@@ -126,6 +126,39 @@ class ReactiveTransformerEncoder(ReactiveTransformerBase):
         return x, torch.stack(hidden_states)
 
 
+class ReactiveTransformerEncoderDetachStm(ReactiveTransformerBase):
+    """
+    Reactive Transformer encoder DetachStm version - reactive transformer encoder that's detaching Short-Term Memory tensors,
+    before processing them in layers (memory cross-attention). Made for Memory-Aware Critic models, to not include memory
+    update gradients in Critic optimization.
+    """
+
+    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> tuple[torch.Tensor, torch.Tensor]:
+        x = super().forward(x)  # apply embeddings
+        if attention_mask is not None:
+            attention_mask = attention_mask.unsqueeze(1).unsqueeze(1).bool()
+
+        hidden_states = []
+        # Process shared layers
+        if self.shared_layers is not None:
+            for i in range(self.num_shared_layers):
+                layer_stm = self.stm(i).detach() # <- Detach STM layer
+                # expand layer STM to batch size, if it's not in batch mode
+                if layer_stm.size(0) == 1:
+                    layer_stm = layer_stm.expand(x.size(0), -1, -1)
+                x = self.shared_layers[i](x, layer_stm, mask=attention_mask)
+                hidden_states.append(x)
+        # Process own layers
+        for i in range(self.num_own_layers):
+            layer_stm = self.stm(i).detach() # <- Detach STM layer
+            # expand layer STM to batch size, if it's not in batch mode
+            if layer_stm.size(0) == 1:
+                layer_stm = layer_stm.expand(x.size(0), -1, -1)
+            x = self.layers[i](x, layer_stm, mask=attention_mask)
+            hidden_states.append(x)
+        return x, torch.stack(hidden_states)
+
+
 class ClassicTransformerBase(nn.Module):
     """Base class for Classic Transformer models - common logic for both decoders and encoders."""
 

{rxnn-0.2.31.dist-info → rxnn-0.2.33.dist-info}/METADATA

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

{rxnn-0.2.31.dist-info → rxnn-0.2.33.dist-info}/RECORD

@@ -7,33 +7,33 @@ rxnn/experimental/moe.py,sha256=jHZ1QhpWiVQOswVpFmuH7b2IUOPf0Uuf-I2Ddwsd7Us,6140
 rxnn/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/memory/attention.py,sha256=PXVBZQYNsRraZh7QDBgUOdPy3lTI8B0d8CzduojBjG0,1747
 rxnn/memory/norm.py,sha256=E98jOQEuIOFFhlkvS8s4fFN-D4tLO6vaOqnObv1oVmA,6592
-rxnn/memory/stm.py,sha256=IH_3INw7FdI013t56ui3Zq9GPUq-k3HeZGjx6BerS4g,3888
+rxnn/memory/stm.py,sha256=SSfc-RL9FE-RLkmOEkLB-9Rb00ZXbMLbsAEPdpIW89o,3851
 rxnn/rxt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/rxt/models.py,sha256=r8wZeeNTC2VAhiiNe4y7LrbnB4wjFu_cupKiGkpdgjI,13002
+rxnn/rxt/models.py,sha256=CzFELVv5-ybAwl1s1ptpmwM7wdJ07M4jaT1-I8PYrR0,13999
 rxnn/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/training/base.py,sha256=TGz_37RfI1qLI31GNRV5rLowW1kAHnJwqPm7DNfLfe4,11730
 rxnn/training/bml.py,sha256=hw6gLpLkGvqLzxIvBg4MvCc5r8cHpEm2RDyh7nH6CtE,16914
 rxnn/training/callbacks.py,sha256=p72lbzFAmFjpcUvyy4aUB3qd53I8C6Sk5w9nQvsKgTk,35852
 rxnn/training/dataset.py,sha256=7hTilFWPpqUEc6zNcMqBPjxFKxCfvTKKF3E8tVlwccQ,51250
 rxnn/training/ddp.py,sha256=VsNBjn3cY-uUj8hbsW7oKvb0_ZKnXnJ2KgObm-Mr9i4,836
-rxnn/training/models.py,sha256=2KhNT7yx0AgUke4nmsFqzQKx_YYp78QvsLWYZjWeUgQ,6812
-rxnn/training/mrl.py,sha256=Aimiiqf_4p6dp5Ty9pY9VwetySBS_OFpCQlcVHVkO4Q,55124
-rxnn/training/reward.py,sha256=7MTVdNm5HnWmt6zFDi3TAYmnVSL_-24riOoY2F7z4x8,11290
-rxnn/training/rl.py,sha256=eL3C0yryiNBgl_xb-D-5dyYUtK4V4-K4t3a60x5ir28,5142
+rxnn/training/models.py,sha256=8FV5eZx1HxtqRSgikwfKoB_bNhPuMYyNi0uSXB65-M4,7223
+rxnn/training/mrl.py,sha256=1pYzjXI17FDZGPTVpmbaBvMYpB-a6SLv-84RHXA4JEA,55142
+rxnn/training/reward.py,sha256=B7nerPk9eNAv2i7umtNF88tVQVwijNNrchIrEITGHKk,11623
+rxnn/training/rl.py,sha256=ckx1nlzIGZBabzwZNRj4isvHqRZwg0y0jGOT-SN6KZc,5841
 rxnn/training/scheduler.py,sha256=LcjU35mEwz2U5x3U6tLfeeYlBqMxbFSxYzJYuXkWbSY,1408
 rxnn/training/tokenizer.py,sha256=umaLByMBx_NMrQElA45HLm9gkuzyKWDTFaKVd-CjXl0,8344
 rxnn/training/utils.py,sha256=Bw8nZLKIt7NQpUVCYkb_79kWKChVFOYgYXwODo4SvNc,5718
 rxnn/transformers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/transformers/attention.py,sha256=d0Igo1Nrn76BphbHrzekiKJfT3RCy4iSXSB6FLAOTgc,16247
-rxnn/transformers/ff.py,sha256=jJnuBDsnnX5uYC_WZH8cXAYrMnz0P-iX7MwcPivjRtI,2533
+rxnn/transformers/ff.py,sha256=WDjO-H9XWInoWnUnxiseIH6Kx5GlHP0zGJygwhcb1gc,2589
 rxnn/transformers/layers.py,sha256=l0bXmhN7KOkCw0KTVLixWSo9Op4SesGabWJ4R4EQBMY,7988
 rxnn/transformers/mask.py,sha256=J0cfLVLt3SzS2ra3KcY4khrkhI975Dw4CjpUi3Sn25s,419
-rxnn/transformers/models.py,sha256=QwVxYN9DrKllEpOiFoAx4CiThOWafeTa-OAY7L6gN0Y,8929
+rxnn/transformers/models.py,sha256=hey6tFN9gmLfWCZLjtl_9OcvIjGpWLI1IDeVnr5y8YM,10583
 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.31.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.31.dist-info/METADATA,sha256=zxD2qPHL_QrFH1bYZrMv4odbXE4B_YIVEpGDzV2MYEI,25960
-rxnn-0.2.31.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.31.dist-info/RECORD,,
+rxnn-0.2.33.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.33.dist-info/METADATA,sha256=im17irb58IYMXOzMXE6QaSPF31Akx0iYS4ay-aRqA9Q,25960
+rxnn-0.2.33.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.33.dist-info/RECORD,,