rxnn 0.2.18__py3-none-any.whl → 0.2.20__py3-none-any.whl

rxnn/training/models.py

@@ -1,6 +1,7 @@
 import torch
 import torch.nn as nn
 from enum import Enum
+from typing import Literal
 from huggingface_hub import PyTorchModelHubMixin
 from ..transformers.models import ReactiveTransformerEncoder, ReactiveTransformerDecoder
 
@@ -74,23 +75,27 @@ class MrlActorModel(nn.Module):
         self.decoder = decoder
         self.memory_attention = memory_attention
 
-    def freeze_components(self):
+    def freeze_components(self, stage: Literal['update', 'fetch', 'both'] = 'both'):
         """Freeze encoder/decoder except memory-related layers."""
         if self.encoder.freeze_without_memory is not None:
             self.encoder.freeze_without_memory()
+            if stage == 'update':
+                self.encoder.freeze_memory()
         else:
             for param in self.encoder.parameters():
                 param.requires_grad = False
-            self.encoder.model.trainable_cross_attention_(True)
+            self.encoder.model.trainable_cross_attention_(True if stage != 'update' else False)
         if self.decoder.freeze_without_memory is not None:
             self.decoder.freeze_without_memory()
+            if stage == 'update':
+                self.decoder.freeze_memory()
         else:
             for param in self.decoder.parameters():
                 param.requires_grad = False
-            self.decoder.model.trainable_cross_attention_(True)
+            self.decoder.model.trainable_cross_attention_(True if stage != 'update' else False)
         # Unfreeze memory attention
         for param in self.memory_attention.parameters():
-            param.requires_grad = True
+            param.requires_grad = True if stage != 'fetch' else False
 
     def unfreeze_components(self):
         """Unfreeze all components after initial training."""

rxnn/training/mrl.py

@@ -3,7 +3,7 @@ from torch.utils.data import DataLoader, DistributedSampler
 from torch.utils.tensorboard import SummaryWriter
 import torch.distributed as dist
 from torch.nn.parallel import DistributedDataParallel
-from typing import Optional, TypedDict
+from typing import Optional, TypedDict, Union
 from enum import Enum
 import random, os
 from ..transformers.sampler import BatchSampler
@@ -37,10 +37,15 @@ class CurriculumConfig(TypedDict):
     eval_dataset: Optional[MrlCurriculumDataset]
     callbacks: Optional[list[MrlTrainerCallback]]
     strategy: MrlStrategy
-    unfreeze_epoch: Optional[int]
+    unfreeze_epoch: Optional[Union[int, tuple[int, int, int]]]
     random_resets: Optional[bool]
     random_resets_from: Optional[int]
     random_resets_ratio: Optional[float]
+    reward_model: Optional[MrlRewardModel]
+    lr: Optional[float]
+    critic_lr: Optional[float]
+    weight_decay: Optional[float]
+    critic_weight_decay: Optional[float]
 
 
 class SamplerConfig(TypedDict):
