Repository: microsoft/archai (language: python)

Issue:
pr(archai): Implements a faster Hugging Face dataset provider.
# Pull Request Template

## Description

- Implements `FastHfDatasetProvider`, which supports loading, encoding and caching Hugging Face Hub datasets. It uses either memory map or shared memory to encode the data and NumPy arrays to save to cache, which dramatically speeds up the processing time.
- Cached data is not dependent on sequence length and may be re-utilized without having to re-encode.

## Changes

- [X] New feature (non-breaking change which adds functionality)

**Configuration**:
* Archai version: 1.0.0
* Environment: Python 3.7/3.8
* OS: Windows 11

## Final checks:

- [X] Follows guidelines of project
- [X] Self-review of code
- [X] Commented code in hard-to-understand areas
- [X] Corresponding changes to the documentation
- [X] Changes generate no new warnings
- [X] Dependent changes have already been merged
- [X] Checked code correction any misspellings


Discussion / hints:

Fix patch:
diff --git a/.vscode/launch.json b/.vscode/launch.json
index d1dff8a01..1b8058581 100644
--- a/.vscode/launch.json
+++ b/.vscode/launch.json
@@ -209,7 +209,7 @@
             "request": "launch",
             "program": "${cwd}/scripts/supergraph/main.py",
             "console": "integratedTerminal",
