rxnn 0.1.52__py3-none-any.whl → 0.1.54__py3-none-any.whl

rxnn/training/base.py

@@ -50,6 +50,10 @@ class BaseTrainer(ABC):
         self.target_field_name = target_field_name
         self.total_tokens = 0
         self.total_steps = 0
+        self.validation_steps = 0
+        self.total_validation_steps = 0
+        self.epoch_steps = 0
+        self.current_epoch = 0
         self.gradient_accumulation_steps = gradient_accumulation_steps
         self.accumulated_loss = 0.0
         self.optimizer_step_count = 0
@@ -108,8 +112,10 @@ class BaseTrainer(ABC):
         scaler = torch.amp.GradScaler() if self.use_amp else None
 
         self.model.train()
-        for epoch in range(epochs):
+        for epoch in range(self.current_epoch, self.current_epoch + epochs):
             if self.is_running:
+                self.current_epoch = epoch
+                self.epoch_steps = 0
                 if train_sampler is not None:
                     train_sampler.set_epoch(epoch)
                 self._run_epoch(dataloader, epoch, optimizer, batch_size, scaler=scaler, scheduler=scheduler)
@@ -142,6 +148,7 @@ class BaseTrainer(ABC):
                     callback.on_batch_start(self.model, batch_idx, batch)
                 if self.get_batch_size(batch) == batch_size:
                     self.total_steps += 1
+                    self.epoch_steps = batch_idx
                     loss = self.train_step(batch, batch_idx)
                     orig_loss = loss.item()
                     self.accumulated_loss += orig_loss
@@ -174,25 +181,28 @@ class BaseTrainer(ABC):
                             self.writer.add_scalar(
                                 'Loss/train',
                                 loss_item,
-                                epoch * len(dataloader) + batch_idx
+                                self.total_steps,
                             )
                             self.writer.add_scalar(
-                                'Loss per epoch/train',
+                                'Loss/train last epoch',
                                 loss_item,
                                 batch_idx
                             )
                             self.writer.add_scalar(
                                 'Perplexity/train',
                                 torch.exp(torch.tensor(loss_item)),
-                                epoch * len(dataloader) + batch_idx
+                                self.total_steps,
                             )
                         self.accumulated_loss = 0.0
                         self.optimizer_step_count = 0
 
                     if self.writer:
                         self.total_tokens += batch['attention_mask'].sum().item()
-                        self.writer.add_scalar('Processed tokens', self.total_tokens,
-                                               epoch * len(dataloader) + batch_idx)
+                        self.writer.add_scalar(
+                            'Processed tokens',
+                            self.total_tokens,
+                            self.total_steps
+                        )
 
                     for callback in self.callbacks:
                         should_stop = callback.on_batch_end(self.model, batch_idx, orig_loss, batch)
@@ -200,6 +210,7 @@ class BaseTrainer(ABC):
                             self.is_running = False
 
         if self.validation_dataset:
+            self.validation_steps = 0
             val_loss, val_metrics = self.validate(batch_size)
             val_loss_tensor = torch.tensor(val_loss).to(self.device)
             if self.use_ddp:
@@ -270,6 +281,8 @@ class BaseTrainer(ABC):
         with torch.no_grad():
             for batch in val_dataloader:
                 if self.get_batch_size(batch) == batch_size:
+                    self.validation_steps += 1
+                    self.total_validation_steps += 1
                     loss, outputs = self.valid_step(batch)
                     val_loss += loss.item()
 

rxnn/training/bml.py

@@ -91,8 +91,8 @@ class MLMTrainer(BaseTrainer):
                 self.writer.add_scalar('Router aux loss/Train', router_loss.item(), self.total_steps)
                 self.writer.add_scalar('Model loss/Train', main_loss.item(), self.total_steps)
             else:
-                self.writer.add_scalar('Router aux loss/Valid', router_loss.item(), self.total_steps)
-                self.writer.add_scalar('Model loss/Valid', main_loss.item(), self.total_steps)
+                self.writer.add_scalar('Router aux loss/Valid', router_loss.item(), self.total_validation_steps)
+                self.writer.add_scalar('Model loss/Valid', main_loss.item(), self.total_validation_steps)
 
         return loss
 
@@ -106,14 +106,25 @@ class MLMTrainer(BaseTrainer):
         with torch.no_grad():
             for batch in val_dataloader:
                 if self.get_batch_size(batch) == batch_size:
+                    self.total_validation_steps += 1
+                    self.validation_steps += 1
                     loss, logits = self.valid_step(batch)
                     val_loss += loss
+                    if self.writer is not None:
+                        self.writer.add_scalar('Loss/Valid total', loss.item(), self.total_validation_steps)
+                        self.writer.add_scalar('Perplexity/Valid', torch.exp(loss).item(), self.total_validation_steps)
+
                     labels = batch[self.target_field_name].to(self.device)
                     valid_indices = labels != -100
                     if valid_indices.any():
                         preds = logits.argmax(-1)
