rxnn 0.1.83__py3-none-any.whl → 0.2.0__py3-none-any.whl

rxnn/training/reward.py

@@ -0,0 +1,111 @@
+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,
+            cos_factor: float = 0.5,
+            cos_ref_factor: float = 0.5,
+            cos_saved_factor: float = 0.5,
+            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,
+    ):
+        self.shared_embedding = shared_embedding.to(device)
+        self.device = device
+        self.bleu_with_saved_data = bleu_with_saved_data
+
+        if not allow_not_summing_factors:
+            assert bleu_factor + cos_factor == 1.0
+            assert cos_ref_factor + cos_saved_factor == 1.0
+            assert neg_bleu_factor + neg_cos_factor == 1.0
+            assert neg_cos_ref_factor + neg_cos_saved_factor == 1.0
+            assert neg_bleu_ref_factor + neg_bleu_saved_factor == 1.0
+
+        self.bleu_factor = bleu_factor
+        self.cos_factor = cos_factor
+        self.cos_ref_factor = cos_ref_factor
+        self.cos_saved_factor = cos_saved_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
+
+    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:
+        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))
+            saved_bleu = 1 - saved_bleu
+
+            return (self.neg_bleu_ref_factor * ref_bleu + self.neg_bleu_saved_factor * saved_bleu) / 2
+        else:
+            return sentence_bleu([reference], generated, weights=(0.25, 0.25, 0.25, 0.25))
+
+    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))
+        gen_and_ref = F.cosine_similarity(generated_emb, self._sequence_embedding(reference))
+        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)
+
+        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 __call__(
+            self,
+            generated: TokenizedDict,
+            reference: TokenizedDict,
+            saved_data: TokenizedDict,
+            mode: MrlRewardMode = MrlRewardMode.STANDARD
+    ) -> list[float]:
+        if mode == MrlRewardMode.STANDARD or 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) + self.cos_factor * cosine).tolist()
+        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) + self.neg_cos_factor * cosine).tolist()
+

rxnn/training/rl.py

@@ -0,0 +1,69 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from abc import ABC, abstractmethod
+from typing import TypedDict
+
+
+class RlAlgorithm(ABC):
+    def __init__(self):
+        super(RlAlgorithm, self).__init__()
+        self.critic_loss = nn.MSELoss()
+
+    @abstractmethod
+    def policy_loss(self, input_ids: torch.Tensor, logits: torch.Tensor, old_log_probs: torch.Tensor, advantages: torch.Tensor) -> torch.Tensor:
+        pass
+
+    @abstractmethod
+    def calculate_advantages(self, rewards: torch.Tensor, values: torch.Tensor) -> torch.Tensor:
+        pass
+
+    def critic_loss(self, rewards: torch.Tensor, values: torch.Tensor) -> torch.Tensor:
+        return self.critic_loss(rewards, values)
+
+class PPOConfig(TypedDict):
+    gae_gamma: float
+    gae_lambda: float
+    clip_eps: float
+
+class PPOAlgorithm(RlAlgorithm):
+    def __init__(self, config: PPOConfig):
+        super(PPOAlgorithm, self).__init__()
+
+        # PPO Config
+        self.gae_gamma = config.get('gae_gamma', 0.99)
+        self.gae_lambda = config.get('gae_lambda', 0.95)
+        self.clip_eps = config.get('clip_eps', 0.2)
+
+    def policy_loss(self, input_ids: torch.Tensor, logits: torch.Tensor, old_log_probs: torch.Tensor, advantages: torch.Tensor) -> torch.Tensor:
+        # a) Get new log probs
+        new_probs = F.log_softmax(logits, dim=-1)
+        new_log_probs = new_probs.gather(-1, input_ids.unsqueeze(-1)).squeeze(-1)
+
+        # b) Calculate ratio
+        ratio = (new_log_probs - old_log_probs).exp()
+
+        # c) Clipped surrogate loss
+        surr1 = ratio * advantages
+        surr2 = torch.clamp(ratio, 1.0 - self.clip_eps, 1.0 + self.clip_eps) * advantages
+        policy_loss = -torch.min(surr1, surr2).mean()
+
+        # d) Entropy bonus
+        entropy = -torch.sum(new_probs * new_probs.exp(), dim=-1).mean()
+        policy_loss -= 0.01 * entropy
+
+        return policy_loss
+
+    def _compute_gae(self, rewards: torch.Tensor, values: torch.Tensor, next_value: torch.Tensor) -> torch.Tensor:
+        advantages = torch.zeros_like(rewards, device=values.device)
+        last_advantage = 0
+        for t in reversed(range(rewards.size(0))):
+            delta = rewards[t] + self.gae_gamma * next_value - values[t]
+            advantages[t] = delta + self.gae_gamma * self.gae_lambda * last_advantage
+            last_advantage = advantages[t]
+        return advantages
+
+    def calculate_advantages(self, rewards: torch.Tensor, values: torch.Tensor) -> torch.Tensor:
+        advantages = self._compute_gae(rewards, values[:-1], values[-1])
+        normalized_advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
+        return normalized_advantages