-            "args": ["--no-search", "--algos", "manual", "--datasets", "imagenet"]
+            "args": ["--no-search", "--algos", "manual"]
         },
         {
             "name": "Resnet-Full",
diff --git a/archai/common/apex_utils.py b/archai/common/apex_utils.py
index 5a5e79a70..b3a1dd8c8 100644
--- a/archai/common/apex_utils.py
+++ b/archai/common/apex_utils.py
@@ -172,7 +172,7 @@ def _set_ranks(self, conf_gpu_ids:str)->None:
     def is_mixed(self)->bool:
         return self._enabled and self._mixed_prec_enabled
     def is_dist(self)->bool:
-        return self._enabled and self._distributed_enabled
+        return self._enabled and self._distributed_enabled and self.world_size > 1
     def is_master(self)->bool:
         return self.global_rank == 0
     def is_ray(self)->bool:
diff --git a/archai/datasets/nlp/fast_hf_dataset_provider.py b/archai/datasets/nlp/fast_hf_dataset_provider.py
new file mode 100644
index 000000000..08756af12
--- /dev/null
+++ b/archai/datasets/nlp/fast_hf_dataset_provider.py
@@ -0,0 +1,314 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+from __future__ import annotations
+
+import json
+import pickle
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, List, Mapping, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from datasets import load_dataset
+from datasets.dataset_dict import DatasetDict
+from overrides import overrides
+from transformers import AutoTokenizer, DataCollatorForLanguageModeling
+
+from archai.api.dataset_provider import DatasetProvider
+from archai.common.ordered_dict_logger import OrderedDictLogger
+from archai.datasets.nlp.fast_hf_dataset_provider_utils import (
+    FastHfDataset,
+    SHMArray,
+    process_with_memory_map_files,
+    process_with_shared_memory,
+    xor,
+)
+from archai.datasets.nlp.hf_dataset_provider_utils import tokenize_concatenated_dataset
+
+logger = OrderedDictLogger(source=__name__)
+
+if sys.version_info.major == 3 and sys.version_info.minor >= 8:
+    ALLOW_SHARED_MEMORY = True
+else:
+    logger.warn("Shared memory is not available in Python < 3.8.")
+    ALLOW_SHARED_MEMORY = False
+
+
+class FastHfDatasetProvider(DatasetProvider):
+    """Fast Hugging Face-based dataset provider."""
+
+    def __init__(
+        self,
+        train_file: str,
+        validation_file: str,
+        test_file: str,
+        tokenizer: Optional[AutoTokenizer] = None,
+    ) -> None:
+        """Initialize Fast Hugging Face-based dataset provider.
+
+        Args:
+            train_file: Path to the training array file (.npy).
+            validation_file: Path to the validation array file (.npy).
+            test_file: Path to the test array file (.npy).
+            tokenizer: Instance of tokenizer to use.
+
+        """
+
+        super().__init__()
+
+        self.train_file = train_file
+        self.validation_file = validation_file
+        self.test_file = test_file
+        self.tokenizer = tokenizer
+
+    @staticmethod
+    def _encode_dataset(
+        dataset_dict: DatasetDict,
+        tokenizer: AutoTokenizer,
+        mapping_column_name: List[str],
+        use_eos_token: bool,
+        dtype: np.dtype,
+        num_workers: int,
+    ) -> DatasetDict:
+        logger.info("Encoding dataset ...")
+        logger.info(f"Number of workers: {num_workers} | EOS token: {use_eos_token}")
+
+        column_names = dataset_dict["train"].column_names
+        encoded_dataset_dict = dataset_dict.map(
+            tokenize_concatenated_dataset,
+            fn_kwargs={
+                "tokenizer": tokenizer,
+                "mapping_column_name": mapping_column_name,
+                "use_eos_token": use_eos_token,
+                "dtype": dtype,
+            },
+            batched=True,
+            num_proc=num_workers,
+            remove_columns=column_names,
+        )
+
+        return encoded_dataset_dict
+
+    @staticmethod
+    def _process_dataset_to_memory(
+        dataset_dict: DatasetDict, cache_dir: str, dtype: np.dtype, num_workers: int, use_shared_memory: int
+    ) -> Dict[str, Union[SHMArray, np.ndarray]]:
+        logger.info("Processing dataset to memory ...")
+        logger.info(f"Number of workers: {num_workers} | Shared memory: {use_shared_memory}")
+
+        if use_shared_memory:
+            return process_with_shared_memory(dataset_dict, dtype, num_proc=num_workers)
+
+        return process_with_memory_map_files(dataset_dict, cache_dir, dtype, num_proc=num_workers)
+
+    @staticmethod
+    def _save_dataset(
+        dataset_dict: Dict[str, Union[SHMArray, np.ndarray]],
+        tokenizer: AutoTokenizer,
+        cache_dir: str,
+        use_shared_memory: bool,
+    ) -> Tuple[Path, Path, Path]:
+        logger.info(f"Saving dataset to: {cache_dir}")
+
+        cache_files = ()
+        for split, dataset in dataset_dict.items():
+            np.save(cache_dir / f"{split}.npy", dataset)
+
+            # If using shared memory, dataset needs to have its shared memory
+            # unlinked to prevent memory leak
+            if use_shared_memory:
+                dataset.shm.unlink()
+
+            # If not using shared memory, dataset needs to have its memory map
+            # closed to prevent an additional .bin file
+            if not use_shared_memory:
+                dataset._mmap.close()
+                Path(cache_dir / f"{split}.bin").unlink()
+
+            cache_files += (cache_dir / f"{split}.npy",)
+
+        with open(cache_dir / "tokenizer.pkl", "wb") as f:
+            pickle.dump(tokenizer, f)
+
+        return cache_files
+
+    @classmethod
+    def from_hub(
+        cls: FastHfDatasetProvider,
+        dataset_name: str,
+        dataset_config_name: Optional[str] = None,
+        data_dir: Optional[str] = None,
+        tokenizer: Optional[AutoTokenizer] = None,
+        tokenizer_name: Optional[str] = None,
+        tokenizer_max_length: Optional[int] = None,
+        mapping_column_name: Optional[List[str]] = None,
+        validation_split: Optional[float] = 0.0,
+        seed: Optional[int] = 42,
+        num_workers: Optional[int] = 1,
+        use_eos_token: Optional[bool] = True,
+        use_shared_memory: Optional[bool] = True,
+        cache_dir: Optional[str] = "cache",
+    ) -> FastHfDatasetProvider:
+        """Load a dataset provider by downloading and encoding data from Hugging Face Hub.
+
+        Args:
+            dataset_name: Name of the dataset.
+            dataset_config_name: Name of the dataset configuration.
+            data_dir: Path to the data directory.
+            tokenizer: Instance of tokenizer to use.
+            tokenizer_name: Name of the tokenizer, if `tokenizer` has not been passed.
+            tokenizer_max_length: Maximum length of the tokenized sequences.
+            mapping_column_name: The columns in `dataset` that should be tokenized.
+            validation_split: Fraction of the dataset to use for validation.
+            seed: Random seed.
+            num_workers: Number of workers to use for encoding.
+            use_eos_token: Whether to use EOS token to separate sequences.
+            use_shared_memory: Whether to use shared memory for caching.
+            cache_dir: Root path to the cache directory.
+
+        Returns:
+            Dataset provider.
+
+        """
+
+        assert xor(tokenizer, tokenizer_name), "`tokenizer` and `tokenizer_name` are mutually exclusive."
+        tokenizer = tokenizer or AutoTokenizer.from_pretrained(tokenizer_name)
+        if tokenizer_max_length:
+            logger.warn(f"New maximum length set for the tokenizer: {tokenizer_max_length}.")
+            tokenizer.model_max_length = tokenizer_max_length
+
+        dtype = np.uint16 if tokenizer.vocab_size < 64 * 1024 else np.int32
+        use_shared_memory = use_shared_memory and ALLOW_SHARED_MEMORY
+
+        cache_dir = Path(cache_dir)
+        if cache_dir.is_dir():
+            logger.warn(f"Cache: {cache_dir} already exists and will be overritten.")
+        cache_dir.mkdir(parents=True, exist_ok=True)
+
+        # Ensure that the downloaded dataset is always a dictionary
+        logger.info("Downloading dataset ...")
+        hub_dataset_dict = load_dataset(dataset_name, name=dataset_config_name, data_dir=data_dir)
+        if not isinstance(hub_dataset_dict, DatasetDict):
+            hub_dataset_dict = DatasetDict({"train": hub_dataset_dict})
+
+        # Ensure that `validation` and `test` splits are available
+        if "validation" not in hub_dataset_dict:
+            logger.info("Creating validation split ...")
+
+            validation_split = validation_split or 0.1
+            tmp_dataset_dict = hub_dataset_dict["train"].train_test_split(
+                test_size=validation_split, shuffle=True, seed=seed
+            )
+            hub_dataset_dict["train"] = tmp_dataset_dict["train"]
+            hub_dataset_dict["validation"] = tmp_dataset_dict["test"]
+
+        if "test" not in hub_dataset_dict:
+            logger.info("Creating test split ...")
+
+            tmp_dataset_dict = hub_dataset_dict["validation"].train_test_split(test_size=0.25, shuffle=True, seed=seed)
+            hub_dataset_dict["validation"] = tmp_dataset_dict["train"]
+            hub_dataset_dict["test"] = tmp_dataset_dict["test"]
+
+        encoded_dataset_dict = FastHfDatasetProvider._encode_dataset(
+            hub_dataset_dict, tokenizer, mapping_column_name, use_eos_token, dtype, num_workers
+        )
+        processed_dataset_dict = FastHfDatasetProvider._process_dataset_to_memory(
+            encoded_dataset_dict, cache_dir, dtype, num_workers, use_shared_memory
+        )
+
+        cache_files = FastHfDatasetProvider._save_dataset(
+            processed_dataset_dict, tokenizer, cache_dir, use_shared_memory
+        )
+
+        with open(cache_dir / "config.json", "w") as f:
+            json.dump(
+                {
+                    "dataset_name": dataset_name,
+                    "dataset_config_name": dataset_config_name,
+                    "data_dir": data_dir,
+                    "tokenizer": {
+                        "name_or_path": tokenizer.name_or_path,
+                        "model_max_length": tokenizer.model_max_length,
+                    },
+                    "mapping_column_name": mapping_column_name or ["text"],
+                    "validation_split": validation_split,
+                    "seed": seed,
+                    "use_eos_token": use_eos_token,
+                },
+                f,
+            )
+
+        return FastHfDatasetProvider(*cache_files, tokenizer=tokenizer)
+
+    @classmethod
+    def from_cache(cls: FastHfDatasetProvider, cache_dir: str) -> FastHfDatasetProvider:
+        """Load a dataset provider from a cache directory.
+
+        Args:
+            cache_dir: Path to the cache directory.
+
+        Returns:
+            Dataset provider.
+
+        """
+
+        logger.info(f"Loading dataset from: {cache_dir}")
+
+        cache_dir = Path(cache_dir)
+        cache_train_file = cache_dir / "train.npy"
+        cache_validation_file = cache_dir / "validation.npy"
+        cache_test_file = cache_dir / "test.npy"
+
+        tokenizer_file = cache_dir / "tokenizer.pkl"
+        if not tokenizer_file.is_file():
+            logger.warn(f"Could not find tokenizer in {cache_dir}.")
+            tokenizer = None
+        else:
+            with open(tokenizer_file, "rb") as f:
+                tokenizer = pickle.load(f)
+
+        return FastHfDatasetProvider(cache_train_file, cache_validation_file, cache_test_file, tokenizer=tokenizer)
+
+    @overrides
+    def get_train_dataset(self, seq_len: Optional[int] = 1) -> FastHfDataset:
+        input_ids = np.load(self.train_file, mmap_mode="r")
+
+        return FastHfDataset(input_ids, seq_len=seq_len)
+
+    @overrides
+    def get_val_dataset(self, seq_len: Optional[int] = 1) -> FastHfDataset:
+        input_ids = np.load(self.validation_file, mmap_mode="r")
+
+        return FastHfDataset(input_ids, seq_len=seq_len)
+
+    @overrides
+    def get_test_dataset(self, seq_len: Optional[int] = 1) -> FastHfDataset:
+        input_ids = np.load(self.test_file, mmap_mode="r")
+
+        return FastHfDataset(input_ids, seq_len=seq_len)
+
+
+@dataclass
+class FastDataCollatorForLanguageModeling(DataCollatorForLanguageModeling):
+    """Language modeling data collator compatible with FastHfDataset.
+
+    Args:
+        use_shifted_labels: Whether to use the original labels (shifted) or the non-shifted labels.
+
+    """
+
+    use_shifted_labels: bool = False
+
+    def torch_call(self, examples: List[Union[List[int], Any, Dict[str, Any]]]) -> Dict[str, Any]:
+        if isinstance(examples[0], Mapping):
+            return super().torch_call(examples)
+
+        batch = super().torch_call([example[0] for example in examples])
+
+        if self.use_shifted_labels:
+            batch["labels"] = torch.stack([example[1] for example in examples], dim=0)
+
+        return batch
diff --git a/archai/datasets/nlp/fast_hf_dataset_provider_utils.py b/archai/datasets/nlp/fast_hf_dataset_provider_utils.py
new file mode 100644
index 000000000..23a810c4e
--- /dev/null
+++ b/archai/datasets/nlp/fast_hf_dataset_provider_utils.py
@@ -0,0 +1,188 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+#
+# Copyright (c) Hazy Research.
+# Licensed under the BSD-3-Clause license.
+# https://github.com/HazyResearch/flash-attention/blob/main/training/src/datamodules
+
+from __future__ import annotations
+
+import math
+import mmap
+import sys
+from typing import Any, Dict, Optional, Tuple
+
+import numpy as np
+import torch
+from datasets.dataset_dict import DatasetDict
+from torch.utils.data import Dataset
+
+# `multiprocessing.shared_memory` is only available in Python 3.8+`
+if sys.version_info.major == 3 and sys.version_info.minor >= 8:
+    from multiprocessing.shared_memory import SharedMemory
+
+
+class FastHfDataset(Dataset):
+    """Fast Hugging Face dataset."""
+
+    def __init__(self, input_ids: torch.Tensor, seq_len: Optional[int] = 1) -> None:
+        """Initialize the dataset.
+
+        Args:
+            input_ids: Tensor with the inputs (encoded data).
+            seq_len: Sequence length.
+
+        """
+
+        super().__init__()
+
+        self.n_input_ids = ((len(input_ids) - 1) // seq_len) * seq_len + 1
+        self.seq_len = seq_len
+
+        # `input_ids` should not be sliced since they could be memory mapped
+        self.input_ids = input_ids
+        self.n_sequences = math.ceil((self.n_input_ids - 1) / self.seq_len)
+
+    def __len__(self) -> int:
+        return self.n_sequences
+
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        start_idx = idx * self.seq_len
+        seq_len = min(self.seq_len, self.n_input_ids - 1 - start_idx)
+
+        input_ids = torch.as_tensor(self.input_ids[start_idx : (start_idx + seq_len + 1)].astype(np.int64))
+        labels = input_ids[1:].clone()
+
+        return input_ids[:-1], labels
+
+
+class SHMArray(np.ndarray):
+    """Numpy array compatible with SharedMemory from `multiprocessing.shared_memory`.
+
+    Reference:
+        https://numpy.org/doc/stable/user/basics.subclassing.html#slightly-more-realistic-example-attribute-added-to-existing-array
+
+    """
+
+    def __new__(cls: SHMArray, input_array: np.ndarray, shm: Optional[SharedMemory] = None) -> SHMArray:
+        obj = np.asarray(input_array).view(cls)
+        obj.shm = shm
+
+        return obj
+
+    def __array_finalize__(self, obj: SHMArray) -> None:
+        if obj is None:
+            return
+
+        self.shm = getattr(obj, "shm", None)
+
+
+def process_with_shared_memory(
+    dataset_dict: DatasetDict, dtype: np.dtype, num_proc: Optional[int] = 1
+) -> Dict[str, SHMArray]:
+    """Process the dataset with a shared memory.
+
+    Args:
+        dataset_dict: Dataset dictionary.
+        dtype: Numpy data type.
+        num_proc: Number of processes.
+
+    Returns:
+        Dictionary with shared memory-processed datasets.
+
+    """
+
+    def _process_with_shared_memory(example: Dict[str, Any], name, length: int) -> None:
+        shared_memory = SharedMemory(name=name)
+
+        shared_memory_array = np.ndarray((length,), dtype=dtype, buffer=shared_memory.buf)
+        start_idx = example["offset"] - len(example["input_ids"])
+        shared_memory_array[start_idx : example["offset"]] = example["input_ids"]
+
+        shared_memory.close()
+
+    processed_dataset_dict = {}
+    for name, ds in dataset_dict.items():
+        dataset_dict[name] = ds.add_column("offset", np.cumsum(ds["length"]))
+        length = dataset_dict[name][-1]["offset"]
+
+        shared_memory = SharedMemory(create=True, size=length * np.dtype(dtype).itemsize)
+        shared_memory_name = shared_memory.name
+
+        dataset_dict[name].map(
+            _process_with_shared_memory,
+            fn_kwargs={"name": shared_memory_name, "length": length},
+            batched=False,
+            num_proc=num_proc,
+        )
+
+        shared_memory_array = np.ndarray((length,), dtype=dtype, buffer=shared_memory.buf)
+        processed_dataset_dict[name] = SHMArray(shared_memory_array, shm=shared_memory)
+
+    return processed_dataset_dict
+
+
+def process_with_memory_map_files(
+    dataset_dict: DatasetDict, cache_dir: str, dtype: np.dtype, num_proc: Optional[int] = 1
+) -> Dict[str, np.ndarray]:
+    """Process the dataset with memory map files.
+
+    Args:
+        dataset_dict: Dataset dictionary.
+        cache_dir: Cache directory.
+        dtype: Numpy data type.
+        num_proc: Number of processes.
+
+    Returns:
+        Dictionary with memory map file-processed datasets.
+
+    """
+
+    def _process_with_memory_map_files(example: Dict[str, Any], file_path: str) -> None:
+        with open(file_path, "r+b") as f:
+            memory_map = mmap.mmap(f.fileno(), 0)
+
+            start_idx = example["offset"] - len(example["input_ids"])
+            length = len(example["input_ids"])
+
+            memory_map_array = np.ndarray(
+                (length,), dtype=dtype, buffer=memory_map, offset=np.dtype(dtype).itemsize * start_idx
+            )
+            memory_map_array[:] = example["input_ids"]
+
+            memory_map.flush()
+
+    processed_dataset_dict = {}
+    for split, dataset in dataset_dict.items():
+        dataset_dict[split] = dataset.add_column("offset", np.cumsum(dataset["length"]))
+        length = dataset_dict[split][-1]["offset"]
+
+        file_path = cache_dir / f"{split}.bin"
+        with open(file_path.as_posix(), "wb") as f:
+            f.truncate(length * np.dtype(dtype).itemsize)
+
+        dataset_dict[split].map(
+            _process_with_memory_map_files,
+            fn_kwargs={"file_path": file_path},
+            batched=False,
+            num_proc=num_proc,
+        )
+
+        processed_dataset_dict[split] = np.memmap(file_path, dtype=dtype, mode="r", shape=(length,))
+
+    return processed_dataset_dict
+
+
+def xor(p: Any, q: Any) -> bool:
+    """Implements the logical XOR operator.
+
+    Args:
+        p: Any instance that may act as `True` or `False`.
+        q: Any instance that may act as `True` or `False`.
+
+    Returns:
+        Logical value.
+
+    """
+
+    return (p and not q) or (not p and q)
diff --git a/archai/datasets/nlp/hf_dataset_provider.py b/archai/datasets/nlp/hf_dataset_provider.py
index 60c1d77fc..b60b78cb6 100644
--- a/archai/datasets/nlp/hf_dataset_provider.py
+++ b/archai/datasets/nlp/hf_dataset_provider.py
@@ -24,8 +24,8 @@ class HfHubDatasetProvider(DatasetProvider):
 
     def __init__(
         self,
-        dataset: Optional[str] = "wikitext",
-        subset: Optional[str] = "wikitext-2-raw-v1",
+        dataset_name: str,
+        dataset_config_name: Optional[str] = None,
         data_dir: Optional[str] = None,
         data_files: Optional[Union[str, List[str], Dict[str, Union[str, List[str]]]]] = None,
         cache_dir: Optional[str] = None,
@@ -34,8 +34,8 @@ def __init__(
         """Initialize Hugging Face Hub dataset provider.
 
         Args:
-            dataset: Name of the dataset.
-            subset: Name of the dataset configuration.
+            dataset_name: Name of the dataset.
+            dataset_config_name: Name of the dataset configuration.
             data_dir: Path to the data directory.
             data_files: Path(s) to the data file(s).
             cache_dir: Path to the read/write cache directory.
@@ -45,8 +45,8 @@ def __init__(
 
         super().__init__()
 
-        self.dataset = dataset
-        self.subset = subset
+        self.dataset_name = dataset_name
+        self.dataset_config_name = dataset_config_name
         self.data_dir = data_dir
         self.data_files = data_files
         self.cache_dir = cache_dir
@@ -60,8 +60,8 @@ def get_dataset(
         streaming: Optional[bool] = False,
     ) -> Union[Dataset, DatasetDict, IterableDataset, IterableDatasetDict]:
         return hf_load_dataset(
-            self.dataset,
-            name=self.subset,
+            self.dataset_name,
+            name=self.dataset_config_name,
             data_dir=self.data_dir,
             data_files=self.data_files,
             split=split,
diff --git a/archai/datasets/nlp/hf_dataset_provider_utils.py b/archai/datasets/nlp/hf_dataset_provider_utils.py
index 6ab7f5b35..ca359ca7c 100644
--- a/archai/datasets/nlp/hf_dataset_provider_utils.py
+++ b/archai/datasets/nlp/hf_dataset_provider_utils.py
@@ -5,6 +5,7 @@
 from itertools import chain
 from typing import Any, Callable, Dict, List, Optional, Union
 
+import numpy as np
 from datasets.arrow_dataset import Dataset
 from datasets.dataset_dict import DatasetDict, IterableDatasetDict
 from datasets.download.download_manager import DownloadMode
@@ -35,12 +36,12 @@ def should_refresh_cache(refresh: bool) -> DownloadMode:
 
 
 def tokenize_dataset(
-    examples: List[str],
+    examples: Dict[str, List[str]],
     tokenizer: Optional[AutoTokenizer] = None,
     mapping_column_name: Optional[List[str]] = None,
+    use_eos_token: Optional[bool] = False,
     truncate: Optional[Union[bool, str]] = True,
     padding: Optional[Union[bool, str]] = "max_length",
-    **kwargs,
 ) -> Dict[str, Any]:
     """Tokenize a list of examples using a specified tokenizer.
 
@@ -48,6 +49,7 @@ def tokenize_dataset(
         examples: A list of examples to be tokenized.
         tokenizer: The tokenizer to use.
         mapping_column_name: The columns in `examples` that should be tokenized.
+        use_eos_token: Whether to append the EOS token to each example.
         truncate: Whether truncation should be applied.
         padding: Whether padding should be applied.
 
@@ -56,20 +58,60 @@ def tokenize_dataset(
 
     """
 
+    def _add_eos_token(examples: List[str]) -> List[str]:
+        return [example + tokenizer.eos_token if example else example for example in examples]
+
     if mapping_column_name is None:
         mapping_column_name = ["text"]
 
-    examples_mapping = tuple(examples[column_name] for column_name in mapping_column_name)
+    examples_mapping = tuple(
+        _add_eos_token(examples[column_name]) if use_eos_token else examples[column_name]
+        for column_name in mapping_column_name
+    )
 
     return tokenizer(*examples_mapping, truncation=truncate, padding=padding)
 
 
+def tokenize_concatenated_dataset(
+    examples: Dict[str, List[str]],
+    tokenizer: Optional[AutoTokenizer] = None,
+    mapping_column_name: Optional[List[str]] = None,
+    use_eos_token: Optional[bool] = False,
+    dtype: Optional[np.dtype] = None,
+) -> Dict[str, Any]:
+    """Tokenize a list of examples using a specified tokenizer and
+    with concatenated batches (no truncation nor padding).
+
+    Args:
+        examples: A list of examples to be tokenized.
+        tokenizer: The tokenizer to use.
+        mapping_column_name: The columns in `examples` that should be tokenized.
+        use_eos_token: Whether to append the EOS token to each example.
+        dtype: Numpy data type of the tokenized examples.
+
+    Returns:
+        Concatenated tokenized examples.
+
+    """
+
+    examples = tokenize_dataset(
+        examples,
+        tokenizer=tokenizer,
+        mapping_column_name=mapping_column_name,
+        use_eos_token=use_eos_token,
+        truncate=False,
+        padding=False,
+    )
+    tokenized_examples = np.fromiter(chain(*examples["input_ids"]), dtype=dtype)
+
+    return {"input_ids": [tokenized_examples], "length": [len(tokenized_examples)]}
+
+
 def tokenize_contiguous_dataset(
-    examples: List[str],
+    examples: Dict[str, List[str]],
     tokenizer: Optional[AutoTokenizer] = None,
     mapping_column_name: Optional[List[str]] = None,
     model_max_length: Optional[int] = 1024,
-    **kwargs,
 ) -> Dict[str, Any]:
     """Tokenize a list of examples using a specified tokenizer and
     with contiguous-length batches (no truncation nor padding).
@@ -104,12 +146,11 @@ def tokenize_contiguous_dataset(
 
 
 def tokenize_nsp_dataset(
-    examples: List[str],
+    examples: Dict[str, List[str]],
     tokenizer: Optional[AutoTokenizer] = None,
     mapping_column_name: Optional[List[str]] = None,
     truncate: Optional[Union[bool, str]] = True,
     padding: Optional[Union[bool, str]] = "max_length",
-    **kwargs,
 ) -> Dict[str, Any]:
     """Tokenize a list of examples using a specified tokenizer and
     with next-sentence prediction (NSP).
diff --git a/archai/supergraph/datasets/providers/imagenet_provider.py b/archai/supergraph/datasets/providers/imagenet_provider.py
index 64fb66ebb..1bb3b0345 100644
--- a/archai/supergraph/datasets/providers/imagenet_provider.py
+++ b/archai/supergraph/datasets/providers/imagenet_provider.py
@@ -4,7 +4,6 @@
 import os
 
 from overrides import overrides
-from PIL import Image
 from torchvision import datasets
 from torchvision.transforms import transforms
 
@@ -59,7 +58,7 @@ def get_transforms(self, img_size:ImgSize)->tuple:
         transform_train = transforms.Compose([
             transforms.RandomResizedCrop(224,
                 scale=(0.08, 1.0), # TODO: these two params are normally not specified
-                interpolation=Image.BICUBIC),
+                interpolation=transforms.InterpolationMode.BICUBIC),
             transforms.RandomHorizontalFlip(),
             transforms.ColorJitter(
                 brightness=0.4,
diff --git a/archai/supergraph/utils/trainer.py b/archai/supergraph/utils/trainer.py
index d7ff8be82..d77975332 100644
--- a/archai/supergraph/utils/trainer.py
+++ b/archai/supergraph/utils/trainer.py
@@ -294,7 +294,8 @@ def _train_epoch(self, train_dl: DataLoader)->None:
 
                 loss_sum += loss_c.item() * len(logits_c)
                 loss_count += len(logits_c)
-                logits_chunks.append(logits_c.detach().cpu())               # pyright: ignore[reportGeneralTypeIssues]
+                # TODO: cannot place on CPU if it was half precision but should we somehow?
+                logits_chunks.append(logits_c.detach())               # pyright: ignore[reportGeneralTypeIssues]
 
             # TODO: original darts clips alphas as well but pt.darts doesn't
             self._apex.clip_grad(self._grad_clip, self.model, self._multi_optim)
@@ -304,7 +305,8 @@ def _train_epoch(self, train_dl: DataLoader)->None:
             # TODO: we possibly need to sync so all replicas are upto date
             self._apex.sync_devices()
 
-            self.post_step(x, y,
+            # TODO: we need to put y on GPU because logits are on GPU. Is this good idea from GPU mem perspective?
+            self.post_step(x, y.to(self.get_device(), non_blocking=True),
                            ml_utils.join_chunks(logits_chunks),
                            torch.tensor(loss_sum/loss_count),
                            steps)
diff --git a/docs/getting_started/notebooks/nlp.rst b/docs/getting_started/notebooks/nlp.rst
index 23d2a750e..fc6170be3 100644
--- a/docs/getting_started/notebooks/nlp.rst
+++ b/docs/getting_started/notebooks/nlp.rst
@@ -4,6 +4,7 @@ Natural Language Processing
 .. toctree::
    :maxdepth: 2
 
+   Fast HF Dataset Provider <nlp/fast_hf_dataset_provider.ipynb>
    HF Dataset Provider <nlp/hf_dataset_provider.ipynb>
    HF Trainer <nlp/hf_trainer.ipynb>
    NVIDIA Dataset Provider <nlp/nvidia_dataset_provider.ipynb>
diff --git a/docs/getting_started/notebooks/nlp/.gitignore b/docs/getting_started/notebooks/nlp/.gitignore
index fd22d4310..cffd14f91 100644
--- a/docs/getting_started/notebooks/nlp/.gitignore
+++ b/docs/getting_started/notebooks/nlp/.gitignore
@@ -2,7 +2,6 @@
 # Licensed under the MIT license.
 
 # Ignore notebook output
-hf-codegen/
-
-# Ignore local data created by `NvidiaDatasetProvider`
+cache/
 dataroot/
+hf-codegen/
diff --git a/docs/getting_started/notebooks/nlp/fast_hf_dataset_provider.ipynb b/docs/getting_started/notebooks/nlp/fast_hf_dataset_provider.ipynb
new file mode 100644
index 000000000..f9709a79c
--- /dev/null
+++ b/docs/getting_started/notebooks/nlp/fast_hf_dataset_provider.ipynb
@@ -0,0 +1,636 @@
+{
+ "cells": [
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Rapid NLP-based Data Creation/Encoding using Memory Mapping\n",
+    "\n",
+    "In this notebook, we will use a fast dataset provider-based abstraction that interfaces with Hugging Face's `datasets` (and has been created by HazyResearch). The key advantage of this approach is the use of either shared memory or memory maps in Python to accelerate the caching process. Furthermore, the dataset is cached as a contiguous numpy array, enabling manipulation of data with any sequence length. This feature eliminates the need for re-encoding data for multiple lengths, streamlining the data processing pipeline."
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Instantiating the Provider\n",
+    "\n",
+    "The first step is to instantiate the `FastHfDatasetProvider`, which loads and encodes the dataset. A set of arguments can be passed to its constructor according to the user's needs:\n",
+    "\n",
+    "* `dataset_name`: Name of the dataset.\n",
+    "* `dataset_config_name`: Name of the dataset configuration.\n",
+    "* `data_dir`: Path to the data directory.\n",
+    "* `tokenizer`: Instance of tokenizer to use.\n",
+    "* `tokenizer_name`: Name of the tokenizer, if `tokenizer` has not been passed.\n",
+    "* `tokenizer_max_length`: Maximum length of the tokenized sequences.\n",
+    "* `mapping_column_name`: The columns in `dataset` that should be tokenized.\n",
+    "* `validation_split`: Fraction of the dataset to use for validation.\n",
+    "* `seed`: Random seed.\n",
+    "* `num_workers`: Number of workers to use for encoding.\n",
+    "* `use_eos_token`: Whether to use EOS token to separate sequences.\n",
+    "* `use_shared_memory`: Whether to use shared memory for caching.\n",
+    "* `cache_dir`: Root path to the cache directory."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2023-02-16 14:19:39,981 - archai.datasets.nlp.fast_hf_dataset_provider — WARNING —  Shared memory is not available in Python < 3.8.\n",
+      "2023-02-16 14:19:42,645 - archai.datasets.nlp.fast_hf_dataset_provider — WARNING —  New maximum length set for the tokenizer: 512.\n",
+      "2023-02-16 14:19:42,649 - archai.datasets.nlp.fast_hf_dataset_provider — INFO —  Downloading dataset ...\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Found cached dataset glue (C:/Users/gderosa/.cache/huggingface/datasets/glue/sst2/1.0.0/dacbe3125aa31d7f70367a07a8a9e72a5a0bfeb5fc42e75c9db75b96da6053ad)\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "d8ef210940ce4773a1675389b5061a9c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/3 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2023-02-16 14:19:46,295 - archai.datasets.nlp.fast_hf_dataset_provider — INFO —  Encoding dataset ...\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Loading cached processed dataset at C:\\Users\\gderosa\\.cache\\huggingface\\datasets\\glue\\sst2\\1.0.0\\dacbe3125aa31d7f70367a07a8a9e72a5a0bfeb5fc42e75c9db75b96da6053ad\\cache-4b268755324077d2.arrow\n",
+      "Loading cached processed dataset at C:\\Users\\gderosa\\.cache\\huggingface\\datasets\\glue\\sst2\\1.0.0\\dacbe3125aa31d7f70367a07a8a9e72a5a0bfeb5fc42e75c9db75b96da6053ad\\cache-3313fdb7893e6922.arrow\n",
+      "Loading cached processed dataset at C:\\Users\\gderosa\\.cache\\huggingface\\datasets\\glue\\sst2\\1.0.0\\dacbe3125aa31d7f70367a07a8a9e72a5a0bfeb5fc42e75c9db75b96da6053ad\\cache-252c3cce9291d908.arrow\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "ef674e6378e8470ead207f24e0dfabed",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/68 [00:00<?, ?ex/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f977234e870643c9bce445768f2dbaa7",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/1 [00:00<?, ?ex/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "b5c7bec1080c4009baf45a6f7f2f265e",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/2 [00:00<?, ?ex/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2023-02-16 14:19:47,798 - archai.datasets.nlp.fast_hf_dataset_provider — INFO —  Saving dataset to: cache\\glue\\sst2\\3ba1dcb015d7664117d16b85af90496896fc74ef\n",
+      "(tensor([24717,   649,  3200,   507,   422,   262, 21694,  4991,   220, 50256,\n",
+      "         3642,  1299,   645, 20868,   837,   691,  2248,  1850,   308,  3775,\n",
+      "          220, 50256,  5562, 10408,   663,  3435,   290, 48556,  1223,  2138,\n",
+      "         4950,   546,  1692,  3450,   220, 50256,  2787,  1299, 15950, 11378,\n",
+      "          284,  3520,   262,   976,  3690,   220, 50256,   261,   262,  5290,\n",
+      "        15827,    12,  1659,    12,  1169,    12,  1008,  9310, 35478, 20954,\n",
+      "          262, 28303,   714, 47478,   469,   510,   220, 50256,  5562,   705,\n",
+      "           82,  1290,  1165, 15444,   284, 17004,   884, 31194,  3513,   220,\n",
+      "        50256, 26567,  2536,   689,   326,   262,  3437,   286,   884,   289,\n",
+      "        31777,  2512, 30181,   355, 29408,  1830,   460,   991,  1210,   503,\n",
+      "          257,  1402,   837,  2614,  2646,   351,   281,  7016,  3355,   404,\n",
+      "          764,   220, 50256,  1659,   473,    84,   948,   220, 50256,    64,\n",
+      "        19095, 17280,    12,  1941,    12,   727,   705,    82, 26781, 19518,\n",
+      "          220, 50256,   533,   517,  7744,  1807,   832,   621,   287,   749,\n",
+      "         4600,   826,    12, 28973,   705,  7328,   220, 50256,  2188,   274,\n",
+      "          284, 12986, 20428,   220, 50256,  1640,   883,  3807, 31006,   508,\n",
+      "        13121,   326,  4600,   484,   466,   299,   470,   787,  6918,   588,\n",
+      "          484,   973,   284,  7471,   220, 50256,  1169,   636,   810,  2147,\n",
+      "          705,    82,  5836,   837,   220, 50256, 43439,   703,  2089,   428,\n",
+      "         3807,   373,   220, 50256,    75,   437,   617, 16247,   284,   257,\n",
+      "        13526,  1621,   220, 50256,  1169,  6000, 17245,   220, 50256, 36673,\n",
+      "         3807,   220, 50256,  4480,   465,  6678,  4430,   290, 11800,   774,\n",
+      "          220, 50256,   445,   917,   415,  3721,   220, 50256,  2032, 27428,\n",
+      "          318,  2029,   477,   546,   257,  1862,  2415,   705,    82,  1986,\n",
+      "          837,   290,   416, 13092,   281, 14549,  3025,  1986,  4493,   326,\n",
+      "         2415,   705,    82, 17188,   290,   614, 23400,   837,   340, 31137,\n",
+      "          764,   220, 50256,  4853,   874,   262,  2656,   290,   287,   617,\n",
+      "         2842,   772,   731,  1010,   340,   220, 50256,   361,  1997,   837,\n",
+      "          766,   340,   329,   479,  5757,  2042,   837,   508, 11665,   510,\n",
+      "          257,  6388,   355,   257, 26821, 14314, 10086, 46410,  3706, 11841,\n",
+      "        19317,   764,   220, 50256,    64,  8212,   319,   534,  1986,   220,\n",
+      "        50256,  8988,   422,   262, 14802,   837, 26329, 44139, 13289,   220,\n",
+      "        50256,  1069, 32838, 46136,   306,  3684, 16948,   290,  6028, 17049,\n",
+      "          555,   398,  5109,   220, 50256,   268, 30486,   416,   281, 25007,\n",
+      "         9404,  7668,  3350,   286, 47792, 13747,   220, 50256,  4758,  2063,\n",
+      "          286, 10441, 12254,   318,  4785,  1058,   262,   636,   810,  2147,\n",
+      "          705,    82,  5836,   837,   393,   262,   636,   810,  1223,   705,\n",
+      "           82,  5836,   220, 50256,   259,   995, 22041,   220, 50256,   548,\n",
+      "          922, 11681,  5559,   220, 50256,  1169,  7110,   318,  2147,   475,\n",
+      "        36741,  6816, 35478, 20954,   422,   923,   284,  5461,   837,   220,\n",
+      "        50256,  1169,  2223,   318,   336,   346,  1513,   220, 50256,   261,\n",
+      "          477, 43386,   220, 50256, 10594,  1064,  1310,   286,  1393,   287,\n",
+      "          428,  2646,   837,   543,   318,  1690, 32950,    88,   290, 13455,\n",
+      "        13134,   220, 50256,  1525,  1290,   262,  5290,  3807,   286,   262,\n",
+      "          614,   220, 50256, 48937,   832,   837,   220, 50256,  3549,   621,\n",
+      "         1194,  7559,  1266,   582, 10148, 17271,   416, 44889,   257,  7505,\n",
+      "         3690,   428,  8258,  2646,   220, 50256,   270,   705,    82,   546,\n",
+      "         2428,   749,  6490,   423,   284,  1986,   287,  4845,   290,  1312,\n",
+      "          892,   326,   705,    82,   644,  1312,  8288,   546,   340,  1377,\n",
+      "          262,  1103,  2428, 29779,  1022,   262, 14397,   290, 14897, 22992,\n",
+      "          220, 50256, 11718,   274,   220, 50256,   672, 35260,   284,   262,\n",
+      "         6224,   286]), tensor([  649,  3200,   507,   422,   262, 21694,  4991,   220, 50256,  3642,\n",
+      "         1299,   645, 20868,   837,   691,  2248,  1850,   308,  3775,   220,\n",
+      "        50256,  5562, 10408,   663,  3435,   290, 48556,  1223,  2138,  4950,\n",
+      "          546,  1692,  3450,   220, 50256,  2787,  1299, 15950, 11378,   284,\n",
+      "         3520,   262,   976,  3690,   220, 50256,   261,   262,  5290, 15827,\n",
+      "           12,  1659,    12,  1169,    12,  1008,  9310, 35478, 20954,   262,\n",
+      "        28303,   714, 47478,   469,   510,   220, 50256,  5562,   705,    82,\n",
+      "         1290,  1165, 15444,   284, 17004,   884, 31194,  3513,   220, 50256,\n",
+      "        26567,  2536,   689,   326,   262,  3437,   286,   884,   289, 31777,\n",
+      "         2512, 30181,   355, 29408,  1830,   460,   991,  1210,   503,   257,\n",
+      "         1402,   837,  2614,  2646,   351,   281,  7016,  3355,   404,   764,\n",
+      "          220, 50256,  1659,   473,    84,   948,   220, 50256,    64, 19095,\n",
+      "        17280,    12,  1941,    12,   727,   705,    82, 26781, 19518,   220,\n",
+      "        50256,   533,   517,  7744,  1807,   832,   621,   287,   749,  4600,\n",
+      "          826,    12, 28973,   705,  7328,   220, 50256,  2188,   274,   284,\n",
+      "        12986, 20428,   220, 50256,  1640,   883,  3807, 31006,   508, 13121,\n",
+      "          326,  4600,   484,   466,   299,   470,   787,  6918,   588,   484,\n",
+      "          973,   284,  7471,   220, 50256,  1169,   636,   810,  2147,   705,\n",
+      "           82,  5836,   837,   220, 50256, 43439,   703,  2089,   428,  3807,\n",
+      "          373,   220, 50256,    75,   437,   617, 16247,   284,   257, 13526,\n",
+      "         1621,   220, 50256,  1169,  6000, 17245,   220, 50256, 36673,  3807,\n",
+      "          220, 50256,  4480,   465,  6678,  4430,   290, 11800,   774,   220,\n",
+      "        50256,   445,   917,   415,  3721,   220, 50256,  2032, 27428,   318,\n",
+      "         2029,   477,   546,   257,  1862,  2415,   705,    82,  1986,   837,\n",
+      "          290,   416, 13092,   281, 14549,  3025,  1986,  4493,   326,  2415,\n",
+      "          705,    82, 17188,   290,   614, 23400,   837,   340, 31137,   764,\n",
+      "          220, 50256,  4853,   874,   262,  2656,   290,   287,   617,  2842,\n",
+      "          772,   731,  1010,   340,   220, 50256,   361,  1997,   837,   766,\n",
+      "          340,   329,   479,  5757,  2042,   837,   508, 11665,   510,   257,\n",
+      "         6388,   355,   257, 26821, 14314, 10086, 46410,  3706, 11841, 19317,\n",
+      "          764,   220, 50256,    64,  8212,   319,   534,  1986,   220, 50256,\n",
+      "         8988,   422,   262, 14802,   837, 26329, 44139, 13289,   220, 50256,\n",
+      "         1069, 32838, 46136,   306,  3684, 16948,   290,  6028, 17049,   555,\n",
+      "          398,  5109,   220, 50256,   268, 30486,   416,   281, 25007,  9404,\n",
+      "         7668,  3350,   286, 47792, 13747,   220, 50256,  4758,  2063,   286,\n",
+      "        10441, 12254,   318,  4785,  1058,   262,   636,   810,  2147,   705,\n",
+      "           82,  5836,   837,   393,   262,   636,   810,  1223,   705,    82,\n",
+      "         5836,   220, 50256,   259,   995, 22041,   220, 50256,   548,   922,\n",
+      "        11681,  5559,   220, 50256,  1169,  7110,   318,  2147,   475, 36741,\n",
+      "         6816, 35478, 20954,   422,   923,   284,  5461,   837,   220, 50256,\n",
+      "         1169,  2223,   318,   336,   346,  1513,   220, 50256,   261,   477,\n",
+      "        43386,   220, 50256, 10594,  1064,  1310,   286,  1393,   287,   428,\n",
+      "         2646,   837,   543,   318,  1690, 32950,    88,   290, 13455, 13134,\n",
+      "          220, 50256,  1525,  1290,   262,  5290,  3807,   286,   262,   614,\n",
+      "          220, 50256, 48937,   832,   837,   220, 50256,  3549,   621,  1194,\n",
+      "         7559,  1266,   582, 10148, 17271,   416, 44889,   257,  7505,  3690,\n",
+      "          428,  8258,  2646,   220, 50256,   270,   705,    82,   546,  2428,\n",
+      "          749,  6490,   423,   284,  1986,   287,  4845,   290,  1312,   892,\n",
+      "          326,   705,    82,   644,  1312,  8288,   546,   340,  1377,   262,\n",
+      "         1103,  2428, 29779,  1022,   262, 14397,   290, 14897, 22992,   220,\n",
+      "        50256, 11718,   274,   220, 50256,   672, 35260,   284,   262,  6224,\n",
+      "          286,   428])) (tensor([  270,   705,    82,   257, 23332,   290,  1690, 13891,  7002,   764,\n",
+      "          220, 50256,   403,  2704,  8589,  4420, 30942,   290, 12111,   220,\n",
+      "        50256, 47205,   514,   284,  2911,   326,   299, 16617,   318, 24357,\n",
+      "          284, 21030,   257,  1688,  3451,   355,   257,  5068,  1865, 47602,\n",
+      "        26479,   764,   220, 50256,  1169,  7205,   837, 24138,   837,  2647,\n",
+      "          837, 13483, 45501,   290,  2128,   389,   477, 34328,  1813,   262,\n",
+      "         3227,   705,    82, 38132,   567,  1957,   274,   764,   220, 50256,\n",
+      "          270,   705,    82,  3105,  1377,   845,   837,   845,  3105,   764,\n",
+      "          220, 50256, 16670, 49699,   351, 14733,   290,   257,  1178, 39980,\n",
+      "         4135, 18105,   837,   262,  2646,   318,   257, 23056,   306,  2726,\n",
+      "          804,   379,  1862,  1466,   764,   220, 50256,    64,  3360, 32460,\n",
+      "         2646,   764,   220, 50256,   273,  1804,   938,   614,   705,    82,\n",
+      "         5704,   351,   534,   409,    12, 22095,   764,   220, 50256,  5832,\n",
+      "          466,   299,   470,   423,   284,   760,   546,  2647,   284,  9144,\n",
+      "          262,  2646,   705,    82,  2562,  5146, 13516,   286, 10997,   290,\n",
+      "        19661,   764,   220, 50256,   259,  3446,  9919,  2431,   837,   749,\n",
+      "          286,   543,  3804,   355,  6364,   355,   611,  1312,   705,    67,\n",
+      "          587,  5586, 12105,   319,   281, 45329, 29680,   837, 10451,  6885,\n",
+      "        30895,   422, 37276,   284, 13665,  2584,   284, 10517, 26876,   764,\n",
+      "          220, 50256,  1169, 46814,  2890, 13289,   286,   262,  5983,  1394,\n",
+      "          262,  2646, 22804,   290,  1394,   262,  5386, 40112,  1513,   764,\n",
+      "          220, 50256,   270,  2753,   257,  6283,  1611,   286, 37296,  1272,\n",
+      "          284,  7030,   262, 18054,   286,   686,  4835,   329,  1706,   837,\n",
+      "          281,   710,   285,   451,    64,   837,   304,  1018,  1734, 35783,\n",
+      "          837,   290,   842,  1292,    67, 11555, 30686,  1559,   477,   287,\n",
+      "          262,   976,  3807,   764,   220, 50256,   986,   262,  2646, 21046,\n",
+      "          422,   257,  3092,   286, 14733,   357,  1223,  2622,   284,  5236,\n",
+      "          503,   262,  3685,  1267,  2644,   220, 50256,   732,  6808,   329,\n",
+      "          357,   537,  3301,   290,   279,  2518,  1267,   837,   772,   588,\n",
+      "          606,   837,   996,  3737,   340,   705,    82,   281,  9942,  5699,\n",
+      "          284, 26246,   764,   220, 50256, 10197,  9961,  3296,   481,   749,\n",
+      "         1884,   407,  1064,   644,   484,   705,   260,  6095,   351,  5876,\n",
+      "          790,  1110,  2162,   262,  3807, 16523,  1111,  5636,  2171,   290,\n",
+      "        14733,   764,   220, 50256,    64, 18857,   837,  1029,    12, 45564,\n",
+      "          863, 10530,   422,   773,   544,   326, 33954,   271,  3973, 32067,\n",
+      "         2647,   837,  9280,   837,  3496,   837,   290,  1029, 10512,   764,\n",
+      "          220, 50256,  1169, 10825,   389,  8246,   290,   481,  5587,   257,\n",
+      "        16384,   351,  2687,   508,   705,    82,  1683,   550,  1641, 14649,\n",
+      "          764,   220, 50256,  3885,  4364,   256,   265,   280,   468,   257,\n",
+      "        47868,   329, 10868,  9176,   326,  7842,  1958,   607, 23077, 20024,\n",
+      "          837,   290,   287,   428,  4187,   378, 48718, 10997,   837,   673,\n",
+      "          705,    82,   355,  3329,    12,  4743,   652,   409, 18478,   415,\n",
+      "          355,   673,   373,   287,   716,  2634, 14485,   764,   220, 50256,\n",
+      "          986,   262,  3807,   318,   655,   257,  8631,  1468,  9234,   764,\n",
+      "          220, 50256,   259,   663,  1266,  7188,   837, 22960,   257,  2089,\n",
+      "         1029,  1524,  3227,   286, 35537,   837,  1231,  4414,   286,  3496,\n",
+      "          764,   220, 50256,    79,   931,  5116,  2753,   281, 37959,   804,\n",
+      "          379,   262, 29442,   286,  1964, 29409,   837,   475,   340,   857,\n",
+      "          523,   351,   884,   281, 30690,  8216,   326,   345,  1239,   760,\n",
+      "          618, 14733,  5645,   290, 13574,  6140,   764,   220, 50256,  1169,\n",
+      "         4686,  7940,   375, 20374,   329,  1528,   532,   428,   655,  2936,\n",
+      "          588,   340]), tensor([  705,    82,   257, 23332,   290,  1690, 13891,  7002,   764,   220,\n",
+      "        50256,   403,  2704,  8589,  4420, 30942,   290, 12111,   220, 50256,\n",
+      "        47205,   514,   284,  2911,   326,   299, 16617,   318, 24357,   284,\n",
+      "        21030,   257,  1688,  3451,   355,   257,  5068,  1865, 47602, 26479,\n",
+      "          764,   220, 50256,  1169,  7205,   837, 24138,   837,  2647,   837,\n",
+      "        13483, 45501,   290,  2128,   389,   477, 34328,  1813,   262,  3227,\n",
+      "          705,    82, 38132,   567,  1957,   274,   764,   220, 50256,   270,\n",
+      "          705,    82,  3105,  1377,   845,   837,   845,  3105,   764,   220,\n",
+      "        50256, 16670, 49699,   351, 14733,   290,   257,  1178, 39980,  4135,\n",
+      "        18105,   837,   262,  2646,   318,   257, 23056,   306,  2726,   804,\n",
+      "          379,  1862,  1466,   764,   220, 50256,    64,  3360, 32460,  2646,\n",
+      "          764,   220, 50256,   273,  1804,   938,   614,   705,    82,  5704,\n",
+      "          351,   534,   409,    12, 22095,   764,   220, 50256,  5832,   466,\n",
+      "          299,   470,   423,   284,   760,   546,  2647,   284,  9144,   262,\n",
+      "         2646,   705,    82,  2562,  5146, 13516,   286, 10997,   290, 19661,\n",
+      "          764,   220, 50256,   259,  3446,  9919,  2431,   837,   749,   286,\n",
+      "          543,  3804,   355,  6364,   355,   611,  1312,   705,    67,   587,\n",
+      "         5586, 12105,   319,   281, 45329, 29680,   837, 10451,  6885, 30895,\n",
+      "          422, 37276,   284, 13665,  2584,   284, 10517, 26876,   764,   220,\n",
+      "        50256,  1169, 46814,  2890, 13289,   286,   262,  5983,  1394,   262,\n",
+      "         2646, 22804,   290,  1394,   262,  5386, 40112,  1513,   764,   220,\n",
+      "        50256,   270,  2753,   257,  6283,  1611,   286, 37296,  1272,   284,\n",
+      "         7030,   262, 18054,   286,   686,  4835,   329,  1706,   837,   281,\n",
+      "          710,   285,   451,    64,   837,   304,  1018,  1734, 35783,   837,\n",
+      "          290,   842,  1292,    67, 11555, 30686,  1559,   477,   287,   262,\n",
+      "          976,  3807,   764,   220, 50256,   986,   262,  2646, 21046,   422,\n",
+      "          257,  3092,   286, 14733,   357,  1223,  2622,   284,  5236,   503,\n",
+      "          262,  3685,  1267,  2644,   220, 50256,   732,  6808,   329,   357,\n",
+      "          537,  3301,   290,   279,  2518,  1267,   837,   772,   588,   606,\n",
+      "          837,   996,  3737,   340,   705,    82,   281,  9942,  5699,   284,\n",
+      "        26246,   764,   220, 50256, 10197,  9961,  3296,   481,   749,  1884,\n",
+      "          407,  1064,   644,   484,   705,   260,  6095,   351,  5876,   790,\n",
+      "         1110,  2162,   262,  3807, 16523,  1111,  5636,  2171,   290, 14733,\n",
+      "          764,   220, 50256,    64, 18857,   837,  1029,    12, 45564,   863,\n",
+      "        10530,   422,   773,   544,   326, 33954,   271,  3973, 32067,  2647,\n",
+      "          837,  9280,   837,  3496,   837,   290,  1029, 10512,   764,   220,\n",
+      "        50256,  1169, 10825,   389,  8246,   290,   481,  5587,   257, 16384,\n",
+      "          351,  2687,   508,   705,    82,  1683,   550,  1641, 14649,   764,\n",
+      "          220, 50256,  3885,  4364,   256,   265,   280,   468,   257, 47868,\n",
+      "          329, 10868,  9176,   326,  7842,  1958,   607, 23077, 20024,   837,\n",
+      "          290,   287,   428,  4187,   378, 48718, 10997,   837,   673,   705,\n",
+      "           82,   355,  3329,    12,  4743,   652,   409, 18478,   415,   355,\n",
+      "          673,   373,   287,   716,  2634, 14485,   764,   220, 50256,   986,\n",
+      "          262,  3807,   318,   655,   257,  8631,  1468,  9234,   764,   220,\n",
+      "        50256,   259,   663,  1266,  7188,   837, 22960,   257,  2089,  1029,\n",
+      "         1524,  3227,   286, 35537,   837,  1231,  4414,   286,  3496,   764,\n",
+      "          220, 50256,    79,   931,  5116,  2753,   281, 37959,   804,   379,\n",
+      "          262, 29442,   286,  1964, 29409,   837,   475,   340,   857,   523,\n",
+      "          351,   884,   281, 30690,  8216,   326,   345,  1239,   760,   618,\n",
+      "        14733,  5645,   290, 13574,  6140,   764,   220, 50256,  1169,  4686,\n",
+      "         7940,   375, 20374,   329,  1528,   532,   428,   655,  2936,   588,\n",
+      "          340,   750]))\n"
+     ]
+    }
+   ],
+   "source": [
+    "from archai.datasets.nlp.fast_hf_dataset_provider import FastHfDatasetProvider\n",
+    "\n",
+    "# The provider will automatically download the dataset and tokenizer, encode\n",
+    "# the dataset and cache it for future use\n",
+    "dataset_provider = FastHfDatasetProvider(\n",
+    "    \"glue\",\n",
+    "    dataset_config_name=\"sst2\",\n",
+    "    tokenizer_name=\"gpt2\",\n",
+    "    tokenizer_max_length=512,\n",
+    "    mapping_column_name=[\"sentence\"],\n",
+    "    use_shared_memory=False\n",
+    ")\n",
+    "\n",
+    "# (inputs, labels) can be retrieved with any sequence length\n",
+    "train_dataset = dataset_provider.get_train_dataset(seq_len=512)\n",
+    "val_dataset = dataset_provider.get_val_dataset(seq_len=512)\n",
+    "print(train_dataset[0], val_dataset[0])"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Loading from Cache\n",
+    "\n",
+    "After loading and encoding the dataset for the first time, a cache will be created with a unique fingerprint (identifier) based on its configuration. The cached is composed by the following files:\n",
+    "\n",
+    "* `config.json`: Dataset provider configuration (used to re-create the object when loaded from cache).\n",
+    "* `tokenizer.pkl`: Tokenizer used to encode the data (also re-created when loaded from cache).\n",
+    "* `train.npy`: Training tokens (inputs and labels).\n",
+    "* `validation.npy`: Validation tokens (inputs and labels).\n",
+    "* `test.npy`: Testing tokens (inputs and labels).\n",
+    "\n",
+    "The `FastHfDatasetProvider` class provides a `from_cache` method which can be used to re-instantiate the cached dataset provider, in case the user wants to re-use in different places."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "(tensor([24717,   649,  3200,   507,   422,   262, 21694,  4991,   220, 50256,\n",
+      "         3642,  1299,   645, 20868,   837,   691,  2248,  1850,   308,  3775,\n",
+      "          220, 50256,  5562, 10408,   663,  3435,   290, 48556,  1223,  2138,\n",
+      "         4950,   546,  1692,  3450,   220, 50256,  2787,  1299, 15950, 11378,\n",
+      "          284,  3520,   262,   976,  3690,   220, 50256,   261,   262,  5290,\n",
+      "        15827,    12,  1659,    12,  1169,    12,  1008,  9310, 35478, 20954,\n",
+      "          262, 28303,   714, 47478,   469,   510,   220, 50256,  5562,   705,\n",
+      "           82,  1290,  1165, 15444,   284, 17004,   884, 31194,  3513,   220,\n",
+      "        50256, 26567,  2536,   689,   326,   262,  3437,   286,   884,   289,\n",
+      "        31777,  2512, 30181,   355, 29408,  1830,   460,   991,  1210,   503,\n",
+      "          257,  1402,   837,  2614,  2646,   351,   281,  7016,  3355,   404,\n",
+      "          764,   220, 50256,  1659,   473,    84,   948,   220, 50256,    64,\n",
+      "        19095, 17280,    12,  1941,    12,   727,   705,    82, 26781, 19518,\n",
+      "          220, 50256,   533,   517,  7744,  1807,   832,   621,   287,   749,\n",
+      "         4600,   826,    12, 28973,   705,  7328,   220, 50256,  2188,   274,\n",
+      "          284, 12986, 20428,   220, 50256,  1640,   883,  3807, 31006,   508,\n",
+      "        13121,   326,  4600,   484,   466,   299,   470,   787,  6918,   588,\n",
+      "          484,   973,   284,  7471,   220, 50256,  1169,   636,   810,  2147,\n",
+      "          705,    82,  5836,   837,   220, 50256, 43439,   703,  2089,   428,\n",
+      "         3807,   373,   220, 50256,    75,   437,   617, 16247,   284,   257,\n",
+      "        13526,  1621,   220, 50256,  1169,  6000, 17245,   220, 50256, 36673,\n",
+      "         3807,   220, 50256,  4480,   465,  6678,  4430,   290, 11800,   774,\n",
+      "          220, 50256,   445,   917,   415,  3721,   220, 50256,  2032, 27428,\n",
+      "          318,  2029,   477,   546,   257,  1862,  2415,   705,    82,  1986,\n",
+      "          837,   290,   416, 13092,   281, 14549,  3025,  1986,  4493,   326,\n",
+      "         2415,   705,    82, 17188,   290,   614, 23400,   837,   340, 31137,\n",
+      "          764,   220, 50256,  4853,   874,   262,  2656,   290,   287,   617,\n",
+      "         2842,   772,   731,  1010,   340,   220, 50256,   361,  1997,   837,\n",
+      "          766,   340,   329,   479,  5757,  2042,   837,   508, 11665,   510,\n",
+      "          257,  6388,   355,   257, 26821, 14314, 10086, 46410,  3706, 11841,\n",
+      "        19317,   764,   220, 50256,    64,  8212,   319,   534,  1986,   220,\n",
+      "        50256,  8988,   422,   262, 14802,   837, 26329, 44139, 13289,   220,\n",
+      "        50256,  1069, 32838, 46136,   306,  3684, 16948,   290,  6028, 17049,\n",
+      "          555,   398,  5109,   220, 50256,   268, 30486,   416,   281, 25007,\n",
+      "         9404,  7668,  3350,   286, 47792, 13747,   220, 50256,  4758,  2063,\n",
+      "          286, 10441, 12254,   318,  4785,  1058,   262,   636,   810,  2147,\n",
+      "          705,    82,  5836,   837,   393,   262,   636,   810,  1223,   705,\n",
+      "           82,  5836,   220, 50256,   259,   995, 22041,   220, 50256,   548,\n",
+      "          922, 11681,  5559,   220, 50256,  1169,  7110,   318,  2147,   475,\n",
+      "        36741,  6816, 35478, 20954,   422,   923,   284,  5461,   837,   220,\n",
+      "        50256,  1169,  2223,   318,   336,   346,  1513,   220, 50256,   261,\n",
+      "          477, 43386,   220, 50256, 10594,  1064,  1310,   286,  1393,   287,\n",
+      "          428,  2646,   837,   543,   318,  1690, 32950,    88,   290, 13455,\n",
+      "        13134,   220, 50256,  1525,  1290,   262,  5290,  3807,   286,   262,\n",
+      "          614,   220, 50256, 48937,   832,   837,   220, 50256,  3549,   621,\n",
+      "         1194,  7559,  1266,   582, 10148, 17271,   416, 44889,   257,  7505,\n",
+      "         3690,   428,  8258,  2646,   220, 50256,   270,   705,    82,   546,\n",
+      "         2428,   749,  6490,   423,   284,  1986,   287,  4845,   290,  1312,\n",
+      "          892,   326,   705,    82,   644,  1312,  8288,   546,   340,  1377,\n",
+      "          262,  1103,  2428, 29779,  1022,   262, 14397,   290, 14897, 22992,\n",
+      "          220, 50256, 11718,   274,   220, 50256,   672, 35260,   284,   262,\n",
+      "         6224,   286]), tensor([  649,  3200,   507,   422,   262, 21694,  4991,   220, 50256,  3642,\n",
+      "         1299,   645, 20868,   837,   691,  2248,  1850,   308,  3775,   220,\n",
+      "        50256,  5562, 10408,   663,  3435,   290, 48556,  1223,  2138,  4950,\n",
+      "          546,  1692,  3450,   220, 50256,  2787,  1299, 15950, 11378,   284,\n",
+      "         3520,   262,   976,  3690,   220, 50256,   261,   262,  5290, 15827,\n",
+      "           12,  1659,    12,  1169,    12,  1008,  9310, 35478, 20954,   262,\n",
+      "        28303,   714, 47478,   469,   510,   220, 50256,  5562,   705,    82,\n",
+      "         1290,  1165, 15444,   284, 17004,   884, 31194,  3513,   220, 50256,\n",
+      "        26567,  2536,   689,   326,   262,  3437,   286,   884,   289, 31777,\n",
+      "         2512, 30181,   355, 29408,  1830,   460,   991,  1210,   503,   257,\n",
+      "         1402,   837,  2614,  2646,   351,   281,  7016,  3355,   404,   764,\n",
+      "          220, 50256,  1659,   473,    84,   948,   220, 50256,    64, 19095,\n",
+      "        17280,    12,  1941,    12,   727,   705,    82, 26781, 19518,   220,\n",
+      "        50256,   533,   517,  7744,  1807,   832,   621,   287,   749,  4600,\n",
+      "          826,    12, 28973,   705,  7328,   220, 50256,  2188,   274,   284,\n",
+      "        12986, 20428,   220, 50256,  1640,   883,  3807, 31006,   508, 13121,\n",
+      "          326,  4600,   484,   466,   299,   470,   787,  6918,   588,   484,\n",
+      "          973,   284,  7471,   220, 50256,  1169,   636,   810,  2147,   705,\n",
+      "           82,  5836,   837,   220, 50256, 43439,   703,  2089,   428,  3807,\n",
+      "          373,   220, 50256,    75,   437,   617, 16247,   284,   257, 13526,\n",
+      "         1621,   220, 50256,  1169,  6000, 17245,   220, 50256, 36673,  3807,\n",
+      "          220, 50256,  4480,   465,  6678,  4430,   290, 11800,   774,   220,\n",
+      "        50256,   445,   917,   415,  3721,   220, 50256,  2032, 27428,   318,\n",
+      "         2029,   477,   546,   257,  1862,  2415,   705,    82,  1986,   837,\n",
+      "          290,   416, 13092,   281, 14549,  3025,  1986,  4493,   326,  2415,\n",
+      "          705,    82, 17188,   290,   614, 23400,   837,   340, 31137,   764,\n",
+      "          220, 50256,  4853,   874,   262,  2656,   290,   287,   617,  2842,\n",
+      "          772,   731,  1010,   340,   220, 50256,   361,  1997,   837,   766,\n",
+      "          340,   329,   479,  5757,  2042,   837,   508, 11665,   510,   257,\n",
+      "         6388,   355,   257, 26821, 14314, 10086, 46410,  3706, 11841, 19317,\n",
+      "          764,   220, 50256,    64,  8212,   319,   534,  1986,   220, 50256,\n",
+      "         8988,   422,   262, 14802,   837, 26329, 44139, 13289,   220, 50256,\n",
+      "         1069, 32838, 46136,   306,  3684, 16948,   290,  6028, 17049,   555,\n",
+      "          398,  5109,   220, 50256,   268, 30486,   416,   281, 25007,  9404,\n",
+      "         7668,  3350,   286, 47792, 13747,   220, 50256,  4758,  2063,   286,\n",
+      "        10441, 12254,   318,  4785,  1058,   262,   636,   810,  2147,   705,\n",
+      "           82,  5836,   837,   393,   262,   636,   810,  1223,   705,    82,\n",
+      "         5836,   220, 50256,   259,   995, 22041,   220, 50256,   548,   922,\n",
+      "        11681,  5559,   220, 50256,  1169,  7110,   318,  2147,   475, 36741,\n",
+      "         6816, 35478, 20954,   422,   923,   284,  5461,   837,   220, 50256,\n",
+      "         1169,  2223,   318,   336,   346,  1513,   220, 50256,   261,   477,\n",
+      "        43386,   220, 50256, 10594,  1064,  1310,   286,  1393,   287,   428,\n",
+      "         2646,   837,   543,   318,  1690, 32950,    88,   290, 13455, 13134,\n",
+      "          220, 50256,  1525,  1290,   262,  5290,  3807,   286,   262,   614,\n",
+      "          220, 50256, 48937,   832,   837,   220, 50256,  3549,   621,  1194,\n",
+      "         7559,  1266,   582, 10148, 17271,   416, 44889,   257,  7505,  3690,\n",
+      "          428,  8258,  2646,   220, 50256,   270,   705,    82,   546,  2428,\n",
+      "          749,  6490,   423,   284,  1986,   287,  4845,   290,  1312,   892,\n",
+      "          326,   705,    82,   644,  1312,  8288,   546,   340,  1377,   262,\n",
+      "         1103,  2428, 29779,  1022,   262, 14397,   290, 14897, 22992,   220,\n",
+      "        50256, 11718,   274,   220, 50256,   672, 35260,   284,   262,  6224,\n",
+      "          286,   428])) (tensor([  270,   705,    82,   257, 23332,   290,  1690, 13891,  7002,   764,\n",
+      "          220, 50256,   403,  2704,  8589,  4420, 30942,   290, 12111,   220,\n",
+      "        50256, 47205,   514,   284,  2911,   326,   299, 16617,   318, 24357,\n",
+      "          284, 21030,   257,  1688,  3451,   355,   257,  5068,  1865, 47602,\n",
+      "        26479,   764,   220, 50256,  1169,  7205,   837, 24138,   837,  2647,\n",
+      "          837, 13483, 45501,   290,  2128,   389,   477, 34328,  1813,   262,\n",
+      "         3227,   705,    82, 38132,   567,  1957,   274,   764,   220, 50256,\n",
+      "          270,   705,    82,  3105,  1377,   845,   837,   845,  3105,   764,\n",
+      "          220, 50256, 16670, 49699,   351, 14733,   290,   257,  1178, 39980,\n",
+      "         4135, 18105,   837,   262,  2646,   318,   257, 23056,   306,  2726,\n",
+      "          804,   379,  1862,  1466,   764,   220, 50256,    64,  3360, 32460,\n",
+      "         2646,   764,   220, 50256,   273,  1804,   938,   614,   705,    82,\n",
+      "         5704,   351,   534,   409,    12, 22095,   764,   220, 50256,  5832,\n",
+      "          466,   299,   470,   423,   284,   760,   546,  2647,   284,  9144,\n",
+      "          262,  2646,   705,    82,  2562,  5146, 13516,   286, 10997,   290,\n",
+      "        19661,   764,   220, 50256,   259,  3446,  9919,  2431,   837,   749,\n",
+      "          286,   543,  3804,   355,  6364,   355,   611,  1312,   705,    67,\n",
+      "          587,  5586, 12105,   319,   281, 45329, 29680,   837, 10451,  6885,\n",
+      "        30895,   422, 37276,   284, 13665,  2584,   284, 10517, 26876,   764,\n",
+      "          220, 50256,  1169, 46814,  2890, 13289,   286,   262,  5983,  1394,\n",
+      "          262,  2646, 22804,   290,  1394,   262,  5386, 40112,  1513,   764,\n",
+      "          220, 50256,   270,  2753,   257,  6283,  1611,   286, 37296,  1272,\n",
+      "          284,  7030,   262, 18054,   286,   686,  4835,   329,  1706,   837,\n",
+      "          281,   710,   285,   451,    64,   837,   304,  1018,  1734, 35783,\n",
+      "          837,   290,   842,  1292,    67, 11555, 30686,  1559,   477,   287,\n",
+      "          262,   976,  3807,   764,   220, 50256,   986,   262,  2646, 21046,\n",
+      "          422,   257,  3092,   286, 14733,   357,  1223,  2622,   284,  5236,\n",
+      "          503,   262,  3685,  1267,  2644,   220, 50256,   732,  6808,   329,\n",
+      "          357,   537,  3301,   290,   279,  2518,  1267,   837,   772,   588,\n",
+      "          606,   837,   996,  3737,   340,   705,    82,   281,  9942,  5699,\n",
+      "          284, 26246,   764,   220, 50256, 10197,  9961,  3296,   481,   749,\n",
+      "         1884,   407,  1064,   644,   484,   705,   260,  6095,   351,  5876,\n",
+      "          790,  1110,  2162,   262,  3807, 16523,  1111,  5636,  2171,   290,\n",
+      "        14733,   764,   220, 50256,    64, 18857,   837,  1029,    12, 45564,\n",
+      "          863, 10530,   422,   773,   544,   326, 33954,   271,  3973, 32067,\n",
+      "         2647,   837,  9280,   837,  3496,   837,   290,  1029, 10512,   764,\n",
+      "          220, 50256,  1169, 10825,   389,  8246,   290,   481,  5587,   257,\n",
+      "        16384,   351,  2687,   508,   705,    82,  1683,   550,  1641, 14649,\n",
+      "          764,   220, 50256,  3885,  4364,   256,   265,   280,   468,   257,\n",
+      "        47868,   329, 10868,  9176,   326,  7842,  1958,   607, 23077, 20024,\n",
+      "          837,   290,   287,   428,  4187,   378, 48718, 10997,   837,   673,\n",
+      "          705,    82,   355,  3329,    12,  4743,   652,   409, 18478,   415,\n",
+      "          355,   673,   373,   287,   716,  2634, 14485,   764,   220, 50256,\n",
+      "          986,   262,  3807,   318,   655,   257,  8631,  1468,  9234,   764,\n",
+      "          220, 50256,   259,   663,  1266,  7188,   837, 22960,   257,  2089,\n",
+      "         1029,  1524,  3227,   286, 35537,   837,  1231,  4414,   286,  3496,\n",
+      "          764,   220, 50256,    79,   931,  5116,  2753,   281, 37959,   804,\n",
+      "          379,   262, 29442,   286,  1964, 29409,   837,   475,   340,   857,\n",
+      "          523,   351,   884,   281, 30690,  8216,   326,   345,  1239,   760,\n",
+      "          618, 14733,  5645,   290, 13574,  6140,   764,   220, 50256,  1169,\n",
+      "         4686,  7940,   375, 20374,   329,  1528,   532,   428,   655,  2936,\n",
+      "          588,   340]), tensor([  705,    82,   257, 23332,   290,  1690, 13891,  7002,   764,   220,\n",
+      "        50256,   403,  2704,  8589,  4420, 30942,   290, 12111,   220, 50256,\n",
+      "        47205,   514,   284,  2911,   326,   299, 16617,   318, 24357,   284,\n",
+      "        21030,   257,  1688,  3451,   355,   257,  5068,  1865, 47602, 26479,\n",
+      "          764,   220, 50256,  1169,  7205,   837, 24138,   837,  2647,   837,\n",
+      "        13483, 45501,   290,  2128,   389,   477, 34328,  1813,   262,  3227,\n",
+      "          705,    82, 38132,   567,  1957,   274,   764,   220, 50256,   270,\n",
+      "          705,    82,  3105,  1377,   845,   837,   845,  3105,   764,   220,\n",
+      "        50256, 16670, 49699,   351, 14733,   290,   257,  1178, 39980,  4135,\n",
+      "        18105,   837,   262,  2646,   318,   257, 23056,   306,  2726,   804,\n",
+      "          379,  1862,  1466,   764,   220, 50256,    64,  3360, 32460,  2646,\n",
+      "          764,   220, 50256,   273,  1804,   938,   614,   705,    82,  5704,\n",
+      "          351,   534,   409,    12, 22095,   764,   220, 50256,  5832,   466,\n",
+      "          299,   470,   423,   284,   760,   546,  2647,   284,  9144,   262,\n",
+      "         2646,   705,    82,  2562,  5146, 13516,   286, 10997,   290, 19661,\n",
+      "          764,   220, 50256,   259,  3446,  9919,  2431,   837,   749,   286,\n",
+      "          543,  3804,   355,  6364,   355,   611,  1312,   705,    67,   587,\n",
+      "         5586, 12105,   319,   281, 45329, 29680,   837, 10451,  6885, 30895,\n",
+      "          422, 37276,   284, 13665,  2584,   284, 10517, 26876,   764,   220,\n",
+      "        50256,  1169, 46814,  2890, 13289,   286,   262,  5983,  1394,   262,\n",
+      "         2646, 22804,   290,  1394,   262,  5386, 40112,  1513,   764,   220,\n",
+      "        50256,   270,  2753,   257,  6283,  1611,   286, 37296,  1272,   284,\n",
+      "         7030,   262, 18054,   286,   686,  4835,   329,  1706,   837,   281,\n",
+      "          710,   285,   451,    64,   837,   304,  1018,  1734, 35783,   837,\n",
+      "          290,   842,  1292,    67, 11555, 30686,  1559,   477,   287,   262,\n",
+      "          976,  3807,   764,   220, 50256,   986,   262,  2646, 21046,   422,\n",
+      "          257,  3092,   286, 14733,   357,  1223,  2622,   284,  5236,   503,\n",
+      "          262,  3685,  1267,  2644,   220, 50256,   732,  6808,   329,   357,\n",
+      "          537,  3301,   290,   279,  2518,  1267,   837,   772,   588,   606,\n",
+      "          837,   996,  3737,   340,   705,    82,   281,  9942,  5699,   284,\n",
+      "        26246,   764,   220, 50256, 10197,  9961,  3296,   481,   749,  1884,\n",
+      "          407,  1064,   644,   484,   705,   260,  6095,   351,  5876,   790,\n",
+      "         1110,  2162,   262,  3807, 16523,  1111,  5636,  2171,   290, 14733,\n",
+      "          764,   220, 50256,    64, 18857,   837,  1029,    12, 45564,   863,\n",
+      "        10530,   422,   773,   544,   326, 33954,   271,  3973, 32067,  2647,\n",
+      "          837,  9280,   837,  3496,   837,   290,  1029, 10512,   764,   220,\n",
+      "        50256,  1169, 10825,   389,  8246,   290,   481,  5587,   257, 16384,\n",
+      "          351,  2687,   508,   705,    82,  1683,   550,  1641, 14649,   764,\n",
+      "          220, 50256,  3885,  4364,   256,   265,   280,   468,   257, 47868,\n",
+      "          329, 10868,  9176,   326,  7842,  1958,   607, 23077, 20024,   837,\n",
+      "          290,   287,   428,  4187,   378, 48718, 10997,   837,   673,   705,\n",
+      "           82,   355,  3329,    12,  4743,   652,   409, 18478,   415,   355,\n",
+      "          673,   373,   287,   716,  2634, 14485,   764,   220, 50256,   986,\n",
+      "          262,  3807,   318,   655,   257,  8631,  1468,  9234,   764,   220,\n",
+      "        50256,   259,   663,  1266,  7188,   837, 22960,   257,  2089,  1029,\n",
+      "         1524,  3227,   286, 35537,   837,  1231,  4414,   286,  3496,   764,\n",
+      "          220, 50256,    79,   931,  5116,  2753,   281, 37959,   804,   379,\n",
+      "          262, 29442,   286,  1964, 29409,   837,   475,   340,   857,   523,\n",
+      "          351,   884,   281, 30690,  8216,   326,   345,  1239,   760,   618,\n",
+      "        14733,  5645,   290, 13574,  6140,   764,   220, 50256,  1169,  4686,\n",
+      "         7940,   375, 20374,   329,  1528,   532,   428,   655,  2936,   588,\n",
+      "          340,   750]))\n"
+     ]
+    }
+   ],
+   "source": [
+    "# The caching mechanism automatically saves `config.json` and `tokenizer.pkl`,\n",
+    "# which are used to recreate the provider when calling `from_cache` method\n",
+    "dataset_provider = FastHfDatasetProvider.from_cache(\"cache/glue/sst2/3ba1dcb015d7664117d16b85af90496896fc74ef\")\n",
+    "\n",
+    "train_dataset = dataset_provider.get_train_dataset(seq_len=512)\n",
+    "val_dataset = dataset_provider.get_val_dataset(seq_len=512)\n",
+    "print(train_dataset[0], val_dataset[0])"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "archai",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.16"
+  },
+  "orig_nbformat": 4,
+  "vscode": {
+   "interpreter": {
+    "hash": "b2be478cf8a2d9a6a1293b022e8589530f7ec0d0340a3a36da6068ef3d344086"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/docs/getting_started/notebooks/nlp/hf_dataset_provider.ipynb b/docs/getting_started/notebooks/nlp/hf_dataset_provider.ipynb
index 7075c88c3..d4d761ad0 100644
--- a/docs/getting_started/notebooks/nlp/hf_dataset_provider.ipynb
+++ b/docs/getting_started/notebooks/nlp/hf_dataset_provider.ipynb
@@ -19,10 +19,9 @@
     "\n",
     "The first step is to create an instance of the `HfHubDatasetProvider`, which offers pre-loads the dataset and offers three methods to retrieve them: `get_train_dataset()`, `get_val_dataset()` and `get_test_dataset()`.\n",
     "\n",
-    "Additionally, a set of arguments can be passed to its constructor according to the user's needs:\n",
+    "Additionally, a set of additional arguments can be passed to its constructor according to the user's needs:\n",
     "\n",
-    "* `dataset`: Name of the dataset.\n",
-    "* `subset`: Name of the dataset configuration.\n",
+    "* `dataset_config_name`: Name of the dataset configuration.\n",
     "* `data_dir`: Path to the data directory.\n",
     "* `data_files`: Path(s) to the data file(s).\n",
     "* `cache_dir`: Path to the read/write cache directory.\n",
@@ -31,7 +30,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [
     {
@@ -59,7 +58,7 @@
    "source": [
     "from archai.datasets.nlp.hf_dataset_provider import HfHubDatasetProvider\n",
     "\n",
-    "dataset_provider = HfHubDatasetProvider(dataset=\"glue\", subset=\"sst2\")\n",
+    "dataset_provider = HfHubDatasetProvider(\"glue\", dataset_config_name=\"sst2\")\n",
     "\n",
     "# When loading `train_dataset`, we will override the split argument to only load 1%\n",
     "# of the data and speed up its encoding\n",
@@ -88,27 +87,20 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [
     {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "191c63e1183b4fe58d462afb6c582bb8",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/1 [00:00<?, ?ba/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Loading cached processed dataset at C:\\Users\\gderosa\\.cache\\huggingface\\datasets\\glue\\sst2\\1.0.0\\dacbe3125aa31d7f70367a07a8a9e72a5a0bfeb5fc42e75c9db75b96da6053ad\\cache-c989f437f7c0d7ad.arrow\n"
+     ]
     },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1217b8d1d9f84f0d9dd737f2d1024876",
+       "model_id": "f7156f9c30db4c8193e6417f84fc693f",
        "version_major": 2,
        "version_minor": 0
       },
diff --git a/docs/getting_started/notebooks/nlp/hf_trainer.ipynb b/docs/getting_started/notebooks/nlp/hf_trainer.ipynb
index baa5c663d..c0c40b421 100644
--- a/docs/getting_started/notebooks/nlp/hf_trainer.ipynb
+++ b/docs/getting_started/notebooks/nlp/hf_trainer.ipynb
@@ -29,119 +29,13 @@
    "execution_count": 1,
    "metadata": {},
    "outputs": [
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "953c7cc4e2d3496b96cc137ede4cbac4",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Downloading (…)okenizer_config.json:   0%|          | 0.00/240 [00:00<?, ?B/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "c:\\Users\\gderosa\\Anaconda3\\envs\\archai\\lib\\site-packages\\huggingface_hub\\file_download.py:129: UserWarning: `huggingface_hub` cache-system uses symlinks by default to efficiently store duplicated files but your machine does not support them in C:\\Users\\gderosa\\.cache\\huggingface\\hub. Caching files will still work but in a degraded version that might require more space on your disk. This warning can be disabled by setting the `HF_HUB_DISABLE_SYMLINKS_WARNING` environment variable. For more details, see https://huggingface.co/docs/huggingface_hub/how-to-cache#limitations.\n",
-      "To support symlinks on Windows, you either need to activate Developer Mode or to run Python as an administrator. In order to see activate developer mode, see this article: https://docs.microsoft.com/en-us/windows/apps/get-started/enable-your-device-for-development\n",
-      "  warnings.warn(message)\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b15f1853cc0f4a7bb0d4109701f2b501",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Downloading (…)olve/main/vocab.json:   0%|          | 0.00/798k [00:00<?, ?B/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b2ec836b48a9477c87ab3c6f18c7ca74",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Downloading (…)olve/main/merges.txt:   0%|          | 0.00/456k [00:00<?, ?B/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "29202e3813f34423a120a63da4b2f520",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Downloading (…)/main/tokenizer.json:   0%|          | 0.00/2.11M [00:00<?, ?B/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "62bd358138da4c83a2519af22702bb0b",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Downloading (…)in/added_tokens.json:   0%|          | 0.00/1.00k [00:00<?, ?B/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "0adf3904d47a4cbcaf10e4b18d596c62",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Downloading (…)cial_tokens_map.json:   0%|          | 0.00/90.0 [00:00<?, ?B/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Found cached dataset wikitext (C:/Users/gderosa/.cache/huggingface/datasets/wikitext/wikitext-103-raw-v1/1.0.0/a241db52902eaf2c6aa732210bead40c090019a499ceb13bcbfa3f8ab646a126)\n"
+      "Found cached dataset wikitext (C:/Users/gderosa/.cache/huggingface/datasets/wikitext/wikitext-103-raw-v1/1.0.0/a241db52902eaf2c6aa732210bead40c090019a499ceb13bcbfa3f8ab646a126)\n",
+      "Loading cached processed dataset at C:\\Users\\gderosa\\.cache\\huggingface\\datasets\\wikitext\\wikitext-103-raw-v1\\1.0.0\\a241db52902eaf2c6aa732210bead40c090019a499ceb13bcbfa3f8ab646a126\\cache-04d7ff93d438ade6.arrow\n"
      ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "8cddeee2ae9841f58e9def82424fdb29",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/19 [00:00<?, ?ba/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
     }
    ],
    "source": [
@@ -153,7 +47,7 @@
     "tokenizer.pad_token = tokenizer.eos_token\n",
     "\n",
     "collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)\n",
-    "dataset_provider = HfHubDatasetProvider(dataset=\"wikitext\", subset=\"wikitext-103-raw-v1\")\n",
+    "dataset_provider = HfHubDatasetProvider(\"wikitext\", dataset_config_name=\"wikitext-103-raw-v1\")\n",
     "\n",
     "# When loading `train_dataset`, we will override the split argument to only load 1%\n",
     "# of the data and speed up its encoding\n",
@@ -228,7 +122,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "cd9c4615466546a5b6180ad7da7e7b87",
+       "model_id": "8600bd55454d42bd9f1e824f41e8e532",
        "version_major": 2,
        "version_minor": 0
       },
@@ -257,14 +151,32 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "{'loss': 11.0746, 'learning_rate': 0.0, 'epoch': 0.0}\n",
-      "{'train_runtime': 13.891, 'train_samples_per_second': 0.072, 'train_steps_per_second': 0.072, 'train_loss': 11.074588775634766, 'epoch': 0.0}\n"
+      "Attempted to log scalar metric loss:\n",
+      "11.0034\n",
+      "Attempted to log scalar metric learning_rate:\n",
+      "0.0\n",
+      "Attempted to log scalar metric epoch:\n",
+      "0.0\n",
+      "{'loss': 11.0034, 'learning_rate': 0.0, 'epoch': 0.0}\n",
+      "Attempted to log scalar metric train_runtime:\n",
+      "10.6794\n",
+      "Attempted to log scalar metric train_samples_per_second:\n",
+      "0.094\n",
+      "Attempted to log scalar metric train_steps_per_second:\n",
+      "0.094\n",
+      "Attempted to log scalar metric total_flos:\n",
+      "759874922496.0\n",
+      "Attempted to log scalar metric train_loss:\n",
+      "11.003437995910645\n",
+      "Attempted to log scalar metric epoch:\n",
+      "0.0\n",
+      "{'train_runtime': 10.6794, 'train_samples_per_second': 0.094, 'train_steps_per_second': 0.094, 'train_loss': 11.003437995910645, 'epoch': 0.0}\n"
      ]
     },
     {
      "data": {
       "text/plain": [
-       "TrainOutput(global_step=1, training_loss=11.074588775634766, metrics={'train_runtime': 13.891, 'train_samples_per_second': 0.072, 'train_steps_per_second': 0.072, 'train_loss': 11.074588775634766, 'epoch': 0.0})"
+       "TrainOutput(global_step=1, training_loss=11.003437995910645, metrics={'train_runtime': 10.6794, 'train_samples_per_second': 0.094, 'train_steps_per_second': 0.094, 'train_loss': 11.003437995910645, 'epoch': 0.0})"
       ]
      },
      "execution_count": 3,
diff --git a/scripts/create_dataroot_symlink.bat b/scripts/create_dataroot_symlink.bat
new file mode 100644
index 000000000..c351d340e
--- /dev/null
+++ b/scripts/create_dataroot_symlink.bat
@@ -0,0 +1,2 @@
+REM Creates symbolic link to datasets folder, so that Archai can find the datasets
+mklink /j %USERPROFILE%\dataroot E:\datasets
\ No newline at end of file
diff --git a/scripts/supergraph/download_datasets/img_stats.py b/scripts/supergraph/download_datasets/img_stats.py
index e0e9e3e89..da877ad70 100644
--- a/scripts/supergraph/download_datasets/img_stats.py
+++ b/scripts/supergraph/download_datasets/img_stats.py
@@ -4,11 +4,18 @@
 import itertools
 import pathlib
 from concurrent.futures import ThreadPoolExecutor
+import subprocess
+import sys
 from threading import Lock
 
 import numpy as np
 from PIL import Image
-from runstats import Statistics
+
+try:
+    from runstats import Statistics
+except:
+    subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'runstats'])
+    from runstats import Statistics
 
 
 class ImageStats:
diff --git a/scripts/supergraph/download_datasets/sport8_install.py b/scripts/supergraph/download_datasets/sport8_install.py
index 763b2e558..da0e2d841 100644
--- a/scripts/supergraph/download_datasets/sport8_install.py
+++ b/scripts/supergraph/download_datasets/sport8_install.py
@@ -6,11 +6,19 @@
 """
 
 import os
+import sys
+import subprocess
 from collections import defaultdict
 from typing import Dict, List
 
 import dataset_utils
-import pyunpack
+
+try:
+    import pyunpack
+except:
+    subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'pyunpack'])
+    import pyunpack
+    
 from torchvision.datasets import utils as tvutils
 
 from archai.common import utils
diff --git a/scripts/supergraph/reports/exprep.py b/scripts/supergraph/reports/exprep.py
index 457c54cdf..866217137 100644
--- a/scripts/supergraph/reports/exprep.py
+++ b/scripts/supergraph/reports/exprep.py
@@ -7,11 +7,18 @@
 import re
 from collections import OrderedDict
 from inspect import getsourcefile
+import subprocess
+import sys
 from typing import Dict, Iterator, List, Tuple
 
 import matplotlib
 import yaml
-from runstats import Statistics
+
+try:
+    from runstats import Statistics
+except:
+    subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'runstats'])
+    from runstats import Statistics
 
 matplotlib.use("Agg")
 
diff --git a/scripts/supergraph/reports/old_logs.py b/scripts/supergraph/reports/old_logs.py
index 2f96cafe8..26d55fda9 100644
--- a/scripts/supergraph/reports/old_logs.py
+++ b/scripts/supergraph/reports/old_logs.py
@@ -4,8 +4,14 @@
 import argparse
 import os
 import pathlib
-
-from runstats import Statistics
+import subprocess
+import sys
+
+try:
+    from runstats import Statistics
+except:
+    subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'runstats'])
+    from runstats import Statistics
 
 
 def main():
diff --git a/scripts/trainers/hf_train_codegen.py b/scripts/trainers/hf_train_codegen.py
index 44db68a60..175494928 100644
--- a/scripts/trainers/hf_train_codegen.py
+++ b/scripts/trainers/hf_train_codegen.py
@@ -62,7 +62,7 @@ def parse_args() -> argparse.Namespace:
 
     collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
 
-    dataset_provider = HfHubDatasetProvider(dataset=args.dataset_name, subset=args.dataset_config_name)
+    dataset_provider = HfHubDatasetProvider(args.dataset_name, dataset_config_name=args.dataset_config_name)
     train_dataset = dataset_provider.get_train_dataset()
     eval_dataset = dataset_provider.get_val_dataset()
 
diff --git a/scripts/trainers/hf_train_opt.py b/scripts/trainers/hf_train_opt.py
new file mode 100644
index 000000000..a1f117c73
--- /dev/null
+++ b/scripts/trainers/hf_train_opt.py
@@ -0,0 +1,97 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+import argparse
+
+from transformers import AutoTokenizer, OPTConfig, OPTForCausalLM, TrainingArguments
+
+from archai.datasets.nlp.fast_hf_dataset_provider import (
+    FastDataCollatorForLanguageModeling,
+    FastHfDatasetProvider,
+)
+from archai.trainers.nlp.hf_trainer import HfTrainer
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Trains an OPT using fast data loading and the Hugging Face trainer.")
+
+    parser.add_argument(
+        "-dn",
+        "--dataset_name",
+        type=str,
+        default="wikitext",
+        help="Name of the dataset to use (via the datasets library).",
+    )
+
+    parser.add_argument(
+        "-dcn",
+        "--dataset_config_name",
+        type=str,
+        default="wikitext-103-raw-v1",
+        help="Configuration name of the dataset to use (via the datasets library).",
+    )
+
+    parser.add_argument("-seq", "--seq_len", type=int, default=192, help="Sequence length.")
+
+    parser.add_argument("-ls", "--logging_steps", type=int, default=10, help="Number of steps between logs.")
+
+    parser.add_argument("-es", "--eval_steps", type=int, default=100, help="Number of steps between evaluations.")
+
+    parser.add_argument("-bsz", "--per_device_train_batch_size", type=int, default=64, help="Batch size per device.")
+
+    parser.add_argument("-lr", "--learning_rate", type=float, default=0.01, help="Learning rate.")
+
+    parser.add_argument("-wd", "--weight_decay", type=float, default=0.0, help="Weight decay.")
+
+    parser.add_argument("-n", "--max_steps", type=int, default=1, help="Maximum number of steps.")
+
+    args = parser.parse_args()
+
+    return args
+
+
+if __name__ == "__main__":
+    args = parse_args()
+
+    tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m", model_max_length=args.seq_len)
+    collator = FastDataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
+    dataset_provider = FastHfDatasetProvider.from_hub(
+        args.dataset_name,
+        dataset_config_name=args.dataset_config_name,
+        tokenizer=tokenizer,
+    )
+
+    train_dataset = dataset_provider.get_train_dataset(seq_len=args.seq_len)
+    eval_dataset = dataset_provider.get_val_dataset(seq_len=args.seq_len)
+
+    config = OPTConfig(
+        n_positions=args.seq_len,
+        hidden_size=768,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        vocab_size=50272,
+    )
+    model = OPTForCausalLM(config=config)
+
+    print(f"Total parameters: {sum(p.numel() for p in model.parameters())}")
+
+    training_args = TrainingArguments(
+        "hf-opt",
+        evaluation_strategy="steps",
+        logging_steps=args.logging_steps,
+        eval_steps=args.eval_steps,
+        per_device_train_batch_size=args.per_device_train_batch_size,
+        learning_rate=args.learning_rate,
+        weight_decay=args.weight_decay,
+        max_steps=args.max_steps,
+    )
+    trainer = HfTrainer(
+        model=model,
+        args=training_args,
+        data_collator=collator,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+    )
+
+    trainer.train()
diff --git a/setup.py b/setup.py
index 3e18ef9d5..95775e394 100644
--- a/setup.py
+++ b/setup.py
@@ -23,10 +23,8 @@
     "psutil",
     "pytest",
     "pytorch-lightning",
-    "pyunpack",
     "pyyaml",
     "ray>=1.0.0",
-    "runstats>=2.0.0",
     "scikit-learn",
     "seaborn",
     "send2trash>=1.8.0",
@@ -86,10 +84,8 @@ def filter_dependencies(*pkgs):
     "overrides",
     "plotly",
     "psutil",
-    "pyunpack",
     "pyyaml",
     "ray",
-    "runstats",
     "seaborn",
     "send2trash",
     "statopt",


Test patch:
diff --git a/tests/datasets/nlp/test_hf_dataset_provider.py b/tests/datasets/nlp/test_hf_dataset_provider.py
index 4f024001c..935bc79aa 100644
--- a/tests/datasets/nlp/test_hf_dataset_provider.py
+++ b/tests/datasets/nlp/test_hf_dataset_provider.py
@@ -10,7 +10,7 @@
 
 
 def test_hf_hub_dataset_provider():
-    dataset_provider = HfHubDatasetProvider("glue", "sst2")
+    dataset_provider = HfHubDatasetProvider("glue", dataset_config_name="sst2")
 
     # Assert that we can individually load training, validation and test datasets
     train_dataset = dataset_provider.get_train_dataset()
@@ -24,7 +24,7 @@ def test_hf_hub_dataset_provider():
 
 
 def test_hf_disk_dataset_provider():
-    dataset_hub_provider = HfHubDatasetProvider("glue", "sst2")
+    dataset_hub_provider = HfHubDatasetProvider("glue", dataset_config_name="sst2")
     train_dataset = dataset_hub_provider.get_train_dataset()
     train_dataset.save_to_disk("dataroot")