-                        correct += (preds[valid_indices] == labels[valid_indices]).sum()
-                        total += valid_indices.sum()
+                        batch_correct = (preds[valid_indices] == labels[valid_indices]).sum()
+                        batch_total = valid_indices.sum()
+                        batch_acc = (batch_correct / batch_total * 100) if total > 0 else torch.tensor(0.0).to(self.device)
+                        if self.writer is not None:
+                            self.writer.add_scalar('Accuracy/Valid total', batch_acc.item(), self.total_validation_steps)
+                        correct += batch_correct
+                        total += batch_total
 
         avg_loss = (val_loss / len(val_dataloader)).item()
         acc = (correct / total * 100) if total > 0 else torch.tensor(0.0).to(self.device)
@@ -197,15 +208,25 @@ class AutoregressiveTrainer(BaseTrainer):
         with torch.no_grad():
             for batch in val_dataloader:
                 if self.get_batch_size(batch) == batch_size:
+                    self.total_validation_steps += 1
+                    self.validation_steps += 1
                     loss, logits = self.valid_step(batch)
                     val_loss += loss
+                    if self.writer is not None:
+                        self.writer.add_scalar('Loss/Valid total', loss.item(), self.total_validation_steps)
+                        self.writer.add_scalar('Perplexity/Valid', torch.exp(loss).item(), self.total_validation_steps)
                     shifted_logits = logits[:, :-1].contiguous()
                     shifted_targets = batch[self.target_field_name][:, 1:].to(self.device).contiguous()
                     valid_indices = shifted_targets != -100
                     if valid_indices.any():
                         preds = shifted_logits.argmax(-1)
-                        correct += (preds[valid_indices] == shifted_targets[valid_indices]).sum()
-                        total += valid_indices.sum()
+                        batch_correct = (preds[valid_indices] == shifted_targets[valid_indices]).sum()
+                        batch_total = valid_indices.sum()
+                        batch_acc = (batch_correct / batch_total * 100) if total > 0 else torch.tensor(0.0).to(self.device)
+                        if self.writer is not None:
+                            self.writer.add_scalar('Accuracy/Valid total', batch_acc.item(), self.total_validation_steps)
+                        correct += batch_correct
+                        total += batch_total
 
         avg_loss = (val_loss / len(val_dataloader)).item()
         acc = (correct / total * 100) if total > 0 else torch.tensor(0.0).to(self.device)

rxnn/training/dataset.py

@@ -1,7 +1,8 @@
 import torch
 from torch.utils.data import Dataset
-from datasets import Dataset as HfDataset
-from transformers import PreTrainedTokenizer
+from datasets import Dataset as HfDataset, load_dataset, concatenate_datasets
+from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
+from .tokenizer import load_tokenizer_from_hf_hub
 
 from typing import Union
 
@@ -10,10 +11,9 @@ class BaseDataset(Dataset):
     def __init__(
             self,
             texts: Union[list[str], HfDataset],
-            tokenizer: PreTrainedTokenizer,
+            tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
             max_seq_len: int = 1024,
             hf_field: str = 'text',
-            merge_short_from: int = None,
             *args,
             **kwargs
     ):
@@ -22,7 +22,6 @@ class BaseDataset(Dataset):
         self.max_seq_len = max_seq_len
         self.texts = texts
         self.hf_field = hf_field
-        self.merge_short_from = merge_short_from
 
     def get_tokenized_text(self, idx: int):
         if isinstance(self.texts, list):
@@ -45,6 +44,93 @@ class BaseDataset(Dataset):
 
         return inputs
 
