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

rxnn/training/base.py

@@ -39,6 +39,8 @@ class BaseTrainer(ABC):
         self.optimizer = optimizer
         self.dataset = dataset
         self.callbacks = callbacks or []
+        if log_dir and not os.path.exists(log_dir):
+            os.makedirs(log_dir)
         self.writer = SummaryWriter(log_dir) if log_dir else None
         self.use_ddp = use_ddp
         self.use_amp = use_amp

rxnn/training/callbacks.py

@@ -499,3 +499,24 @@ class JointModelSaveCallback(TrainerCallback):
             if not self.mlm_mode:
                 self._save_final(model.decoder, 'decoder', hub_id=self.hub_model_decoder)
             self._save_final(model.mlm_head, 'head', hub_id=self.hub_model_head)
+
+class EarlyStoppageCallback(TrainerCallback):
+  def __init__(self, num_plateau_epochs: int = 3) -> None:
+      super().__init__()
+      self.num_plateau_epochs = num_plateau_epochs
+      self.best_loss = 9999.0
+      self.best_loss_epoch = 0
+
+  def on_validation_end(
+        self,
+        model: torch.nn.Module,
+        epoch: int,
+        val_loss: float,
+        val_metrics: dict
+  ):
+    if val_loss < self.best_loss:
+      self.best_loss = val_loss
+      self.best_loss_epoch = epoch
+    elif epoch - self.best_loss_epoch > self.num_plateau_epochs:
+      return True
+    return None

rxnn/training/dataset.py

@@ -314,6 +314,8 @@ class JointLMDataset(BaseDataset):
     def __getitem__(self, idx: int) -> dict[str, dict[str, torch.Tensor]]:
         inputs = self.get_tokenized_text(idx)
         encoder_input_ids = inputs['input_ids'][0]
+        if self.is_pre_tokenized:
+            encoder_input_ids = encoder_input_ids.clone()
         attention_mask = inputs['attention_mask'][0]
 
         decoder_input_ids = encoder_input_ids.clone()
@@ -361,6 +363,8 @@ class MaskedLMDataset(BaseDataset):
         inputs = self.get_tokenized_text(idx)
 
         input_ids = inputs['input_ids'][0]
+        if self.is_pre_tokenized:
+            input_ids = input_ids.clone()
         attention_mask = inputs['attention_mask'][0]
         labels = input_ids.clone()
 

rxnn/training/scheduler.py