@@ -90,6 +95,7 @@ class MRLTrainer:
         """
         self.actor = actor
         self.critic = critic
+        self.shared_reward_model = reward
         self.reward = reward
         self.device = device
         self.max_seq_len = config.get('max_seq_len', 256)
@@ -117,17 +123,15 @@ class MRLTrainer:
         self.use_amp = use_amp
         self.dtype = dtype
 
+        self.base_optim_config = {
+            'lr': config.get('lr', 3e-4),
+            'critic_lr': config.get('critic_lr', 1e-4),
+            'weight_decay': config.get('weight_decay', 0.01),
+            'critic_weight_decay': config.get('critic_weight_decay', 0.01),
+        }
+
         # Optimizers
-        self.optimizer = torch.optim.AdamW(
-            self.actor.unique_parameters(),
-            lr=config.get("lr", 3e-4),
-            weight_decay=config.get("weight_decay", 0.01),
-        )
-        self.critic_optimizer = torch.optim.AdamW(
-            self.critic.parameters(),
-            lr=config.get("critic_lr", 1e-4),
-            weight_decay=config.get("critic_weight_decay", 0.01),
-        )
+        self.optimizer, self.critic_optimizer = self._init_optimizers(**self.base_optim_config)
 
         self.scaler = torch.amp.GradScaler() if self.use_amp else None
         self.critic_scaler = torch.amp.GradScaler() if self.use_amp else None
@@ -154,6 +158,21 @@ class MRLTrainer:
         self.global_epoch = 0
         self.global_epochs_count = 0
 
+    def _init_optimizers(self, lr: float, critic_lr: float, weight_decay: float, critic_weight_decay: float):
+        optimizer = torch.optim.AdamW(
+            self.actor.unique_parameters(),
+            lr=lr,
+            weight_decay=weight_decay,
+        )
+
+        critic_optimizer = torch.optim.AdamW(
+            self.critic.parameters(),
+            lr=critic_lr,
+            weight_decay=critic_weight_decay,
+        )
+        return optimizer, critic_optimizer
+
+
     def _init_steps(self):
         return {
             'collect': 0,
@@ -221,21 +240,29 @@ class MRLTrainer:
 
         return generated_answer, log_probs
 
+    def _calculate_reward(self, generated: TokenizedDict, reference: TokenizedDict,
+                          saved_query: TokenizedDict, saved_answer: TokenizedDict,
+                          mode: MrlRewardMode = MrlRewardMode.STANDARD,
+                          prev_data: tuple[TokenizedDict, TokenizedDict] = None):
+        saved_interaction = smart_concat(saved_query, saved_answer, max_length=self.max_seq_len,
+                                         pad_token_id=self.pad_token_id)
+        prev_data = smart_concat(prev_data[0], prev_data[1], self.max_seq_len,
+                                 self.pad_token_id) if prev_data is not None else None
+        return self.reward(generated, reference, saved_interaction, mode=mode, prev_data=prev_data), saved_interaction
+
     def compute_reward(self, generated: TokenizedDict, reference: TokenizedDict,
                        saved_data: tuple[TokenizedDict, TokenizedDict], mode: MrlRewardMode = MrlRewardMode.STANDARD,
-                       eval_mode: bool = False) -> list[float]:
+                       eval_mode: bool = False, prev_data: tuple[TokenizedDict, TokenizedDict] = None) -> list[float]:
         """Compute reward based on memory retention (e.g., BLEU-4)."""
         saved_query, saved_answer = saved_data
         # 1. Concat saved (previous) interaction and calculate reward using generated sequence, reference and saved data - with autocast on/off
         if self.use_amp:
             with torch.amp.autocast(device_type=self.device.type, dtype=self.dtype):
-                saved_interaction = smart_concat(saved_query, saved_answer, max_length=self.max_seq_len,
-                                                 pad_token_id=self.pad_token_id)
-                reward = self.reward(generated, reference, saved_interaction, mode=mode)
+                reward, saved_interaction = self._calculate_reward(generated, reference, saved_query, saved_answer,
+                                                                   mode=mode, prev_data=prev_data)
         else:
-            saved_interaction = smart_concat(saved_query, saved_answer, max_length=self.max_seq_len,
-                                             pad_token_id=self.pad_token_id)
-            reward = self.reward(generated, reference, saved_interaction, mode=mode)
+            reward, saved_interaction = self._calculate_reward(generated, reference, saved_query, saved_answer,
+                                                               mode=mode, prev_data=prev_data)
 
         # 2. Run 'on reward' callbacks
         for cb in self.callbacks:
@@ -289,22 +316,27 @@ class MRLTrainer:
                     # state from existing one, instead of new random one)
                     reset_done = self.reset_stm()
 
-                    # 4. Get first batch of interactions (data to save) and follow-up interactions for current episode, based on curriculum step
+                    # 4. Reset reward prev data running mean - it's calculated for multi-step retention, we have to reset it before episode
+                    self.reward.reset_running_mean()
+
+                    # 5. Get first batch of interactions (data to save) and follow-up interactions for current episode, based on curriculum step
                     first_query, first_answer, interactions = batch['query'], batch['answer'], batch['interactions']
                     interactions = interactions[:self.curriculum_steps]
                     interactions_len = len(interactions)
-                    # 5. Encode and update STM with data to save from 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))
 
-                    # 6. Save first interaction as data to save (for trajectory state)
+                    # 7. Save first interaction as data to save (for trajectory state)
                     query, answer = first_query, first_answer
 
-                    # 7. Run training strategy for follow-up interactions
+                    # 8. Run training strategy for follow-up interactions
                     episode_steps = []
                     episode_rewards = []
 
+                    prev_interaction = None
+
                     for i, interaction in enumerate(interactions):
-                        # 8. Generate batch of answers based on batch of follow-up queries
+                        # 9. Generate batch of answers based on batch of follow-up queries
                         next_query = self._move_batch(interaction['query'])
                         generated_answer, log_probs = self.generate_answer(next_query)
 
@@ -312,7 +344,7 @@ class MRLTrainer:
 
                         detached_answer = self._cpu_detach(generated_answer)  # detach and keep states on CPU
 
-                        # 9. Depending on strategy compute reward
+                        # 10. Depending on strategy compute reward
                         if self.strategy == MrlStrategy.LONG_RANGE_STRATEGY and i == 0:
                             # a) long-range - first interaction - change topic - negative reward (it shouldn't include saved data)
                             reward = self.compute_reward(detached_answer, interaction['answer'], (query, answer),
@@ -320,18 +352,19 @@ class MRLTrainer:
                         elif self.strategy == MrlStrategy.LONG_RANGE_STRATEGY and is_last_interaction:
                             # b) long-range - last interaction - first interaction topic - long-range reward (it should include content from first interaction)
                             reward = self.compute_reward(detached_answer, interaction['answer'],
-                                                         (first_query, first_answer), mode=MrlRewardMode.LONG_RANGE)
+                                                         (first_query, first_answer), mode=MrlRewardMode.LONG_RANGE,
+                                                         prev_data=prev_interaction)
                         else:
                             # c) standard reward - generated answer should include some content from previous interaction (saved data), like reference answer
                             reward = self.compute_reward(detached_answer, interaction['answer'], (query, answer),
-                                                         mode=MrlRewardMode.STANDARD)
+                                                         mode=MrlRewardMode.STANDARD, prev_data=prev_interaction)
 
-                        # 10. Update STM with generated response (except last interaction, it's not needed)
+                        # 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
 
-                        # 11. Store trajectory step
+                        # 12. Store trajectory step
                         trajectory: MrlTrajectoryStep = {
                             'state': (query, answer, interaction['query']),
                             'action': detached_answer,
@@ -342,10 +375,12 @@ class MRLTrainer:
                         episode_steps.append(trajectory)
                         episode_rewards.append(reward)
 
-                        # 12. Set current interaction query and generated answer (batches), as saved data for next interaction
+                        # 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
 
-                    # 13. Append full batched episode (number of steps depends on curriculum stage) to trajectories
+                    # 14. Append full batched episode (number of steps depends on curriculum stage) to trajectories
                     episode_trajectory: MrlTrajectoryEpisode = {
                         'reset_stm': reset_done,
                         'steps': episode_steps,
@@ -356,7 +391,7 @@ class MRLTrainer:
 
                     self._collect_writer(mean_episode_reward, epoch)
 
-                    # 14. Run "on episode collected" callbacks
+                    # 15. Run "on episode collected" callbacks
                     for cb in self.callbacks:
                         cb.on_episode_collected(self.actor, batch_idx, episode_trajectory, mean_episode_reward)
 
@@ -595,63 +630,70 @@ class MRLTrainer:
         for batch in dataloader:
             with torch.no_grad():
                 if batch['query']['input_ids'].size(0) == batch_size:
-                  self._increment_steps('eval')
-                  # 3. Reset STM with random resets ratio
-                  self.reset_stm()
-
-                  # 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))
-
-                  # 6. Save follow-up interactions len and first query and answer as previous one for iteration
-                  interactions_len = len(interactions)
-                  query, answer = first_query, first_answer
-                  episode_reward = torch.tensor(0.0).to(self.device)
-                  episode_interactions = torch.tensor(0).to(self.device)
-                  # 7. Run all follow-up interactions
-                  for i, interaction in enumerate(interactions):
-                      # 8. Generate batch of answers
-                      next_query = self._move_batch(interaction['query'])
-                      generated_answer, _ = self.generate_answer(next_query)
-
-                      is_last_interaction = (i + 1) == interactions_len
-
-                      detached_answer = self._cpu_detach(generated_answer)
-
-                      # 9. Depending on current strategy and step, compute reward
-                      if self.strategy == MrlStrategy.LONG_RANGE_STRATEGY and i == 0:
-                          reward = self.compute_reward(detached_answer, interaction['answer'], (query, answer),
-                                                      mode=MrlRewardMode.NEGATIVE, eval_mode=True)
-                      elif self.strategy == MrlStrategy.LONG_RANGE_STRATEGY and is_last_interaction:
-                          reward = self.compute_reward(detached_answer, interaction['answer'],
-                                                      (first_query, first_answer), mode=MrlRewardMode.LONG_RANGE,
-                                                      eval_mode=True)
-                      else:
-                          reward = self.compute_reward(detached_answer, interaction['answer'], (query, answer),
-                                                      mode=MrlRewardMode.STANDARD, eval_mode=True)
-
-                      # 10. Encode and update memory for the next interaction
-                      if not is_last_interaction:
-                          self.encode_and_update_stm(next_query, generated_answer)
-
-                      # 11. Accumulate rewards
-                      step_reward = torch.tensor(reward).mean().to(self.device)
-                      # total
-                      total_reward += step_reward
-                      count += 1
-                      # episode
-                      episode_reward += step_reward
-                      episode_interactions += 1
-                      # 12. Save previous interaction
-                      query, answer = interaction['query'], 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)
-
-                  # 14. Run "on eval episode end" callbacks
-                  for cb in self.callbacks:
-                      cb.on_eval_episode_end(self.actor, epoch, self.epoch_step['eval'], avg_episode_reward)
+                    self._increment_steps('eval')
+                    # 3. Reset STM with random resets ratio and reward model running mean
+                    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))
+
+                    # 6. Save follow-up interactions len and first query and answer as previous one for iteration
+                    interactions_len = len(interactions)
+                    query, answer = first_query, first_answer
+                    episode_reward = torch.tensor(0.0).to(self.device)
+                    episode_interactions = torch.tensor(0).to(self.device)
+
+                    prev_interaction = None
+
+                    # 7. Run all follow-up interactions
+                    for i, interaction in enumerate(interactions):
+                        # 8. Generate batch of answers
+                        next_query = self._move_batch(interaction['query'])
+                        generated_answer, _ = self.generate_answer(next_query)
+
+                        is_last_interaction = (i + 1) == interactions_len
+
+                        detached_answer = self._cpu_detach(generated_answer)
+
+                        # 9. Depending on current strategy and step, compute reward
+                        if self.strategy == MrlStrategy.LONG_RANGE_STRATEGY and i == 0:
+                            reward = self.compute_reward(detached_answer, interaction['answer'], (query, answer),
+                                                         mode=MrlRewardMode.NEGATIVE, eval_mode=True)
+                        elif self.strategy == MrlStrategy.LONG_RANGE_STRATEGY and is_last_interaction:
+                            reward = self.compute_reward(detached_answer, interaction['answer'],
+                                                         (first_query, first_answer), mode=MrlRewardMode.LONG_RANGE,
+                                                         eval_mode=True, prev_data=prev_interaction)
+                        else:
+                            reward = self.compute_reward(detached_answer, interaction['answer'], (query, answer),
+                                                         mode=MrlRewardMode.STANDARD, eval_mode=True,
+                                                         prev_data=prev_interaction)
+
+                        # 10. Encode and update memory for the next interaction
+                        if not is_last_interaction:
+                            self.encode_and_update_stm(next_query, generated_answer)
+
+                        # 11. Accumulate rewards
+                        step_reward = torch.tensor(reward).mean().to(self.device)
+                        # total
+                        total_reward += step_reward
+                        count += 1
+                        # episode
+                        episode_reward += step_reward
+                        episode_interactions += 1
+                        # 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
+                    avg_episode_reward = (episode_reward / episode_interactions).item()
+                    # 13. Run eval TensorBoard writer with average episode reward
+                    self._eval_writer(avg_episode_reward, epoch)
+
+                    # 14. Run "on eval episode end" callbacks
+                    for cb in self.callbacks:
+                        cb.on_eval_episode_end(self.actor, epoch, self.epoch_step['eval'], avg_episode_reward)
 
         # 15. Calculate average reward
         if self.use_ddp:
@@ -679,6 +721,14 @@ class MRLTrainer:
                                     self.shared_callbacks)  # trainer callbacks for current curriculum stage
         self.strategy = config.get('strategy',
                                    MrlStrategy.MULTI_STEP_STRATEGY)  # MRL strategy for given curriculum stage
+        self.reward = config.get('reward_model', self.shared_reward_model)  # MRL Reward Model for curriculum stage
+        if config['lr'] is not None or config['critic_lr'] is not None or config['weight_decay'] is not None or config['critic_weight_decay'] is not None:
+            self.optimizer, self.critic_optimizer = self._init_optimizers(
+                lr=config['lr'] or self.base_optim_config['lr'],
+                critic_lr=config['critic_lr'] or self.base_optim_config['critic_lr'],
+                weight_decay=config['weight_decay'] or self.base_optim_config['weight_decay'],
+                critic_weight_decay=config['critic_weight_decay'] or self.base_optim_config['critic_weight_decay']
+            )
 
         # 2. Get epochs and random resets configs
         epochs = config.get('epochs', 5)  # number of epochs for current stage
@@ -720,7 +770,11 @@ class MRLTrainer:
 
             # 4. Freeze all components except memory attention and memory cross-attention layers in decoder/encoder
             if unfreeze_epoch != 0:
-                self.actor.freeze_components()
+                is_staged_unfreeze = isinstance(unfreeze_epoch, tuple)
+                if is_staged_unfreeze:
+                    self.actor.freeze_components('update')
+                else:
+                    self.actor.freeze_components()
 
             # 5. Setup train DataLoader
             if self.use_ddp:
@@ -761,8 +815,18 @@ class MRLTrainer:
                     self.random_resets_ratio = 1.0
 
                 # 11. Unfreeze all components before selected epoch
-                if epoch == unfreeze_epoch:
-                    self.actor.unfreeze_components()
+                is_staged_unfreeze = isinstance(unfreeze_epoch, tuple)
+                if is_staged_unfreeze:
+                    fetch_epoch, both_epoch, all_epoch = unfreeze_epoch
+                    if epoch == fetch_epoch:
+                        self.actor.freeze_components('fetch')
+                    elif epoch == both_epoch:
+                        self.actor.freeze_components('both')
+                    elif epoch == all_epoch:
+                        self.actor.unfreeze_components()
+                else:
+                    if epoch == unfreeze_epoch:
+                        self.actor.unfreeze_components()
 
                 # 12. Set epoch for distributed sampler
                 if train_sampler is not None:
@@ -805,4 +869,3 @@ class MRLTrainer:
         # 21. Close writer
         if self.writer:
             self.writer.close()
-

rxnn/training/reward.py

@@ -11,6 +11,7 @@ class MrlRewardMode(Enum):
     NEGATIVE = 2
     LONG_RANGE = 3
 
+
 class MrlRewardModel:
     def __init__(
             self,
@@ -18,9 +19,14 @@ class MrlRewardModel:
             device: torch.device,
             bleu_with_saved_data: bool = False,
             bleu_factor: float = 0.5,
+            bleu_ref_factor: float = 0.5,
+            bleu_saved_factor: float = 0.5,
             cos_factor: float = 0.5,
             cos_ref_factor: float = 0.5,
             cos_saved_factor: float = 0.5,
+            multi_cos_ref_factor: float = 0.3,
+            multi_cos_saved_factor: float = 0.5,
+            multi_cos_running_mean_factor: float = 0.2,
             neg_bleu_factor: Optional[float] = None,
             neg_cos_factor: Optional[float] = None,
             neg_cos_ref_factor: Optional[float] = None,
@@ -28,45 +34,88 @@ class MrlRewardModel:
             neg_bleu_ref_factor: float = 0.5,
             neg_bleu_saved_factor: float = 0.5,
             allow_not_summing_factors: bool = False,
+            reward_len: bool = False,
+            neg_reward_len: bool = False,
+            max_rewarded_len: int = None,
+            len_factor: int = None,
+            use_running_mean: bool = True,
+            running_mean_decay: float = 0.2,
+            bleu_saved_weights: tuple = (0.5, 0.5),
+            bleu_ref_weights: tuple = (0.5, 0.5),
+            rewards_scale: float = 1.0,
     ):
         self.shared_embedding = shared_embedding.to(device)
         self.device = device
         self.bleu_with_saved_data = bleu_with_saved_data
 
         self.bleu_factor = bleu_factor
+        self.bleu_ref_factor = bleu_ref_factor
+        self.bleu_saved_factor = bleu_saved_factor
         self.cos_factor = cos_factor
         self.cos_ref_factor = cos_ref_factor
         self.cos_saved_factor = cos_saved_factor
+        self.multi_cos_ref_factor = multi_cos_ref_factor
+        self.multi_cos_saved_factor = multi_cos_saved_factor
+        self.multi_cos_running_mean_factor = multi_cos_running_mean_factor
         self.neg_bleu_factor = neg_bleu_factor if neg_bleu_factor is not None else bleu_factor
         self.neg_cos_factor = neg_cos_factor if neg_cos_factor is not None else cos_factor
         self.neg_cos_ref_factor = neg_cos_ref_factor if neg_cos_ref_factor is not None else cos_ref_factor
         self.neg_cos_saved_factor = neg_cos_saved_factor if neg_cos_saved_factor is not None else cos_saved_factor
         self.neg_bleu_ref_factor = neg_bleu_ref_factor
         self.neg_bleu_saved_factor = neg_bleu_saved_factor
+        self.reward_len = reward_len
+        self.neg_reward_len = neg_reward_len
+        self.max_rewarded_len = max_rewarded_len
+        self.len_factor = len_factor
+        self.use_running_mean = use_running_mean
+        self.running_mean_decay = running_mean_decay
+        self.bleu_ref_weights = bleu_ref_weights
+        self.bleu_saved_weights = bleu_saved_weights
+        self.rewards_scale = rewards_scale
+
+        self.prev_data_running_mean = None
 
         if not allow_not_summing_factors:
-            assert self.bleu_factor + self.cos_factor == 1.0
-            assert self.cos_ref_factor + self.cos_saved_factor == 1.0
-            assert self.neg_bleu_factor + self.neg_cos_factor == 1.0
-            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
+            if reward_len:
+                assert self.bleu_factor + self.cos_factor + self.len_factor == 1.0
+                assert self.neg_bleu_factor + self.neg_cos_factor + self.len_factor == 1.0
+                assert self.multi_cos_ref_factor + self.multi_cos_saved_factor + self.multi_cos_running_mean_factor == 1.0
+                assert self.bleu_ref_factor + self.bleu_saved_factor == 1.0
+                assert self.cos_ref_factor + self.cos_saved_factor == 1.0
+                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
+            else:
+                assert self.bleu_factor + self.cos_factor == 1.0
+                assert self.bleu_ref_factor + self.bleu_saved_factor == 1.0
+                assert self.cos_ref_factor + self.cos_saved_factor == 1.0
+                assert self.multi_cos_ref_factor + self.multi_cos_saved_factor + self.multi_cos_running_mean_factor == 1.0
+                assert self.neg_bleu_factor + self.neg_cos_factor == 1.0
+                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
-        refs = [reference, saved_data] if self.bleu_with_saved_data else [reference]
-        return sentence_bleu(refs, generated, weights=(0.25, 0.25, 0.25, 0.25))
 
-    def _negative_sentence_bleu(self, generated: torch.Tensor, reference: torch.Tensor, saved_data: torch.Tensor) -> float:
+        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)
+            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, generated: torch.Tensor, reference: torch.Tensor,
+                                saved_data: torch.Tensor) -> float:
         from nltk.translate.bleu_score import sentence_bleu
 
         if self.bleu_with_saved_data:
-            ref_bleu = sentence_bleu([reference], generated, weights=(0.25, 0.25, 0.25, 0.25))
-            saved_bleu = sentence_bleu([saved_data], generated, weights=(0.25, 0.25, 0.25))
+            ref_bleu = sentence_bleu([reference], generated, weights=self.bleu_ref_weights)
+            saved_bleu = sentence_bleu([saved_data], generated, weights=self.bleu_saved_weights)
             saved_bleu = 1 - saved_bleu
 
-            return (self.neg_bleu_ref_factor * ref_bleu + self.neg_bleu_saved_factor * saved_bleu) / 2
+            return self.neg_bleu_ref_factor * ref_bleu + self.neg_bleu_saved_factor * saved_bleu
         else:
-            return sentence_bleu([reference], generated, weights=(0.25, 0.25, 0.25, 0.25))
+            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)
@@ -79,33 +128,81 @@ class MrlRewardModel:
     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, self._sequence_embedding(saved_data))
-        gen_and_ref = F.cosine_similarity(generated_emb, self._sequence_embedding(reference))
+        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
 
-    def batch_cosine(self, generated: torch.Tensor, reference: torch.Tensor, saved_data: torch.Tensor) -> torch.Tensor:
-        gen_and_saved, gen_and_ref = self._cosine_sim(generated, reference, saved_data)
+    def _cosine_sim_running_mean(self, generated: torch.Tensor, reference: torch.Tensor, saved_data: torch.Tensor):
+        generated_emb = self._sequence_embedding(generated)
 
-        return self.cos_saved_factor * gen_and_saved + self.cos_ref_factor * gen_and_ref
+        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
+
+    def batch_cosine(self, generated: torch.Tensor, reference: torch.Tensor, saved_data: torch.Tensor,
+                     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)
+        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
+        else:
+            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: torch.Tensor, reference: torch.Tensor,
+                        saved_data: torch.Tensor) -> 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):
+        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)
+        return len_reward
+
+    def reset_running_mean(self):
+        self.prev_data_running_mean = None
+
+    def init_running_mean(self, prev_data: torch.Tensor):
+        self.prev_data_running_mean = self._sequence_embedding(prev_data)
+
+    def update_running_mean(self, prev_data: torch.Tensor):
+        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 __call__(
             self,
             generated: TokenizedDict,
             reference: TokenizedDict,
             saved_data: TokenizedDict,
+            prev_data: TokenizedDict = None,
             mode: MrlRewardMode = MrlRewardMode.STANDARD
     ) -> list[float]:
-        if mode == MrlRewardMode.STANDARD or mode == MrlRewardMode.LONG_RANGE:
+        if prev_data is not None:
+            if self.prev_data_running_mean is None:
+                self.init_running_mean(prev_data['input_ids'])
+            else:
+                self.update_running_mean(prev_data['input_ids'])
+
+        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)
+            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'])
-            return (self.bleu_factor * torch.tensor(bleu, device=self.device) + self.cos_factor * cosine).tolist()
+            cosine = self.batch_cosine(generated['input_ids'], reference['input_ids'], saved_data['input_ids'],
+                                       negative_running_mean=prev_data is not None)
+            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'])
-            return (self.neg_bleu_factor * torch.tensor(bleu, device=self.device) + self.neg_cos_factor * cosine).tolist()
+            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
+        return rewards.tolist()

{rxnn-0.2.18.dist-info → rxnn-0.2.20.dist-info}/METADATA

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

{rxnn-0.2.18.dist-info → rxnn-0.2.20.dist-info}/RECORD

@@ -15,9 +15,9 @@ rxnn/training/base.py,sha256=_xik1GXE4RJ_nxwqLQ1ccXA5pRtBCi-jL-jeRFBdHBU,11851
 rxnn/training/bml.py,sha256=FJszaQXOLx2ZHBa1CQpyMrG8i4Kj14E-gzDAEK_Ei5k,17272
 rxnn/training/callbacks.py,sha256=-N0MQPpZQaUWCINdTOsjul4bDGbGr2JgQBqOIXBLS6o,35053
 rxnn/training/dataset.py,sha256=7hTilFWPpqUEc6zNcMqBPjxFKxCfvTKKF3E8tVlwccQ,51250
-rxnn/training/models.py,sha256=renPa5YH443XNTMFI-YTCwi5vNp3QzwF5UXedNd5hDk,5187
-rxnn/training/mrl.py,sha256=Ar2g-vjqTq_4qLKc4L1Ai0j2LX-x98dmsx_VaWVV-Es,39448
-rxnn/training/reward.py,sha256=C0ToTz-u-L-qyBd2yJ1HlvVPS110OChYj9ZhD6iSSMU,5654
+rxnn/training/models.py,sha256=qXfD3_97T9z724NN4myjzrpX6-jYA9Igl266ZwtJCtc,5519
+rxnn/training/mrl.py,sha256=zk4m1JFuX0y82J0tG2XkY0Pz6Uy2did9cngOXqR9lMk,43326
+rxnn/training/reward.py,sha256=7MTVdNm5HnWmt6zFDi3TAYmnVSL_-24riOoY2F7z4x8,11290
 rxnn/training/rl.py,sha256=DHFwnPUlnq2JVj6CS6DwifnC_eMeBAUVp36UCAWNMis,3934
 rxnn/training/scheduler.py,sha256=LcjU35mEwz2U5x3U6tLfeeYlBqMxbFSxYzJYuXkWbSY,1408
 rxnn/training/tokenizer.py,sha256=umaLByMBx_NMrQElA45HLm9gkuzyKWDTFaKVd-CjXl0,8344
@@ -32,7 +32,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.18.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.18.dist-info/METADATA,sha256=_hGNlaH_rclBfQdzA7tCFhkI-RZPiK5tNBM8tjUsbWQ,25960
-rxnn-0.2.18.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.18.dist-info/RECORD,,
+rxnn-0.2.20.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.20.dist-info/METADATA,sha256=mgimK5GvI27RapfLjhlIdBwgfVdKoMA5Ig5yVxfeYIw,25960
+rxnn-0.2.20.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.20.dist-info/RECORD,,