+    @classmethod
+    def from_hf_hub(
+            cls,
+            dataset_id: str,
+            subset: str = None,
+            split: str = 'train',
+            target_field: str = 'text',
+            tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast] = None,
+            tokenizer_hub_id: str = None,
+            max_seq_len: int = 1024,
+            load_kwargs: dict = None,
+            load_tokenizer_kwargs: dict = None,
+            **kwargs
+    ):
+        """
+        Load dataset from HuggingFace Hub and convert it to RxNN training dataset.
+
+        One of the `tokenizer` or `tokenizer_hub_id` args must be provided. If both are provided, `tokenizer` will be used.
+
+        Args:
+            dataset_id (str): Hub dataset repository name
+            subset (str): Dataset subset
+            split (str): Dataset split (default: "train")
+            target_field (str): Name of dataset field used for training (default: "text")
+            tokenizer (PreTrainedTokenizer): HuggingFace Tokenizer used for training (default: None)
+            tokenizer_hub_id (str): HuggingFace Hub ID of tokenizer to load (default: None)
+            max_seq_len (int): Maximum sequence length for training (default: 1024)
+            load_kwargs (dict): Additional args for HuggingFace API load_dataset function
+            load_tokenizer_kwargs (dict): Additional args for loading tokenizer from HuggingFace API with `huggingface_hub.hf_hub_download`
+            **kwargs: Additional args for RxNN Dataset class
+        """
+        assert tokenizer is not None or tokenizer_hub_id is not None, "One of the `tokenizer` or `tokenizer_hub_id` args must be provided."
+
+        if tokenizer is None:
+            tokenizer = load_tokenizer_from_hf_hub(tokenizer_hub_id, **load_tokenizer_kwargs)
+
+        hf_dataset = load_dataset(dataset_id, subset, split=split, **load_kwargs)
+
+        return cls(hf_dataset, tokenizer, max_seq_len=max_seq_len, hf_field=target_field, **kwargs)
+
+    @classmethod
+    def concat_from_hf_hub(
+            cls,
+            dataset_ids: tuple[str],
+            subsets: tuple[str] = None,
+            split: str = 'train',
+            target_field: str = 'text',
+            tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast] = None,
+            tokenizer_hub_id: str = None,
+            max_seq_len: int = 1024,
+            load_kwargs: dict = None,
+            load_tokenizer_kwargs: dict = None,
+            **kwargs
+    ):
+        """
+        Load and concatenate multiple datasets from HuggingFace Hub and convert them to RxNN training dataset.
+        All datasets should use the same split and target field. If it's not the case, just use `load_dataset` and pass the
+        result to RxNN dataset constructor directly.
+
+        One of the `tokenizer` or `tokenizer_hub_id` args must be provided. If both are provided, `tokenizer` will be used.
+
+        Args:
+            dataset_ids (tuple[str]): Hub dataset repository names
+            subsets (tuple[str]): Dataset subsets (default: None)
+            split (str): Dataset split (default: "train")
+            target_field (str): Name of dataset field used for training (default: "text")
+            tokenizer (PreTrainedTokenizer): HuggingFace Tokenizer used for training (default: None)
+            tokenizer_hub_id (str): HuggingFace Hub ID of tokenizer to load (default: None)
+            max_seq_len (int): Maximum sequence length for training (default: 1024)
+            load_kwargs (dict): Additional args for HuggingFace API load_dataset function
+            load_tokenizer_kwargs (dict): Additional args for loading tokenizer from HuggingFace API with `huggingface_hub.hf_hub_download`
+            **kwargs: Additional args for RxNN Dataset class
+        """
+        assert tokenizer is not None or tokenizer_hub_id is not None, "One of the `tokenizer` or `tokenizer_hub_id` args must be provided."
+
+        if tokenizer is None:
+            tokenizer = load_tokenizer_from_hf_hub(tokenizer_hub_id, **load_tokenizer_kwargs)
+
+        hf_datasets = [
+            load_dataset(dataset_id, subset, split=split, **load_kwargs) for dataset_id, subset in zip(dataset_ids, subsets)
+        ] if subsets is not None else [
+            load_dataset(dataset_id, split=split, **load_kwargs) for dataset_id in dataset_ids
+        ]
+        hf_dataset = concatenate_datasets(hf_datasets)
+
+        return cls(hf_dataset, tokenizer, max_seq_len=max_seq_len, hf_field=target_field, **kwargs)
+
 
 class JointLMDataset(BaseDataset):
     def __init__(

rxnn/training/tokenizer.py

@@ -206,3 +206,9 @@ class TokenizerTrainer:
         trainer = cls()
         trainer.load(tokenizer_file)
         return trainer
+
+def load_tokenizer_from_hf_hub(repo_id: str, **kwargs) -> PreTrainedTokenizerFast:
+    return TokenizerTrainer.from_pretrained(repo_id, **kwargs).get_hf_tokenizer()
+
+def load_tokenizer_from_file(path: str) -> PreTrainedTokenizerFast:
+    return TokenizerTrainer.hf_tokenizer_from_file(path)

rxnn/transformers/moe.py

@@ -23,20 +23,21 @@ class MoeRouter(nn.Module):
 
     def calculate_aux_loss(self, top_k_indices: torch.Tensor, probs: torch.Tensor) -> torch.Tensor:
         # Get shapes
-        B, S, K = top_k_indices.shape  # Batch, Sequence length, Top-K
+        T, K = top_k_indices.shape  # Batch, Sequence length, Top-K
 
         # 1. Compute expert selection mask (one-hot encoded)
         expert_mask = F.one_hot(top_k_indices, self.num_experts).float()  # (B, S, K, E)
 
         # 2. Total number of times each expert is selected
-        expert_usage = expert_mask.sum(dim=(0, 1, 2))  # (E,)
+        expert_usage = expert_mask.sum(dim=(0, 1))  # (E,)
 
         # 3. Fraction of tokens assigned to each expert
-        total_tokens = B * S * K
-        fraction_expert = expert_usage / total_tokens  # (E,)
+        total_selections = T * K
+        fraction_expert = expert_usage / (total_selections + 1e-6)  # (E,)
 
         # 4. Sum of probabilities for each expert's selected tokens
-        sum_probs = (probs.unsqueeze(-1) * expert_mask).sum(dim=(0, 1, 2))  # (E,)
+        probs_expanded = probs.unsqueeze(1).expand(-1, K, -1)  # (B_K, K, E)
+        sum_probs = (probs_expanded * expert_mask).sum(dim=(0, 1))
 
         # 5. Average probability per expert (avoid division by zero)
         avg_probs = sum_probs / expert_usage.clamp(min=1e-6)  # (E,)

{rxnn-0.1.52.dist-info → rxnn-0.1.54.dist-info}/METADATA

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

{rxnn-0.1.52.dist-info → rxnn-0.1.54.dist-info}/RECORD

@@ -9,23 +9,23 @@ rxnn/memory/stm.py,sha256=EsD8slSP4_9dLuq6aFPDmuFe8PWilxh90so5Z3nm-ig,2057
 rxnn/rxt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/rxt/models.py,sha256=INTFeNcqzAsjyWhNtbBHL4Tx7tYDsaQHgm72tf6u20M,6918
 rxnn/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/training/base.py,sha256=QD8uS14jSyR5Y_8BgCaBQTKpsarerU3lyufsWsCq_6o,11227
-rxnn/training/bml.py,sha256=HtxSzI-WcpRclAuIccF_WoTZ24KzH5ZWKe8SnWgjjm4,17581
+rxnn/training/base.py,sha256=gEWASLSuWR8UF8b2e-DYqkBZ1lBx0VsIm4kGf9eWSHM,11678
+rxnn/training/bml.py,sha256=S1ZaXTybzeJH7uVFamCr4TPl2bLyZ5xmn_lSsjThTiM,19162
 rxnn/training/callbacks.py,sha256=_YfMKY_eFdc-tubhO9nYH2PXDZDQwlSI74FVOoCXpQg,22108
-rxnn/training/dataset.py,sha256=vQ5mDF3bA0HXya474n4D4iL8Mn3AEpJukgzFNVkxjGU,5106
+rxnn/training/dataset.py,sha256=JQuWSUdT5AnsrG6M_EsewoU6uroVHhg4K715nbtDx8A,9643
 rxnn/training/scheduler.py,sha256=ow6oALzWjWQmHSpcJEjv6tg4g4CDMvr73TypxfcefMc,712
-rxnn/training/tokenizer.py,sha256=4Y41f07uo2KPA_7bp3FCcwGKbXoS2hsckOoXUsXfQxY,8052
+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
 rxnn/transformers/ff.py,sha256=jJnuBDsnnX5uYC_WZH8cXAYrMnz0P-iX7MwcPivjRtI,2533
 rxnn/transformers/layers.py,sha256=OX8CsFY9A7uqH1SLwyexR_5BNlwheYrJHCGXjF8Q7HU,7186
 rxnn/transformers/mask.py,sha256=J0cfLVLt3SzS2ra3KcY4khrkhI975Dw4CjpUi3Sn25s,419
 rxnn/transformers/models.py,sha256=QFzBrOR7tDp9d_T0HoIukBMfEbLxsCictV5p3e2ilxg,7552
-rxnn/transformers/moe.py,sha256=88-w4cQhYNcebdq4zBsdkaoFa4VxJi1LFXDKAAkfVLk,5791
+rxnn/transformers/moe.py,sha256=j6jEx6Ip0zttlUZKKn82azxo95lkLZs-H2GLSMD88hY,5859
 rxnn/transformers/positional.py,sha256=2l38RS0Dini3f6Z3LUHr3XwWzg1UK7fO2C6wazWDAYU,4292
 rxnn/transformers/sampler.py,sha256=poWBpxg1iuK5gEJtxHkk5VVfS9V48hs2Olqdhy_Gw8c,6548
 rxnn/utils.py,sha256=d5U8i5ukovgDyqiycc2AoxObTz_eF_bgo2MKvdtJ98s,467
-rxnn-0.1.52.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.1.52.dist-info/METADATA,sha256=aae9Bt0SpsDgugeHY-7Bi6SN3wWhXneD3Kbz1NMtxJo,16627
-rxnn-0.1.52.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-rxnn-0.1.52.dist-info/RECORD,,
+rxnn-0.1.54.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.1.54.dist-info/METADATA,sha256=FF9XlvOeROGLpVR5pHuuceoeXTzbMNJhEusmQdfPTD0,16627
+rxnn-0.1.54.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+rxnn-0.1.54.dist-info/RECORD,,