@@ -17,3 +17,21 @@ def get_transformer_lr_scheduler(
         return LambdaLR(optimizer, lr_lambda)
     else:
         return CosineAnnealingLR(optimizer, T_max=num_training_steps)
+
+def calculate_steps(
+        dataset_size: int,
+        epochs: int,
+        batch_size: int,
+        warmup_ratio: float = 0.0,
+        num_workers: int = 1,
+        gradient_accumulation_steps: int = 1,
+        verbose: bool = True,
+):
+    steps_per_epoch = int((dataset_size / batch_size - 1) // num_workers)
+    total_steps = int((epochs * steps_per_epoch) / gradient_accumulation_steps)
+    warmup_steps = int(warmup_ratio * total_steps)
+    if verbose:
+        print(f'Total steps: {total_steps}')
+        print(f'Warmup steps: {warmup_steps}')
+        print(f'Total steps per epoch: {steps_per_epoch}')
+    return { 'total': total_steps, 'warmup': warmup_steps, 'epoch': steps_per_epoch}

rxnn/utils.py

@@ -1,6 +1,9 @@
+import random, gc
 import torch
+import numpy as np
 
 def human_format(num: int):
+    """Format numbers to human-readable format."""
     num = float('{:.3g}'.format(num))
     magnitude = 0
     while abs(num) >= 1000:
@@ -10,5 +13,23 @@ def human_format(num: int):
 
 
 def get_model_size(model: torch.nn.Module):
+    """Calculate all models parameters with requires_grad param set as True"""
     trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
-    return f'Model params {human_format(trainable_params)}'
+    return f'Model params {human_format(trainable_params)}'
+
+def set_random_seed(seed: int):
+    """
+    Set random seed for reproducibility.
+
+    Applied on 3 libs: PyTorch, Numpy and random
+
+    seed (int): Random seed value
+    """
+    torch.random.manual_seed(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+
+def cache_clean():
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()

{rxnn-0.1.82.dist-info → rxnn-0.1.83.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: rxnn
-Version: 0.1.82
+Version: 0.1.83
 Summary: RxNN: Reactive Neural Networks Platform
 License: Apache-2.0
 Keywords: deep-learning,ai,machine-learning
@@ -29,7 +29,7 @@ Description-Content-Type: text/markdown
 # Reactive AI - RxNN
 ## Reactive Neural Networks Platform
 
-RxNN is AI/DeepLearning development platform made for Reactive Neural Networks and Event-driven AI, introduced by Reactive AI.
+RxNN is AI/Deep Learning development platform made for Reactive Neural Networks and Event-driven AI, introduced by Reactive AI.
 
 ## Reactive Neural Networks and Event-driven AI
 Reactive neural networks (RxNN) are a new family of memory-augmented neural networks that combine classical deep learning
@@ -75,6 +75,172 @@ released with next versions of **RxNN** framework:
 - 1.x.x: Multimodal reactive models (could be released earlier, depending on progress)
 - 2.0.0: Real-Time Vision Reactor - Worker class models
 - x.x.x: ...and more!
+
+## Usage
+**RxNN** is made to train models based on reactive architectures, as well as transformer language models. Current version
+is based on PyTorch and HuggingFace libraries (Transformers/Datasets/Tokenizer/Hub), and is integrated with [HuggingFace Hub](https://hugginface.co)
+and [TensorBoard](https://github.com/tensorflow/tensorboard).
+
+> We are also planning a version for **TensorFlow**, more info soon
+
+### Install library and dependencies
+- RxNN and required deps: `pip install rxnn torch transformers tokenizers huggingface_hub`
+- Datasets are required only for training: `pip install datasets`
+- TensorBoard is optional: `pip install tensorboard`
+- [Flash Attention](https://github.com/Dao-AILab/flash-attention) is recommended for faster training/inference (required for models with explicit `use_flash_attention=True`) - check its separate [installation guide](#installing-flash-attention)
+- **NumPy** should be installed too: `pip install numpy`
+
+> ### Installing Flash Attention
+> Installing `flash-attn` could be very frustrating and may take hours (with standard method), only to result in some incompatibility
+> error. Fortunately, the prebuilt versions could be downloaded from GitHub and installed just in seconds. However, you should choose
+> the compatible version based on:
+> - Python version
+> - CUDA version
+> - PyTorch version (2.7 is currently not supported)
+> - ABI
+>
+> #### Steps
+> 1. Choose your version from [https://github.com/Dao-AILab/flash-attention/releases](https://github.com/Dao-AILab/flash-attention/releases)
+> 2. Download prebuilt release, in example: `wget https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp312-cp312-linux_x86_64.whl`
+> 3. Install it, in example: `pip install --no-dependencies --upgrade flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp312-cp312-linux_x86_64.whl`
+> 4. Verify: `flash_attn.__version__` (an incorrect version will cause the error when importing)
+> 
+> #### Note on `use_flash_attention` option in models/layers
+> Explicit `use_flash_attention` option is made to enable direct calls to `flash_attn_func` without using **PyTorch** `scaled_dot_product_attention`. Even
+> if it's set to `False`, when `flash-attn` library is installed, **PyTorch** will try to use it implicitly through _SDPA backend_. It's better to set it
+> to `False` and use automatically, because of better compatibility. Explicit options could be used for research
+
+### Modules
+**RxNN** framework has multiple modules with models, layers, training and inference tools, made for complete development
+of _reactive models_, and could be also used for regular **Transformers**.
+
+#### Transformers
+Transformers module includes classes for models and layers. It includes **Reactive Transformers** as well as **Classic Transformers**
+
+Submodules:
+- `rxnn.transformers.attention` - basic, most common attention layers - `MultiHeadAttention`, `GroupedQueryAttention` and `MultiQueryAttention`
+  - additional attention layers, especially `SparseQueryAttention` could be found in `rxnn.experimental.attention` module
+  - `SparseQueryAttention` will be moved to `rxnn.transformers.attention` in 0.2.x version
+- `rxnn.transformers.positional` - positional encoding layers - `RotaryPositionalEmbedding` and legacy ones - `AbsolutePositionalEmbedding`/`RelativePositionalEmbedding`
+- `rxnn.transformers.ff` - dense feed forward layers, including gated layers (_SwiGLU_, etc.) - `FeedForward` & `GatedFeedForward` (recommended)
+- `rxnn.transformers.moe` - Mixture-of-Experts feed forward layers - `MoeFeedForward` & `GatedMoeFeedForward` (recommended)
+- `rxnn.transformer.layers` - complete reactive/classic transformer layers - `ReactiveTransformerLayer` & `ClassicTransformerLayer`
+- `rxnn.transformer.models` - reactive/classic transformer models - `ReactiveTransformerEncoder`, `ReactiveTransformerDecoder` & `ClassicTransformerEncoder`, `ClassicTransformerDecoder`
+- `rxnn.transformer.sampler` - samplers for reactive models (Sampler is the integral part of reactive architectures) - `Sampler` & `SampleDecoder`
+
+In **RxNN** models are initialized in declarative style by class composition, but then they are wrapped in imperative classes,
+to be compatible with HuggingFace **JSON** config. In example:
+
+```python
+from typing import TypedDict
+import torch
+import torch.nn as nn
+from huggingface_hub import PyTorchModelHubMixin
+from rxnn.transformers.attention import GroupedQueryAttention
+from rxnn.transformers.positional import RotaryPositionalEmbedding
+from rxnn.transformers.layers import ReactiveTransformerLayer
+from rxnn.transformers.models import ReactiveTransformerDecoder
+from rxnn.memory.stm import ShortTermMemory
+
+class YourReactiveTransformerConfig(TypedDict):
+    num_layers: int
+    vocab_size: int
+    embed_dim: int
+    ff_dim: int
+    att_heads: int
+    seq_len: int
+    stm_size: int
+    att_groups: int
+    cross_att_groups: int
+
+
+class YourReactiveTransformerDecoder(nn.Module, PyTorchModelHubMixin):
+    def __init__(
+            self,
+            config: YourReactiveTransformerConfig,
+            **kwargs
+    ):
+        super(YourReactiveTransformerDecoder, self).__init__(**kwargs)
+
+        embedding = nn.Embedding(config['vocab_size'], config['embed_dim'])
+        rope = RotaryPositionalEmbedding(config['embed_dim'] // config['att_heads'], config['seq_len'])
+        stm = ShortTermMemory(config['num_layers'], config['embed_dim'], config['stm_size'])
+
+        self.model = ReactiveTransformerDecoder(
+            stm=stm,
+            embedding=embedding,
+            own_layers=nn.ModuleList([
+                ReactiveTransformerLayer(
+                    config['embed_dim'],
+                    config['ff_dim'],
+                    use_gated=True,
+                    use_moe=False,
+                    ff_activation=nn.GELU(),
+                    ff_dropout=0.1,
+                    use_rms_norm=True,
+                    self_attention=GroupedQueryAttention(
+                        config['embed_dim'],
+                        config['att_heads'],
+                        config['att_groups'],
+                        rope=rope,
+                        dropout=0.1,
+                        max_seq_len=config['seq_len'],
+                        is_causal=True,
+                    ),
+                    memory_cross_attention=GroupedQueryAttention(
+                        config['embed_dim'],
+                        config['att_heads'],
+                        config['att_groups'],
+                        rope=rope,
+                        dropout=0.1,
+                        max_seq_len=config['seq_len'],
+                        is_causal=True,
+                        rope_only_for_query=True
+                    ),
+                ) for _ in range(config['num_layers'])
+            ])
+        )
+    
+    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None):
+        return self.model(x, attention_mask=attention_mask)
+```
+
+#### Memory
+The _memory_ module includes **Short-Term Memory** and layers responsible for its update. In future versions it will also
+include **Long-Term Memory**.
+
+The main `ShortTermMemory` class is located in `rxnn.memory.stm` module - the usage example is in Transformers module description.
+
+Other submodules are connected to **Memory Attention** and will be described in 0.2.x version, after MRL
+
+#### Training
+Training module includes **Trainers** for different training stages of reactive models and shared training utils.
+
+Submodules:
+- `rxnn.training.tokenizer` - custom Trainer for **HuggingFace** `tokenizers` and utils to load tokenizer from Hub
+  - Tokenizer could be loaded from Hub with `load_tokenizer_from_hf_hub(repo_id)`
+- `rxnn.training.dataset` - datasets for different training stages:
+  - `MaskedLMDataset` & `AutoregressiveLMDataset` are made for base models pre-training
+  - `EncoderSftDataset` & `DecoderSftDataset` are made for Interaction Supervised Fine-Tuning for reactive models
+  - `MrlCurriculumDataset` is the dataset for single MRL Curriculum step
+  - `MrlDatasets` is wrapping MRL datasets for all curriculum steps
+  - each dataset has `from_hf_hub` class method to load dataset from Hub
+  - they have also `concat_from_hf_hub` class method to load multiple Hub datasets into single training dataset
+  - if dataset has no validation/test split, each dataset has `get_subset(subset_size, from_start=False)` method - it
+    returns new subset and modifying existing one - i.e. `valid_dataset = train_dataset.get_subset(0.1)`
+  - for concatenated datasets, validation/test split could be created with `concat_from_hf_hub_with_subset` - it cuts the
+    same percentage of each loaded dataset
+- `rxnn.training.callbacks` contain Trainer callbacks, for different kind of utils (more info below)
+- `rxnn.training.scheduler` includes learning rate scheduler for training
+- `rxnn.training.bml` - Base Model Learning module with Trainers for pre-training and fine-tuning
+- `rxnn.training.mrl` - Memory Reinforcement Learning module with Trainers for MRL (from 0.2.x)
+- `rxnn.training.rxrlhf` - Reinforcement Learning from Human Feedback for Reactive Models module (from 0.3.x)
+- `rxnn.training.brl` - Behavioral Reinforcement Learning module (Reactor / from 0.7.x
+
+##### Base Model Learning
+Docs in progress
+
+
                                  Apache License
                            Version 2.0, January 2004
                         http://www.apache.org/licenses/

{rxnn-0.1.82.dist-info → rxnn-0.1.83.dist-info}/RECORD

@@ -9,11 +9,11 @@ rxnn/memory/stm.py,sha256=EsD8slSP4_9dLuq6aFPDmuFe8PWilxh90so5Z3nm-ig,2057
 rxnn/rxt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/rxt/models.py,sha256=iUlSvdXrD1NVzZFmdC55qp4_3xoJj31FC40BGgYlf4Q,8763
 rxnn/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/training/base.py,sha256=xPMA2Bg9-oUZvSZg67ls2p7Gk9pZ9IHUiIJwUzSe2K8,11766
+rxnn/training/base.py,sha256=_xik1GXE4RJ_nxwqLQ1ccXA5pRtBCi-jL-jeRFBdHBU,11851
 rxnn/training/bml.py,sha256=S1ZaXTybzeJH7uVFamCr4TPl2bLyZ5xmn_lSsjThTiM,19162
-rxnn/training/callbacks.py,sha256=_YfMKY_eFdc-tubhO9nYH2PXDZDQwlSI74FVOoCXpQg,22108
-rxnn/training/dataset.py,sha256=xI7bbARRWifunVX6HakCroSFqkM401BQmxfsf9pDeY4,35621
-rxnn/training/scheduler.py,sha256=ow6oALzWjWQmHSpcJEjv6tg4g4CDMvr73TypxfcefMc,712
+rxnn/training/callbacks.py,sha256=xcU3W6_OsIEDTFTbN7S3uIWyGqLulbUWZMpW0aIXmF4,22699
+rxnn/training/dataset.py,sha256=XEDmOwD8v0c9u0QCk7I3xZShKaMtBDwYlfK1ofu6A1E,35789
+rxnn/training/scheduler.py,sha256=LcjU35mEwz2U5x3U6tLfeeYlBqMxbFSxYzJYuXkWbSY,1408
 rxnn/training/tokenizer.py,sha256=umaLByMBx_NMrQElA45HLm9gkuzyKWDTFaKVd-CjXl0,8344
 rxnn/transformers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/transformers/attention.py,sha256=dC0UmC-_kjX8US6Sf0Fi5zw5kJ-P6orH3JDHeBB5gI8,15695
@@ -24,8 +24,8 @@ rxnn/transformers/models.py,sha256=xbnn3FTNZFhaqq9A0XEM12ie_WL_58pPeq0qFXIgve0,7
 rxnn/transformers/moe.py,sha256=j6jEx6Ip0zttlUZKKn82azxo95lkLZs-H2GLSMD88hY,5859
 rxnn/transformers/positional.py,sha256=ge-kaS6WnWnPGnWVp25ZK5bVkmhBUNCaELaN2rN_fSY,4097
 rxnn/transformers/sampler.py,sha256=poWBpxg1iuK5gEJtxHkk5VVfS9V48hs2Olqdhy_Gw8c,6548
-rxnn/utils.py,sha256=d5U8i5ukovgDyqiycc2AoxObTz_eF_bgo2MKvdtJ98s,467
-rxnn-0.1.82.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.1.82.dist-info/METADATA,sha256=xhip3_H9uGKIHKfyTnR0vk_a9zr0TzTIr8buNIiDUQY,16589
-rxnn-0.1.82.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.1.82.dist-info/RECORD,,
+rxnn/utils.py,sha256=ihb6OTyDtPiocB_lOvnq7eOkjjpCkgs8wxvXUBNQ7mM,996
+rxnn-0.1.83.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.1.83.dist-info/METADATA,sha256=AhGTqWM9mvBzDRWliKeTRySDAL2cXXTYefRL_HGJN_Q,25930
+rxnn-0.1.83.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.1.83.dist-info/RECORD,,