chore: import upstream snapshot with attribution

labml_nn/utils/__init__.py

@@ -0,0 +1,65 @@
+"""
+---
+title: Utilities
+summary: A bunch of utility functions and classes
+---
+
+# Utilities
+"""
+
+import copy
+from torch import nn
+from torch.utils.data import Dataset, IterableDataset
+
+
+
+def clone_module_list(module: nn.Module, n: int) -> nn.ModuleList:
+    """
+    ## Clone Module
+
+    Make a `nn.ModuleList` with clones of a given module
+    """
+    return nn.ModuleList([copy.deepcopy(module) for _ in range(n)])
+
+
+def cycle_dataloader(data_loader):
+    """
+    <a id="cycle_dataloader"></a>
+
+    ## Cycle Data Loader
+
+    Infinite loader that recycles the data loader after each epoch
+    """
+    while True:
+        for batch in data_loader:
+            yield batch
+
+
+class MapStyleDataset(Dataset):
+    """
+    <a id="map_style_dataset"></a>
+
+    ## Map Style Dataset
+
+    This converts an [`IterableDataset`](https://pytorch.org/docs/stable/data.html#torch.utils.data.IterableDataset)
+    to a [map-style dataset](https://pytorch.org/docs/stable/data.html#map-style-datasets)
+    so that we can shuffle the dataset.
+
+    *This only works when the dataset size is small and can be held in memory.*
+    """
+
+    def __init__(self, dataset: IterableDataset):
+        # Load the data to memory
+        self.data = [d for d in dataset]
+
+    def __getitem__(self, idx: int):
+        """Get a sample by index"""
+        return self.data[idx]
+
+    def __iter__(self):
+        """Create an iterator"""
+        return iter(self.data)
+
+    def __len__(self):
+        """Size of the dataset"""
+        return len(self.data)

labml_nn/utils/tokenizer.py

@@ -0,0 +1,50 @@
+from typing import Callable
+
+from labml.configs import BaseConfigs, option
+
+
+class TokenizerConfigs(BaseConfigs):
+    """
+    <a id="TokenizerConfigs"></a>
+
+    ## Tokenizer Configurations
+    """
+
+    tokenizer: Callable = 'character'
+
+    def __init__(self):
+        super().__init__(_primary='tokenizer')
+
+
+@option(TokenizerConfigs.tokenizer)
+def basic_english():
+    """
+    ### Basic  english tokenizer
+
+    We use character level tokenizer in this experiment.
+    You can switch by setting,
+
+    ```
+    'tokenizer': 'basic_english'
+    ```
+
+    in the configurations dictionary when starting the experiment.
+
+    """
+    from torchtext.data import get_tokenizer
+    return get_tokenizer('basic_english')
+
+
+def character_tokenizer(x: str):
+    """
+    ### Character level tokenizer
+    """
+    return list(x)
+
+
+@option(TokenizerConfigs.tokenizer)
+def character():
+    """
+    Character level tokenizer configuration
+    """
+    return character_tokenizer