rxnn 0.2.31__tar.gz → 0.2.33__tar.gz

{rxnn-0.2.31 → rxnn-0.2.33}/PKG-INFO

@@ -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 → rxnn-0.2.33}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "rxnn"
-version = "0.2.31"
+version = "0.2.33"
 description = "RxNN: Reactive Neural Networks Platform"
 
 license = "Apache-2.0"

{rxnn-0.2.31 → rxnn-0.2.33}/src/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-0.2.31 → rxnn-0.2.33}/src/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-0.2.31 → rxnn-0.2.33}/src/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-0.2.31 → rxnn-0.2.33}/src/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-0.2.31 → rxnn-0.2.33}/src/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-0.2.31 → rxnn-0.2.33}/src/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-0.2.31 → rxnn-0.2.33}/src/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-0.2.31 → rxnn-0.2.33}/src/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."""