rxnn/training/utils.py

@@ -0,0 +1,148 @@
+import torch
+from typing import TypedDict
+
+class SpecialTokenIds(TypedDict):
+    bos: int
+    eos: int
+    pad: int
+
+class TokenizedDict(TypedDict):
+    input_ids: torch.Tensor
+    attention_mask: torch.Tensor
+
+def smart_concat_critic_states(
+        prev_query: TokenizedDict,
+        prev_answer: TokenizedDict,
+        next_query: TokenizedDict,
+        max_length: int,
+        pad_token_id: int
+) -> TokenizedDict:
+    """
+    Smart vectorized concatenation of MRL critic states - previous interaction (query and answer) and next query.
+    It creates a batch of critic input sequences from previous query, previous answer and next query batches.
+    Used in MRL to concatenate critic states in correct format.
+
+    All the concatenated sequences (batches) are padded to the same max length, but the result should have two times
+    longer max length. Single max length is made to fit single query and answer, but here we have additional next query,
+    so we are using 2x longer sequence for safety.
+
+    Args:
+        prev_query (TokenizedDict): Batch of tokenized queries with attention masks from previous interaction
+        prev_answer (TokenizedDict): Batch of tokenized answers with attention masks from previous interaction
+        next_query (TokenizedDict): Batch of tokenized queries with attention masks from next interaction
+        max_length (int): Max length of result sequence.
+        pad_token_id (int): Index of padding token
+    """
+    device = prev_query['input_ids'].device
+    batch_size = prev_query['input_ids'].size(0)
+
+    # Get input dimensions
+    query_max_len = prev_query['input_ids'].size(1)
+    answer_max_len = prev_answer['input_ids'].size(1)
+    next_q_max_len = next_query['input_ids'].size(1)
+
+    # Get actual lengths using attention masks
+    query_lens = prev_query['attention_mask'].sum(dim=1)
+    answer_lens = prev_answer['attention_mask'].sum(dim=1)
+    next_query_lens = next_query['attention_mask'].sum(dim=1)
+
+    # Calculate positions and boundaries
+    positions = torch.arange(max_length, device=device).expand(batch_size, -1)
+    section1_end = query_lens.unsqueeze(1)
+    section2_end = section1_end + answer_lens.unsqueeze(1)
+    section3_end = section2_end + next_query_lens.unsqueeze(1)
+
+    # Create masks for each section
+    mask_prev = positions < section1_end
+    mask_answer = (positions >= section1_end) & (positions < section2_end)
+    mask_next = (positions >= section2_end) & (positions < section3_end)
+
+    # Build combined tensor
+    combined_ids = torch.full((batch_size, max_length), pad_token_id, device=device)
+
+    # 1. Fill previous query section (with input length clamping)
+    query_indices = positions.clamp(max=query_max_len - 1)
+    combined_ids = torch.where(
+        mask_prev,
+        prev_query['input_ids'].gather(1, query_indices),
+        combined_ids
+    )
+
+    # 2. Fill answer section (with answer length clamping)
+    answer_pos = (positions - section1_end).clamp(min=0, max=answer_max_len - 1)
+    combined_ids = torch.where(
+        mask_answer,
+        prev_answer['input_ids'].gather(1, answer_pos),
+        combined_ids
+    )
+
+    # 3. Fill next query section (with next query length clamping)
+    next_q_pos = (positions - section2_end).clamp(min=0, max=next_q_max_len - 1)
+    combined_ids = torch.where(
+        mask_next,
+        next_query['input_ids'].gather(1, next_q_pos),
+        combined_ids
+    )
+
+    # Create attention mask
+    combined_mask = (positions < section3_end).long()
+
+    return {
+        'input_ids': combined_ids,
+        'attention_mask': combined_mask
+    }
+
+def smart_concat(query: TokenizedDict, answer: TokenizedDict, max_length: int, pad_token_id: int) -> TokenizedDict:
+    """
+    Smart vectorized concatenation of interaction parts - query and answer. It creates
+    batch of interactions from query and answer batches. Used in MRL to concatenate data
+    to encode and update memory.
+
+    Query and answer sequences are padded to the same max length, and the result also has
+    the same length.
+
+    Args:
+        query (TokenizedDict): Batch of tokenized queries with attention masks
+        answer (TokenizedDict): Batch of tokenized answers with attention masks
+        max_length (int): Max length of each sequence - query, answer and result.
+        pad_token_id (int): Index of padding token
+    """
+    device = query['input_ids'].device
+    batch_size = query['input_ids'].size(0)
+
+    # Get actual lengths from attention masks
+    query_lens = query['attention_mask'].sum(dim=1)
+    answer_lens = answer['attention_mask'].sum(dim=1)
+
+    # Create combined length tensor
+    combined_lens = torch.minimum(query_lens + answer_lens,
+                                  torch.full_like(query_lens, max_length))
+
+    # Create position indices [batch_size, max_length]
+    positions = torch.arange(max_length, device=device).expand(batch_size, -1)
+
+    # Create mask for query/answer parts
+    query_mask = positions < query_lens.unsqueeze(1)
+    answer_mask = (positions >= query_lens.unsqueeze(1)) & (positions < combined_lens.unsqueeze(1))
+
+    # Calculate answer positions with overflow protection
+    answer_pos = (positions - query_lens.unsqueeze(1)).clamp(min=0)
+
+    # Build combined_ids using vectorized where
+    combined_ids = torch.where(
+        query_mask,
+        query['input_ids'].gather(1, torch.minimum(positions, query_lens.unsqueeze(1) - 1)),
+        torch.where(
+            answer_mask,
+            answer['input_ids'].gather(1, answer_pos),
+            query['input_ids'].new_full((1,), pad_token_id)
+        )
+    )
+
+    # Build attention mask
+    combined_mask = (positions < combined_lens.unsqueeze(1)).long()
+
+    return {
+        'input_ids': combined_ids,
+        'attention_mask': combined_mask
+    }

