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

rxnn/memory/norm.py

@@ -163,7 +163,7 @@ def init_memory_norm(
     init_scale: float = 1.0,
     per_dim_scale: bool = False,
 ) -> nn.Module:
-    assert norm_type in ['layer', 'rms', 'adaptive', 'positional']
+    assert norm_type in ['layer', 'rms', 'adaptive', 'positional', 'classic-rms']
     if norm_type == 'layer':
         return MemoryLayerNorm(dim, use_gate, init_scale, init_gate)
     elif norm_type == 'rms':
@@ -172,4 +172,6 @@ def init_memory_norm(
         return AdaptiveRMSMemoryNorm(dim, use_gate, decay, init_scale, init_gate)
     elif norm_type == 'positional':
         return AdaptivePositionalMemoryNorm(num_slots, dim, decay, use_scale, use_gate, init_gate, per_dim_scale)
+    elif norm_type == 'classic-rms':
+        return nn.RMSNorm(dim)
     return MemoryLayerNorm(dim, use_gate, init_scale, init_gate)

rxnn/training/mrl.py

@@ -35,6 +35,7 @@ class MrlConfig(TypedDict):
     moe_aux_loss_scale: Optional[float]
     freeze_embeddings: Optional[bool]
     embedding_lr: Optional[float]
+    use_memory_warmup: Optional[bool]
 
 
 class MrlStrategy(Enum):
@@ -70,6 +71,7 @@ class CurriculumConfig(TypedDict):
     update_epochs: Optional[int]
     freeze_embeddings: Optional[bool]
     embedding_lr: Optional[float]
+    teacher_forcing: Optional[bool]
 
 
 class SamplerConfig(TypedDict):
@@ -136,6 +138,7 @@ class MRLTrainer:
         self.moe_aux_loss_scale = config.get('moe_aux_loss_scale', 0.01)
         self.shared_freeze_embeddings = config.get('freeze_embeddings', False)
         self.freeze_embeddings = self.shared_freeze_embeddings
+        self.use_memory_warmup = config.get('use_memory_warmup', False)
         # Internal update epochs config
         self.shared_update_epochs = config.get('update_epochs', 10)
         self.update_epochs = self.shared_update_epochs
@@ -213,6 +216,7 @@ class MRLTrainer:
         self.callbacks = []
         self.global_epoch = 0
         self.global_epochs_count = 0
+        self.teacher_forcing = False
 
     def _init_optimizers(
             self,
@@ -381,6 +385,11 @@ class MRLTrainer:
             self.writer.add_scalar(f'Collect/episode reward (steps: {self.curriculum_steps})', avg_reward,
                                    self.stage_step['collect'])
 
+    def memory_warmup(self, query: TokenizedDict, answer: TokenizedDict):
+        if self.use_memory_warmup:
+            with torch.no_grad():
+                self.encode_and_update_stm(query, answer)
+
     def collect_trajectories(self, dataloader: DataLoader, epoch: int, batch_size: int) -> list[MrlTrajectoryEpisode]:
         """Collect trajectories for PPO for current curriculum step."""
         # 1. Init trajectories list
@@ -402,8 +411,13 @@ class MRLTrainer:
                     first_query, first_answer, interactions = batch['query'], batch['answer'], batch['interactions']
                     interactions = interactions[:self.curriculum_steps]
                     interactions_len = len(interactions)
+
+                    first_interaction = self._move_multiple_batches(first_query, first_answer)
+
+                    if reset_done:
+                        self.memory_warmup(*first_interaction)
                     # 6. Encode and update STM with data to save from first interaction
-                    self.encode_and_update_stm(*self._move_multiple_batches(first_query, first_answer))
+                    self.encode_and_update_stm(*first_interaction)
 
                     # 7. Save first interaction as data to save (for trajectory state)
                     query, answer = first_query, first_answer
@@ -440,8 +454,10 @@ class MRLTrainer:
 
                         # 11. Update STM with generated response (except last interaction, it's not needed)
                         if not is_last_interaction:
-                            self.encode_and_update_stm(next_query,
-                                                       generated_answer)  # update with generated_answer on GPU
+                            self.encode_and_update_stm(
+                                next_query,
+                                self._move_batch(interaction['answer']) if self.teacher_forcing else generated_answer
+                            )  # update with generated_answer on GPU
 
                         # 12. Store trajectory step
                         trajectory: MrlTrajectoryStep = {
@@ -458,7 +474,7 @@ class MRLTrainer:
                         # 13. Set previous and current interaction query and generated answer (batches), as saved data for next interaction
                         if not (self.strategy == MrlStrategy.LONG_RANGE_STRATEGY and i == 0):
                             prev_interaction = (query, answer)
-                        query, answer = interaction['query'], detached_answer
+                        query, answer = interaction['query'], (interaction['answer'] if self.teacher_forcing else detached_answer)
 
                     # 14. Append full batched episode (number of steps depends on curriculum stage) to trajectories
                     episode_trajectory: MrlTrajectoryEpisode = {
@@ -649,6 +665,9 @@ class MRLTrainer:
 
                 self.actor.clone_reset_memory()
 
+                if should_reset_stm and step_idx == 0:
+                    self.memory_warmup(query, answer)
+
                 # 7. In memory aware critic version, encode and update STM before critic update, to include its gradients in critic loss too
                 if self.memory_aware_critic:
                     self.encode_and_update_stm(query, answer)
@@ -798,13 +817,16 @@ class MRLTrainer:
                 if batch['query']['input_ids'].size(0) == batch_size:
                     self._increment_steps('eval')
                     # 3. Reset STM with random resets ratio and reward model running mean
-                    self.reset_stm()
+                    reset_stm = self.reset_stm()
                     self.reward.reset_running_mean()
 
                     # 4. Get batches for first queries, answers and all follow-up interactions
                     first_query, first_answer, interactions = batch['query'], batch['answer'], batch['interactions']
                     # 5. Encode and update STM with initial interactions (batch)
-                    self.encode_and_update_stm(*self._move_multiple_batches(first_query, first_answer))
+                    first_interaction = self._move_multiple_batches(first_query, first_answer)
+                    if reset_stm:
+                        self.memory_warmup(*first_interaction)
+                    self.encode_and_update_stm(*first_interaction)
 
                     # 6. Save follow-up interactions len and first query and answer as previous one for iteration
                     interactions_len = len(interactions)
@@ -839,7 +861,10 @@ class MRLTrainer:
 
                         # 10. Encode and update memory for the next interaction
                         if not is_last_interaction:
-                            self.encode_and_update_stm(next_query, generated_answer)
+                            self.encode_and_update_stm(
+                                next_query,
+                                self._move_batch(interaction['answer']) if self.teacher_forcing else generated_answer
+                            )
 
                         # 11. Accumulate rewards
                         step_reward = torch.tensor(reward).mean().to(self.device)
@@ -852,7 +877,7 @@ class MRLTrainer:
                         # 12. Save previous interaction
                         if not (self.strategy == MrlStrategy.LONG_RANGE_STRATEGY and i == 0):
                             prev_interaction = (query, answer)
-                        query, answer = interaction['query'], detached_answer
+                        query, answer = interaction['query'], (interaction['answer'] if self.teacher_forcing else detached_answer)
                     avg_episode_reward = (episode_reward / episode_interactions).item()
                     # 13. Run eval TensorBoard writer with average episode reward
                     self._eval_writer(avg_episode_reward, epoch)
@@ -982,8 +1007,7 @@ class MRLTrainer:
         self.reward = config.get('reward_model', self.shared_reward_model)  # MRL Reward Model for curriculum stage
         self.update_epochs = config.get('update_epochs', self.shared_update_epochs)  # Internal update epochs
         self.freeze_embeddings = config.get('freeze_embeddings', self.shared_freeze_embeddings)
-
-
+        self.teacher_forcing = config.get('teacher_forcing', False)
 
         def has_param(field: OptimField) -> bool:
             return field in config and config[field] is not None

rxnn/training/reward.py

@@ -1,6 +1,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
 from enum import Enum
 from typing import Optional
 from .utils import TokenizedDict
@@ -37,6 +38,7 @@ class MrlRewardModel:
             reward_len: bool = False,
             neg_reward_len: bool = False,
             max_rewarded_len: int = None,
+            target_len_as_ref: bool = False,
             len_factor: int = None,
             use_running_mean: bool = True,
             running_mean_decay: float = 0.2,
@@ -44,6 +46,7 @@ class MrlRewardModel:
             bleu_ref_weights: tuple = (0.5, 0.5),
             tanh_reward_scale: bool = False,
             rewards_scale: float = 1.0,
+            debug_mode: int = 0,
     ):
         self.shared_embedding = shared_embedding.to(device)
         self.device = device
@@ -67,6 +70,7 @@ class MrlRewardModel:
         self.reward_len = reward_len
         self.neg_reward_len = neg_reward_len
         self.max_rewarded_len = max_rewarded_len
+        self.target_len_as_ref = target_len_as_ref
         self.len_factor = len_factor
         self.use_running_mean = use_running_mean
         self.running_mean_decay = running_mean_decay
@@ -74,6 +78,8 @@ class MrlRewardModel:
         self.bleu_saved_weights = bleu_saved_weights
         self.tanh_reward_scale = tanh_reward_scale
         self.rewards_scale = rewards_scale
+        self.bleu_smoothing = SmoothingFunction().method4
+        self.debug_mode = debug_mode
 
         self.prev_data_running_mean = None
 
@@ -95,59 +101,133 @@ class MrlRewardModel:
                 assert self.neg_cos_ref_factor + self.neg_cos_saved_factor == 1.0
                 assert self.neg_bleu_ref_factor + self.neg_bleu_saved_factor == 1.0
 
-    def _sentence_bleu(self, generated: torch.Tensor, reference: torch.Tensor, saved_data: torch.Tensor) -> float:
-        from nltk.translate.bleu_score import sentence_bleu
+    def _sentence_bleu(self, input_ids: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
+                       masks: tuple[torch.Tensor, torch.Tensor, torch.Tensor]) -> float:
+        generated, reference, saved_data = input_ids
+        generated_mask, reference_mask, saved_data_mask = masks
+
+        generated = generated.tolist()[:generated_mask.sum().item()]
+        reference = reference.tolist()[:reference_mask.sum().item()]
+        saved_data = saved_data.tolist()[:saved_data_mask.sum().item()]
+
+        if self.debug_mode == 2:
+            print('LENS: ', (len(generated), len(reference), len(saved_data)))
 
         if self.bleu_with_saved_data:
-            ref_bleu = sentence_bleu([reference], generated, weights=self.bleu_ref_weights)
-            saved_bleu = sentence_bleu([saved_data], generated, weights=self.bleu_saved_weights)
+            ref_bleu = sentence_bleu([reference], generated, weights=self.bleu_ref_weights,
+                                     smoothing_function=self.bleu_smoothing)
+            saved_bleu = sentence_bleu([saved_data], generated, weights=self.bleu_saved_weights,
+                                       smoothing_function=self.bleu_smoothing)
+            if self.debug_mode == 2:
+                print('REF BLEU: ', ref_bleu)
+                print('SAVED BLEU: ', saved_bleu)
+
             return self.bleu_ref_factor * ref_bleu + self.bleu_saved_factor * saved_bleu
         else:
             return sentence_bleu([reference], generated, weights=self.bleu_ref_weights)
 
+    def _negative_sentence_bleu(self, input_ids: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
+                                masks: tuple[torch.Tensor, torch.Tensor, torch.Tensor]) -> float:
+        generated, reference, saved_data = input_ids
+        generated_mask, reference_mask, saved_data_mask = masks
+
+        generated = generated.tolist()[:generated_mask.sum().item()]
+        reference = reference.tolist()[:reference_mask.sum().item()]
+        saved_data = saved_data.tolist()[:saved_data_mask.sum().item()]
 
-    def _negative_sentence_bleu(self, generated: torch.Tensor, reference: torch.Tensor,
-                                saved_data: torch.Tensor) -> float:
-        from nltk.translate.bleu_score import sentence_bleu
+        if self.debug_mode == 2:
+            print('LENS: ', (len(generated), len(reference), len(saved_data)))
 
         if self.bleu_with_saved_data:
-            ref_bleu = sentence_bleu([reference], generated, weights=self.bleu_ref_weights)
-            saved_bleu = sentence_bleu([saved_data], generated, weights=self.bleu_saved_weights)
+            ref_bleu = sentence_bleu([reference], generated, weights=self.bleu_ref_weights,
+                                     smoothing_function=self.bleu_smoothing)
+            saved_bleu = sentence_bleu([saved_data], generated, weights=self.bleu_saved_weights,
+                                       smoothing_function=self.bleu_smoothing)
             saved_bleu = 1 - saved_bleu
 
+            if self.debug_mode == 2:
+                print('REF BLEU: ', ref_bleu)
+                print('SAVED BLEU: ', saved_bleu)
+
             return self.neg_bleu_ref_factor * ref_bleu + self.neg_bleu_saved_factor * saved_bleu
         else:
             return sentence_bleu([reference], generated, weights=self.bleu_ref_weights)
 
-    def batch_bleu(self, generated: torch.Tensor, reference: torch.Tensor, saved_data: torch.Tensor) -> list[float]:
-        batch_size = generated.size(0)
-        return [self._sentence_bleu(generated[i], reference[i], saved_data[i]) for i in range(batch_size)]
+    def batch_bleu(self, generated: TokenizedDict, reference: TokenizedDict, saved_data: TokenizedDict) -> list[float]:
+        batch_size = generated['input_ids'].size(0)
+
+        return [
+            self._sentence_bleu(
+                input_ids=(generated['input_ids'][i], reference['input_ids'][i], saved_data['input_ids'][i]),
+                masks=(generated['attention_mask'][i], reference['attention_mask'][i], saved_data['attention_mask'][i])
+            ) for i in range(batch_size)
+        ]
+
+    def negative_bleu(self, generated: TokenizedDict, reference: TokenizedDict, saved_data: TokenizedDict) -> list[
+        float]:
+        batch_size = generated['input_ids'].size(0)
+
+        return [
+            self._negative_sentence_bleu(
+                input_ids=(generated['input_ids'][i], reference['input_ids'][i], saved_data['input_ids'][i]),
+                masks=(generated['attention_mask'][i], reference['attention_mask'][i], saved_data['attention_mask'][i])
+            ) for i in range(batch_size)
+        ]
+
+    def _sequence_embedding(self, sequence: TokenizedDict) -> torch.Tensor:
+        input_ids = sequence['input_ids']
+        attention_mask = sequence['attention_mask']
+
+        # Get embeddings
+        embeddings = self.shared_embedding(input_ids.to(self.device))
+
+        # Apply attention mask
+        mask_expanded = attention_mask.unsqueeze(-1).to(self.device)
+        masked_embeddings = embeddings * mask_expanded
+
+        # Compute mean with masking
+        sum_embeddings = torch.sum(masked_embeddings, dim=1)
+        token_counts = torch.sum(mask_expanded, dim=1)
+        token_counts = torch.clamp(token_counts, min=1e-8)  # Avoid division by zero
+
+        return sum_embeddings / token_counts
 
-    def _sequence_embedding(self, sequence: torch.Tensor) -> torch.Tensor:
-        embedding = self.shared_embedding(sequence.to(self.device))
-        return embedding.mean(dim=1)
+    def _cosine_sim(self, generated: TokenizedDict, reference: TokenizedDict, saved_data: TokenizedDict):
+        generated_emb = F.normalize(self._sequence_embedding(generated), dim=-1)
+        saved_data_emb = F.normalize(self._sequence_embedding(saved_data), dim=-1)
+        reference_emb = F.normalize(self._sequence_embedding(reference), dim=-1)
 
-    def _cosine_sim(self, generated: torch.Tensor, reference: torch.Tensor, saved_data: torch.Tensor):
-        generated_emb = self._sequence_embedding(generated)
+        gen_and_saved = F.cosine_similarity(generated_emb, saved_data_emb, dim=1)
+        gen_and_ref = F.cosine_similarity(generated_emb, reference_emb, dim=1)
 
-        gen_and_saved = (F.cosine_similarity(generated_emb, self._sequence_embedding(saved_data)) + 1) / 2
-        gen_and_ref = (F.cosine_similarity(generated_emb, self._sequence_embedding(reference)) + 1) / 2
-        return gen_and_saved, gen_and_ref
+        if self.debug_mode >= 1:
+            print('GEN AND SAVED: ', gen_and_saved.mean())
+            print('GEN AND REF: ', gen_and_ref.mean())
+        return torch.clamp(gen_and_saved, min=0), torch.clamp(gen_and_ref, min=0)
 
-    def _cosine_sim_running_mean(self, generated: torch.Tensor, reference: torch.Tensor, saved_data: torch.Tensor):
-        generated_emb = self._sequence_embedding(generated)
+    def _cosine_sim_running_mean(self, generated: TokenizedDict, reference: TokenizedDict, saved_data: TokenizedDict):
+        generated_emb = F.normalize(self._sequence_embedding(generated), dim=-1)
+        saved_data_emb = F.normalize(self._sequence_embedding(saved_data), dim=-1)
+        reference_emb = F.normalize(self._sequence_embedding(reference), dim=-1)
+        running_emb = F.normalize(self.prev_data_running_mean, dim=-1)
 
-        gen_and_saved = (F.cosine_similarity(generated_emb, self._sequence_embedding(saved_data)) + 1) / 2
-        gen_and_ref = (F.cosine_similarity(generated_emb, self._sequence_embedding(reference)) + 1) / 2
-        gen_and_mean = (F.cosine_similarity(generated_emb, self.prev_data_running_mean) + 1) / 2
-        return gen_and_saved, gen_and_ref, gen_and_mean
+        gen_and_saved = F.cosine_similarity(generated_emb, saved_data_emb, dim=1)
+        gen_and_ref = F.cosine_similarity(generated_emb, reference_emb, dim=1)
+        gen_and_mean = F.cosine_similarity(generated_emb, running_emb, dim=1)
 
-    def batch_cosine(self, generated: torch.Tensor, reference: torch.Tensor, saved_data: torch.Tensor,
+        if self.debug_mode >= 1:
+            print('GEN AND SAVED: ', gen_and_saved.mean())
+            print('GEN AND REF: ', gen_and_ref.mean())
+            print('GEN AND MEAN: ', gen_and_mean.mean())
+
+        return torch.clamp(gen_and_saved, min=0), torch.clamp(gen_and_ref, min=0), torch.clamp(gen_and_mean, min=0)
+
+    def batch_cosine(self, generated: TokenizedDict, reference: TokenizedDict, saved_data: TokenizedDict,
                      include_running_mean: bool = False, negative_running_mean: bool = False) -> torch.Tensor:
         if self.use_running_mean and negative_running_mean:
             gen_and_saved, gen_and_ref, gen_and_mean = self._cosine_sim_running_mean(generated, reference, saved_data)
             return self.multi_cos_saved_factor * gen_and_saved + self.multi_cos_ref_factor * gen_and_ref + self.multi_cos_saved_factor * (
-                        1 - gen_and_mean)
+                    1 - gen_and_mean)
         elif self.use_running_mean and include_running_mean:
             gen_and_saved, gen_and_ref, gen_and_mean = self._cosine_sim_running_mean(generated, reference, saved_data)
             return self.multi_cos_saved_factor * gen_and_saved + self.multi_cos_ref_factor * gen_and_ref + self.multi_cos_saved_factor * gen_and_mean
@@ -155,25 +235,26 @@ class MrlRewardModel:
             gen_and_saved, gen_and_ref = self._cosine_sim(generated, reference, saved_data)
             return self.cos_saved_factor * gen_and_saved + self.cos_ref_factor * gen_and_ref
 
-    def negative_cosine(self, generated: torch.Tensor, reference: torch.Tensor,
-                        saved_data: torch.Tensor) -> torch.Tensor:
+    def negative_cosine(self, generated: TokenizedDict, reference: TokenizedDict,
+                        saved_data: TokenizedDict) -> torch.Tensor:
         gen_and_saved, gen_and_ref = self._cosine_sim(generated, reference, saved_data)
 
         return self.neg_cos_saved_factor * (1 - gen_and_saved) + self.neg_cos_ref_factor * gen_and_ref
 
-    def len_reward(self, generated: TokenizedDict):
+    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)
-        neg_lens = self.max_rewarded_len / lens if self.neg_reward_len else 1.0
-        len_reward = torch.where(lens >= self.max_rewarded_len, neg_lens, lens / self.max_rewarded_len)
+        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
 
     def reset_running_mean(self):
         self.prev_data_running_mean = None
 
-    def init_running_mean(self, prev_data: torch.Tensor):
+    def init_running_mean(self, prev_data: TokenizedDict):
         self.prev_data_running_mean = self._sequence_embedding(prev_data)
 
-    def update_running_mean(self, prev_data: torch.Tensor):
+    def update_running_mean(self, prev_data: TokenizedDict):
         self.prev_data_running_mean = (1 - self.running_mean_decay) * self._sequence_embedding(
             prev_data) + self.running_mean_decay * self.prev_data_running_mean
 
@@ -193,24 +274,50 @@ class MrlRewardModel:
     ) -> list[float]:
         if prev_data is not None:
             if self.prev_data_running_mean is None:
-                self.init_running_mean(prev_data['input_ids'])
+                self.init_running_mean(prev_data)
             else:
-                self.update_running_mean(prev_data['input_ids'])
+                self.update_running_mean(prev_data)
 
         if mode == MrlRewardMode.STANDARD:
-            bleu = self.batch_bleu(generated['input_ids'], reference['input_ids'], saved_data['input_ids'])
-            cosine = self.batch_cosine(generated['input_ids'], reference['input_ids'], saved_data['input_ids'],
-                                       include_running_mean=prev_data is not None)
+            bleu = self.batch_bleu(generated, reference, saved_data)
+            cosine = self.batch_cosine(generated, reference, saved_data, include_running_mean=prev_data is not None)
+
+            if self.debug_mode >= 1:
+                print('STANDARD MODE')
+                print('BLEU: ', sum(bleu) / len(bleu))
+                print('COSINE: ', sum(cosine) / len(cosine))
+
             sim_rewards = self.bleu_factor * torch.tensor(bleu, device=self.device) + self.cos_factor * cosine
         elif mode == MrlRewardMode.LONG_RANGE:
-            bleu = self.batch_bleu(generated['input_ids'], reference['input_ids'], saved_data['input_ids'])
-            cosine = self.batch_cosine(generated['input_ids'], reference['input_ids'], saved_data['input_ids'],
+            bleu = self.batch_bleu(generated, reference, saved_data)
+            cosine = self.batch_cosine(generated, reference, saved_data,
                                        negative_running_mean=prev_data is not None)
+
+            if self.debug_mode >= 1:
+                print('LONG MODE')
+                print('BLEU: ', sum(bleu) / len(bleu))
+                print('COSINE: ', sum(cosine) / len(cosine))
+
             sim_rewards = self.bleu_factor * torch.tensor(bleu, device=self.device) + self.cos_factor * cosine
         else:
-            bleu = self.batch_bleu(generated['input_ids'], reference['input_ids'], saved_data['input_ids'])
-            cosine = self.negative_cosine(generated['input_ids'], reference['input_ids'], saved_data['input_ids'])
+            bleu = self.negative_bleu(generated, reference, saved_data)
+            cosine = self.negative_cosine(generated, reference, saved_data)
+
+            if self.debug_mode >= 1:
+                print('NEGATIVE MODE')
+                print('BLEU: ', sum(bleu) / len(bleu))
+                print('COSINE: ', sum(cosine) / len(cosine))
+
             sim_rewards = self.neg_bleu_factor * torch.tensor(bleu, device=self.device) + self.neg_cos_factor * cosine
 
-        rewards = self._pre_scale_rewards(sim_rewards + self.len_factor * self.len_reward(generated) if self.reward_len else sim_rewards) * self.rewards_scale
+        if self.reward_len:
+            len_reward = self.len_reward(generated, reference)
+
+            if self.debug_mode >= 1:
+                print('REWARD LEN: ', (len_reward.sum() / len_reward.size(0)).item())
+
+            rewards = self._pre_scale_rewards(sim_rewards + self.len_factor * len_reward) * self.rewards_scale
+        else:
+            rewards = self._pre_scale_rewards(sim_rewards) * self.rewards_scale
+
         return rewards.tolist()

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: rxnn
-Version: 0.2.47
+Version: 0.2.49
 Summary: RxNN: Reactive Neural Networks Platform
 License: Apache-2.0
 Keywords: deep-learning,ai,machine-learning
@@ -15,6 +15,7 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Requires-Dist: datasets (>=3.5.0,<4.0.0)
 Requires-Dist: huggingface-hub (>=0.30.0,<0.31.0)
+Requires-Dist: nltk (>=3.9.1,<4.0.0)
 Requires-Dist: tensorboard (>=2.19.0,<3.0.0)
 Requires-Dist: tokenizers (>=0.21.0,<0.22.0)
 Requires-Dist: torch (>=2.6.0,<3.0.0)

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

@@ -6,7 +6,7 @@ rxnn/experimental/models.py,sha256=foBo0n0ufvBnfIdJomiEg3CuSOiWSt-q5ako7vzYxx4,4
 rxnn/experimental/moe.py,sha256=jHZ1QhpWiVQOswVpFmuH7b2IUOPf0Uuf-I2Ddwsd7Us,6140
 rxnn/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/memory/attention.py,sha256=kan6UNPTjLfO7zKNp92hGooldgWPi3li_2-_L5xiErs,2784
-rxnn/memory/norm.py,sha256=E98jOQEuIOFFhlkvS8s4fFN-D4tLO6vaOqnObv1oVmA,6592
+rxnn/memory/norm.py,sha256=cVjjhCLqR5K6-321SP_ObG17y-ddlcTJeCTXvW4vpk0,6675
 rxnn/memory/stm.py,sha256=jv57gsH9XW19sLbxpRDqsp1yfsii_4Ef4Ncr_ztk-i4,3937
 rxnn/rxt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/rxt/models.py,sha256=new_YXLe9vfIBPX-pmFRoV523d7yCjEgfTY06EaH3Ms,14605
@@ -17,8 +17,8 @@ rxnn/training/callbacks.py,sha256=rS8leuVFPVVfE5Zc8DMkUZhRIPN-vpPbUjowXE5TSBw,36
 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/mrl.py,sha256=VXwRJ4wQtE0OoRsrsjYlWa2toTvHjoBJ_kril3EiK_A,59811
-rxnn/training/reward.py,sha256=B7nerPk9eNAv2i7umtNF88tVQVwijNNrchIrEITGHKk,11623
+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/scheduler.py,sha256=LcjU35mEwz2U5x3U6tLfeeYlBqMxbFSxYzJYuXkWbSY,1408
 rxnn/training/tokenizer.py,sha256=umaLByMBx_NMrQElA45HLm9gkuzyKWDTFaKVd-CjXl0,8344
@@ -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.47.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.47.dist-info/METADATA,sha256=OqRYFY68bnqQXdXfBNboYLAmXRmojMmR1YFUVQa4Jgo,25960
-rxnn-0.2.47.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.47.dist-info/RECORD,,
+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,,