rxnn 0.2.47__tar.gz → 0.2.49__tar.gz

{rxnn-0.2.47 → rxnn-0.2.49}/PKG-INFO

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

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "rxnn"
-version = "0.2.47"
+version = "0.2.49"
 description = "RxNN: Reactive Neural Networks Platform"
 
 license = "Apache-2.0"
@@ -30,4 +30,4 @@ datasets = "^3.5.0"
 tokenizers = "^0.21.0"
 huggingface-hub = "^0.30.0"
 tensorboard = "^2.19.0"
-
+nltk = "^3.9.1"

{rxnn-0.2.47 → rxnn-0.2.49}/src/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-0.2.47 → rxnn-0.2.49}/src/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-0.2.49/src/rxnn/training/reward.py

@@ -0,0 +1,323 @@
+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
+
+
+class MrlRewardMode(Enum):
+    STANDARD = 1
+    NEGATIVE = 2
+    LONG_RANGE = 3
+
+
+class MrlRewardModel:
+    def __init__(
+            self,
+            shared_embedding: nn.Embedding,
+            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,
+            neg_cos_saved_factor: Optional[float] = None,
+            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,
+            target_len_as_ref: bool = False,
+            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),
+            tanh_reward_scale: bool = False,
+            rewards_scale: float = 1.0,
+            debug_mode: int = 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.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
+        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.bleu_smoothing = SmoothingFunction().method4
+        self.debug_mode = debug_mode
+
+        self.prev_data_running_mean = None
+
+        if not allow_not_summing_factors:
+            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, 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,
+                                     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()]
+
+        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,
+                                     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: 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 _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)
+
+        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)
+
+        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: 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, 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)
+
+        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)
+        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: 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, 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 = 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: TokenizedDict):
+        self.prev_data_running_mean = self._sequence_embedding(prev_data)
+
+    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
+
+    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,
+            reference: TokenizedDict,
+            saved_data: TokenizedDict,
+            prev_data: TokenizedDict = None,
+            mode: MrlRewardMode = MrlRewardMode.STANDARD
+    ) -> list[float]:
+        if prev_data is not None:
+            if self.prev_data_running_mean is None:
+                self.init_running_mean(prev_data)
+            else:
+                self.update_running_mean(prev_data)
+
+        if mode == MrlRewardMode.STANDARD:
+            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, 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.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
+
+        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/src/rxnn/training/reward.py

@@ -1,216 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from enum import Enum
-from typing import Optional
-from .utils import TokenizedDict
-
-
-class MrlRewardMode(Enum):
-    STANDARD = 1
-    NEGATIVE = 2
-    LONG_RANGE = 3
-
-
-class MrlRewardModel:
-    def __init__(
-            self,
-            shared_embedding: nn.Embedding,
-            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,
-            neg_cos_saved_factor: Optional[float] = None,
-            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),
-            tanh_reward_scale: bool = False,
-            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.tanh_reward_scale = tanh_reward_scale
-        self.rewards_scale = rewards_scale
-
-        self.prev_data_running_mean = None
-
-        if not allow_not_summing_factors:
-            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
-
-        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=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
-        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 _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: 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)) + 1) / 2
-        gen_and_ref = (F.cosine_similarity(generated_emb, self._sequence_embedding(reference)) + 1) / 2
-        return gen_and_saved, gen_and_ref
-
-    def _cosine_sim_running_mean(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)) + 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:
-        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 _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,
-            reference: TokenizedDict,
-            saved_data: TokenizedDict,
-            prev_data: TokenizedDict = None,
-            mode: MrlRewardMode = MrlRewardMode.STANDARD
-    ) -> list[float]:
-        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'],
-                                       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'])
-            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
-        return rewards.tolist()