rxnn/transformers/attention.py

@@ -16,6 +16,7 @@ class MultiHeadAttention(nn.Module):
             dropout: float = 0.0,
             rope: RotaryPositionalEmbedding = None,
             rope_only_for_query: bool = False,
+            rope_only_for_keys: bool = False,
             use_relative_embeddings: bool = False,
             max_seq_len: int = 1024,
             use_flash_attention: bool = True,
@@ -37,10 +38,12 @@ class MultiHeadAttention(nn.Module):
             self.rel_embed = RelativePositionalEmbedding(max_seq_len, embed_dim // num_heads)
             self.rope = None
             self.rope_only_for_query = False
+            self.rope_only_for_keys = False
         else:
             self.rel_embed = None
             self.rope = rope
             self.rope_only_for_query = rope_only_for_query
+            self.rope_only_for_keys = rope_only_for_keys
         self.dropout = nn.Dropout(dropout)
         self._init_q(embed_dim)
         self._init_kv(embed_dim)
@@ -70,6 +73,8 @@ class MultiHeadAttention(nn.Module):
         if self.rope is not None:
             if self.rope_only_for_query:
                 q = self.rope.forward_one(q)
+            elif self.rope_only_for_keys:
+                k = self.rope.forward_one(k)
             else:
                 q, k = self.rope(q, k)
         return q, k
@@ -192,6 +197,7 @@ class GroupedQueryAttention(MultiHeadAttention):
             k = self.k_proj(key).view(b, -1, self.num_groups, head_dim).transpose(1, 2)
             v = self.v_proj(value).view(b, -1, self.num_groups, head_dim).transpose(1, 2)
         else:
+            # Relative embedding version is not working without this strange mapping - it will be removed in next versions
             group_heads = self.num_heads // self.num_groups
 
             # Process Q
@@ -289,6 +295,7 @@ def init_attention(
         dropout: float = 0.0,
         rope: RotaryPositionalEmbedding = None,
         rope_only_for_query: bool = False,
+        rope_only_for_keys: bool = False,
         use_relative_embeddings: bool = False,
         max_seq_len: int = 1024,
         use_flash_attention: bool = False,
@@ -308,6 +315,7 @@ def init_attention(
             use_relative_embeddings=use_relative_embeddings,
             max_seq_len=max_seq_len,
             rope_only_for_query=rope_only_for_query,
+            rope_only_for_keys=rope_only_for_keys,
             use_flash_attention=use_flash_attention,
             is_causal=is_causal,
             use_bias=use_bias,
@@ -321,6 +329,7 @@ def init_attention(
             use_relative_embeddings=use_relative_embeddings,
             max_seq_len=max_seq_len,
             rope_only_for_query=rope_only_for_query,
+            rope_only_for_keys=rope_only_for_keys,
             use_flash_attention=use_flash_attention,
             is_causal=is_causal,
             use_bias=use_bias,
@@ -334,6 +343,7 @@ def init_attention(
             use_relative_embeddings=use_relative_embeddings,
             max_seq_len=max_seq_len,
             rope_only_for_query=rope_only_for_query,
+            rope_only_for_keys=rope_only_for_keys,
             use_flash_attention=use_flash_attention,
             is_causal=is_causal,
             use_bias=use_bias,

rxnn/transformers/layers.py

@@ -61,6 +61,12 @@ class ReactiveTransformerLayer(nn.Module):
         for param in self.memory_cross_attention.parameters():
             param.requires_grad_(is_trainable)
 
+    def update_max_len(self, max_seq_len: int):
+        if self.attention.rope is not None:
+            self.attention.rope.update_max_len(max_seq_len)
+        if self.memory_cross_attention.rope is not None:
+            self.memory_cross_attention.rope.update_max_len(max_seq_len)
+
     def moe_router_loss(self):
         ff_router_loss = self.ff.router_loss() if self.use_moe else None
         att_router_loss = None

rxnn/transformers/models.py

@@ -72,11 +72,17 @@ class ReactiveTransformerDecoder(ReactiveTransformerBase):
         # Process shared layers
         if self.shared_layers is not None:
             for i in range(self.num_shared_layers):
-                layer_stm = self.stm(i).expand(x.size(0), -1, -1)
+                layer_stm = self.stm(i)
+                # expand layer STM to batch size, if it's not in batch mode
+                if layer_stm.size(0) == 1:
+                    layer_stm = layer_stm.expand(x.size(0), -1, -1)
                 x = self.shared_layers[i](x, layer_stm, mask=mask)
         # Process own layers
         for i in range(self.num_own_layers):
-            layer_stm = self.stm(i).expand(x.size(0), -1, -1)
+            layer_stm = self.stm(i)
+            # expand layer STM to batch size, if it's not in batch mode
+            if layer_stm.size(0) == 1:
+                layer_stm = layer_stm.expand(x.size(0), -1, -1)
             x = self.layers[i](x, layer_stm, mask=mask)
         return self.head(x)
 
@@ -93,12 +99,18 @@ class ReactiveTransformerEncoder(ReactiveTransformerBase):
         # Process shared layers
         if self.shared_layers is not None:
             for i in range(self.num_shared_layers):
-                layer_stm = self.stm(i).expand(x.size(0), -1, -1)
+                layer_stm = self.stm(i)
+                # expand layer STM to batch size, if it's not in batch mode
+                if layer_stm.size(0) == 1:
+                    layer_stm = layer_stm.expand(x.size(0), -1, -1)
                 x = self.shared_layers[i](x, layer_stm, mask=attention_mask)
                 hidden_states.append(x)
         # Process own layers
         for i in range(self.num_own_layers):
-            layer_stm = self.stm(i).expand(x.size(0), -1, -1)
+            layer_stm = self.stm(i)
+            # expand layer STM to batch size, if it's not in batch mode
+            if layer_stm.size(0) == 1:
+                layer_stm = layer_stm.expand(x.size(0), -1, -1)
             x = self.layers[i](x, layer_stm, mask=attention_mask)
             hidden_states.append(x)
         return x, torch.stack(hidden_states)

rxnn/transformers/positional.py

@@ -18,6 +18,11 @@ class RotaryPositionalEmbedding(nn.Module):
         freqs = torch.einsum('i,j->ij', t, self.inv_freq)
         self.register_buffer('cache', freqs)
 
+    def update_max_len(self, max_seq_len: int):
+        self.max_seq_len = max_seq_len
+        t = torch.arange(max_seq_len).type_as(self.inv_freq)
+        freqs = torch.einsum('i,j->ij', t, self.inv_freq)
+        self.cache = freqs
 
     def forward(self, q: torch.Tensor, k: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         seq_len = q.size(-2)
@@ -42,6 +47,8 @@ class RotaryPositionalEmbedding(nn.Module):
         return q_embed
 
     def _prepare_freqs(self, seq_len: int) -> torch.Tensor:
+        if seq_len > self.max_seq_len:
+            self.update_max_len(seq_len)
         return self.cache[:seq_len][None, None, :, :]
 
     def _rotate(self, x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: