Repository: malteos/llm-datasets (language: python)

Issue:
Added informing condition to the license
Just added the "informing condition" to the license class.


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
new file mode 100644
index 0000000..e692190
--- /dev/null
+++ b/.github/workflows/docs.yml
@@ -0,0 +1,30 @@
+name: docs
+on:
+  push:
+    branches:
+      - master
+      - main
+      - dev
+permissions:
+  contents: write
+jobs:
+  deploy:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Configure Git Credentials
+        run: |
+          git config user.name github-actions[bot]
+          git config user.email 41898282+github-actions[bot]@users.noreply.github.com
+      - uses: actions/setup-python@v4
+        with:
+          python-version: 3.x
+      - run: echo "cache_id=$(date --utc '+%V')" >> $GITHUB_ENV
+      - uses: actions/cache@v3
+        with:
+          key: mkdocs-material-${{ env.cache_id }}
+          path: .cache
+          restore-keys: |
+            mkdocs-material-
+      - run: pip install -r ./requirements/docs.txt
+      - run: mkdocs gh-deploy --force
\ No newline at end of file
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 227bcb1..1afe847 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -3,9 +3,9 @@ repos:
     rev: 23.9.1
     hooks:
       - id: black
-  - repo: https://github.com/PyCQA/flake8
-    rev:  3.9.2
-    hooks:
-      - id: flake8
-        entry: pflake8
-        additional_dependencies: [pyproject-flake8]
\ No newline at end of file
+  # - repo: https://github.com/PyCQA/flake8
+  #   rev:  3.9.2
+  #   hooks:
+  #     - id: flake8
+  #       entry: pflake8
+  #       additional_dependencies: [pyproject-flake8]
\ No newline at end of file
diff --git a/.vscode/launch.json b/.vscode/launch.json
new file mode 100644
index 0000000..b6910c6
--- /dev/null
+++ b/.vscode/launch.json
@@ -0,0 +1,64 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Python: Current File",
+            "type": "python",
+            "request": "launch",
+            "program": "${file}",
+            "console": "integratedTerminal",
+            "justMyCode": false,
+            "env": {
+                "PYTHONPATH": "${env:PYTHONPATH}:${workspaceFolder}${pathSeparator}src${pathSeparator}"
+            }
+        },
+        {
+            "name": "Python: Current File - custom env",
+            "type": "python",
+            "request": "launch",
+            "program": "${file}",
+            "console": "integratedTerminal",
+            "justMyCode": false,
+            "env": {
+                "PYTHONPATH": "${env:PYTHONPATH}:${workspaceFolder}${pathSeparator}src${pathSeparator}:/data/experiments/mostendorff/internal-lm-datasets/src"
+            }
+        },
+        // {
+        //     "name": "Debug Tests",
+        //     "type": "python",
+        //     "request": "launch",
+        //     "purpose": [
+        //         "debug-test"
+        //     ],
+        //     "module": "pytest",
+        //     "console": "integratedTerminal",
+        //     "justMyCode": false,
+        //     "env": {
+        //         "_PYTEST_RAISE": "1"
+        //     },
+        // },
+        {
+            "name": "Debug Tests: Current File",
+            "type": "python",
+            "request": "launch",
+            // "purpose": [
+            //     "debug-test"
+            // ],
+            "module": "pytest",
+            "console": "integratedTerminal",
+            "justMyCode": true,
+            "args": [
+                "${file}",
+                "-s",
+                // "--log-cli",
+                "--log-cli-level=INFO",
+            ],
+            "env": {
+                "_PYTEST_RAISE": "1"
+            },
+        }
+    ]
+}
\ No newline at end of file
diff --git a/.vscode/settings.json b/.vscode/settings.json
index cd30724..b33fc0f 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -4,16 +4,17 @@
     "files.trimTrailingWhitespace": true,
     "files.autoSave": "onFocusChange",
     "git.autofetch": true,
-    "[jsonc]": {
-        "editor.defaultFormatter": "vscode.json-language-features"
-    },
-    "[python]": {
-        "editor.defaultFormatter": "ms-python.black-formatter"
-    },
-    "python.defaultInterpreterPath": "/usr/local/bin/python",
-    "python.formatting.provider": "black",
+    // "[jsonc]": {
+    //     "editor.defaultFormatter": "vscode.json-language-features"
+    // },
+    // "[python]": {
+    //     "editor.defaultFormatter": "ms-python.black-formatter"
+    // },
+    // "python.languageServer": "Pylance",
+    "python.defaultInterpreterPath": "/usr/bin/python",
+    // "python.formatting.provider": "black",
     "python.testing.unittestEnabled": false,
-    "python.testing.pytestEnabled": true,
+    // "python.testing.pytestEnabled": true,
     "pylint.args": [
         "--rcfile=pyproject.toml"
     ],
@@ -26,8 +27,35 @@
     "isort.args": [
         "--settings-path=pyproject.toml"
     ],
-    "python.linting.enabled": true,
-    "python.linting.flake8Enabled": true,
-    "python.linting.lintOnSave": true,
-    "python.linting.pylintEnabled": true,
+    // "python.linting.enabled": true,
+    // "python.linting.flake8Enabled": true,
+    // "python.linting.lintOnSave": true,
+    // "python.linting.pylintEnabled": true,
+    "workbench.colorCustomizations": {
+        "activityBar.activeBackground": "#79d9bc",
+        "activityBar.background": "#79d9bc",
+        "activityBar.foreground": "#15202b",
+        "activityBar.inactiveForeground": "#15202b99",
+        "activityBarBadge.background": "#b46ad5",
+        "activityBarBadge.foreground": "#15202b",
+        "commandCenter.border": "#15202b99",
+        "sash.hoverBorder": "#79d9bc",
+        "statusBar.background": "#51cea8",
+        "statusBar.foreground": "#15202b",
+        "statusBarItem.hoverBackground": "#34b890",
+        "statusBarItem.remoteBackground": "#51cea8",
+        "statusBarItem.remoteForeground": "#15202b",
+        "titleBar.activeBackground": "#51cea8",
+        "titleBar.activeForeground": "#15202b",
+        "titleBar.inactiveBackground": "#51cea899",
+        "titleBar.inactiveForeground": "#15202b99"
+    },
+    "peacock.remoteColor": "#51cea8",
+    "python.testing.pytestArgs": [
+        "tests"
+    ],
+    "[python]": {
+        "editor.defaultFormatter": "ms-python.black-formatter"
+    },
+    "python.formatting.provider": "none",
 }
\ No newline at end of file
diff --git a/README.md b/README.md
index db2e966..500b39e 100644
--- a/README.md
+++ b/README.md
@@ -1,31 +1,43 @@
 # lm-datasets
 
-<img align="left" src="https://github.com/malteos/lm-datasets/raw/main/images/A_colorful_parrot_sitting_on_a_pile_of_books__whit-removebg-preview.png" height="200" />
+<img align="left" src="https://github.com/malteos/lm-datasets/raw/main/docs/images/A_colorful_parrot_sitting_on_a_pile_of_books__whit-removebg-preview.png" height="200" />
 
 ![](https://img.shields.io/pypi/l/lm-datasets?style=flat-square)
 [![PRs Welcome](https://img.shields.io/badge/PRs-welcome-brightgreen.svg?style=flat-square)](https://makeapullrequest.com)
 
 **lm-datasets is a collection of datasets for language model training including scripts for downloading, preprocesssing, and sampling.**
 
+The documentation is available [here](https://malteos.github.io/lm-datasets/).
 
-## Installation
+## Quick start
+
+### Installation
+
+Install the `lm-datasets` package with [pip](https://pypi.org/project/lm-datasets/):
 
 ```bash
 pip install lm-datasets
 ```
 
-## Usage
+In order to keep the package minimal by default, `lm-datasets` comes with optional dependencies useful for some use cases.
+For example, if you want to have the text extraction for all available datasets, run:
+
+```bash
+pip install lm-datasets[datasets]
+```
+
+### Download and text extraction
 
 To download and extract the plain-text of one or more datasets, run the following command:
 
 ```bash
-python -m lm_datasets.extract_plaintext $DATASET_ID $OUTPUT_DIR
+lm_datasets extract_text $DATASET_ID $OUTPUT_DIR
 ```
 
 By default, output is saved as JSONL files. To change the output format, you can use the `--output_format` argument as below:
 
 ```bash
-python -m lm_datasets.extract_plaintext $DATASET_ID $OUTPUT_DIR --output_format parquet  --output_compression zstd
+lm_datasets extract_text $DATASET_ID $OUTPUT_DIR --output_format parquet  --output_compression zstd
 ```
 
 ### Available datasets
@@ -33,7 +45,7 @@ python -m lm_datasets.extract_plaintext $DATASET_ID $OUTPUT_DIR --output_format
 A list or table with all available datasets can be print with the follow command:
 
 ```bash
-python -m lm_datasets.print_stats --print_output md
+lm_datasets print_stats --print_output md
 ```
 #### Token count by language
 
@@ -108,7 +120,6 @@ python -m lm_datasets.print_stats --print_output md
 | macocu                           | 23 B     |
 | marcell_legislative_subcorpus_v2 | 31 M     |
 | norwegian_cc                     | 5 B      |
-| opengptx                         | 26 B     |
 | openlegaldata                    | 10 B     |
 | oscar                            | 9 T      |
 | oscar_opengptx                   | 245 B    |
@@ -145,10 +156,15 @@ python -m lm_datasets.print_stats --print_output md
 ### Dataset viewer
 
 We provide a Web-based application through streamlit to browse all datasets and their contained text content.
-To start the app, run the following command:
+To start the app, first clone this repository, install dependencies, and run the following command:
 
 ```bash
-streamlit viewer/app.py $RAW_DATASETS_DIR $PROCESSED_DATASET_DIR
+# clone is needed since streamlit does not support apps from modules yet
+git clone https://github.com/malteos/lm-datasets.git
+
+streamlit run src/lm_datasets/viewer/app.py -- \
+    --raw_datasets_dir=$RAW_DATASETS_DIR \
+    --output_dir=$PROCESSED_DATASET_DIR
 ```
 
 
@@ -156,6 +172,9 @@ streamlit viewer/app.py $RAW_DATASETS_DIR $PROCESSED_DATASET_DIR
 
 ### Setup environment
 
+To setup, your local development environment we recommend conda and cloning the repository.
+The repository also includes settings and launch scripts for VSCode.
+
 ```bash
 git clone git@github.com:malteos/lm-datasets.git
 cd lm-datasets
@@ -166,6 +185,11 @@ conda activate lm-datasets
 pip install -r requirements.txt
 ```
 
+Alternatively, you can install the Python package directly from the dev branch:
+
+```bash
+pip install git+https://github.com/malteos/lm-datasets.git@dev
+```
 
 ### Install the pre-commit hooks
 
@@ -188,6 +212,11 @@ The tests can be executed with:
 pytest --doctest-modules --cov-report term --cov=lm_datasets
 ```
 
+## Acknowledgements
+
+The work on the lm-datasets software is partially funded by the [German Federal Ministry for Economic Affairs and Climate Action (BMWK)](https://www.bmwk.de/Navigation/EN/Home/home.html)
+through the project [OpenGPT-X](https://opengpt-x.de/en/) (project no. 68GX21007D).
+
 ## License
 
 Apache 2.0
diff --git a/docs/README.md b/docs/README.md
new file mode 100644
index 0000000..40641cf
--- /dev/null
+++ b/docs/README.md
@@ -0,0 +1,17 @@
+# Documentation
+
+The documentation uses [mkdocs](https://www.mkdocs.org).
+
+## Commands
+
+* `mkdocs serve` - Start the live-reloading docs server.
+* `mkdocs serve -a 0.0.0.0:8000` - Server accessible over network.
+* `mkdocs build` - Build the documentation site.
+* `mkdocs -h` - Print help message and exit.
+
+## Project layout
+
+    mkdocs.yml    # The configuration file.
+    docs/
+        index.md  # The documentation homepage.
+        ...       # Other markdown pages, images and other files.
\ No newline at end of file
diff --git a/docs/add-your-own-data.md b/docs/add-your-own-data.md
new file mode 100644
index 0000000..a41a027
--- /dev/null
+++ b/docs/add-your-own-data.md
@@ -0,0 +1,59 @@
+# Integrate a custom dataset
+
+## Write a dataset class
+
+The first step for adding a new dataset is write a new dataset class.
+If your data comes from a common source such as Huggingface, you can build upon existing abstractions.
+For example, Huggingface datasets only needed to specify some metadata like dataset ID, title etc. and the column where the textual data can be extracted from (by default `text` column):
+
+```python
+# my_datasets/pg19.py
+
+from lm_datasets.datasets.hf_dataset import HFDataset
+from lm_datasets.datasets.base import License, Availability
+
+class PG19Dataset(HFDataset):
+    DATASET_ID = "pg19"
+    TITLE = "Project Gutenberg books published before 1919"
+    HOMEPAGE = "https://huggingface.co/datasets/pg19"
+    LICENSE = License("Apache License Version 2.0 (or public domain?)", url="https://www.apache.org/licenses/LICENSE-2.0.html")
+    CITATION = """@article{raecompressive2019,
+        author = {Rae, Jack W and Potapenko, Anna and Jayakumar, Siddhant M and
+                    Hillier, Chloe and Lillicrap, Timothy P},
+        title = {Compressive Transformers for Long-Range Sequence Modelling},
+        journal = {arXiv preprint},
+        url = {https://arxiv.org/abs/1911.05507},
+        year = {2019},
+        }
+        """  # noqa
+    AVAILIBILITY = Availability.DIRECT_DOWNLOAD
+
+    HF_DATASET_ID = "pg19"
+    HF_DATASET_SPLIT = "train"
+    streaming = True
+    text_column_name = "text"
+    title_column_name = "short_book_title"
+```
+
+## Register new dataset classes
+
+Each dataset class needs to be registered with `lm-dataset` such that the commands know what classes are available.
+This can be done by making a new Python module with a `get_registered_dataset_classes` method that returns a list of dataset classes:
+
+```python
+# my_datasets/dataset_registry.py
+from my_datasets.pg19 import PG19Dataset
+
+def get_registered_dataset_classes():
+    return [
+        PG19Dataset,
+    ]
+```
+
+## Load registry in commands
+
+To load the registerd datasets in the pipeline commands, you need to specify the `--extra_dataset_registries` argument:
+
+```bash
+lm_datasets compose ... -extra_dataset_registries=my_datasets.dataset_registry
+```
\ No newline at end of file
diff --git a/docs/api/base_dataset.md b/docs/api/base_dataset.md
new file mode 100644
index 0000000..952be2e
--- /dev/null
+++ b/docs/api/base_dataset.md
@@ -0,0 +1,3 @@
+# BaseDataset
+
+::: lm_datasets.datasets.base.BaseDataset
\ No newline at end of file
diff --git a/docs/api/hf_dataset.md b/docs/api/hf_dataset.md
new file mode 100644
index 0000000..610a86e
--- /dev/null
+++ b/docs/api/hf_dataset.md
@@ -0,0 +1,3 @@
+# HFDataset
+
+::: lm_datasets.datasets.hf_dataset.HFDataset
\ No newline at end of file
diff --git a/docs/api/jsonl_dataset.md b/docs/api/jsonl_dataset.md
new file mode 100644
index 0000000..7d157e2
--- /dev/null
+++ b/docs/api/jsonl_dataset.md
@@ -0,0 +1,3 @@
+# BaseDataset
+
+::: lm_datasets.datasets.jsonl_dataset.JSONLDataset
\ No newline at end of file
diff --git a/docs/compose-train-validation-data.md b/docs/compose-train-validation-data.md
new file mode 100644
index 0000000..5d1b09a
--- /dev/null
+++ b/docs/compose-train-validation-data.md
@@ -0,0 +1,12 @@
+# Compose dataset
+
+The pipeline step that produces the final training or validation set is the `compose` step.
+Before you run this command, you should specify in the [config](config-files.md) files what datasets should be selected and how they should be sampled.
+
+```bash
+lm-datasets compose –-split=train –-configs=my_dataset.yaml \
+	--text_data_dir=/data/my_text_data \
+	--composed_data_dir=/data/my_composed_data/train/
+```
+
+Depending on the your system (especially IO-speed) and dataset size this step can take a substantial amount of time (> 24 hours for a 1T token dataset).
diff --git a/docs/config-files.md b/docs/config-files.md
new file mode 100644
index 0000000..82177f6
--- /dev/null
+++ b/docs/config-files.md
@@ -0,0 +1,64 @@
+# Config Files
+
+`lm-datasets` allows you to specific general settings through config files so you do not need to specific always the same command line arguments.
+Several commands support passing the `--configs` argument which should point to one or more YAML-files on your file system. For example, the text extraction command:
+
+```bash
+lm_datasets extract_text ... --configs $PATH_TO_YAML_CONFIG_FILE
+```
+
+## Specifing local paths
+
+In the config files, you can store for example system specific settings like the local paths, where the raw dataset files are located:
+
+```yaml
+# ./examples/lm_datasets_configs/my_system.yaml
+local_dirs_by_source_id:
+  redpajama: /my_system_specific_data_directory/redpajama
+```
+
+The [RedPajama dataset](https://huggingface.co/datasets/togethercomputer/RedPajama-Data-1T) requires the manual download prior to the text extraction.
+With the above config, we tell the extraction command the path where we downloaded the RedPajama data by providing the config file:
+
+```bash
+lm_datasets extract_text redpajama_book --configs ./examples/lm_datasets_configs/my_system.yaml
+```
+
+## Dataset selection and sampling
+
+The configuration files are also needed for specifying the final dataset composition, including the selection of the datasets and their sampling.
+The following examples shows a config for an Italian dataset:
+
+```yaml
+# ./examples/lm_datasets_configs/italian_data.yaml
+
+# a fixed random seed for shuffling etc.
+seed: 0
+
+selected_dataset_ids:
+  # italian subsets
+  - itwac
+  - eurlex_it
+  - wikipedia_20231101_it
+  - wikibooks_it
+  - wikinews_it
+  - colossal_oscar_2023-23_it
+  - parlamint_it
+
+# down-sample webcrawled + up-sampled high quality
+sampling_factor_by_source_id:
+  colossal_oscar: 0.1
+
+sampling_factor_by_dataset_id:
+  itwac: 0.5
+  eurlex_it: 2
+  wikipedia_20231101_it: 3
+```
+
+To use this config, provide the path in the `--configs` argument:
+
+```bash
+# compose final dataset
+lm_datasets compose ... --configs ./examples/lm_datasets_configs/italian_data.yaml
+
+```
\ No newline at end of file
diff --git a/docs/extract-text-data.md b/docs/extract-text-data.md
new file mode 100644
index 0000000..661ea72
--- /dev/null
+++ b/docs/extract-text-data.md
@@ -0,0 +1,14 @@
+
+### Download and text extraction
+
+To download and extract the plain-text of one or more datasets, run the following command:
+
+```bash
+lm_datasets extract_text $DATASET_ID $OUTPUT_DIR
+```
+
+By default, output is saved as JSONL files. To change the output format, you can use the `--output_format` argument as below:
+
+```bash
+lm_datasets extract_text $DATASET_ID $OUTPUT_DIR --output_format parquet  --output_compression zstd
+```
diff --git a/docs/getting-started.md b/docs/getting-started.md
new file mode 100644
index 0000000..a9a9124
--- /dev/null
+++ b/docs/getting-started.md
@@ -0,0 +1,63 @@
+# Getting Started
+
+## Installation
+
+Install the `lm-datasets` package with [pip](https://pypi.org/project/lm-datasets/):
+
+```bash
+pip install lm-datasets
+```
+
+In order to keep the package minimal by default, `lm-datasets` comes with optional dependencies useful for some use cases.
+For example, if you want to have the text extraction for all available datasets, run:
+
+```bash
+pip install lm-datasets[datasets]
+```
+
+## Quick start
+
+### Download and text extraction
+
+To download and extract the plain-text of one or more datasets, run the following command:
+
+```bash
+lm_datasets extract_text $DATASET_ID $OUTPUT_DIR
+```
+
+By default, output is saved as JSONL files. To change the output format, you can use the `--output_format` argument as below:
+
+```bash
+lm_datasets extract_text $DATASET_ID $OUTPUT_DIR --output_format parquet  --output_compression zstd
+```
+
+### Available datasets
+
+A list or table with all available datasets can be print with the follow command:
+
+```bash
+lm_datasets print_stats --print_output md
+```
+
+### Pipeline commands
+
+```
+usage: lm-datasets <command> [<args>]
+
+positional arguments:
+  {chunkify,collect_metrics,compose,convert_parquet_to_jsonl,extract_text,hf_upload,print_stats,shuffle,train_tokenizer}
+                        lm-datasets command helpers
+    chunkify            Split the individual datasets into equally-sized file chunks (based on bytes or rows)
+    collect_metrics     Collect metrics (token count etc.) from extracted texts
+    compose             Compose the final train/validation set based on the individual datasets
+    convert_parquet_to_jsonl
+                        Convert Parquet files to JSONL
+    extract_text        Extract text from raw datasets
+    hf_upload           Upload files or directories to Huggingface Hub.
+    print_stats         Print dataset statistics as CSV, Markdown, ...
+    shuffle             Shuffle the individual datasets on the file-chunk level (no global shuffle!)
+    train_tokenizer     Train a tokenizer (only: sentencepiece supproted)
+
+options:
+  -h, --help            show this help message and exit
+```
\ No newline at end of file
diff --git a/images/A_colorful_parrot_sitting_on_a_pile_of_books__whit-removebg-preview.png b/docs/images/A_colorful_parrot_sitting_on_a_pile_of_books__whit-removebg-preview.png
similarity index 100%
rename from images/A_colorful_parrot_sitting_on_a_pile_of_books__whit-removebg-preview.png
rename to docs/images/A_colorful_parrot_sitting_on_a_pile_of_books__whit-removebg-preview.png
diff --git a/docs/images/data-schema.svg b/docs/images/data-schema.svg
new file mode 100644
index 0000000..0c5e36f
--- /dev/null
+++ b/docs/images/data-schema.svg
@@ -0,0 +1,1 @@
+<svg width="2210" height="1526" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" xml:space="preserve" overflow="hidden"><g transform="translate(-1067 -605)"><g><path d="M1155.5 959.288C1155.5 940.075 1171.07 924.5 1190.29 924.5L2081.71 924.5C2100.92 924.5 2116.5 940.075 2116.5 959.288L2116.5 1926.71C2116.5 1945.92 2100.92 1961.5 2081.71 1961.5L1190.29 1961.5C1171.07 1961.5 1155.5 1945.92 1155.5 1926.71Z" stroke="#172C51" stroke-width="6.875" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#DEEBF7" fill-rule="evenodd" fill-opacity="1"/><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="700" font-stretch="normal" font-size="64" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1532.34 1013)">Dataset</text><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="italic" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1263 1081)">A specific dataset implementation </text><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="italic" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1321.87 1147)">(e.g., Wikipedia 2023</text><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="italic" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1792.08 1147)">-</text><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="italic" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1810.41 1147)">01 de)</text><path d="M1161.5 657.676C1161.5 631.621 1182.62 610.5 1208.68 610.5L2075.32 610.5C2101.38 610.5 2122.5 631.621 2122.5 657.676L2122.5 845.324C2122.5 871.379 2101.38 892.5 2075.32 892.5L1208.68 892.5C1182.62 892.5 1161.5 871.379 1161.5 845.324Z" stroke="#172C51" stroke-width="6.875" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#DEEBF7" fill-rule="evenodd" fill-opacity="1"/><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="700" font-stretch="normal" font-size="64" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1435.65 709)">Dataset Source</text><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="italic" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1306.97 777)">A common source for different </text><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="italic" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1361.51 843)">datasets (e.g., Wikipedia)</text><path d="M1208.5 1726.67C1208.5 1718.29 1215.29 1711.5 1223.67 1711.5L1608.33 1711.5C1616.71 1711.5 1623.5 1718.29 1623.5 1726.67L1623.5 1787.33C1623.5 1795.71 1616.71 1802.5 1608.33 1802.5L1223.67 1802.5C1215.29 1802.5 1208.5 1795.71 1208.5 1787.33Z" stroke="#172C51" stroke-width="4.58333" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#F2F2F2" fill-rule="evenodd" fill-opacity="1"/><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1250.76 1777)">JSONL dataset</text><path d="M1648.5 1729.67C1648.5 1721.29 1655.29 1714.5 1663.67 1714.5L2053.33 1714.5C2061.71 1714.5 2068.5 1721.29 2068.5 1729.67L2068.5 1790.33C2068.5 1798.71 2061.71 1805.5 2053.33 1805.5L1663.67 1805.5C1655.29 1805.5 1648.5 1798.71 1648.5 1790.33Z" stroke="#172C51" stroke-width="4.58333" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#F2F2F2" fill-rule="evenodd" fill-opacity="1"/><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1739.83 1780)">HF dataset</text><path d="M1258.5 1838.67C1258.5 1830.29 1265.29 1823.5 1273.67 1823.5L1696.33 1823.5C1704.71 1823.5 1711.5 1830.29 1711.5 1838.67L1711.5 1899.33C1711.5 1907.71 1704.71 1914.5 1696.33 1914.5L1273.67 1914.5C1265.29 1914.5 1258.5 1907.71 1258.5 1899.33Z" stroke="#172C51" stroke-width="4.58333" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#F2F2F2" fill-rule="evenodd" fill-opacity="1"/><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1316.74 1889)">Parquet dataset</text><rect x="1260.5" y="1217.5" width="318" height="91.0001" stroke="#172C51" stroke-width="4.58333" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#7F7F7F" fill-opacity="1"/><text fill="#FFFFFF" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1299.75 1282)">Languages</text><rect x="1286.5" y="1329.5" width="250" height="89.9998" stroke="#172C51" stroke-width="4.58333" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#7F7F7F" fill-opacity="1"/><text fill="#FFFFFF" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1325.02 1394)">License</text><rect x="1300.5" y="1443.5" width="203" height="91.0001" stroke="#172C51" stroke-width="4.58333" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#7F7F7F" fill-opacity="1"/><text fill="#FFFFFF" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1334.33 1509)">Genre</text><rect x="1565.5" y="1329.5" width="390" height="89.9998" stroke="#172C51" stroke-width="4.58333" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#7F7F7F" fill-opacity="1"/><text fill="#FFFFFF" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1604.6 1394)">Quality issues</text><rect x="1610.5" y="1214.5" width="339" height="91.0001" stroke="#172C51" stroke-width="4.58333" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#7F7F7F" fill-opacity="1"/><text fill="#FFFFFF" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1649.89 1279)">Availability</text><rect x="1523.5" y="1443.5" width="246" height="91.0001" stroke="#172C51" stroke-width="4.58333" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#7F7F7F" fill-opacity="1"/><text fill="#FFFFFF" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1557.35 1509)">Citation</text><rect x="1790.5" y="1443.5" width="125" height="91.0001" stroke="#172C51" stroke-width="4.58333" stroke-linecap="square" stroke-linejoin="miter" stroke-miterlimit="8" stroke-dasharray="13.75 4.58333" stroke-opacity="1" fill="#7F7F7F" fill-opacity="1"/><text fill="#FFFFFF" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1824.93 1509)">…</text><path d="M2310.5 1485.84C2310.5 1460.25 2331.25 1439.5 2356.84 1439.5L3224.16 1439.5C3249.75 1439.5 3270.5 1460.25 3270.5 1485.84L3270.5 1670.16C3270.5 1695.75 3249.75 1716.5 3224.16 1716.5L2356.84 1716.5C2331.25 1716.5 2310.5 1695.75 2310.5 1670.16Z" stroke="#172C51" stroke-width="6.875" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#DEEBF7" fill-rule="evenodd" fill-opacity="1"/><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="700" font-stretch="normal" font-size="64" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 2562.18 1531)">Dataset Registry</text><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="italic" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 2503.17 1620)">List of all default datasets</text><path d="M1762.5 1839.67C1762.5 1831.29 1769.29 1824.5 1777.67 1824.5L1962.33 1824.5C1970.71 1824.5 1977.5 1831.29 1977.5 1839.67L1977.5 1900.33C1977.5 1908.71 1970.71 1915.5 1962.33 1915.5L1777.67 1915.5C1769.29 1915.5 1762.5 1908.71 1762.5 1900.33Z" stroke="#172C51" stroke-width="4.58333" stroke-linecap="square" stroke-linejoin="miter" stroke-miterlimit="8" stroke-dasharray="13.75 4.58333" stroke-opacity="1" fill="#F2F2F2" fill-rule="evenodd" fill-opacity="1"/><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="50" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1844.47 1888)">…</text><path d="M2309.5 1906.32C2309.5 1883.78 2327.78 1865.5 2350.32 1865.5L3229.68 1865.5C3252.23 1865.5 3270.5 1883.78 3270.5 1906.32L3270.5 2068.68C3270.5 2091.23 3252.23 2109.5 3229.68 2109.5L2350.32 2109.5C2327.78 2109.5 2309.5 2091.23 2309.5 2068.68Z" stroke="#000000" stroke-width="4.58333" stroke-linecap="square" stroke-linejoin="miter" stroke-miterlimit="8" stroke-dasharray="13.75 4.58333" stroke-opacity="1" fill="#F2F2F2" fill-rule="evenodd" fill-opacity="1"/><text fill="#7F7F7F" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="700" font-stretch="normal" font-size="64" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 2446.8 1955)">Custom Dataset Registry</text><text fill="#7F7F7F" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="italic" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 2537.62 2044)">List of custom datasets</text><path d="M2309.5 635.962C2309.5 621.347 2321.35 609.5 2335.96 609.5L3243.04 609.5C3257.65 609.5 3269.5 621.347 3269.5 635.962L3269.5 1314.04C3269.5 1328.65 3257.65 1340.5 3243.04 1340.5L2335.96 1340.5C2321.35 1340.5 2309.5 1328.65 2309.5 1314.04Z" stroke="#172C51" stroke-width="6.875" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#DEEBF7" fill-rule="evenodd" fill-opacity="1"/><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="700" font-stretch="normal" font-size="64" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 2696.08 695)">Config</text><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="italic" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 2468.01 763)">A reproducible configuration </text><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="italic" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 2582.34 829)">of the final dataset</text><rect x="2786.5" y="985.5" width="315" height="91.0001" stroke="#172C51" stroke-width="4.58333" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#7F7F7F" fill-opacity="1"/><text fill="#FFFFFF" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 2889.89 1050)">Seed</text><rect x="2454.5" y="874.5" width="647" height="92.0001" stroke="#172C51" stroke-width="4.58333" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#7F7F7F" fill-opacity="1"/><text fill="#FFFFFF" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 2589.51 940)">Selected datasets</text><rect x="2454.5" y="984.5" width="315" height="92.0001" stroke="#172C51" stroke-width="4.58333" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#7F7F7F" fill-opacity="1"/><text fill="#FFFFFF" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 2505.82 1050)">Sampling</text><rect x="2940.5" y="1214.5" width="161" height="92.0001" stroke="#172C51" stroke-width="4.58333" stroke-linecap="square" stroke-linejoin="miter" stroke-miterlimit="8" stroke-dasharray="13.75 4.58333" stroke-opacity="1" fill="#7F7F7F" fill-opacity="1"/><text fill="#FFFFFF" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 2993.18 1280)">…</text><rect x="2454.5" y="1096.5" width="315" height="91.0001" stroke="#172C51" stroke-width="4.58333" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#7F7F7F" fill-opacity="1"/><text fill="#FFFFFF" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 2498.7 1162)">Split sizes</text><rect x="1229.5" y="1564.5" width="390" height="91.0001" stroke="#172C51" stroke-width="4.58333" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#7F7F7F" fill-opacity="1"/><text fill="#FFFFFF" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1287.5 1630)">Token count</text><rect x="1641.5" y="1564.5" width="390" height="91.0001" stroke="#172C51" stroke-width="4.58333" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#7F7F7F" fill-opacity="1"/><text fill="#FFFFFF" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1715.57 1630)">Byte count</text><path d="M0-3.4375 91.2022-3.4375 91.2022 695.102 28.1814 695.102 28.1814 688.227 87.7647 688.227 84.3272 691.664 84.3272 0 87.7647 3.4375 0 3.4375ZM32.7647 705.414 5.2647 691.664 32.7647 677.914Z" fill="#000000" fill-rule="nonzero" fill-opacity="1" transform="matrix(-1 1.22465e-16 1.22465e-16 1 1160.76 751.5)"/><path d="M2116.5 1440.06 2216.61 1440.06 2216.61 1578.37 2213.18 1574.93 2286.94 1574.93 2286.94 1581.8 2209.74 1581.8 2209.74 1443.5 2213.18 1446.94 2116.5 1446.94ZM2282.35 1564.62 2309.85 1578.37 2282.35 1592.12Z" fill="#000000" fill-rule="nonzero" fill-opacity="1"/><rect x="2786.5" y="1096.5" width="315" height="91.0001" stroke="#172C51" stroke-width="4.58333" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#7F7F7F" fill-opacity="1"/><text fill="#FFFFFF" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 2826.3 1162)">Data paths</text><path d="M1212.5 1993.5 2039 1993.5 2095.5 2050 2039 2106.5 1212.5 2106.5Z" stroke="#172C51" stroke-width="6.875" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#2F5597" fill-rule="evenodd" fill-opacity="1"/><text fill="#FFFFFF" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="64" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1323.52 2073)">Extract text from dataset</text><path d="M3.43748-0.0118472 3.87169 125.974-3.00327 125.997-3.43748 0.0118472ZM14.1683 121.355 0.513189 148.902-13.3315 121.45Z" fill="#000000" fill-rule="nonzero" fill-opacity="1" transform="matrix(1 0 0 -1 2789.5 1865.4)"/><path d="M2792.94 1340.46 2793.73 1416.98 2786.85 1417.05 2786.06 1340.54ZM2803.99 1412.29 2790.53 1439.93 2776.49 1412.58Z" fill="#000000" fill-rule="nonzero" fill-opacity="1"/><rect x="2454.5" y="1216.5" width="469" height="92.0001" stroke="#172C51" stroke-width="4.58333" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#7F7F7F" fill-opacity="1"/><text fill="#FFFFFF" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 2508.33 1282)">Filter thresholds</text><path d="M0 0 0.000360892 31.2259" stroke="#000000" stroke-width="6.875" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="none" fill-rule="evenodd" transform="matrix(-1 0 0 1 1626.5 1961.5)"/></g></g></svg>
\ No newline at end of file
diff --git a/docs/images/pipeline.svg b/docs/images/pipeline.svg
new file mode 100644
index 0000000..923ecfe
--- /dev/null
+++ b/docs/images/pipeline.svg
@@ -0,0 +1,1 @@
+<svg width="4042" height="1146" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" xml:space="preserve" overflow="hidden"><g transform="translate(-238 -916)"><g><path d="M2497.5 957.169C2497.5 938.574 2512.57 923.5 2531.17 923.5L3100.83 923.5C3119.43 923.5 3134.5 938.574 3134.5 957.169L3134.5 1327.83C3134.5 1346.43 3119.43 1361.5 3100.83 1361.5L2531.17 1361.5C2512.57 1361.5 2497.5 1346.43 2497.5 1327.83Z" stroke="#172C51" stroke-width="6.875" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#DEEBF7" fill-rule="evenodd" fill-opacity="1"/><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="700" font-stretch="normal" font-size="64" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 2677.8 1011)">Data Split</text><path d="M3577.5 957.169C3577.5 938.574 3592.57 923.5 3611.17 923.5L4089.83 923.5C4108.43 923.5 4123.5 938.574 4123.5 957.169L4123.5 1327.83C4123.5 1346.43 4108.43 1361.5 4089.83 1361.5L3611.17 1361.5C3592.57 1361.5 3577.5 1346.43 3577.5 1327.83Z" stroke="#172C51" stroke-width="6.875" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#DEEBF7" fill-rule="evenodd" fill-opacity="1"/><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="700" font-stretch="normal" font-size="64" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 3636.3 1011)">Tokenized Data</text><path d="M3634.5 1081.67C3634.5 1072.19 3642.19 1064.5 3651.67 1064.5L4062.33 1064.5C4071.81 1064.5 4079.5 1072.19 4079.5 1081.67L4079.5 1150.33C4079.5 1159.81 4071.81 1167.5 4062.33 1167.5L3651.67 1167.5C3642.19 1167.5 3634.5 1159.81 3634.5 1150.33Z" stroke="#172C51" stroke-width="4.58333" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#F2F2F2" fill-rule="evenodd" fill-opacity="1"/><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 3698.68 1135)">Megatron</text><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 3912.95 1135)">-</text><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 3931.29 1135)">LM</text><path d="M3636.5 1213.67C3636.5 1204.19 3644.19 1196.5 3653.67 1196.5L4062.33 1196.5C4071.81 1196.5 4079.5 1204.19 4079.5 1213.67L4079.5 1282.33C4079.5 1291.81 4071.81 1299.5 4062.33 1299.5L3653.67 1299.5C3644.19 1299.5 3636.5 1291.81 3636.5 1282.33Z" stroke="#172C51" stroke-width="4.58333" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#F2F2F2" fill-rule="evenodd" fill-opacity="1"/><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 3710.1 1267)">Transformers</text><path d="M2532.5 1084.33C2532.5 1073.38 2541.38 1064.5 2552.33 1064.5L3074.67 1064.5C3085.62 1064.5 3094.5 1073.38 3094.5 1084.33L3094.5 1163.67C3094.5 1174.62 3085.62 1183.5 3074.67 1183.5L2552.33 1183.5C2541.38 1183.5 2532.5 1174.62 2532.5 1163.67Z" stroke="#172C51" stroke-width="4.58333" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#F2F2F2" fill-rule="evenodd" fill-opacity="1"/><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 2635.13 1144)">Training dataset</text><path d="M1408.5 970.834C1408.5 946.349 1428.35 926.5 1452.83 926.5L1988.17 926.5C2012.65 926.5 2032.5 946.349 2032.5 970.834L2032.5 1148.17C2032.5 1172.65 2012.65 1192.5 1988.17 1192.5L1452.83 1192.5C1428.35 1192.5 1408.5 1172.65 1408.5 1148.17Z" stroke="#172C51" stroke-width="6.875" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#DEEBF7" fill-rule="evenodd" fill-opacity="1"/><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="700" font-stretch="normal" font-size="64" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1586.74 1017)">Text Data</text><path d="M2532.5 1228.83C2532.5 1217.6 2541.6 1208.5 2552.83 1208.5L3068.17 1208.5C3079.4 1208.5 3088.5 1217.6 3088.5 1228.83L3088.5 1310.17C3088.5 1321.4 3079.4 1330.5 3068.17 1330.5L2552.83 1330.5C2541.6 1330.5 2532.5 1321.4 2532.5 1310.17Z" stroke="#172C51" stroke-width="4.58333" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#F2F2F2" fill-rule="evenodd" fill-opacity="1"/><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 2612.54 1289)">Validation dataset</text><path d="M1439.5 1072.33C1439.5 1061.38 1448.38 1052.5 1459.33 1052.5L1696.67 1052.5C1707.62 1052.5 1716.5 1061.38 1716.5 1072.33L1716.5 1151.67C1716.5 1162.62 1707.62 1171.5 1696.67 1171.5L1459.33 1171.5C1448.38 1171.5 1439.5 1162.62 1439.5 1151.67Z" stroke="#172C51" stroke-width="4.58333" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#F2F2F2" fill-rule="evenodd" fill-opacity="1"/><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1495.03 1131)">JSONL</text><path d="M1728.5 1074.33C1728.5 1063.38 1737.38 1054.5 1748.33 1054.5L1985.67 1054.5C1996.62 1054.5 2005.5 1063.38 2005.5 1074.33L2005.5 1153.67C2005.5 1164.62 1996.62 1173.5 1985.67 1173.5L1748.33 1173.5C1737.38 1173.5 1728.5 1164.62 1728.5 1153.67Z" stroke="#172C51" stroke-width="4.58333" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#F2F2F2" fill-rule="evenodd" fill-opacity="1"/><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="55" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1781.86 1133)">Parquet</text><path d="M915.5 984.5 1311.5 984.5 1379.5 1052.5 1311.5 1120.5 915.5 1120.5Z" stroke="#172C51" stroke-width="6.875" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#2F5597" fill-rule="evenodd" fill-opacity="1"/><text fill="#FFFFFF" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="73" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 947.833 1079)">Extract text</text><path d="M459.5 963.773C459.5 943.188 476.188 926.5 496.773 926.5L849.227 926.5C869.812 926.5 886.5 943.188 886.5 963.773L886.5 1661.23C886.5 1681.81 869.812 1698.5 849.227 1698.5L496.773 1698.5C476.188 1698.5 459.5 1681.81 459.5 1661.23Z" stroke="#172C51" stroke-width="6.875" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#DEEBF7" fill-rule="evenodd" fill-opacity="1"/><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="700" font-stretch="normal" font-size="64" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 535.878 1015)">Raw Data</text><path d="M497.5 1332.67C497.5 1324.29 504.29 1317.5 512.667 1317.5L833.333 1317.5C841.709 1317.5 848.5 1324.29 848.5 1332.67L848.5 1393.33C848.5 1401.71 841.709 1408.5 833.333 1408.5L512.667 1408.5C504.29 1408.5 497.5 1401.71 497.5 1393.33Z" stroke="#172C51" stroke-width="4.58333" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#F2F2F2" fill-rule="evenodd" fill-opacity="1"/><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="50" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 548.647 1381)">HF Datasets</text><path d="M493.5 1217.67C493.5 1209.29 500.291 1202.5 508.667 1202.5L829.333 1202.5C837.71 1202.5 844.5 1209.29 844.5 1217.67L844.5 1278.33C844.5 1286.71 837.71 1293.5 829.333 1293.5L508.667 1293.5C500.291 1293.5 493.5 1286.71 493.5 1278.33Z" stroke="#172C51" stroke-width="4.58333" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#F2F2F2" fill-rule="evenodd" fill-opacity="1"/><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="50" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 620.992 1267)">Web</text><path d="M497.5 1450.67C497.5 1442.29 504.29 1435.5 512.667 1435.5L833.333 1435.5C841.709 1435.5 848.5 1442.29 848.5 1450.67L848.5 1511.33C848.5 1519.71 841.709 1526.5 833.333 1526.5L512.667 1526.5C504.29 1526.5 497.5 1519.71 497.5 1511.33Z" stroke="#172C51" stroke-width="4.58333" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#F2F2F2" fill-rule="evenodd" fill-opacity="1"/><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="50" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 623.126 1499)">ELG</text><path d="M492.5 1099.67C492.5 1091.29 499.291 1084.5 507.667 1084.5L828.333 1084.5C836.71 1084.5 843.5 1091.29 843.5 1099.67L843.5 1160.33C843.5 1168.71 836.71 1175.5 828.333 1175.5L507.667 1175.5C499.291 1175.5 492.5 1168.71 492.5 1160.33Z" stroke="#172C51" stroke-width="4.58333" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#F2F2F2" fill-rule="evenodd" fill-opacity="1"/><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="50" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 561.139 1149)">Local files</text><path d="M497.5 1574.67C497.5 1566.29 504.29 1559.5 512.667 1559.5L833.333 1559.5C841.709 1559.5 848.5 1566.29 848.5 1574.67L848.5 1635.33C848.5 1643.71 841.709 1650.5 833.333 1650.5L512.667 1650.5C504.29 1650.5 497.5 1643.71 497.5 1635.33Z" stroke="#000000" stroke-width="4.58333" stroke-linecap="square" stroke-linejoin="miter" stroke-miterlimit="8" stroke-dasharray="13.75 4.58333" stroke-opacity="1" fill="#F2F2F2" fill-rule="evenodd" fill-opacity="1"/><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="50" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 647.189 1623)">…</text><path d="M1464.5 1225.5 1899.5 1225.5 1967.5 1293.5 1899.5 1361.5 1464.5 1361.5Z" stroke="#172C51" stroke-width="6.875" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#2F5597" fill-rule="evenodd" fill-opacity="1"/><text fill="#FFFFFF" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="73" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1497.24 1320)">Shuffle</text><path d="M1464.5 1404.5 1899.5 1404.5 1967.5 1472.5 1899.5 1540.5 1464.5 1540.5Z" stroke="#172C51" stroke-width="6.875" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#2F5597" fill-rule="evenodd" fill-opacity="1"/><text fill="#FFFFFF" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="73" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1497.24 1499)">Split by size</text><path d="M2062.5 979.5 2401 979.5 2469.5 1048 2401 1116.5 2062.5 1116.5Z" stroke="#172C51" stroke-width="6.875" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#2F5597" fill-rule="evenodd" fill-opacity="1"/><text fill="#FFFFFF" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="73" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 2094.58 1074)">Compose</text><path d="M3162.5 978.5 3481.5 978.5 3549.5 1046.5 3481.5 1114.5 3162.5 1114.5Z" stroke="#172C51" stroke-width="6.875" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#2F5597" fill-rule="evenodd" fill-opacity="1"/><text fill="#FFFFFF" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="73" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 3195.06 1073)">Tokenize</text><path d="M1464.5 1580.5 1899.5 1580.5 1967.5 1648.5 1899.5 1716.5 1464.5 1716.5Z" stroke="#172C51" stroke-width="6.875" stroke-linecap="butt" stroke-linejoin="miter" stroke-miterlimit="8" stroke-opacity="1" fill="#2F5597" fill-rule="evenodd" fill-opacity="1"/><text fill="#FFFFFF" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="normal" font-variant="normal" font-weight="400" font-stretch="normal" font-size="73" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1497.24 1675)">Count tokens</text><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="italic" font-variant="normal" font-weight="400" font-stretch="normal" font-size="64" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 316.229 1845)">The raw data can originate </text><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="italic" font-variant="normal" font-weight="400" font-stretch="normal" font-size="64" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 292.579 1922)">from different sources and be </text><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="italic" font-variant="normal" font-weight="400" font-stretch="normal" font-size="64" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 427.375 1999)">in diverse formats.</text><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="italic" font-variant="normal" font-weight="400" font-stretch="normal" font-size="64" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1432.22 1845)">The intermediate data </text><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="italic" font-variant="normal" font-weight="400" font-stretch="normal" font-size="64" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1413.7 1922)">representation contains </text><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="italic" font-variant="normal" font-weight="400" font-stretch="normal" font-size="64" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 1383.25 1999)">primarily the “plaintext”.</text><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="italic" font-variant="normal" font-weight="400" font-stretch="normal" font-size="64" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 2506.2 1850)">Based on the config, the </text><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="italic" font-variant="normal" font-weight="400" font-stretch="normal" font-size="64" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 2423.42 1927)">training and validation data is </text><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="italic" font-variant="normal" font-weight="400" font-stretch="normal" font-size="64" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 2532.55 2004)">selected and sampled.</text><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="italic" font-variant="normal" font-weight="400" font-stretch="normal" font-size="64" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 3514.49 1850)">The tokenized data can be </text><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="italic" font-variant="normal" font-weight="400" font-stretch="normal" font-size="64" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 3530.57 1927)">directly used by common </text><text fill="#000000" fill-opacity="1" font-family="Times New Roman,Times New Roman_MSFontService,sans-serif" font-style="italic" font-variant="normal" font-weight="400" font-stretch="normal" font-size="64" text-anchor="start" direction="ltr" writing-mode="lr-tb" unicode-bidi="normal" text-decoration="none" transform="matrix(1 0 0 1 3491.84 2004)">model training frameworks.</text></g></g></svg>
\ No newline at end of file
diff --git a/docs/index.md b/docs/index.md
new file mode 100644
index 0000000..56a709a
--- /dev/null
+++ b/docs/index.md
@@ -0,0 +1,8 @@
+# lm-datasets
+
+**lm-datasets** is a collection of datasets for language model training including scripts for downloading, preprocesssing, and sampling.
+
+- [Getting started](getting-started.md)
+- [Config files](config-files.md)
+- [Add your own data](add-your-own-data.md)
+
diff --git a/docs/overview.md b/docs/overview.md
new file mode 100644
index 0000000..d66a814
--- /dev/null
+++ b/docs/overview.md
@@ -0,0 +1,11 @@
+# Framework Overview
+
+## Data schema
+
+![Schema](images/data-schema.svg)
+
+## Pipeline
+
+![Pipeline](images/pipeline.svg)
+
+
diff --git a/docs/related-work.md b/docs/related-work.md
new file mode 100644
index 0000000..aef5792
--- /dev/null
+++ b/docs/related-work.md
@@ -0,0 +1,6 @@
+# Related Work
+
+This framework is heavily inspired by other open source projects:
+
+- [Huggingface Datasets](https://github.com/huggingface/datasets)
+- [TensorFlow Datasets](https://www.tensorflow.org/datasets)
diff --git a/examples/custom_datasets/README.md b/examples/custom_datasets/README.md
new file mode 100644
index 0000000..e69de29
diff --git a/examples/lm_datasets_configs/README.md b/examples/lm_datasets_configs/README.md
new file mode 100644
index 0000000..e69de29
diff --git a/examples/lm_datasets_configs/italian_data.yaml b/examples/lm_datasets_configs/italian_data.yaml
new file mode 100644
index 0000000..72c167b
--- /dev/null
+++ b/examples/lm_datasets_configs/italian_data.yaml
@@ -0,0 +1,23 @@
+# Italian data
+
+# a fixed random seed for shuffling etc.
+seed: 0
+
+selected_dataset_ids:
+  # italian subsets
+  - itwac
+  - eurlex_it
+  - wikipedia_20231101_it
+  - wikibooks_it
+  - wikinews_it
+  - colossal_oscar_2023-23_it
+  - parlamint_it
+
+# down-sample webcrawled + up-sampled high quality
+sampling_factor_by_source_id:
+  colossal_oscar: 0.1
+
+sampling_factor_by_dataset_id:
+  itwac: 0.5
+  eurlex_it: 2
+  wikipedia_20231101_it: 3
diff --git a/examples/lm_datasets_configs/my_system.yaml b/examples/lm_datasets_configs/my_system.yaml
new file mode 100644
index 0000000..f3a2e9d
--- /dev/null
+++ b/examples/lm_datasets_configs/my_system.yaml
@@ -0,0 +1,2 @@
+local_dirs_by_source_id:
+  redpajama: /my_system_specific_data_directory/redpajama
\ No newline at end of file
diff --git a/mkdocs.yml b/mkdocs.yml
new file mode 100644
index 0000000..4f6c27c
--- /dev/null
+++ b/mkdocs.yml
@@ -0,0 +1,40 @@
+site_name: "lm-datasets: Documentation"
+site_url: https://github.com/malteos/lm-datasets/
+
+site_description: "Documentation of the lm-datasets framework."
+site_author: "Malte Ostendorff and lm-datasets contributors"
+docs_dir: docs/
+repo_name: "GitHub"
+repo_url: "https://github.com/malteos/lm-datasets/"
+
+nav:
+  - Home: index.md
+  - "Getting started": getting-started.md
+  - "Framework overview": overview.md
+  - "Config files": config-files.md
+  - "Extract text data": extract-text-data.md
+  - "Adding your own data": add-your-own-data.md
+  - "Compose training and validation dataset": compose-train-validation-data.md
+
+  - "Related work": related-work.md
+  - "API reference":
+    - "BaseDataset": api/base_dataset.md
+    - "HFDataset": api/hf_dataset.md
+    - "JSONLDataset": api/jsonl_dataset.md
+
+
+
+theme:
+  name: "material"
+  # logo: "images/apple-touch-icon.png"
+
+plugins:
+  - mkdocstrings:
+      default_handler: python
+      handlers:
+        python:
+          paths: [src]
+  - search
+  - exclude:
+      glob:
+        - README.md
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index 2d4d986..b0cb27f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name =  "lm-datasets"
-version = "0.0.1"
+version = "0.0.2"
 authors = [
     {name = "Malte Ostendorff", email = "malte.ostendorff@dfki.de"},
 ]
@@ -27,27 +27,57 @@ dependencies = [
     "smart-open",
     "wget",
     "zstandard",
+    "polars",
+    "pyarrow>=14.0.1",
 ]
 license = {text = "Apache 2.0"}
 
+[project.scripts]
+lm-datasets = "lm_datasets.commands.lm_datasets_cli:main"
+
 [project.optional-dependencies]
-test = [
-    "black==23.3.0",
-    "check-manifest==0.49",
-    "flake8-bugbear==23.5.9",
-    "flake8-docstrings",
-    "flake8-formatter_junit_xml",
+datasets = [
+    "wget==3.2",
+    "mwparserfromhell>=0.6.4",
+    "gensim>=4.3.0",
+    "zstandard>=0.20.0",
+    "treelib>=1.6.4",
+    "conllu>=4.5.3",
+    "prevert>=1.0.2",
+    "translate-toolkit>=3.9.1",
+    "folia>=2.5.8",
+    "estnltk>=1.7.1",
+    "Jinja2"
+]
+dev = [
+    "pre-commit",
+    "black",
     "flake8",
     "flake8-pyproject",
-    "pre-commit==3.3.1",
-    "pylint==2.17.4",
-    "pylint_junit",
-    "pytest-cov==4.0.0",
-    "pytest-mock<3.10.1",
+    "pytest",
+    "pytest-cov",
     "pytest-runner",
-    "pytest==7.3.1",
-    "pytest-github-actions-annotate-failures",
-    "shellcheck-py==0.9.0.2"
+    "coverage",
+    "pylint",
+    "twine",
+    "jupyter"
+]
+docs = [
+    "mkdocs",
+    "mkdocs-material",
+    "mkdocs-exclude",
+    "mkdocstrings",
+    "mkdocstrings[python]"
+]
+viewer = [
+    "streamlit",
+    "ngrok"
+]
+all = [
+    "lm-datasets[datasets]",
+    "lm-datasets[dev]",
+    "lm-datasets[docs]",
+    "lm-datasets[viewer]"
 ]
 
 [project.urls]
@@ -73,6 +103,8 @@ exclude = [
     ".eggs",
     ".git",
     ".tox",
+    ".github",
+    ".vscode",
     "nssm",
     "obj",
     "out",
@@ -82,6 +114,7 @@ exclude = [
     "swagger_client",
     "dist",
     "images",
+    "docs",
 ]
 ignore = [
     "E722",
diff --git a/requirements.txt b/requirements.txt
index 35821a1..c2436de 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,78 +1,5 @@
-parfive==2.0.2
-pandas>=1.5.3
-smart-open>=6.3.0
-# for s3 support
-smart_open[s3]
-
-requests>=2.28.2
-beautifulsoup4>=4.11.2
-lxml>=4.9.2
-pyyaml
-
-# Wikimedia dumps
-wget==3.2
-mwparserfromhell>=0.6.4
-gensim==4.3.0
-
-# OSCAR
-zstandard>=0.20.0
-
-# HF
-transformers>=4.26.1
-datasets>=2.10.1
-tokenizers>=0.13.2
-huggingface-hub==0.14.0rc1
-
-tensorboard>=2.12.0
-jupyter
-
-# oasst
-treelib>=1.6.4
-
-# CONLL
-conllu>=4.5.3
-
-# MaCoCu corpora
-prevert>=1.0.2
-
-# translation datasets
-translate-toolkit>=3.9.1
-
-# sonar corpus
-folia==2.5.8
-
-# google drive downloads
-gdown>=4.7.1
-
-# same as docker image
-torch==2.0.0
-
-# TEI parser
-# git+https://github.com/standoff-nlp/standoffconverter.git
-
-# estonian
-estnltk==1.7.1
-
-# https://github.com/miurahr/py7zr
-#py7zr>=0.20.5
-
-# encodings
-chardet>=5.1.0
-
-# dev
-pre-commit
-black
-flake8
-flake8-pyproject
-pytest
-pytest-cov
-coverage
-pylint
-twine
-
-# parquet
-polars>=0.18.13
-
-# dataset viewer
-streamlit
-ngrok
\ No newline at end of file
+-r requirements/base.txt
+-r requirements/datasets.txt
+-r requirements/dev.txt
+-r requirements/viewer.txt
+-r requirements/docs.txt
\ No newline at end of file
diff --git a/requirements/base.txt b/requirements/base.txt
new file mode 100644
index 0000000..4249b5a
--- /dev/null
+++ b/requirements/base.txt
@@ -0,0 +1,36 @@
+parfive==2.0.2
+pandas>=1.5.3
+smart-open>=6.3.0
+# for s3 support
+smart_open[s3]
+
+# parquet
+polars>=0.18.13
+pyarrow>=14.0.1
+
+requests>=2.28.2
+beautifulsoup4>=4.11.2
+lxml>=4.9.2
+pyyaml
+
+# google drive downloads
+gdown>=4.7.1
+
+# HF
+transformers>=4.35.2
+datasets>=2.10.1
+tokenizers>=0.13.2
+huggingface-hub>=0.14.0rc1
+accelerate
+
+tensorboard>=2.12.0
+
+# same as docker image
+torch==2.0.0
+
+# https://github.com/miurahr/py7zr
+#py7zr>=0.20.5
+
+# encodings
+chardet>=5.1.0
+
diff --git a/requirements/datasets.txt b/requirements/datasets.txt
new file mode 100644
index 0000000..c8f1442
--- /dev/null
+++ b/requirements/datasets.txt
@@ -0,0 +1,32 @@
+
+# Wikimedia dumps
+wget==3.2
+mwparserfromhell>=0.6.4
+gensim==4.3.0
+
+# OSCAR
+zstandard>=0.20.0
+
+# oasst
+treelib>=1.6.4
+
+# CONLL
+conllu>=4.5.3
+
+# MaCoCu corpora
+prevert>=1.0.2
+
+# translation datasets
+translate-toolkit>=3.9.1
+
+# sonar corpus
+folia==2.5.8
+
+# TEI parser
+# git+https://github.com/standoff-nlp/standoffconverter.git
+
+# estonian
+estnltk==1.7.1
+
+# translation templates
+Jinja2
diff --git a/requirements/dev.txt b/requirements/dev.txt
new file mode 100644
index 0000000..5a4bbb6
--- /dev/null
+++ b/requirements/dev.txt
@@ -0,0 +1,14 @@
+# dev
+pre-commit
+black
+flake8
+flake8-pyproject
+flake8-docstrings
+pytest
+pytest-cov
+pytest-runner
+pytest-mock
+coverage
+pylint
+twine
+jupyter
diff --git a/requirements/docs.txt b/requirements/docs.txt
new file mode 100644
index 0000000..7dd0248
--- /dev/null
+++ b/requirements/docs.txt
@@ -0,0 +1,6 @@
+# docs
+mkdocs
+mkdocs-material
+mkdocs-exclude
+mkdocstrings
+mkdocstrings[python]
\ No newline at end of file
diff --git a/requirements/viewer.txt b/requirements/viewer.txt
new file mode 100644
index 0000000..763f089
--- /dev/null
+++ b/requirements/viewer.txt
@@ -0,0 +1,4 @@
+
+# dataset viewer
+streamlit
+ngrok
\ No newline at end of file
diff --git a/src/lm_datasets/__init__.py b/src/lm_datasets/__init__.py
index f102a9c..3b93d0b 100644
--- a/src/lm_datasets/__init__.py
+++ b/src/lm_datasets/__init__.py
@@ -1,1 +1,1 @@
-__version__ = "0.0.1"
+__version__ = "0.0.2"
diff --git a/src/lm_datasets/split_by_size.py b/src/lm_datasets/chunkify_datasets.py
similarity index 68%
rename from src/lm_datasets/split_by_size.py
rename to src/lm_datasets/chunkify_datasets.py
index 06302e2..25a9d2d 100644
--- a/src/lm_datasets/split_by_size.py
+++ b/src/lm_datasets/chunkify_datasets.py
@@ -1,9 +1,8 @@
-import argparse
-
 import os
-import logging
 from pathlib import Path
 
+from lm_datasets.utils.config import Config
+
 
 from .datasets.dataset_registry import get_registered_dataset_classes
 from .datasets.base import BaseDataset
@@ -24,84 +23,23 @@ def iter_parquet(pq_file, batch_size):
         yield batch
 
 
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("datasets", help="Name of datasets to shuffle (comma separated)")
-    parser.add_argument(
-        "--output_dir",
-        default=None,
-        type=str,
-        help="Unshuffled datasets are loaded from this directory",
-    )
-    parser.add_argument(
-        "--output_format",
-        default="parquet",
-        type=str,
-        help="Format of processed dataset",
-    )
-    parser.add_argument(
-        "--max_uncompressed_bytes_per_chunk",
-        default="5GB",
-        type=str,
-        help="Chunks are splitted if they exceed this byte count (<n>, <n>KB, <n>MB, or <n>GB)",
-    )
-    parser.add_argument(
-        "--safety_factor",
-        default=0.975,
-        type=float,
-        help="Max. chunk size is multiplied with this factor (accounts for inaccurate chunk sizes due to batching)",
-    )
-    parser.add_argument(
-        "--output_compression",
-        default=None,
-        type=str,
-        help="""Compression of output (default: compression of existing output; jsonl: "gzip"; """
-        + """parquet: "NONE", "SNAPPY", "GZIP", "BROTLI", "LZ4", "ZSTD")""",
-    )
-    parser.add_argument(
-        "--batch_size",
-        default=16,
-        type=int,
-        help="""Read and write batch size; smaller batch size = more accurate splitts but slower""",
-    )
-    parser.add_argument(
-        "--log_file",
-        default=None,
-        type=str,
-        help="Log file is saved at this path",
-    )
-    parser.add_argument("--verbose", action="store_true", help="Enable verbose logging (log level = debug)")
-    parser.add_argument("--override", action="store_true", help="Override existing output files")
-    parser.add_argument("--rename_original", action="store_true", help="Renames orignal output file")
-    args = parser.parse_args()
-
-    log_handlers = [logging.StreamHandler()]
-
-    if args.log_file:
-        log_handlers.append(logging.FileHandler(args.log_file))
-
-    logging.basicConfig(
-        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
-        datefmt="%Y-%m-%d %H:%M:%S",
-        level=logging.DEBUG if args.verbose else logging.INFO,
-        handlers=log_handlers,
-    )
-    logger = logging.getLogger(__name__)
-
-    if args.output_format != "parquet":
+def chunkify_datasets(config: Config):
+    logger = config.init_logger(__name__)
+
+    if config.output_format != "parquet":
         raise ValueError("Output format is not supported (currently only parquet is supported)")
 
     id_to_dataset_class = {cls.DATASET_ID: cls for cls in get_registered_dataset_classes()}
 
-    datasets_list = args.datasets.split(",")
+    datasets_list = config.datasets.split(",")
 
     if len(datasets_list) == 1 and datasets_list[0] == "all":
         # Get list of all non-dummy datasets
         datasets_list = id_to_dataset_class.keys()
 
-    max_uncompressed_bytes_per_chunk = get_bytes_from_int_or_string(args.max_uncompressed_bytes_per_chunk)
-    max_chunk_bytes_with_safety = int(max_uncompressed_bytes_per_chunk * args.safety_factor)
-    logger.info(f"Max. chunk size with {args.safety_factor=}: {max_chunk_bytes_with_safety:,} bytes")
+    max_uncompressed_bytes_per_chunk = get_bytes_from_int_or_string(config.max_uncompressed_bytes_per_chunk)
+    max_chunk_bytes_with_safety = int(max_uncompressed_bytes_per_chunk * config.safety_factor)
+    logger.info(f"Max. chunk size with {config.safety_factor=}: {max_chunk_bytes_with_safety:,} bytes")
 
     # Iterate over datasets
     for i, dataset_id in enumerate(datasets_list, 1):
@@ -113,9 +51,8 @@ def iter_parquet(pq_file, batch_size):
         logger.info(f"Dataset ID: {dataset_id} ({i} / {len(datasets_list)})")
 
         dataset: BaseDataset = dataset_cls(
-            output_dir=args.output_dir,
-            output_format=args.output_format,
-            # output_compression=args.output_compression,
+            output_dir=config.output_dir,
+            output_format=config.output_format,
         )
 
         if dataset.has_chunked_output_files():
@@ -140,7 +77,7 @@ def iter_parquet(pq_file, batch_size):
             existing_chunk_fps = list(output_dir.glob(existing_chunk_fp_pattern))
 
             if existing_chunk_fps:
-                if args.override:
+                if config.override:
                     logger.info("Removing existing part files ...")
                     for fp in existing_chunk_fps:
                         logger.info(f"Removing {fp} ...")
@@ -153,13 +90,13 @@ def iter_parquet(pq_file, batch_size):
             with open(dataset.get_single_output_file_path(), "rb") as file_handler:
                 # pq_file = open_parquet_file_with_retries(dataset.get_output_file_path(), retries=2)
                 pq_file = pq.ParquetFile(file_handler)
-                batch_iter = iter_parquet(pq_file, args.batch_size)  # iterate over row groups and batches
+                batch_iter = iter_parquet(pq_file, config.batch_size)  # iterate over row groups and batches
 
                 # Auto-detect compression from source file
-                if args.output_compression is None:
+                if config.output_compression is None:
                     compression = pq_file.metadata.row_group(0).column(0).compression
                 else:
-                    compression = get_parquet_compression(args.output_compression)
+                    compression = get_parquet_compression(config.output_compression)
 
                 logger.info(f"Unsplitted file size: {file_stats.st_size:,} bytes ({compression=})")
 
@@ -220,7 +157,7 @@ def iter_parquet(pq_file, batch_size):
                     os.rename(chunk_fp, new_chunk_fp)
 
                 # Keep original: by renaming to ".orignal"
-                if args.rename_original:
+                if config.rename_original:
                     new_path = str(output_file_path).replace(".parquet", ".parquet.original")
                     logger.info(f"Moving original file to {new_path}")
                     os.rename(str(output_file_path), new_path)
diff --git a/src/lm_datasets/collect_metrics.py b/src/lm_datasets/collect_metrics.py
new file mode 100644
index 0000000..8c7f2d1
--- /dev/null
+++ b/src/lm_datasets/collect_metrics.py
@@ -0,0 +1,121 @@
+import json
+
+from pathlib import Path
+
+from .datasets.dataset_registry import (
+    get_dataset_class_by_id,
+    get_registered_dataset_ids,
+)
+from .datasets.base import BaseDataset
+from .utils.config import Config
+
+from transformers import AutoTokenizer
+
+
+def collect_metrics(config: Config):
+    logger = config.init_logger(__name__)
+
+    save_to_path = Path(config.save_to)
+
+    if save_to_path.exists() and not config.override:
+        raise FileExistsError(f"Cannot save stats because path exists already (fix with --override): {save_to_path}")
+
+    datasets_list = config.datasets.split(",")
+
+    if len(datasets_list) == 1:
+        if datasets_list[0] == "all":
+            # Get list of all regsitered datasets
+            datasets_list = get_registered_dataset_ids(config.extra_dataset_registries)
+
+        elif datasets_list[0] == "all_from_source":
+            # Get registered datasets based on source
+            if config.source_id is None:
+                raise ValueError("The argument --source_id must be set.")
+
+            datasets_list = get_registered_dataset_ids(
+                config.extra_dataset_registries, needed_source_id=config.source_id
+            )
+
+    if config.hf_tokenizer_name_or_path:
+        tokenizer = AutoTokenizer.from_pretrained(
+            config.hf_tokenizer_name_or_path,
+            use_fast=True,
+        )
+
+        if not config.texts_limit:
+            raise ValueError("Tokenizer must be used with --texts_limit argument.")
+    else:
+        tokenizer = None
+
+    # Iterate over datasets
+    dataset_id_to_stats = {}
+
+    for i, dataset_id in enumerate(datasets_list, 1):
+        logger.info(f"Dataset ID: {dataset_id} ({i} / {len(datasets_list)})")
+
+        if i <= config.skip_datasets:
+            logger.warning(f"Skip dataset")
+            continue
+
+        dataset_cls = get_dataset_class_by_id(dataset_id, config.extra_dataset_registries)
+        dataset: BaseDataset = dataset_cls(
+            output_dir=config.output_dir,
+            output_format=config.output_format,
+            shuffled_output_dir=config.shuffled_output_dir,
+            config=config,
+        )
+
+        if config.only_selected_datasets and not dataset.is_selected():
+            logger.info("Skip %s (not part of selected datasets or sources)", dataset_id)
+            continue
+
+        total_ws_count = 0
+        total_byte_count = 0
+        texts = []
+
+        for i, text in enumerate(
+            dataset.generate_texts_from_output(shuffled=True, limit=config.texts_limit, shuffle_output_file_paths=True)
+        ):
+            # cast to string
+            text = str(text)
+
+            ws_count = text.count(" ")
+            byte_count = len(text.encode("utf-8"))
+
+            total_ws_count += ws_count
+            total_byte_count += byte_count
+
+            if tokenizer:
+                texts.append(text)
+
+        # Append to stats
+        dataset_id_to_stats[dataset.DATASET_ID] = {
+            "whitespace_count": total_ws_count,
+            "byte_count": total_byte_count,
+        }
+
+        if tokenizer and texts:
+            tokenizer_out = tokenizer(
+                text=texts,
+                add_special_tokens=False,
+                padding=False,
+                truncation=False,
+                return_length=True,
+            )
+            dataset_id_to_stats[dataset.DATASET_ID].update(
+                {
+                    "tokenizer": config.hf_tokenizer_name_or_path,
+                    "tokens_count": sum(tokenizer_out["length"]),
+                }
+            )
+
+        # Save stats to to JSON after each dataset
+        with open(save_to_path, "w") as f:
+            logger.info("Saving stats to %s", save_to_path)
+            json.dump(dataset_id_to_stats, f)
+
+        if config.datasets_limit > 0 and len(dataset_id_to_stats) >= config.datasets_limit:
+            logger.warning("Datasets limit reached")
+            break
+
+    logger.info("Done")
diff --git a/src/lm_datasets/commands/__init__.py b/src/lm_datasets/commands/__init__.py
new file mode 100644
index 0000000..48da6ff
--- /dev/null
+++ b/src/lm_datasets/commands/__init__.py
@@ -0,0 +1,78 @@
+from abc import ABC, abstractmethod
+from argparse import _SubParsersAction
+
+
+class BaseCLICommand(ABC):
+    @staticmethod
+    @abstractmethod
+    def register_subcommand(parser: _SubParsersAction):
+        raise NotImplementedError()
+
+    @abstractmethod
+    def run(self):
+        raise NotImplementedError()
+
+    @staticmethod
+    def add_common_args(
+        parser: _SubParsersAction,
+        raw_datasets_dir=False,
+        output=False,
+        extra_dataset_registries=False,
+        configs=False,
+        required_configs=False,
+        log=False,
+    ):
+        if raw_datasets_dir:
+            parser.add_argument(
+                "--raw_datasets_dir",
+                default=None,
+                type=str,
+                help="Dataset files are read from this directory (needed for `is_downloaded` field)",
+            )
+
+        if output:
+            parser.add_argument(
+                "--output_dir",
+                default=None,
+                type=str,
+                help="Processed dataset are saved in this directory (need for `has_output_file` field)",
+            )
+            parser.add_argument(
+                "--output_format",
+                default="jsonl",
+                type=str,
+                help="Format of processed dataset (jsonl,parquet; need for `has_output_file` field)",
+            )
+
+        if extra_dataset_registries:
+            parser.add_argument(
+                "--extra_dataset_registries",
+                default=None,
+                type=str,
+                help="List of Python packages to load dataset registries",
+            )
+
+        if configs:
+            parser.add_argument(
+                "-c",
+                "--configs",
+                nargs="+",
+                help=(
+                    "Paths to one or more YAML-config files (duplicated settings are override based on config order; "
+                    "cmd args override configs)"
+                ),
+                default=None,
+                dest="config_paths",
+                required=required_configs,
+            )
+
+        if log:
+            parser.add_argument(
+                "--log_file",
+                default=None,
+                type=str,
+                help="Log file is saved at this path",
+            )
+            parser.add_argument("--verbose", action="store_true", help="Enable verbose logging (log level = debug)")
+
+        return parser
diff --git a/src/lm_datasets/commands/chunkify_command.py b/src/lm_datasets/commands/chunkify_command.py
new file mode 100644
index 0000000..4791297
--- /dev/null
+++ b/src/lm_datasets/commands/chunkify_command.py
@@ -0,0 +1,59 @@
+from argparse import Namespace, _SubParsersAction
+
+from lm_datasets.chunkify_datasets import chunkify_datasets
+from lm_datasets.commands import BaseCLICommand
+from lm_datasets.utils.config import Config, get_config_from_paths
+
+
+class ChunkifyCommand(BaseCLICommand):
+    @staticmethod
+    def register_subcommand(parser: _SubParsersAction):
+        subcommand_parser = parser.add_parser(
+            "chunkify", help="Split the individual datasets into equally-sized file chunks (based on bytes or rows)"
+        )
+
+        subcommand_parser.add_argument("datasets", help="Name of datasets to shuffle (comma separated)")
+        subcommand_parser.add_argument(
+            "--max_uncompressed_bytes_per_chunk",
+            default="5GB",
+            type=str,
+            help="Chunks are splitted if they exceed this byte count (<n>, <n>KB, <n>MB, or <n>GB)",
+        )
+        subcommand_parser.add_argument(
+            "--safety_factor",
+            default=0.975,
+            type=float,
+            help="Max. chunk size is multiplied with this factor (accounts for inaccurate chunk sizes due to batching)",
+        )
+        subcommand_parser.add_argument(
+            "--output_compression",
+            default=None,
+            type=str,
+            help="""Compression of output (default: compression of existing output; jsonl: "gzip"; """
+            + """parquet: "NONE", "SNAPPY", "GZIP", "BROTLI", "LZ4", "ZSTD")""",
+        )
+        subcommand_parser.add_argument(
+            "--batch_size",
+            default=16,
+            type=int,
+            help="""Read and write batch size; smaller batch size = more accurate splitts but slower""",
+        )
+        subcommand_parser.add_argument("--override", action="store_true", help="Override existing output files")
+        subcommand_parser.add_argument("--rename_original", action="store_true", help="Renames orignal output file")
+
+        subcommand_parser = BaseCLICommand.add_common_args(
+            subcommand_parser,
+            raw_datasets_dir=False,
+            output=True,
+            extra_dataset_registries=True,
+            configs=True,
+            required_configs=False,
+            log=True,
+        )
+        subcommand_parser.set_defaults(func=ChunkifyCommand)
+
+    def __init__(self, args: Namespace) -> None:
+        self.config: Config = get_config_from_paths(args.config_paths, override=args.__dict__)
+
+    def run(self) -> None:
+        chunkify_datasets(config=self.config)
diff --git a/src/lm_datasets/commands/collect_metrics_command.py b/src/lm_datasets/commands/collect_metrics_command.py
new file mode 100644
index 0000000..0c464dc
--- /dev/null
+++ b/src/lm_datasets/commands/collect_metrics_command.py
@@ -0,0 +1,73 @@
+from argparse import Namespace, _SubParsersAction
+import logging
+
+from lm_datasets.collect_metrics import collect_metrics
+from lm_datasets.commands import BaseCLICommand
+from lm_datasets.utils.config import Config, get_config_from_paths
+
+
+logger = logging.getLogger(__name__)
+
+
+class CollectMetricsCommand(BaseCLICommand):
+    @staticmethod
+    def register_subcommand(parser: _SubParsersAction):
+        subcommand_parser = parser.add_parser(
+            "collect_metrics", help="Collect metrics (token count etc.) from extracted texts"
+        )
+
+        subcommand_parser.add_argument("datasets", help="Name of datasets to shuffle (comma separated)")
+        subcommand_parser.add_argument(
+            "--shuffled_output_dir",
+            default=None,
+            type=str,
+            help="Shuffled dataset are saved in this directory",
+        )
+        subcommand_parser.add_argument(
+            "--save_to",
+            default=None,
+            type=str,
+            help="""Save collected stats to this file path (JSON format)""",
+        )
+        subcommand_parser.add_argument("--override", action="store_true", help="Override existing output files")
+        subcommand_parser.add_argument(
+            "--texts_limit", type=int, default=0, help="Limit number of texts generated for each dataset"
+        )
+        subcommand_parser.add_argument(
+            "--datasets_limit", type=int, default=0, help="Limit number of texts generated for each dataset"
+        )
+        subcommand_parser.add_argument("--skip_datasets", type=int, default=0, help="Skip n datasets before starting")
+        subcommand_parser.add_argument(
+            "--hf_tokenizer_name_or_path",
+            default=None,
+            type=str,
+            help="""Name or path to HF tokenizer (if is set, tokens are counted)""",
+        )
+        subcommand_parser.add_argument(
+            "--source_id",
+            default=None,
+            type=str,
+            help="Filter datasets by source ID (used if `datasets`='all_from_source')",
+        )
+        subcommand_parser.add_argument(
+            "--only_selected_datasets",
+            action="store_true",
+            help="Include only datasets there were explicitly selected (via config)",
+        )
+
+        subcommand_parser = BaseCLICommand.add_common_args(
+            subcommand_parser,
+            raw_datasets_dir=True,
+            output=True,
+            extra_dataset_registries=True,
+            configs=True,
+            required_configs=False,
+            log=True,
+        )
+        subcommand_parser.set_defaults(func=CollectMetricsCommand)
+
+    def __init__(self, args: Namespace) -> None:
+        self.config: Config = get_config_from_paths(args.config_paths, override=args.__dict__)
+
+    def run(self) -> None:
+        collect_metrics(config=self.config)
diff --git a/src/lm_datasets/commands/compose_command.py b/src/lm_datasets/commands/compose_command.py
new file mode 100644
index 0000000..db157cf
--- /dev/null
+++ b/src/lm_datasets/commands/compose_command.py
@@ -0,0 +1,93 @@
+from argparse import Namespace, _SubParsersAction
+import logging
+
+from lm_datasets.compose_dataset import compose_dataset
+from lm_datasets.commands import BaseCLICommand
+from lm_datasets.utils.config import Config, get_config_from_paths
+from lm_datasets.utils.dataset_generator import DatasetSplit
+
+
+logger = logging.getLogger(__name__)
+
+
+class ComposeCommand(BaseCLICommand):
+    @staticmethod
+    def register_subcommand(parser: _SubParsersAction):
+        subcommand_parser = parser.add_parser(
+            "compose", help="Compose the final train/validation set based on the individual datasets"
+        )
+
+        subcommand_parser.add_argument(
+            "--split", type=DatasetSplit, help="Dataset split (full, train, tokenizer_train, validation)"
+        )
+        subcommand_parser.add_argument(
+            "--shuffled_output_dir",
+            help="Shuffled output is saved in this directory (<language code>/<source name>.<jsonl/parquet>)",
+        )
+        subcommand_parser.add_argument(
+            "--composed_dataset_dir",
+            required=True,
+            type=str,
+            help="""Save composed dataset this directory""",
+        )
+        subcommand_parser.add_argument(
+            "--save_dataset_ids", action="store_true", help="Save dataset ID in addition to text field"
+        )
+        subcommand_parser.add_argument(
+            "--limit",
+            type=int,
+            default=0,
+            help="""Limit number of output samples (for debugging)""",
+        )
+        subcommand_parser.add_argument(
+            "--output_batch_size",
+            type=int,
+            default=1000,
+            help="""Write batch size; smaller batch size = more accurate splitts but slower""",
+        )
+        subcommand_parser.add_argument(
+            "--input_batch_size",
+            type=int,
+            default=1000,
+            help="""Reader batch size; smaller batch size = less memory consuption but slower""",
+        )
+        subcommand_parser.add_argument(
+            "--interleave_random_batch_size",
+            type=int,
+            default=100,
+            help="""Datasets are randomly interleaves with this batch size""",
+        )
+        subcommand_parser.add_argument(
+            "--max_output_chunk_uncompressed_bytes",
+            type=str,
+            default="10GB",
+            help="Chunks are splitted if they exceed this byte count (<n>, <n>KB, <n>MB, or <n>GB)",
+        )
+        subcommand_parser.add_argument(
+            "--disable_sampling",
+            action="store_false",
+            dest="use_sampling",
+            help="Disable dataset up/down sampling based on sampling factors (see config file)",
+        )
+        subcommand_parser.add_argument(
+            "--use_separated_validation_sets",
+            action="store_true",
+            help="If enabled, validation set is separated stored on disk (one for each dataset)",
+        )
+
+        subcommand_parser = BaseCLICommand.add_common_args(
+            subcommand_parser,
+            raw_datasets_dir=True,
+            output=True,
+            extra_dataset_registries=True,
+            configs=True,
+            required_configs=True,
+            log=True,
+        )
+        subcommand_parser.set_defaults(func=ComposeCommand)
+
+    def __init__(self, args: Namespace) -> None:
+        self.config: Config = get_config_from_paths(args.config_paths, override=args.__dict__)
+
+    def run(self) -> None:
+        compose_dataset(config=self.config)
diff --git a/src/lm_datasets/commands/convert_parquet_to_jsonl_command.py b/src/lm_datasets/commands/convert_parquet_to_jsonl_command.py
new file mode 100644
index 0000000..251ab62
--- /dev/null
+++ b/src/lm_datasets/commands/convert_parquet_to_jsonl_command.py
@@ -0,0 +1,37 @@
+from argparse import Namespace, _SubParsersAction
+
+from lm_datasets.convert_parquet_to_jsonl import convert_parquet_to_jsonl
+from lm_datasets.commands import BaseCLICommand
+from lm_datasets.utils.config import Config
+
+
+class ConvertParquetToJSONLCommand(BaseCLICommand):
+    @staticmethod
+    def register_subcommand(parser: _SubParsersAction):
+        subcommand_parser = parser.add_parser("convert_parquet_to_jsonl", help="Convert Parquet files to JSONL")
+
+        subcommand_parser.add_argument(
+            "input_dir_or_file", help="Directory containing *.parquet files or the file itself"
+        )
+        subcommand_parser.add_argument("output_dir", help="Directory containing *.parquet files")
+        subcommand_parser.add_argument("--override", action="store_true", help="Override existing output files")
+        subcommand_parser.add_argument("--input_glob", type=str, default="*.parquet", help="Glob pattern")
+
+        subcommand_parser = BaseCLICommand.add_common_args(
+            subcommand_parser,
+            log=True,
+        )
+        subcommand_parser.set_defaults(func=ConvertParquetToJSONLCommand)
+
+    def __init__(self, args: Namespace) -> None:
+        self.config: Config = Config(**args.__dict__)
+
+    def run(self) -> None:
+        self.config.init_logger(__name__)
+
+        convert_parquet_to_jsonl(
+            input_dir_or_file=self.config.input_dir_or_file,
+            output_dir=self.config.output_dir,
+            override=self.config.override,
+            input_glob=self.config.input_glob,
+        )
diff --git a/src/lm_datasets/commands/extract_text_command.py b/src/lm_datasets/commands/extract_text_command.py
new file mode 100644
index 0000000..fc2e290
--- /dev/null
+++ b/src/lm_datasets/commands/extract_text_command.py
@@ -0,0 +1,89 @@
+from argparse import Namespace, _SubParsersAction
+import logging
+
+from lm_datasets.extract_text import extract_text
+from lm_datasets.commands import BaseCLICommand
+from lm_datasets.utils.config import Config, get_config_from_paths
+from lm_datasets.utils.settings import DEFAULT_MIN_TEXT_LENGTH
+
+
+logger = logging.getLogger(__name__)
+
+
+class ExtractTextCommand(BaseCLICommand):
+    @staticmethod
+    def register_subcommand(parser: _SubParsersAction):
+        subcommand_parser = parser.add_parser("extract_text", help="Extract text from raw datasets")
+
+        subcommand_parser.add_argument("datasets", help="Name of datasets to process (comma separated)")
+        subcommand_parser.add_argument(
+            "output_dir",
+            help="Output is saved in this directory (<language code>/<source name>.<jsonl/parquet>)",
+        )
+        subcommand_parser.add_argument("--override", action="store_true", help="Override existing output files")
+        subcommand_parser.add_argument(
+            "--ignore_errors",
+            action="store_true",
+            help="Ignore dataset-level errors (use when processing multiple datasets)",
+        )
+        subcommand_parser.add_argument(
+            "--json_ensure_ascii", action="store_true", help="Escape non-ASCII characters in JSON output"
+        )
+        subcommand_parser.add_argument(
+            "--output_compression",
+            default=None,
+            type=str,
+            help="""Compression of output (jsonl: "gzip"; parquet: "NONE", "SNAPPY", "GZIP", "BROTLI", "LZ4", "ZSTD")""",
+        )
+        subcommand_parser.add_argument(
+            "--output_batch_size",
+            default=1000,
+            type=int,
+            help="""Write batch size; smaller batch size = more accurate splitts but slower""",
+        )
+        subcommand_parser.add_argument(
+            "--max_output_chunk_uncompressed_bytes",
+            default="10GB",
+            type=str,
+            help="Chunks are splitted if they exceed this byte count (<n>, <n>KB, <n>MB, or <n>GB)",
+        )
+        subcommand_parser.add_argument(
+            "--workers",
+            default=1,
+            type=int,
+            help="Number of workers for parallel processing",
+        )
+        subcommand_parser.add_argument("--limit", default=0, type=int, help="Limit dataset size (for debugging)")
+        subcommand_parser.add_argument(
+            "--min_text_length", default=DEFAULT_MIN_TEXT_LENGTH, type=int, help="Text with less length is discarded"
+        )
+        subcommand_parser.add_argument(
+            "--skip_items",
+            default=0,
+            type=int,
+            help="Skip N items (depending on dataset: directories, subsets, files, documents) (for debugging)",
+        )
+        subcommand_parser.add_argument("--hf_auth_token", default=None, type=str, help="HuggingFace auth token")
+        subcommand_parser.add_argument(
+            "--source_id",
+            default=None,
+            type=str,
+            help="Filter datasets by source ID (used if `datasets`='all_from_source')",
+        )
+
+        subcommand_parser = BaseCLICommand.add_common_args(
+            subcommand_parser,
+            raw_datasets_dir=True,
+            output=True,
+            extra_dataset_registries=True,
+            configs=True,
+            required_configs=False,
+            log=True,
+        )
+        subcommand_parser.set_defaults(func=ExtractTextCommand)
+
+    def __init__(self, args: Namespace) -> None:
+        self.config: Config = get_config_from_paths(args.config_paths, override=args.__dict__)
+
+    def run(self) -> None:
+        extract_text(config=self.config)
diff --git a/src/lm_datasets/commands/hf_upload_command.py b/src/lm_datasets/commands/hf_upload_command.py
new file mode 100644
index 0000000..4c269f6
--- /dev/null
+++ b/src/lm_datasets/commands/hf_upload_command.py
@@ -0,0 +1,82 @@
+from argparse import Namespace, _SubParsersAction
+
+from lm_datasets.commands import BaseCLICommand
+from lm_datasets.utils.config import Config
+
+import os
+from pathlib import Path
+from huggingface_hub import HfApi
+from tqdm.auto import tqdm
+
+
+class HFUploadCommand(BaseCLICommand):
+    @staticmethod
+    def register_subcommand(parser: _SubParsersAction):
+        subcommand_parser = parser.add_parser("hf_upload", help="Upload files or directories to Huggingface Hub.")
+
+        subcommand_parser.add_argument("input_path", help="Files from this directory are uploaded")
+        subcommand_parser.add_argument("repo_id", help="HF hub repo ID (e.g., username/dataset_name)")
+        subcommand_parser.add_argument("--path_in_repo", type=str, default=".", help="Path in HF repo")
+        subcommand_parser.add_argument("--repo_type", type=str, default="dataset", help="Path in HF repo")
+        subcommand_parser.add_argument(
+            "--commit_message", type=str, default="Files uploaded", help="Git commit message"
+        )
+        subcommand_parser.add_argument(
+            "--skip_if_exists", action="store_true", help="Skip if a file it already exists in the HF repo"
+        )
+        subcommand_parser.add_argument(
+            "--input_glob", type=str, default=None, help="Glob pattern to selected matching files from input directory"
+        )
+        subcommand_parser = BaseCLICommand.add_common_args(
+            subcommand_parser,
+            log=True,
+        )
+        subcommand_parser.set_defaults(func=HFUploadCommand)
+
+    def __init__(self, args: Namespace) -> None:
+        self.args = args
+        self.config = Config(**args.__dict__)
+
+    def run(self) -> None:
+        logger = self.config.init_logger(__name__)
+        args = self.args
+
+        api = HfApi()
+
+        input_path = args.input_path
+
+        # is directory
+        if os.path.isdir(input_path):
+            if args.input_glob is None:
+                # all files from input directory
+                input_file_names = list(sorted(os.listdir(input_path)))
+
+            else:
+                # select files based on glob pattern
+                input_file_names = [p.name for p in sorted(Path(input_path).glob(args.input_glob))]
+
+        elif os.path.isfile(input_path):
+            # single file upload
+            input_file_names = [Path(input_path).name]
+            input_path = Path(input_path).parent
+
+        repo_type = args.repo_type
+
+        logger.info(f"Starting upload to {args.repo_id} from {input_path}")
+
+        for i, file_name in tqdm(enumerate(input_file_names, 1), total=len(input_file_names), desc="Uploading"):
+            path_in_repo = f"{args.path_in_repo}/{file_name}"
+
+            if args.skip_if_exists:
+                if api.file_exists(repo_id=args.repo_id, filename=path_in_repo, repo_type=repo_type):
+                    logger.info("Skip file because it already exists in the repo: %s", path_in_repo)
+                    continue
+
+            api.upload_file(
+                path_or_fileobj=os.path.join(input_path, file_name),
+                repo_id=args.repo_id,
+                path_in_repo=path_in_repo,
+                commit_message=args.commit_message + f" ({i}/{len(input_file_names)})",
+                repo_type=repo_type,
+            )
+        logger.info("done")
diff --git a/src/lm_datasets/commands/lm_datasets_cli.py b/src/lm_datasets/commands/lm_datasets_cli.py
new file mode 100644
index 0000000..645cdbf
--- /dev/null
+++ b/src/lm_datasets/commands/lm_datasets_cli.py
@@ -0,0 +1,42 @@
+from argparse import ArgumentParser
+
+from lm_datasets.commands.chunkify_command import ChunkifyCommand
+from lm_datasets.commands.collect_metrics_command import CollectMetricsCommand
+from lm_datasets.commands.compose_command import ComposeCommand
+from lm_datasets.commands.convert_parquet_to_jsonl_command import ConvertParquetToJSONLCommand
+from lm_datasets.commands.extract_text_command import ExtractTextCommand
+from lm_datasets.commands.hf_upload_command import HFUploadCommand
+from lm_datasets.commands.print_stats_command import PrintStatsCommand
+from lm_datasets.commands.shuffle_command import ShuffleCommand
+from lm_datasets.commands.train_tokenizer_command import TrainTokenizerCommand
+
+
+def main():
+    parser = ArgumentParser("lm-datasets", usage="lm-datasets <command> [<args>]")
+    commands_parser = parser.add_subparsers(help="lm-datasets command helpers")
+
+    # Register commands
+    ChunkifyCommand.register_subcommand(commands_parser)
+    CollectMetricsCommand.register_subcommand(commands_parser)
+    ComposeCommand.register_subcommand(commands_parser)
+    ConvertParquetToJSONLCommand.register_subcommand(commands_parser)
+    ExtractTextCommand.register_subcommand(commands_parser)
+    HFUploadCommand.register_subcommand(commands_parser)
+    PrintStatsCommand.register_subcommand(commands_parser)
+    ShuffleCommand.register_subcommand(commands_parser)
+    TrainTokenizerCommand.register_subcommand(commands_parser)
+
+    # Let's go
+    args = parser.parse_args()
+
+    if not hasattr(args, "func"):
+        parser.print_help()
+        exit(1)
+
+    # Run
+    service = args.func(args)
+    service.run()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/lm_datasets/commands/print_stats_command.py b/src/lm_datasets/commands/print_stats_command.py
new file mode 100644
index 0000000..a7678e1
--- /dev/null
+++ b/src/lm_datasets/commands/print_stats_command.py
@@ -0,0 +1,76 @@
+from argparse import Namespace, _SubParsersAction
+
+from lm_datasets.print_stats import print_stats
+from lm_datasets.commands import BaseCLICommand
+from lm_datasets.utils.config import Config, get_config_from_paths
+
+
+class PrintStatsCommand(BaseCLICommand):
+    @staticmethod
+    def register_subcommand(parser: _SubParsersAction):
+        subcommand_parser = parser.add_parser("print_stats", help="Print dataset statistics as CSV, Markdown, ...")
+
+        subcommand_parser.add_argument(
+            "--shuffled_output_dir",
+            default=None,
+            type=str,
+            help="Shuffled dataset are saved in this directory",
+        )
+        subcommand_parser.add_argument(
+            "--metrics_dir",
+            default=None,
+            type=str,
+            help="Dataset metrics (whitespace count, byte count, ...) are loaded from this directory (*.json files)",
+        )
+        subcommand_parser.add_argument(
+            "--token_estimation_path",
+            default=None,
+            type=str,
+            help="Path to dataset metrics with token count based on a sample (JSON-file; this is used to estimate the total token count)",
+        )
+        subcommand_parser.add_argument(
+            "--rows_count", action="store_true", help="Extract number of rows from output files"
+        )
+        subcommand_parser.add_argument(
+            "--shuffled_rows_count", action="store_true", help="Extract number of rows from shuffled output files"
+        )
+        subcommand_parser.add_argument(
+            "--output_compression", action="store_true", help="Extract compression from output files"
+        )
+        subcommand_parser.add_argument("--print_format", default="csv", type=str, help="Print format (tsv,csv,md)")
+        subcommand_parser.add_argument("--limit", default=0, type=int, help="Limit the datasets (for debugging)")
+        subcommand_parser.add_argument(
+            "--save_to",
+            default=None,
+            type=str,
+            help="Save output to this path on the disk (default: only print to stdout)",
+        )
+        subcommand_parser.add_argument("--exclude_dummy_datasets", action="store_true", help="Exclude dummy datasets")
+        subcommand_parser.add_argument(
+            "--only_selected_datasets",
+            action="store_true",
+            help="Include only datasets there were explicitly selected (via config)",
+        )
+        subcommand_parser.add_argument(
+            "--extra_columns",
+            default=None,
+            type=str,
+            help="Comma separated list of columns (see AVAILABLE_DATAFRAME_COLUMNS)",
+        )
+
+        subcommand_parser = BaseCLICommand.add_common_args(
+            subcommand_parser,
+            raw_datasets_dir=True,
+            output=True,
+            extra_dataset_registries=True,
+            configs=True,
+            required_configs=False,
+            log=True,
+        )
+        subcommand_parser.set_defaults(func=PrintStatsCommand)
+
+    def __init__(self, args: Namespace) -> None:
+        self.config: Config = get_config_from_paths(args.config_paths, override=args.__dict__)
+
+    def run(self) -> None:
+        print_stats(config=self.config)
diff --git a/src/lm_datasets/commands/shuffle_command.py b/src/lm_datasets/commands/shuffle_command.py
new file mode 100644
index 0000000..7258465
--- /dev/null
+++ b/src/lm_datasets/commands/shuffle_command.py
@@ -0,0 +1,75 @@
+from argparse import Namespace, _SubParsersAction
+
+from lm_datasets.shuffle_datasets import shuffle_datasets
+from lm_datasets.commands import BaseCLICommand
+from lm_datasets.utils.config import Config, get_config_from_paths
+from lm_datasets.utils.settings import DEFAULT_MIN_FILE_SIZE_FOR_BUFFERED_SHUFFLING
+
+
+class ShuffleCommand(BaseCLICommand):
+    @staticmethod
+    def register_subcommand(parser: _SubParsersAction):
+        subcommand_parser = parser.add_parser(
+            "shuffle", help="Shuffle the individual datasets on the file-chunk level (no global shuffle!)"
+        )
+
+        subcommand_parser.add_argument("datasets", help="Name of datasets to shuffle (comma separated)")
+        subcommand_parser.add_argument(
+            "--shuffled_output_dir",
+            default=None,
+            type=str,
+            help="Shuffled dataset are saved in this directory",
+        )
+        subcommand_parser.add_argument(
+            "--output_compression",
+            default=None,
+            type=str,
+            help="""Compression of output (jsonl: "gzip"; parquet: "NONE", "SNAPPY", "GZIP", "BROTLI", "LZ4", "ZSTD")""",
+        )
+        subcommand_parser.add_argument(
+            "--seed",
+            default=0,
+            type=int,
+            help="Random seed",
+        )
+        subcommand_parser.add_argument(
+            "--shuffle_buffer_size",
+            default=1_000_000,
+            type=int,
+            help="Number of items in buffer to be shuffled at once (larger buffer = more memory but better shuffing)",
+        )
+        subcommand_parser.add_argument(
+            "--min_file_size_for_buffered_shuffling",
+            default=DEFAULT_MIN_FILE_SIZE_FOR_BUFFERED_SHUFFLING,
+            type=int,
+            help="Min. file size bytes for buffered shuffling (default: 5GB; set to 0 to disable)",
+        )
+        subcommand_parser.add_argument("--override", action="store_true", help="Override existing output files")
+        subcommand_parser.add_argument(
+            "--skip_large_datasets",
+            action="store_true",
+            help="Skip datasets with bytes > --min_file_size_for_buffered_shuffling",
+        )
+        subcommand_parser.add_argument(
+            "--source_id",
+            default=None,
+            type=str,
+            help="Filter datasets by source ID (used if `datasets`='all_from_source')",
+        )
+
+        subcommand_parser = BaseCLICommand.add_common_args(
+            subcommand_parser,
+            raw_datasets_dir=True,
+            output=True,
+            extra_dataset_registries=True,
+            configs=True,
+            required_configs=False,
+            log=True,
+        )
+        subcommand_parser.set_defaults(func=ShuffleCommand)
+
+    def __init__(self, args: Namespace) -> None:
+        self.config: Config = get_config_from_paths(args.config_paths, override=args.__dict__)
+
+    def run(self) -> None:
+        shuffle_datasets(config=self.config)
diff --git a/src/lm_datasets/commands/train_tokenizer_command.py b/src/lm_datasets/commands/train_tokenizer_command.py
new file mode 100644
index 0000000..c97f444
--- /dev/null
+++ b/src/lm_datasets/commands/train_tokenizer_command.py
@@ -0,0 +1,99 @@
+from argparse import Namespace, _SubParsersAction
+
+from lm_datasets.train_sp_tokenizer import train_sp_tokenizer
+from lm_datasets.commands import BaseCLICommand
+from lm_datasets.utils.config import Config, get_config_from_paths
+
+from lm_datasets.utils.settings import DEFAULT_TOKENIZER_RATIO
+
+
+class TrainTokenizerCommand(BaseCLICommand):
+    @staticmethod
+    def register_subcommand(parser: _SubParsersAction):
+        subcommand_parser = parser.add_parser(
+            "train_tokenizer", help="Train a tokenizer (only: sentencepiece supproted)"
+        )
+
+        subcommand_parser.add_argument(
+            "tokenizer_backend", help="Tokenizer backend (currently only `sentencepiece` supported)"
+        )
+
+        subcommand_parser.add_argument(
+            "--composed_dataset_dir",
+            required=True,
+            type=str,
+            help="""Save composed dataset this directory""",
+        )
+        subcommand_parser.add_argument(
+            "--source_tokenizer_path",
+            default=None,
+            type=str,
+            help="Source tokenizer is loaded from this path (.model file)",
+        )
+        subcommand_parser.add_argument(
+            "--output_tokenizer_path",
+            default=None,
+            type=str,
+            help="SP tokenizer model is saved to this path",
+        )
+        subcommand_parser.add_argument(
+            "--text_field_name",
+            default="text",
+            type=str,
+            help="Text is read from this field name from composed dataset files.",
+        )
+        subcommand_parser.add_argument(
+            "--tokenizer_ratio",
+            default=DEFAULT_TOKENIZER_RATIO,
+            type=float,
+            help="This ratio of the training set is used for the tokenizer training.",
+        )
+        subcommand_parser.add_argument(
+            "--tokenizer_vocab_size",
+            default=256000,
+            type=int,
+            help="Vocab size.",
+        )
+        subcommand_parser.add_argument(
+            "--tokenizer_model_type",
+            default="bpe",
+            type=str,
+            help="Tokenizer model type (options: bpe, unigram, word, char).",
+        )
+        subcommand_parser.add_argument(
+            "--input_batch_size",
+            default=10_000,
+            type=int,
+            help="Input data is read with this batch size",
+        )
+        subcommand_parser.add_argument("--override", action="store_true", help="Override existing output files")
+        subcommand_parser.add_argument(
+            "--sentence_splitting", action="store_true", help="Split into texts into sentences"
+        )
+        subcommand_parser.add_argument(
+            "--workers",
+            default=1,
+            type=int,
+            help="Number of workers for parallel processing",
+        )
+
+        subcommand_parser = BaseCLICommand.add_common_args(
+            subcommand_parser,
+            raw_datasets_dir=False,
+            output=True,
+            extra_dataset_registries=True,
+            configs=True,
+            required_configs=True,
+            log=True,
+        )
+        subcommand_parser.set_defaults(func=TrainTokenizerCommand)
+
+    def __init__(self, args: Namespace) -> None:
+        self.tokenizer_backend = args.tokenizer_backend
+        self.config: Config = get_config_from_paths(args.config_paths, override=args.__dict__)
+
+    def run(self) -> None:
+        if self.tokenizer_backend == "sentencepiece":
+            train_sp_tokenizer(config=self.config)
+        else:
+            raise ValueError(f"Unsupported tokenizer backend: {self.tokenizer_backend }")
diff --git a/src/lm_datasets/compose_dataset.py b/src/lm_datasets/compose_dataset.py
new file mode 100644
index 0000000..72917bc
--- /dev/null
+++ b/src/lm_datasets/compose_dataset.py
@@ -0,0 +1,101 @@
+import json
+from pathlib import Path
+from lm_datasets.io.parquet import save_texts_to_parquet_chunks
+from datetime import datetime
+import pyarrow as pa
+from lm_datasets.utils import get_bytes_from_int_or_string
+
+from lm_datasets.utils.config import Config
+from lm_datasets.utils.dataset_generator import DatasetGenerator, DatasetSplit
+
+
+def compose_dataset(config: Config):
+    logger = config.init_logger(__name__)
+
+    logger.info("Compose dataset")
+
+    max_output_chunk_uncompressed_bytes = get_bytes_from_int_or_string(config.max_output_chunk_uncompressed_bytes)
+    # split_str = config.split
+    save_to_dir = Path(config.composed_dataset_dir)
+    output_compression = "ZSTD"
+    output_text_field = "text"
+    output_dataset_id_field = "dataset_id"
+    output_format = "parquet"
+
+    if not save_to_dir.exists():
+        logger.info("Save directory did not exist, creating it at %s", save_to_dir)
+        save_to_dir.mkdir()
+
+    # Initialize dataset generator (interleaves pre-shuffled datasets)
+    dataset_generator = DatasetGenerator(
+        config,
+        shuffled_output_dir=config.shuffled_output_dir,
+        output_format=output_format,
+        save_to_dir=save_to_dir,
+        split=config.split,
+    )
+    dataset_generator.prepare_datasets(
+        use_sampling=config.use_sampling,
+        print_progress=True,
+    )
+
+    # By default, the only column is the text column
+    parquet_columns = [
+        (output_text_field, pa.string()),
+    ]
+    if config.save_dataset_ids and config.split != DatasetSplit.VALIDATION:
+        # Add `dataset_id` column at in the first place (ID are prepended)
+        parquet_columns.insert(0, (output_dataset_id_field, pa.string()))
+
+    parquet_schema = pa.schema(parquet_columns)
+
+    if config.split == DatasetSplit.VALIDATION and config.use_separated_validation_sets:
+        # Validation splits are saved into separate files (based on dataeset ID)
+        for dataset_id in dataset_generator.list_of_dataset_ids:
+
+            def validation_output_path(part, total_parts=None):
+                return dataset_generator.save_to_path(part, total_parts, "_" + dataset_id)
+
+            save_texts_to_parquet_chunks(
+                dataset_generator.generate_texts_from_single_dataset(
+                    dataset_id,
+                    print_progress=True,
+                    reader_batch_size=config.input_batch_size,
+                ),
+                parquet_schema,
+                output_path_func=validation_output_path,
+                max_chunk_uncompressed_bytes=max_output_chunk_uncompressed_bytes,
+                compression=output_compression,
+                batch_size=config.output_batch_size,
+                print_write_progress=0,
+                limit=config.limit,
+            )
+    else:
+        # Save texts from iterator into a single (chunked) file to disk
+        save_texts_to_parquet_chunks(
+            dataset_generator.generate_texts_from_interleaved_datasets(
+                generate_dataset_ids=config.save_dataset_ids,
+                print_progress=True,
+                random_batch_size=config.interleave_random_batch_size,
+                reader_batch_size=config.input_batch_size,
+            ),
+            parquet_schema,
+            output_path_func=dataset_generator.save_to_path,
+            max_chunk_uncompressed_bytes=max_output_chunk_uncompressed_bytes,
+            compression=output_compression,
+            batch_size=config.output_batch_size,
+            print_write_progress=0,
+            limit=config.limit,
+        )
+
+    # Save stats
+    stats = {
+        "config": config.__dict__,
+        "dataset_id_to_stats": dataset_generator.dataset_id_to_stats,
+        "datetime": str(datetime.now()),
+    }
+    stats_path = save_to_dir / f"{config.split}_stats.json"
+    with open(stats_path, "w") as f:
+        json.dump(stats, f, indent=4)
+
+    logger.info("done")
diff --git a/src/lm_datasets/convert_parquet_to_jsonl.py b/src/lm_datasets/convert_parquet_to_jsonl.py
new file mode 100644
index 0000000..2a7bce1
--- /dev/null
+++ b/src/lm_datasets/convert_parquet_to_jsonl.py
@@ -0,0 +1,35 @@
+import argparse
+from pathlib import Path
+from typing import List
+import polars as pl
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def convert_parquet_to_jsonl(
+    input_dir_or_file: str, output_dir: str, override: bool = False, input_glob="*.parquet"
+) -> List[str]:
+    output_dir_path = Path(output_dir)
+    input_path = Path(input_dir_or_file)
+    input_paths = [input_path] if input_path.is_file() else input_path.glob(input_glob)
+    output_file_paths = []
+
+    logger.debug("Input paths: %s", input_paths)
+
+    for file_path in input_paths:
+        logger.debug(f"Converting {file_path}")
+
+        output_file_path = output_dir_path / file_path.with_suffix(".jsonl").name
+
+        if output_file_path.exists() and not override:
+            logger.warning(f"Skipping because output exists already: {output_file_path}")
+
+        else:
+            df = pl.scan_parquet(file_path, low_memory=True).collect(streaming=True)
+
+            df.write_ndjson(output_file_path)
+
+        output_file_paths.append(output_file_path)
+
+    return output_file_paths
diff --git a/src/lm_datasets/datasets/base.py b/src/lm_datasets/datasets/base.py
index c688471..0c4d3b4 100644
--- a/src/lm_datasets/datasets/base.py
+++ b/src/lm_datasets/datasets/base.py
@@ -3,20 +3,25 @@
 import logging
 import os
 from enum import Enum
+import random
 from urllib.error import HTTPError
-from typing import Iterable, List, Optional, TextIO, Tuple, Union
+from typing import Iterable, List, Literal, Optional, TextIO, Tuple, Type, Union
 
 import wget
 
 import pyarrow as pa
+import pyarrow.parquet as pq
+import polars as pl
 
-from smart_open import open
+from smart_open import open as smart_open
 
 from pathlib import Path
-from lm_datasets.io.parquet import save_texts_to_parquet_chunks
+from lm_datasets.io.parquet import get_selected_row_groups, save_texts_to_parquet_chunks
+from lm_datasets.utils.settings import DEFAULT_MIN_TEXT_LENGTH
 
-from lm_datasets.systems import get_path_by_system
+from lm_datasets.utils.systems import get_path_by_system
 from lm_datasets.utils import get_parquet_compression
+from lm_datasets.utils.config import Config
 
 MILLION = 1_000_000
 BILLION = 1_000_000_000
@@ -43,11 +48,12 @@ def __str__(self):
 class QualityWarning(Enum):
     SHORT_TEXT = "short_text"  # Dataset contains mostly short text (below paragraphs)
     BAD_ENCODING = "bad_encoding"  # Text might bit wrong encoded (utf-8 etc.)
-    BAD_WHITESPACES = (  # Text might contain bad whitespaces, e.g., before punctuation (converting tokens to text)
+    BAD_WHITESPACES = (  # Text might contain bad whitespaces, e.g., before punctuation (converting tokens to text).
         "bad_whitespaces"
     )
-    BAD_LINEBREAKS = "bad_linebreaks"  # Text has missing or too many line breaks
-    BAD_PUNCTUATION = "bad_punctuation"
+    BAD_LINEBREAKS = "bad_linebreaks"  # Text has missing or too many line breaks.
+    BAD_PUNCTUATION = "bad_punctuation"  # Punctuation may be missing or incorrect.
+    BAD_DOCUMENT_SPLITS = "bad_document_splits"  # Extract texts may span accross multi documents or documents are split into multiple texts.
 
     def __str__(self):
         return str(self.value)
@@ -56,6 +62,7 @@ def __str__(self):
 class Genre(Enum):
     LEGAL = "legal"
     SCIENCE = "science"
+    PATENT = "patent"
     MATH = "math"
     NEWS = "news"
     DIALOGUE = "dialogue"
@@ -75,6 +82,42 @@ def __str__(self):
     "No - we would need to spontaneously reach out to the current owners/custodians",
 ]
 
+
+class License(object):
+    """
+    Basic licensing information. Set attributions must be verified. If an attribution is unset, it is unknown.
+
+    See https://choosealicense.com/ and https://creativecommons.org/share-your-work/cclicenses/
+    """
+
+    def __init__(
+        self,
+        name,
+        url: Optional[str] = None,
+        commercial_use: Optional[bool] = None,
+        research_use: Optional[bool] = None,
+        modification: Optional[bool] = None,
+        distribution: Optional[bool] = None,
+        sharealike: Optional[bool] = None,
+        attribution: Optional[bool] = None,
+        derivates: Optional[bool] = None,
+        informing: Optional[bool] = None,
+    ):
+        self.name = name  # name and/or description of license
+        self.url = url  # link to full license text
+        self.commercial_use = commercial_use  # like CC-NC
+        self.research_use = research_use
+        self.modification = modification
+        self.distribution = distribution
+        self.sharealike = sharealike  # like CC-SA
+        self.attribution = attribution  # like CC-BY
+        self.derivates = derivates  # like CC-ND
+        self.informing = informing  # like ACA ID-BY-NC-INF-NORED
+
+    def __str__(self):
+        return f"{self.name} ({self.commercial_use=}; {self.sharealike=})"
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -88,24 +131,28 @@ class BaseDataset(object):
 
     TITLE = None
     DESCRIPTION = None
-    HOMEPAGE = None
-    AVAILIBILITY = None
-    DOWNLOAD_URLS = []
+    HOMEPAGE: Optional[str] = None
+    AVAILIBILITY: Availability = None
+    DOWNLOAD_URLS: List[Union[str, Tuple[str]]] = []
     LOCAL_DIRS = []
     VERSION = None
     DOI = None
+    CITATION = None
 
-    LICENSE = None
+    LICENSE: Optional[Union[str, License]] = None
     PII = None
 
     LANGUAGES = []
 
+    TRANSLATIONS = False
     WEB_CRAWLED = False
-    QUALITY_WARNINGS = []
-    GENRES = []
+    QUALITY_WARNINGS: List[QualityWarning] = []
+    GENRES: List[Genre] = []
+    HAS_OVERLAP_WITH: List[Union[Type, str]] = []
     USED_BY = None
     DUMMY = False
     SINGLE_OUTPUT_FILE = True
+    HAS_PREDEFINED_VALIDATION_SET = False
 
     # Statistics
     TOKENS = None
@@ -124,18 +171,20 @@ def __init__(
         skip_items: int = 0,
         hf_auth_token: str = None,
         print_write_progress: int = 10_000,
-        min_length: int = 250,
+        min_length: Optional[int] = None,
         json_ensure_ascii: bool = False,
         title_delimiter: str = ":\n\n",
         paragraph_delimiter: str = "\n\n",
         sentence_delimiter: str = " ",
-        output_format: str = "jsonl",  # options: jsonl, parquet
+        output_format: Literal["jsonl", "parquet"] = "jsonl",
         output_compression: Optional[
             str
         ] = None,  # jsonl: gzip, parquet: ‘NONE’, ‘SNAPPY’, ‘GZIP’, ‘BROTLI’, ‘LZ4’, ‘ZSTD’
         output_batch_size: int = 1000,
         shuffled_output_dir: Optional[str] = None,
-        max_output_chunk_uncompressed_bytes: int = 0,
+        max_output_chunk_uncompressed_bytes: Optional[int] = None,
+        max_output_chunk_rows: Optional[int] = None,
+        config: Config = None,
     ) -> None:
         self.output_dir = output_dir
         self.raw_datasets_dir = raw_datasets_dir
@@ -146,7 +195,7 @@ def __init__(
         self.skip_items = skip_items
         self.hf_auth_token = hf_auth_token
         self.print_write_progress = print_write_progress
-        self.min_length = min_length
+        self.min_length = min_length if min_length is not None else DEFAULT_MIN_TEXT_LENGTH
         self.json_ensure_ascii = json_ensure_ascii
         self.title_delimiter = title_delimiter
         self.paragraph_delimiter = paragraph_delimiter
@@ -156,6 +205,8 @@ def __init__(
         self.output_batch_size = output_batch_size
         self.shuffled_output_dir = shuffled_output_dir
         self.max_output_chunk_uncompressed_bytes = max_output_chunk_uncompressed_bytes
+        self.max_output_chunk_rows = max_output_chunk_rows
+        self.config = config
 
     def get_source_id(self):
         if self.SOURCE_ID:
@@ -173,6 +224,9 @@ def get_language_code(self, unknown: str = "unknown", mixed: str = "mixed"):
 
         return lang
 
+    def get_output_text_field(self):
+        return self.output_text_field
+
     def has_output_files(self, min_file_size: int = 1, shuffled=False) -> bool:
         return self.has_single_output_file(
             min_file_size=min_file_size, shuffled=shuffled
@@ -181,7 +235,7 @@ def has_output_files(self, min_file_size: int = 1, shuffled=False) -> bool:
     def has_single_output_file(self, min_file_size: int = 1, shuffled=False) -> bool:
         fp = self.get_single_output_file_path(shuffled=shuffled)
 
-        return os.path.exists(fp) and os.stat(fp).st_size >= min_file_size
+        return fp is not None and os.path.exists(fp) and os.stat(fp).st_size >= min_file_size
 
     def has_chunked_output_files(self, min_file_size: int = 1, shuffled=False) -> bool:
         for fp in self.get_chunked_output_file_paths(shuffled=shuffled):
@@ -301,7 +355,7 @@ def save_texts(self, texts: Iterable[str], append: bool = False):
             if append:
                 raise NotImplementedError("Appending is not supported by parquet output format")
 
-            docs_count = self.save_texts_to_parquet(texts)
+            docs_count, saved_chunks = self.save_texts_to_parquet(texts)
 
         else:
             raise ValueError(f"Unsupported output format: {self.output_format}")
@@ -321,6 +375,8 @@ def save_texts_to_parquet(self, texts: Iterable[str], file_path: Optional[str] =
         """
         Save text in parquet (single column schema, in batches)
         """
+        assert self.output_format == "parquet"
+
         if file_path is None:
             file_path = self.get_output_file_paths(single=True)[0]
 
@@ -329,7 +385,7 @@ def save_texts_to_parquet(self, texts: Iterable[str], file_path: Optional[str] =
 
         schema = pa.schema(
             [
-                (self.output_text_field, pa.string()),
+                (self.get_output_text_field(), pa.string()),
             ]
         )
         # Max. chunk size is multiplied with this factor
@@ -342,7 +398,10 @@ def save_texts_to_parquet(self, texts: Iterable[str], file_path: Optional[str] =
         saved_docs, saved_chunks = save_texts_to_parquet_chunks(
             texts=texts,
             schema=schema,
-            max_chunk_uncompressed_bytes=self.max_output_chunk_uncompressed_bytes * safety_factor,
+            max_chunk_uncompressed_bytes=self.max_output_chunk_uncompressed_bytes * safety_factor
+            if self.max_output_chunk_uncompressed_bytes is not None
+            else None,
+            max_chunk_rows=self.max_output_chunk_rows,
             output_path_func=self.get_single_or_chunked_output_file_path,
             compression=get_parquet_compression(self.output_compression),
             batch_size=self.output_batch_size,
@@ -354,7 +413,7 @@ def save_texts_to_parquet(self, texts: Iterable[str], file_path: Optional[str] =
             logger.info("Killing all remaining workers, if any (iterator end)")
             texts.terminate()
 
-        return saved_docs
+        return saved_docs, saved_chunks
 
     def save_texts_to_jsonl(self, texts: Iterable[str], append: bool = False):
         """
@@ -369,12 +428,12 @@ def save_texts_to_jsonl(self, texts: Iterable[str], append: bool = False):
 
         docs_count = 0
 
-        with open(fp, mode) as f:
+        with smart_open(fp, mode) as f:
             for docs_count, text in enumerate(self.filter_texts(texts), 1):
-                f.write(json.dumps({self.output_text_field: text}, ensure_ascii=self.json_ensure_ascii) + "\n")
+                f.write(json.dumps({self.get_output_text_field(): text}, ensure_ascii=self.json_ensure_ascii) + "\n")
 
                 if docs_count > 0 and (docs_count % self.print_write_progress) == 0:
-                    logger.info(f"Writen {docs_count:,} docs ...")
+                    logger.info(f"Written {docs_count:,} docs ...")
 
                 if self.limit > 0 and docs_count >= self.limit:
                     logger.warning(f"Limit reached ({docs_count:,} docs)")
@@ -403,7 +462,14 @@ def get_hf_auth_token(self):
         return None
 
     def get_local_dataset_dir(self):
-        if self.LOCAL_DIRS:
+        if self.config:
+            if self.DATASET_ID in self.config.local_dirs_by_dataset_id:
+                return self.config.local_dirs_by_dataset_id[self.DATASET_ID]
+
+            if self.get_source_id() in self.config.local_dirs_by_source_id:
+                return self.config.local_dirs_by_source_id[self.get_source_id()]
+
+        if self.LOCAL_DIRS:  # TODO deprecated -> use config instead!
             # manually defined dataset directory
             return get_path_by_system(self.LOCAL_DIRS)
         elif self.raw_datasets_dir:
@@ -447,8 +513,10 @@ def get_dataset_file_paths(
         fps = [fp for fp in fps if os.path.isfile(fp)]
 
         if single_file:
-            if len(fps) != 1:
-                raise ValueError(f"Multiple files in download directory but only a single one was expected: {fps}")
+            if len(fps) > 1:
+                raise FileExistsError(f"Multiple files in download directory but only a single one was expected: {fps}")
+            elif len(fps) == 0:
+                raise FileNotFoundError(f"No file found but a single one was expected: {fps}")
 
             return fps[0]
 
@@ -540,8 +608,341 @@ def get_texts_from_conllu_file(self, file_handler: TextIO):
     def get_texts(self) -> Iterable[str]:
         raise NotImplementedError
 
-    def extract_plaintext(self):
-        self.save_texts(self.get_texts())
+    def extract_plaintext(self) -> int:
+        saved_texts_count = self.save_texts(self.get_texts())
 
         if self.counter:
             logger.info(f"Statistics {self.counter}")
+
+        return saved_texts_count
+
+    def get_output_rows_count(self, shuffled: bool = False) -> int:
+        """
+        Read metadata from parquet files and extract number of rows
+        """
+        if self.output_format == "parquet":
+            output_paths = list(self.get_output_file_paths(shuffled=shuffled))
+
+            # Filter for existing
+            output_paths = [output_path for output_path in output_paths if os.path.exists(output_path)]
+
+            if output_paths:
+                rows_count = 0
+
+                for output_path in output_paths:
+                    with open(output_path, "rb") as f:
+                        parquet_file = pq.ParquetFile(
+                            f,
+                            # increased to avoid OSErrors
+                            thrift_string_size_limit=1000000000,  # default: 100000000
+                            thrift_container_size_limit=10000000,  # default: 1000000
+                        )
+                        rows_count += parquet_file.metadata.num_rows
+
+                        logger.debug("Rows = %s in %s", rows_count, output_path)
+
+                return rows_count
+
+            logger.debug("No output files exists for %s", self.DATASET_ID)
+            return -1
+        else:
+            raise ValueError(f"Cannot determine the output rows count with {self.output_format=}")
+
+    def get_compression_from_output_files(self, shuffled=False):
+        """
+        NOTE: Currently only implemented for `parquet` format.
+        """
+        if self.output_format == "parquet":
+            for output_path in self.get_output_file_paths(shuffled=shuffled):
+                if os.path.exists(output_path):
+                    with open(output_path, "rb") as f:
+                        parquet_file = pq.ParquetFile(
+                            f,
+                            # increased to avoid OSErrors
+                            thrift_string_size_limit=1000000000,  # default: 100000000
+                            thrift_container_size_limit=10000000,  # default: 1000000
+                        )
+                        parquet_metadata = parquet_file.metadata
+                        for i in range(parquet_metadata.num_row_groups):
+                            for j in range(parquet_metadata.num_columns):
+                                return parquet_file.metadata.row_group(i).column(j).compression
+
+        return "unknown"
+
+    def generate_texts_from_output(
+        self,
+        shuffled: bool = False,
+        batch_size: Optional[int] = None,
+        limit: int = 0,
+        offset: int = 0,
+        shuffle_output_file_paths: bool = False,
+        reader_implementation: Literal["polars_read_parquet", "pyarrow"] = "pyarrow",
+    ):
+        """
+        A iterator over texts from processed output files.
+        """
+        if batch_size is None:
+            batch_size = self.output_batch_size
+
+        if self.output_format != "parquet":
+            raise ValueError(f"Cannot texts with {self.output_format=}")
+
+        # Check if output files exists and sort them
+        output_paths = [
+            file_path
+            for file_path in sorted(self.get_output_file_paths(shuffled=shuffled))
+            if os.path.exists(file_path)
+        ]
+
+        # Count generated rows
+        rows = 0
+        rows_limit = limit - offset
+
+        # if limit > 0:
+        #     batch_size = min(batch_size, limit)
+
+        # Shuffle output chunks:
+        # This changes the order in that the chunks are read ensure also shuffling on the full dataset level.
+        if shuffle_output_file_paths:
+            random.seed(self.config.seed)  # reset seed to avoid inference by other scripts
+            random.shuffle(output_paths)
+
+        chunk_start = 0
+        chunk_end = None
+
+        if output_paths:
+            for file_path in output_paths:
+                logger.info("Generating text from %s", file_path)
+
+                # PyArrow implementation
+                with open(file_path, "rb") as file_handler:
+                    pq_file = pq.ParquetFile(
+                        file_handler,
+                        # memory_map=False,
+                    )
+                    file_rows_count = pq_file.metadata.num_rows
+
+                    chunk_end = chunk_start + file_rows_count - 1
+
+                    # Should we read from the current chunk?
+                    # Yes, if
+                    # - offset is smaller or equal chunk_start
+                    # (- limit is greater or equal chunk_end) --- limit does not matter
+
+                    # variants
+                    # A) requested rows start in chunk and ends in chunk
+                    # B) requested rows start in chunk but ends in following chunk
+                    # C) requested rows start before chunk and ends in chunk
+                    # D) requested rows start before chunk and ends in following chunk
+
+                    if (
+                        chunk_start <= offset < chunk_end
+                        or offset < chunk_start
+                        and (limit == 0 or chunk_start < limit)
+                    ):
+                        file_offset = max(
+                            0, offset - chunk_start
+                        )  # global offset minus start of current file (current chunk)
+                        file_limit = (
+                            max(0, limit - chunk_start) if limit > 0 else 0  # limit - chunk_start
+                        )  # Length of the slice: global limit minus start of current chunk
+                        # TODO before: limit - chunk_start - file_offset
+
+                        logger.debug(
+                            f"Reading file chunk from %s: file [%s - %s]; global [%s - %s]; chunk [%s - %s]",
+                            file_path,
+                            file_offset,
+                            file_limit,
+                            offset,
+                            limit,
+                            chunk_start,
+                            chunk_end,
+                        )
+                        if reader_implementation == "pyarrow":
+                            # PyArrow implementation with iter_batches
+                            # with open(file_path, "rb") as file_handler:
+                            #     parquet_file = pq.ParquetFile(file_handler)
+
+                            for batch_idx, pq_batch in enumerate(
+                                pq_file.iter_batches(
+                                    columns=[self.get_output_text_field()], batch_size=batch_size, use_threads=False
+                                )
+                            ):
+                                for row_idx, text_column in enumerate(pq_batch.columns[0], batch_idx * batch_size):
+                                    if row_idx >= file_offset:
+                                        if rows_limit > 0 and rows >= rows_limit:
+                                            # break row loop
+                                            logger.debug("break row loop")
+                                            break
+
+                                        # cast to string
+                                        # text = text_column.as_py()
+
+                                        text = text_column
+
+                                        yield text
+                                        rows += 1
+
+                                if rows_limit > 0 and rows >= rows_limit:
+                                    # break batch loop
+                                    logger.debug("break batch loop")
+                                    break
+
+                            # PyArrow implementation with read_row_group
+                            # with open(file_path, "rb") as file_handler:
+                            #     parquet_file = pq.ParquetFile(file_handler)
+
+                            #     # 1. What row groups need to be read?
+                            #     row_groups, group_idx_to_offset_last_row = get_selected_row_groups(
+                            #         parquet_file, file_offset, file_limit
+                            #     )
+                            #     logger.debug("Selected row groups: %s; %s", row_groups, group_idx_to_offset_last_row)
+
+                            #     # 2. Read selected row groups
+                            #     for selected_row_group in row_groups:
+                            #         logger.debug("Read row group: %s", selected_row_group)
+                            #         group_table = parquet_file.read_row_group(
+                            #             selected_row_group, columns=[self.get_output_text_field()]
+                            #         )
+
+                            #         # What offsets and limit? (only if needed)
+                            #         if group_idx_to_offset_last_row is not None:
+                            #             group_offset, _ = group_idx_to_offset_last_row[selected_row_group]
+
+                            #             row_offset = max(0, file_offset - group_offset)
+                            #             logger.debug("Row group: %s; row offset: %s", selected_row_group, row_offset)
+
+                            #         # Iterate over rows
+                            #         for row_idx, text_column in enumerate(group_table.columns[0]):
+                            #             # Skip rows before offset
+                            #             if group_idx_to_offset_last_row is None or row_idx >= row_offset:
+                            #                 if rows_limit > 0 and rows >= rows_limit:
+                            #                     # break row loop
+                            #                     logger.debug("break row loop")
+                            #                     break
+
+                            #                 text = text_column.as_py()  # cast to str
+                            #                 yield text
+                            #                 rows += 1
+
+                            #         if rows_limit > 0 and rows >= rows_limit:
+                            #             # break row group loop
+                            #             logger.debug("break row group loop")
+                            #             break
+
+                        elif reader_implementation == "polars_read_parquet":
+                            # Polars "scan_parquet" implementation: Error "Segmentation fault (core dumped)"
+                            # df = (
+                            #     pl.scan_parquet(file_path, low_memory=True).collect(
+                            #     streaming=True
+                            # ).slice(offset=file_offset, length=file_limit if file_limit != 0 else None)
+                            #     .collect(streaming=True)
+                            # )
+                            # text_column_index = df.columns.index(self.get_output_text_field())
+
+                            df = pl.read_parquet(
+                                file_path, low_memory=True, columns=[self.get_output_text_field()]
+                            ).slice(offset=file_offset, length=file_limit if file_limit != 0 else None)
+                            text_column_index = 0
+
+                            # Iterate over rows
+                            for row in df.iter_rows():
+                                text = str(row[text_column_index])
+                                yield text
+                                rows += 1
+
+                                if rows_limit > 0 and rows >= rows_limit:
+                                    # break row loop
+                                    break
+                            else:
+                                raise ValueError("Invalid `reader_implementation`")
+                    else:
+                        logger.debug("Skip this file because output does not contain the requested rows: %s", file_path)
+
+                    # current_offset += file_rows_count  # TODO +1?
+                    chunk_start = chunk_end + 1  # set start for the next chunk
+
+                if rows_limit > 0 and rows >= rows_limit:
+                    # break file loop
+                    logger.debug("break file loop")
+                    break
+        else:
+            logger.warning("Cannot generate texts because output files do not exist.")
+
+        logger.info(
+            "Texts generated: %s (expected size: %s; offset: %s; limit: %s;)", rows, limit - offset, offset, limit
+        )
+
+    def get_estimated_bytes_from_output(self, shuffled: bool = False, read_first_n_rows: int = 1_000) -> int:
+        """
+        Estimate byte size of output text:
+        - read first N rows of shuffled output files and count their byte size
+        - multiply counted bytes by total number of rows
+        """
+        if not shuffled:
+            raise NotImplementedError
+
+        if self.output_format != "parquet":
+            raise NotImplementedError
+
+        bytes_sum = 0
+        total_rows = 0
+
+        # iterate over output files (use shuffled files for a better estimate)
+        for output_path in self.get_output_file_paths(shuffled=shuffled):
+            if os.path.exists(output_path):
+                # read the first n rows
+                df = pl.scan_parquet(
+                    output_path,
+                    low_memory=True,
+                    n_rows=read_first_n_rows,
+                ).collect(streaming=True)
+                for row in df.iter_rows():
+                    text = str(row[0])
+                    bytes_sum += len(text.encode("utf-8"))  # count the byte size of the text
+
+                # read total row count from metadata
+                with open(output_path, "rb") as f:
+                    parquet_file = pq.ParquetFile(
+                        f,
+                        # increased to avoid OSErrors
+                        thrift_string_size_limit=1000000000,  # default: 100000000
+                        thrift_container_size_limit=10000000,  # default: 1000000
+                    )
+                    total_rows += parquet_file.metadata.num_rows
+
+        # estimated bytes
+        bytes_per_row = bytes_sum / read_first_n_rows
+        total_bytes = int(total_rows * bytes_per_row)
+
+        return total_bytes
+
+    def get_sampling_factor(self) -> float:
+        """
+        Sampling is defined based on dataset ID, source ID, or language.
+        """
+        if self.config:
+            if self.DATASET_ID in self.config.sampling_factor_by_dataset_id:
+                return self.config.sampling_factor_by_dataset_id[self.DATASET_ID]
+
+            if self.get_source_id() in self.config.sampling_factor_by_source_id:
+                return self.config.sampling_factor_by_source_id[self.get_source_id()]
+
+            if self.get_language_code() in self.config.sampling_factor_by_language:
+                return self.config.sampling_factor_by_language[self.get_language_code()]
+
+        return 1.0  # default factor
+
+    def is_selected(self) -> bool:
+        """
+        Is this dataset part of selected datasets or sources?
+        """
+        return (
+            self.DATASET_ID in self.config.selected_dataset_ids
+            or self.get_source_id() in self.config.selected_source_ids
+        )
+
+    def get_shuffled_output_file_path(self, unshuffled_output_file_path: str) -> str:
+        output_file_name = Path(unshuffled_output_file_path).name
+
+        return os.path.join(self.config.shuffled_output_dir, output_file_name.replace(".parquet", ".shuffled.parquet"))
diff --git a/src/lm_datasets/datasets/bg/bgnc_admin_eur.py b/src/lm_datasets/datasets/bg/bgnc_admin_eur.py
index 6909175..e492edf 100644
--- a/src/lm_datasets/datasets/bg/bgnc_admin_eur.py
+++ b/src/lm_datasets/datasets/bg/bgnc_admin_eur.py
@@ -6,6 +6,7 @@
 logger = logging.getLogger(__name__)
 
 
+# deprecated -> use bulnc instead!
 class BGNCAdminEURDataset(BaseDataset):
     """
     Part of Bulgarian National Corpus
@@ -20,11 +21,11 @@ class BGNCAdminEURDataset(BaseDataset):
     AVAILIBILITY = Availability.DIRECT_DOWNLOAD
 
     LANGUAGES = ["bg"]
-
+    HAS_OVERLAP_WITH = [
+        "bulnc",
+    ]
     DOWNLOAD_URLS = ["https://dcl.bas.bg/BulNC-registration/dl.php?dl=feeds/ADMIN_EUR.BG.zip"]
 
-    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/bg/BgNC/admin_eur/ADMIN_EUR.BG"]
-
     BYTES = 257 * MB
 
     def download(self):
diff --git a/src/lm_datasets/datasets/bg/bgnc_news_corpus.py b/src/lm_datasets/datasets/bg/bgnc_news_corpus.py
index 4981a3b..9ab40d3 100644
--- a/src/lm_datasets/datasets/bg/bgnc_news_corpus.py
+++ b/src/lm_datasets/datasets/bg/bgnc_news_corpus.py
@@ -6,6 +6,7 @@
 logger = logging.getLogger(__name__)
 
 
+# deprecated -> use bulnc instead!
 class BGNCNewsCorpusDataset(BaseDataset):
     """
     Part of Bulgarian National Corpus
@@ -20,9 +21,7 @@ class BGNCNewsCorpusDataset(BaseDataset):
     LANGUAGES = ["bg"]
 
     DOWNLOAD_URLS = ["https://dcl.bas.bg/BulNC-registration/dl.php?dl=feeds/JOURNALISM.BG.zip"]
-
-    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/bg/BgNC/news_corpus"]
-
+    HAS_OVERLAP_WITH = ["bulnc"]
     BYTES = 57 * MB
 
     def download(self):
diff --git a/src/lm_datasets/datasets/bg/bulgarian_news.py b/src/lm_datasets/datasets/bg/bulgarian_news.py
index d0351bc..a3c5d32 100644
--- a/src/lm_datasets/datasets/bg/bulgarian_news.py
+++ b/src/lm_datasets/datasets/bg/bulgarian_news.py
@@ -1,4 +1,6 @@
-from lm_datasets.datasets.base import MB, BaseDataset
+import json
+from pathlib import Path
+from lm_datasets.datasets.base import MB, BaseDataset, Availability, License
 
 import logging
 
@@ -7,49 +9,36 @@
 
 class BulgarianNewsDataset(BaseDataset):
     DATASET_ID = "bulgarian_news"
+    TITLE = "Crawl of Bulgarian news websites"
     DOWNLOAD_URLS = ["http://old.dcl.bas.bg/dataset/Bulgarian_news.7z"]
     DESCRIPTION = (
         "The collection was collected by crawling Bulgarian websites in Bulgarian. Text samples are in json format. We"
         " can provide raw tests."
     )
-
+    WEB_CRAWLED = True
     LANGUAGES = ["bg"]
     BYTES = 919 * MB
+    AVAILIBILITY = Availability.ON_REQUEST
+    LICENSE = License("research only")
 
     def decompress(self):
         # 7z x Bulgarian_news.7z
         pass
 
     def get_texts(self):
-        # from py7zr import SevenZipFile
-
         # read from extracted JSON files
-        # fps = self.get_dataset_file_paths(subdirectories=True, needed_suffix=".json")
-
-        # read from 7z
-        # fp = self.get_dataset_file_paths(single_file=True, needed_suffix=".7z")
-        # logger.info(f"Reading from {fp}")
-
-        # with SevenZipFile(fp, "r") as zip:
-        #     fn = "Bulgarian_news/n/nbox.bg/bg_00000634-24dd-4c46-b846-af4e5c3db559.json"
-        #     print(fn)
-
-        #     content = zip.read(fn)[fn].read()
-
-        #     # for fname, bio in zip.read(targets=[fn]).items():
-        #     #     print(f"{fname}: {bio.read(10)}...")
-        #     #     break
-
-        #     # allfiles = zip.getnames()
-
-        # with SevenZipFile(fp, "r") as zip:
-        #     for fn in allfiles:
-        #         if fn.endswith(".json"):
-        #             for fname, bio in zip.read(fn).items():
-        #                 print(f"{fname}: {bio.read(10)}...")
-        #         break
-
-        # with lzma.open(fp) as f:
-        #     file_content = f.read()
-
-        print("x")
+        for i, file_path in enumerate(Path(self.get_local_dataset_dir()).rglob("*.json")):
+            if self.skip_items > 0 and i < self.skip_items:
+                continue
+
+            with open(file_path) as f:
+                try:
+                    doc = json.load(f)
+                    if "bg_a_text" in doc:
+                        text = self.paragraph_delimiter.join(doc["bg_a_text"])
+                        yield text
+                    else:
+                        logger.warning("JSON has no text field: %s", file_path)
+
+                except ValueError:
+                    logger.error("Cannot parse JSON from %s", file_path)
diff --git a/src/lm_datasets/datasets/bg/bulnc.py b/src/lm_datasets/datasets/bg/bulnc.py
index 0dcd2ab..b987ded 100644
--- a/src/lm_datasets/datasets/bg/bulnc.py
+++ b/src/lm_datasets/datasets/bg/bulnc.py
@@ -1,4 +1,5 @@
-from lm_datasets.datasets.base import BaseDataset, GB, Availability
+from pathlib import Path
+from lm_datasets.datasets.base import BaseDataset, GB, Availability, License
 import logging
 
 logger = logging.getLogger(__name__)
@@ -8,7 +9,7 @@ class BulNCDataset(BaseDataset):
     DATASET_ID = "bulnc"
     TITLE = "Bulgarian National Corpus"
     AVAILIBILITY = Availability.ON_REQUEST
-    DOWNLOAD_URLS = ["http://old.dcl.bas.bg/dataset/BulNC.7z"]
+    DOWNLOAD_URLS = ["http://old.dcl.bas.bg/dataset/BulNC.7z"]  # password-protected file!
     DESCRIPTION = (
         "The Bulgarian National Corpus contains a wide range of texts in various sizes, media types (written and "
         "spoken), styles, periods (synchronic and diachronic), and licenses. Each text in the collection is supplied "
@@ -20,15 +21,28 @@ class BulNCDataset(BaseDataset):
         "popular science. The BulNC is not fully downloadable due to the inclusion of copyrighted material. We've "
         "provided a link to a password-protected archive for evaluation."
     )
-
+    AVAILIBILITY = Availability.ON_REQUEST
+    LICENSE = License("research only", sharealike=False)
     LANGUAGES = ["bg"]
     BYTES = 1.8 * GB
 
     def decompress(self):
-        # 7z x BulNC.7z
+        """
+        7z x BulNC.7z
+
+        Folders: 125
+        Files: 256906
+        Size:       13279357395
+        Compressed: 1981942477
+        """
         pass
 
     def get_texts(self):
-        # # read from extracted JSON files
-        # fps = self.get_dataset_file_paths(subdirectories=True, needed_suffix=".json")
-        raise NotImplementedError()
+        # read from extracted TXT files
+        for file_path in Path(self.get_local_dataset_dir()).rglob(
+            "*.txt"
+        ):  # self.get_dataset_file_paths(subdirectories=True, needed_suffix=".txt"):
+            with open(file_path) as f:
+                text = f.read()
+
+                yield text
diff --git a/src/lm_datasets/datasets/code/starcoder.py b/src/lm_datasets/datasets/code/starcoder.py
index b48cffa..21b0334 100644
--- a/src/lm_datasets/datasets/code/starcoder.py
+++ b/src/lm_datasets/datasets/code/starcoder.py
@@ -1,16 +1,122 @@
+import os
 from pathlib import Path
-from lm_datasets.datasets.base import BaseDataset, Availability
+from typing import List, Optional
+from lm_datasets.datasets.base import Availability, License
 from lm_datasets.datasets.hf_dataset import HFDataset
 
+import logging
+
+from lm_datasets.datasets.parquet_dataset import ParquetDataset
+
+
+logger = logging.getLogger(__name__)
+
+# Dataset IDs:
+# starcoder_emacs-lisp,starcoder_literate-haskell,starcoder_shell,starcoder_ada,starcoder_erlang,starcoder_lua,starcoder_smalltalk,starcoder_agda,starcoder_f-sharp,starcoder_makefile,starcoder_solidity,starcoder_alloy,starcoder_fortran,starcoder_maple,starcoder_sparql,starcoder_antlr,starcoder_git-commits-cleaned,starcoder_markdown,starcoder_sql,starcoder_applescript,starcoder_github-issues-filtered-structured,starcoder_mathematica,starcoder_stan,starcoder_assembly,starcoder_glsl,starcoder_matlab,starcoder_standard-ml,starcoder_augeas,starcoder_go,starcoder_ocaml,starcoder_stata,starcoder_awk,starcoder_groovy,starcoder_pascal,starcoder_systemverilog,starcoder_batchfile,starcoder_haskell,starcoder_perl,starcoder_tcl,starcoder_bluespec,starcoder_html,starcoder_php,starcoder_tcsh,starcoder_c,starcoder_idris,starcoder_powershell,starcoder_tex,starcoder_c-sharp,starcoder_isabelle,starcoder_prolog,starcoder_thrift,starcoder_clojure,starcoder_java,starcoder_protocol-buffer,starcoder_typescript,starcoder_cmake,starcoder_java-server-pages,starcoder_python,starcoder_verilog,starcoder_coffeescript,starcoder_javascript,starcoder_r,starcoder_vhdl,starcoder_common-lisp,starcoder_json,starcoder_racket,starcoder_visual-basic,starcoder_cpp,starcoder_julia,starcoder_restructuredtext,starcoder_xslt,starcoder_css,starcoder_jupyter-scripts-dedup-filtered,starcoder_rmarkdown,starcoder_yacc,starcoder_cuda,starcoder_jupyter-structured-clean-dedup,starcoder_ruby,starcoder_yaml,starcoder_dart,starcoder_kotlin,starcoder_rust,starcoder_zig,starcoder_dockerfile,starcoder_lean,starcoder_sas,starcoder_elixir,starcoder_literate-agda,starcoder_scala,starcoder_elm,starcoder_literate-coffeescript,starcoder_scheme
+STARCODER_PROGRAMMING_LANGUAGES = [
+    "emacs-lisp",
+    "literate-haskell",
+    "shell",
+    "ada",
+    "erlang",
+    "lua",
+    "smalltalk",
+    "agda",
+    "f-sharp",
+    "makefile",
+    "solidity",
+    "alloy",
+    "fortran",
+    "maple",
+    "sparql",
+    "antlr",
+    "git-commits-cleaned",
+    "markdown",
+    "sql",
+    "applescript",
+    "github-issues-filtered-structured",
+    "mathematica",
+    "stan",
+    "assembly",
+    "glsl",
+    "matlab",
+    "standard-ml",
+    "augeas",
+    "go",
+    "ocaml",
+    "stata",
+    "awk",
+    "groovy",
+    "pascal",
+    "systemverilog",
+    "batchfile",
+    "haskell",
+    "perl",
+    "tcl",
+    "bluespec",
+    "html",
+    "php",
+    "tcsh",
+    "c",
+    "idris",
+    "powershell",
+    "tex",
+    "c-sharp",
+    "isabelle",
+    "prolog",
+    "thrift",
+    "clojure",
+    "java",
+    "protocol-buffer",
+    "typescript",
+    "cmake",
+    "java-server-pages",
+    "python",
+    "verilog",
+    "coffeescript",
+    "javascript",
+    "r",
+    "vhdl",
+    "common-lisp",
+    "json",
+    "racket",
+    "visual-basic",
+    "cpp",
+    "julia",
+    "restructuredtext",
+    "xslt",
+    "css",
+    "jupyter-scripts-dedup-filtered",
+    "rmarkdown",
+    "yacc",
+    "cuda",
+    "jupyter-structured-clean-dedup",
+    "ruby",
+    "yaml",
+    "dart",
+    "kotlin",
+    "rust",
+    "zig",
+    "dockerfile",
+    "lean",
+    "sas",
+    "elixir",
+    "literate-agda",
+    "scala",
+    "elm",
+    "literate-coffeescript",
+    "scheme",
+]
+
 
 # files already downloaded => use parquet DS instead!
 class StarcoderHFDataset(HFDataset):
     DATASET_ID = "starcoder"
 
-    TITLE = "StarCoder"
+    TITLE = "Starcoder"
     HOMEPAGE = "https://huggingface.co/datasets/bigcode/starcoderdata"
     AVAILIBILITY = Availability.SIGNIN_DOWNLOAD
-
+    LICENSE = "mixed permissive liceses (see https://huggingface.co/datasets/bigcode/the-stack-dedup)"
     LANGUAGES = ["code"]
 
     TOKENS = 250_000_000_000
@@ -22,33 +128,29 @@ class StarcoderHFDataset(HFDataset):
     # remove_columns = ["doc_id", "LICENSE", "uri", "date_built"]
 
 
-class ParquetDataset(BaseDataset):
-    SINGLE_OUTPUT_FILE = False
-
-    def extract_plaintext(self):
-        if self.output_format == "parquet":
-            raise ValueError("Dataset is already in parquet format; no text extraction needed!")
-        else:
-            super().extract_plaintext()
-
-
-class StarcoderDataset(ParquetDataset):
+class StarcoderBaseDataset(ParquetDataset):
     """
-    Output files:
-    /netscratch/ortiz/corpora/starcoder/starcoderdata/<code_lang>/train-<i>-of-<n>.parquet
+    Output files: $LOCAL_DIR/<code_lang>/train-<i>-of-<n>.parquet
     """
 
-    DATASET_ID = "starcoder"
+    DATASET_ID = None
+    PROGRAMMING_LANGUAGE = None
 
-    TITLE = "StarCoder"
+    SOURCE_ID = "starcoder"
+    TITLE = "Starcoder"
     HOMEPAGE = "https://huggingface.co/datasets/bigcode/starcoderdata"
     AVAILIBILITY = Availability.SIGNIN_DOWNLOAD
-
+    LICENSE = License(
+        "mixed permissive liceses",
+        url="https://huggingface.co/datasets/bigcode/the-stack-dedup",
+        attribution=True,
+        commercial_use=True,
+        research_use=True,
+        sharealike=False,
+    )
     LANGUAGES = ["code"]
 
-    TOKENS = 250_000_000_000
-
-    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/starcoder/starcoderdata"]
+    TOKENS = 250_000_000_000 / len(STARCODER_PROGRAMMING_LANGUAGES)  # wrongly assuming uniform distribution
 
     HF_DATASET_ID = "bigcode/starcoderdata"
     HF_DATASET_SPLIT = "train"
@@ -56,5 +158,53 @@ class StarcoderDataset(ParquetDataset):
     text_column_name = "text"
     # remove_columns = ["doc_id", "LICENSE", "uri", "date_built"]
 
-    def get_output_file_paths(self):
-        return Path(self.get_local_dataset_dir()).glob("*.parquet")
+    def get_output_text_field(self):
+        # hard-coded text field in parquet schema
+        return "content"
+
+    def get_output_dir(self, shuffled=False):
+        if shuffled:
+            if self.shuffled_output_dir:
+                return self.shuffled_output_dir
+            raise ValueError("shuffled_output_dir is not set")
+        else:
+            return str(Path(self.get_local_dataset_dir()) / self.PROGRAMMING_LANGUAGE)
+
+    def get_single_output_file_path(self, shuffled=False) -> str:
+        return None
+
+    def has_chunked_output_files(self, **kwargs):
+        return True
+
+    # def get_output_file_paths(self, single=False, chunked=False, shuffled=False):
+    #     if shuffled is not None or single or not chunked:
+    #         logger.warning(f"starcoder data is only provided as it is (no shuffled version, only chunked output files)")
+
+    #     return (Path(self.get_local_dataset_dir()) / self.PROGRAMMING_LANGUAGE).rglob("*.parquet")
+
+    def get_shuffled_output_file_path(self, unshuffled_output_file_path: str) -> str:
+        output_file_name = Path(unshuffled_output_file_path).name
+        shuffled_file_name = output_file_name.replace(".parquet", ".shuffled.parquet")
+
+        return os.path.join(self.config.shuffled_output_dir, f"{self.DATASET_ID}_{shuffled_file_name}")
+
+    def get_chunked_output_file_paths(self, shuffled=False) -> List[str]:
+        if shuffled:
+            # from shuffled output dir
+            return list(Path(self.shuffled_output_dir).glob(f"{self.DATASET_ID}_*.shuffled.parquet"))
+        else:
+            # from starcoder data directory
+            return list((Path(self.get_local_dataset_dir()) / self.PROGRAMMING_LANGUAGE).rglob("*.parquet"))
+
+
+def get_starcoder_class(lang):
+    class StarcoderDataset(StarcoderBaseDataset):
+        DATASET_ID = "starcoder_" + lang
+        PROGRAMMING_LANGUAGE = lang
+        TITEL = f"Starcoder ({lang})"
+
+    return StarcoderDataset
+
+
+def get_auto_starcoder_classes():
+    return [get_starcoder_class(lang) for lang in STARCODER_PROGRAMMING_LANGUAGES]
diff --git a/src/lm_datasets/datasets/cs/cs_en_parallel.py b/src/lm_datasets/datasets/cs/cs_en_parallel.py
index be6b856..756c997 100644
--- a/src/lm_datasets/datasets/cs/cs_en_parallel.py
+++ b/src/lm_datasets/datasets/cs/cs_en_parallel.py
@@ -1,142 +1,93 @@
+from io import BytesIO
 import json
 import logging
 import argparse
 import os
+import tarfile
 from typing import Iterable
 from smart_open import open
 import re
 import gzip
 
-logging.basicConfig(
-    format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
-    datefmt="%Y-%m-%d %H:%M:%S",
-    level=logging.INFO,
-    handlers=[logging.StreamHandler()],
-)
-logger = logging.getLogger(__name__)
-
-# dataset metadata
-DATASET_ID = "cs_en_parallel"
-TITLE = "Czech-English Parallel Corpus 1.0 (CzEng 1.0)"
-HOMEPAGE = "http://hdl.handle.net/11234/1-1458"
-
-# language codes of the dataset (see ./resources/languages.json)
-LANGUAGES = ["cs"]
-
-
-"""
-DOWNLOAD
------------
-
-Instruction
-
-- Run the command on the server:
-
-curl --remote-name-all \
-    https://lindat.mff.cuni.cz/repository/xmlui/bitstream/handle/11234/1-1458/data-plaintext-format.tar
-
-- Extract the archive:
-
-tar -xvf data-plaintext-format.tar
-
-"""
-
-
-def file_iterator(data_path: str) -> Iterable[str]:
-    """
-    Iterates through all files in a given directory and outputs path names
-    """
-
-    for file in os.listdir(data_path):
-        filename = os.path.join(data_path, file)
-        yield filename
+from lm_datasets.datasets.base import BaseDataset, Availability, License
 
 
-def text_extractor(files_path: Iterable[str]) -> Iterable[str]:
-    """
-    Extracts all texts from each file. Each file contains several texts, and one line per sentence.
-    This function concatenates all sentences belonging to one text and then outputs the text.
-    """
-    for input_file in file_iterator(files_path):
-        with gzip.open(input_file, "rt") as inp:
-            lines = inp.readlines()
-            cs_texts = []
-            current_filename = ""
-            for line in lines:
-                rows = line.split("\t")
-                name = rows[0]
-                filename = re.sub("-s[0-9]+", "", name)
-                if current_filename == "":
-                    current_filename = filename
-                elif current_filename != filename:
-                    yield " ".join(cs_texts)
-                    cs_texts = []
-                    current_filename = filename
-                # score = rows[1]
-                cs_text = rows[2]
-                # en_text = rows[3]
-                cs_texts.append(cs_text.strip("\n"))
-            yield " ".join(cs_texts)
+logger = logging.getLogger(__name__)
 
 
-def save_texts_to_jsonl(
-    texts: Iterable[str],
-    output_file_path: str,
-    append: bool = False,
-    limit: int = 0,
-    min_length: int = 0,
-    print_write_progress: int = 1_000,
-):
+class CzechEnglishParallelDataset(BaseDataset):
     """
-    Saves a list or generator of texts into a JSON-line file (each line is a valid JSON object)
+    This dataset is only using the Czech part of the original parallel corpus.
     """
-    mode = "a" if append else "w"
-    output_text_field = "text"
-
-    # Save as JSONL
-    logger.info(f"Writing output to {output_file_path} ({mode=})")
-
-    docs_count = 0
-
-    with open(output_file_path, mode, encoding="utf-8") as f:
-        for i, text in enumerate(texts):
-            if min_length > 0 and len(text) < min_length:
-                # skip because of short text length
-                continue
-
-            f.write(json.dumps({output_text_field: text}, ensure_ascii=False) + "\n")
-            docs_count += 1
-
-            if i > 0 and (i % print_write_progress) == 0:
-                logger.info(f"Writen {i:,} lines ...")
 
-            if limit > 0 and docs_count >= limit:
-                logger.warning(f"Limit reached ({docs_count:,} docs)")
-                break
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("input_path", help="Path to input file or directory")
-    parser.add_argument(
-        "output_path",
-        help="Output is saved in file path (JSONL file format)",
+    DATASET_ID = "cs_en_parallel"
+    TITLE = "Czech-English Parallel Corpus 1.0 (CzEng 1.0)"
+    HOMEPAGE = "http://hdl.handle.net/11234/1-1458"
+    DESCRIPTION = (
+        "CzEng 1.0 is the fourth release of a sentence-parallel Czech-English corpus compiled at the "
+        "Institute of Formal and Applied Linguistics (ÚFAL) freely available for non-commercial research purposes. "
+        "CzEng 1.0 contains 15 million parallel sentences (233 million English and 206 million Czech tokens) from "
+        "seven different types of sources automatically annotated at surface and deep (a- and t-) layers of "
+        "syntactic representation."
     )
-    parser.add_argument("--override", action="store_true", help="Override existing output files")
-    parser.add_argument("--download", action="store_true", help="Download dataset")
-    parser.add_argument("--limit", default=0, type=int, help="Limit dataset size (for debugging)")
-    parser.add_argument("--min_text_length", default=0, type=int, help="Min. text length (shorter texts are discarded)")
-    args = parser.parse_args()
-
-    if os.path.exists(args.output_path) and not args.override:
-        raise FileExistsError(f"Output exists already ({args.output_path}). Fix this error with --override")
-
-    if not os.path.exists(args.input_path):
-        raise FileNotFoundError(f"Input does not exist: {args.input_path}")
-
-    # if args.download:
-    #     download()
-
-    texts = text_extractor(args.input_path)
-
-    save_texts_to_jsonl(texts, args.output_path, limit=args.limit, min_length=args.min_text_length)
+    LICENSE = License(
+        "Attribution-NonCommercial-ShareAlike 3.0 Unported (CC BY-NC-SA 3.0)",
+        url="http://creativecommons.org/licenses/by-nc-sa/3.0/",
+        research_use=True,
+        commercial_use=False,
+        sharealike=True,
+    )
+    LANGUAGES = ["cs"]
+
+    DOWNLOAD_URLS = [
+        "https://lindat.mff.cuni.cz/repository/xmlui/bitstream/handle/11234/1-1458/data-plaintext-format.tar"
+    ]
+    AVAILIBILITY = Availability.DIRECT_DOWNLOAD
+
+    def is_downloaded(self):
+        return len(self.get_dataset_file_paths(needed_suffix=".tar")) == len(self.DOWNLOAD_URLS)
+
+    def get_texts(self):
+        """
+        Extracts all texts from each file. Each file contains several texts, and one line per sentence.
+        This function concatenates all sentences belonging to one text and then outputs the text.
+        """
+        if not self.is_downloaded():
+            self.download()
+
+        # read from tar files
+        tar_file_paths = self.get_dataset_file_paths(needed_suffix=".tar")
+        for tar_fp in tar_file_paths:
+            logger.info(f"Extracting from {tar_fp}")
+
+            with tarfile.open(tar_fp) as tar_f:
+                members = [m for m in tar_f.getmembers() if m.name.endswith(".gz")]
+
+                # Read from gz files
+                for member in members:
+                    logger.info(f"Decompress: {member.name}")
+
+                    member_content = tar_f.extractfile(member).read()
+
+                    with gzip.open(BytesIO(member_content), "rt") as inp:
+                        lines = inp.readlines()
+                        cs_texts = []
+                        current_filename = ""
+                        for line in lines:
+                            rows = line.split("\t")
+                            name = rows[0]
+                            filename = re.sub("-s[0-9]+", "", name)
+                            if current_filename == "":
+                                current_filename = filename
+                            elif current_filename != filename:
+                                yield " ".join(cs_texts)
+                                cs_texts = []
+                                current_filename = filename
+                            # score = rows[1]
+                            cs_text = rows[2]
+                            # en_text = rows[3]
+                            cs_texts.append(cs_text.strip("\n"))
+
+                        text = " ".join(cs_texts)
+                        # print(text)
+                        yield text
diff --git a/src/lm_datasets/datasets/cs/syn_v9.py b/src/lm_datasets/datasets/cs/syn_v9.py
index d612ab4..e217316 100644
--- a/src/lm_datasets/datasets/cs/syn_v9.py
+++ b/src/lm_datasets/datasets/cs/syn_v9.py
@@ -2,23 +2,27 @@
 
 import lzma as xz
 
-from lm_datasets.datasets.base import BaseDataset, Availability
+from lm_datasets.datasets.base import BaseDataset, Availability, License
 
 
 logger = logging.getLogger(__name__)
 
 
 class SynV9Dataset(BaseDataset):
-    """ """
-
     DATASET_ID = "syn_v9"
     TITLE = "SYN v9: large corpus of written Czech"
     HOMEPAGE = "https://lindat.mff.cuni.cz/repository/xmlui/handle/11234/1-4635"
     AVAILIBILITY = Availability.SIGNIN_DOWNLOAD
-
     LANGUAGES = ["cs"]
 
-    LICENSE = "Academic Use"
+    LICENSE = License(
+        "Academic Use - Czech National Corpus (Shuffled Corpus Data)",
+        url="https://lindat.mff.cuni.cz/repository/xmlui/page/license-cnc",
+        research_use=True,
+        commercial_use=False,
+        attribution=True,
+        distribution=False,
+    )
 
     USED_BY = ["https://arxiv.org/abs/2103.13031"]
 
diff --git a/src/lm_datasets/datasets/da/danewsroom.py b/src/lm_datasets/datasets/da/danewsroom.py
index 3b15683..3a9c145 100644
--- a/src/lm_datasets/datasets/da/danewsroom.py
+++ b/src/lm_datasets/datasets/da/danewsroom.py
@@ -3,7 +3,7 @@
 from typing import Iterable
 from smart_open import open
 
-from lm_datasets.datasets.base import GB, BaseDataset, Availability
+from lm_datasets.datasets.base import GB, BaseDataset, Availability, License
 
 logger = logging.getLogger(__name__)
 
@@ -12,10 +12,24 @@ class DANewsroomDataset(BaseDataset):
     DATASET_ID = "danewsroom"
     TITLE = "DaNewsroom: A Large-scale Danish Summarisation Dataset"
     HOMEPAGE = "https://github.com/danielvarab/da-newsroom"  # "https://aclanthology.org/2020.lrec-1.831/"
-
+    CITATION = """@inproceedings{varab-schluter-2020-danewsroom,
+        title = "{D}a{N}ewsroom: A Large-scale {D}anish Summarisation Dataset",
+        author = "Varab, Daniel  and
+        Schluter, Natalie",
+        booktitle = "Proceedings of The 12th Language Resources and Evaluation Conference",
+        month = may,
+        year = "2020",
+        address = "Marseille, France",
+        publisher = "European Language Resources Association",
+        url = "https://www.aclweb.org/anthology/2020.lrec-1.831",
+        pages = "6731--6739",
+        abstract = "Dataset development for automatic summarisation systems is notoriously English-oriented. In this paper we present the first large-scale non-English language dataset specifically curated for automatic summarisation. The document-summary pairs are news articles and manually written summaries in the Danish language. There has previously been no work done to establish a Danish summarisation dataset, nor any published work on the automatic summarisation of Danish. We provide therefore the first automatic summarisation dataset for the Danish language (large-scale or otherwise). To support the comparison of future automatic summarisation systems for Danish, we include system performance on this dataset of strong well-established unsupervised baseline systems, together with an oracle extractive summariser, which is the first account of automatic summarisation system performance for Danish. Finally, we make all code for automatically acquiring the data freely available and make explicit how this technology can easily be adapted in order to acquire automatic summarisation datasets for further languages.",
+        language = "English",
+        ISBN = "979-10-95546-34-4",
+        }
+        """  # noqa
     LANGUAGES = ["da"]
-    # LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/danewsroom"]
-
+    LICENSE = License("research-only (unknown license)", commercial_use=False, research_use=True)
     AVAILIBILITY = Availability.ON_REQUEST
     BYTES = 1.5 * GB
 
diff --git a/src/lm_datasets/datasets/da/danish_gigaword.py b/src/lm_datasets/datasets/da/danish_gigaword.py
index d1d27a7..bfc2256 100644
--- a/src/lm_datasets/datasets/da/danish_gigaword.py
+++ b/src/lm_datasets/datasets/da/danish_gigaword.py
@@ -1,16 +1,19 @@
-import json
-import os
-from lm_datasets.datasets.base import BILLION, Availability, BaseDataset
+from lm_datasets.datasets.base import BILLION, Availability, License
 from lm_datasets.datasets.hf_dataset import HFDataset
 
 
 class DanishGigawordDataset(HFDataset):
     DATASET_ID = "danish_gigaword"
-
-    TITLE = "Danish GiagaWord"
+    TITLE = "Danish GigaWord"
     HOMEPAGE = "https://sprogteknologi.dk/dataset/danish-gigaword"
     AVAILIBILITY = Availability.DIRECT_DOWNLOAD
-
+    LICENSE = License(
+        "CC-BY 4.0, Creative Commons with Attribution",
+        research_use=True,
+        commercial_use=True,
+        attribution=True,
+        sharealike=False,
+    )
     LANGUAGES = ["da"]
 
     DESCRIPTION = (
@@ -19,8 +22,34 @@ class DanishGigawordDataset(HFDataset):
         "",
         "The license is CC-BY 4.0, Creative Commons with Attribution. Owners: ITU; Leon Derczynski, Manuel R. Ciosici",
     )
-    PII = "I have not checked the data source for personally identifiable or sensitive information."
-    LICENSE = "public domain"
+    CITATION = """@inproceedings{stromberg-derczynski-etal-2021-danish,
+    title = "The {D}anish {G}igaword Corpus",
+    author = "Str{\o}mberg-Derczynski, Leon  and
+      Ciosici, Manuel  and
+      Baglini, Rebekah  and
+      Christiansen, Morten H.  and
+      Dalsgaard, Jacob Aarup  and
+      Fusaroli, Riccardo  and
+      Henrichsen, Peter Juel  and
+      Hvingelby, Rasmus  and
+      Kirkedal, Andreas  and
+      Kjeldsen, Alex Speed  and
+      Ladefoged, Claus  and
+      Nielsen, Finn {\AA}rup  and
+      Madsen, Jens  and
+      Petersen, Malte Lau  and
+      Rystr{\o}m, Jonathan Hvithamar  and
+      Varab, Daniel",
+    booktitle = "Proceedings of the 23rd Nordic Conference on Computational Linguistics (NoDaLiDa)",
+    month = may # " 31--2 " # jun,
+    year = "2021",
+    address = "Reykjavik, Iceland (Online)",
+    publisher = {Link{\"o}ping University Electronic Press, Sweden},
+    url = "https://aclanthology.org/2021.nodalida-main.46",
+    pages = "413--421",
+    abstract = "Danish language technology has been hindered by a lack of broad-coverage corpora at the scale modern NLP prefers. This paper describes the Danish Gigaword Corpus, the result of a focused effort to provide a diverse and freely-available one billion word corpus of Danish text. The Danish Gigaword corpus covers a wide array of time periods, domains, speakers{'} socio-economic status, and Danish dialects.",
+    }
+    """  # noqa
 
     # Size: 1 billion words
     TOKENS = 1 * BILLION
@@ -43,63 +72,3 @@ def filter_func(example):
             return example["source"] not in self.excluded_sources
 
         return filter_func
-
-
-# deprecated
-class DanishGigawordOriginalDataset(BaseDataset):
-    """
-    Dataset based on orignal source (not HF version)
-    """
-
-    DATASET_ID = "danish_gigaword"
-    TITLE = "Danish GiagaWord"
-    HOMEPAGE = "https://sprogteknologi.dk/dataset/danish-gigaword"
-    AVAILIBILITY = "Yes - it has a direct download link or links"
-
-    LANGUAGES = ["da"]
-
-    PII = "I have not checked the data source for personally identifiable or sensitive information."
-    LICENSE = "public domain"
-
-    # Size: 1 billion words
-    TOKENS = 1 * BILLION
-
-    DOWNLOAD_URLS = ["https://itu.dk/research/dagw/dagw_v1.0-release.zip"]
-    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/danish_gigaword"]
-
-    def get_dataset_file_path(self):
-        # TODO handle datasets with multiple filles
-        dataset_dir = self.get_local_dataset_dir()
-        files = os.listdir(dataset_dir)
-
-        if len(files) == 1:
-            return os.path.join(dataset_dir, files[0])
-        else:
-            raise ValueError(f"Cannot determine file path. Either none or multiple files in dataset dir: {files=}")
-
-    def get_texts(self):
-        import zipfile
-
-        dataset_file_path = self.get_dataset_file_path()
-        print(dataset_file_path)
-
-        with zipfile.ZipFile(dataset_file_path, "r") as zf:
-            # zf.printdir()
-            zfns = zf.namelist()
-
-            for fn in zfns:
-                if "sektioner" in fn:
-                    print(fn)
-
-                    with zf.open(fn) as content_file:
-                        for json_str in content_file:
-                            doc = json.loads(json_str)
-                            print(doc)
-
-                            yield doc["text"]
-
-                        # content = content_file.read()
-
-                        # for l in myfile:
-                        #     print(l)
-                    # break
diff --git a/src/lm_datasets/datasets/da/danish_parliament_corpus.py b/src/lm_datasets/datasets/da/danish_parliament_corpus.py
index 6ca6f7e..3b8ee36 100644
--- a/src/lm_datasets/datasets/da/danish_parliament_corpus.py
+++ b/src/lm_datasets/datasets/da/danish_parliament_corpus.py
@@ -30,7 +30,8 @@ class DanishParliamentCorpusDataset(BaseDataset):
     LICENSE = "public domain"
 
     # Size:  40,841,226 words
-    TOKENS = 40_841_226
+    # TOKENS = 40_841_226
+    DUMMY = True
 
     DOWNLOAD_URLS = []
-    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/danish_parliament_corpus"]
+    HAS_OVERLAP_WITH = ["danish_gigaword"]
diff --git a/src/lm_datasets/datasets/da/dk_clarin.py b/src/lm_datasets/datasets/da/dk_clarin.py
index 09a5fee..35822c4 100644
--- a/src/lm_datasets/datasets/da/dk_clarin.py
+++ b/src/lm_datasets/datasets/da/dk_clarin.py
@@ -4,7 +4,7 @@
 import zipfile
 
 
-from lm_datasets.datasets.base import BaseDataset, Availability, GB
+from lm_datasets.datasets.base import BaseDataset, Availability, GB, License
 
 logger = logging.getLogger(__name__)
 
@@ -15,7 +15,14 @@ class DKClarinDataset(BaseDataset):
     HOMEPAGE = "https://korpus.dsl.dk/clarin/"
 
     AVAILIBILITY = Availability.SIGNIN_DOWNLOAD
-    LICENSE = "Academic Use; CLARIN-ACA-NC"
+    LICENSE = License(
+        "Academic Use; CLARIN-ACA-NC",
+        url="https://kitwiki.csc.fi/twiki/bin/view/FinCLARIN/ClarinEulaAca?ID=1&AFFIL=EDU&BY=1&NC=1&NORED=1",
+        commercial_use=False,
+        research_use=True,
+        sharealike=False,
+        attribution=True,
+    )
 
     LANGUAGES = ["da"]
 
diff --git a/src/lm_datasets/datasets/dataset_registry.py b/src/lm_datasets/datasets/dataset_registry.py
index 2c0e0eb..be7b4a3 100644
--- a/src/lm_datasets/datasets/dataset_registry.py
+++ b/src/lm_datasets/datasets/dataset_registry.py
@@ -1,11 +1,11 @@
 from typing import List, Optional, Union
 from .multilingual.wikimedia import get_wikimedia_auto_classes
-from .multilingual.oscar_dummy import get_oscar_dummy_classes
 from .multilingual.colossal_oscar import get_colossal_oscar_auto_classes
 from .multilingual.eurlex import get_eurlex_auto_classes
 from .multilingual.legal_mc4 import get_legal_mc4_auto_classes
+from .code.starcoder import get_auto_starcoder_classes
 from .nl.sonar import get_sonar_classes
-
+from .en.pile_of_law import get_pile_of_law_auto_classes
 import importlib
 import logging
 
@@ -47,9 +47,11 @@
     ".en.pes2o.PeS2oDataset",
     ".en.proof_pile.ProofPileDataset",
     ".en.dialogstudio.DialogstudioDataset",
+    ".en.pile_of_law.PileOfLawDataset",
+    ".en.math_amps.MathAMPSDataset",
     # bg
-    ".bg.bgnc_admin_eur.BGNCAdminEURDataset",
-    ".bg.bgnc_news_corpus.BGNCNewsCorpusDataset",
+    # ".bg.bgnc_admin_eur.BGNCAdminEURDataset",  # deprecated -> use bulnc
+    # ".bg.bgnc_news_corpus.BGNCNewsCorpusDataset",  # deprecated -> use bulnc
     ".bg.bulgarian_news.BulgarianNewsDataset",
     ".bg.bulnc.BulNCDataset",
     # de
@@ -78,6 +80,7 @@
     ".sk.sk_laws.SKLawsDataset",
     # cs
     ".cs.syn_v9.SynV9Dataset",
+    ".cs.cs_en_parallel.CzechEnglishParallelDataset",
     # da
     ".da.danish_gigaword.DanishGigawordDataset",
     # danish_parliament_corpus.DanishParliamentCorpusDataset,
@@ -111,8 +114,10 @@
     # et
     ".et.estonian_reference_corpus.EstonianReferenceCorpusDataset",
     ".et.enc.ENC2021Dataset",
+    ".et.ekspress.EkspressDataset",
     # eu
     ".eu.euscrawl.EUSCrawlDataset",
+    ".eu.euscrawl.EUSCrawlFilteredDataset",
     # es
     ".es.spanish_legal.SpanishLegalDataset",
     # ".es.escorpius.ESCorpiusDataset",  # based on CC => use OSCAR instead
@@ -126,8 +131,6 @@
     ".ro.marcell_legislative_subcorpus_v2.MarcellLegislativeSubcorpusV2Dataset",
     # uk
     ".uk.uk_laws.UKLawsDataset",
-    # code
-    ".code.starcoder.StarcoderDataset",
 ]
 
 
@@ -160,17 +163,32 @@ def get_class_by_import_string(
         return import_string_or_cls  # already object, no need to import from string
 
 
-
-def get_registered_dataset_classes(extra_dataset_registries: Optional[Union[str, List[str]]] = None):
-    dataset_classes = (
-        [get_class_by_import_string(clss) for clss in ALL_DATASET_IMPORTS]
-        + get_oscar_dummy_classes()
-        + get_eurlex_auto_classes()
-        + get_legal_mc4_auto_classes()
-        + get_wikimedia_auto_classes()
-        + get_colossal_oscar_auto_classes()
-    )
-
+def get_registered_dataset_classes(
+    extra_dataset_registries: Optional[Union[str, List[str]]] = None,
+    extra_dataset_classes: Optional[List] = None,
+    use_default_registry: bool = True,
+):
+    """
+    Construct list of registered dataset classes
+    """
+    dataset_classes = []
+
+    # Predefined dataset classes (default registry)
+    if use_default_registry:
+        dataset_classes += (
+            [get_class_by_import_string(clss) for clss in ALL_DATASET_IMPORTS]
+            + get_eurlex_auto_classes()
+            + get_legal_mc4_auto_classes()
+            + get_wikimedia_auto_classes()
+            + get_colossal_oscar_auto_classes()
+            + get_auto_starcoder_classes()
+            + get_pile_of_law_auto_classes()
+        )
+
+    if extra_dataset_classes:
+        dataset_classes += extra_dataset_classes
+
+    # Load dataset classes from extra registries
     if extra_dataset_registries:
         if isinstance(extra_dataset_registries, str):
             extra_dataset_registries = extra_dataset_registries.split(",")
@@ -181,11 +199,33 @@ def get_registered_dataset_classes(extra_dataset_registries: Optional[Union[str,
             extra_dataset_registry_package = importlib.import_module(extra_dataset_registry_str)
 
             extra_dataset_registry_getter = getattr(extra_dataset_registry_package, "get_registered_dataset_classes")
-            extra_dataset_classes = extra_dataset_registry_getter()
+            extra_dataset_classes_from_registry = extra_dataset_registry_getter()
 
-            logger.info(f"Extra datasets: {extra_dataset_classes}")
+            logger.debug(f"Extra datasets: {extra_dataset_classes_from_registry}")
 
-            dataset_classes += extra_dataset_classes
+            dataset_classes += extra_dataset_classes_from_registry
 
     return dataset_classes
 
+
+def get_registered_dataset_ids(
+    extra_dataset_registries: Optional[Union[str, List[str]]] = None, needed_source_id: Optional[str] = None
+):
+    return [
+        cls.DATASET_ID
+        for cls in get_registered_dataset_classes(extra_dataset_registries)
+        if needed_source_id is None or cls.SOURCE_ID == needed_source_id
+    ]
+
+
+def get_dataset_class_by_id(dataset_id, extra_dataset_registries: Optional[Union[str, List[str]]] = None):
+    id_to_dataset_class = {cls.DATASET_ID: cls for cls in get_registered_dataset_classes(extra_dataset_registries)}
+
+    if dataset_id in id_to_dataset_class:
+        dataset_cls = id_to_dataset_class[dataset_id]
+    else:
+        raise ValueError(f"Unknown dataset ID: {dataset_id} (available: {id_to_dataset_class.keys()})")
+
+    dataset_cls = id_to_dataset_class[dataset_id]
+
+    return dataset_cls
diff --git a/src/lm_datasets/datasets/de/dewac.py b/src/lm_datasets/datasets/de/dewac.py
index 52995f3..cac9344 100644
--- a/src/lm_datasets/datasets/de/dewac.py
+++ b/src/lm_datasets/datasets/de/dewac.py
@@ -1,6 +1,6 @@
 import logging
 from typing import Iterable
-from lm_datasets.datasets.base import BILLION, Availability, BaseDataset
+from lm_datasets.datasets.base import BILLION, Availability, BaseDataset, License
 from smart_open import open
 
 logger = logging.getLogger(__name__)
@@ -20,8 +20,9 @@ class DEWacDataset(BaseDataset):
     LANGUAGES = ["de"]
     AVAILIBILITY = Availability.ON_REQUEST
     HOMEPAGE = "https://docs.sslmit.unibo.it/doku.php?id=corpora:dewac"
-
+    WEB_CRAWLED = True
     TOKENS = 1.7 * BILLION
+    LICENSE = License("unknown license; likely fair-use / research-only", sharealike=False)
 
     def decompress(self):
         #  7z x dewac_preproc.7z
diff --git a/src/lm_datasets/datasets/de/openlegaldata.py b/src/lm_datasets/datasets/de/openlegaldata.py
index cd87076..2a534bb 100644
--- a/src/lm_datasets/datasets/de/openlegaldata.py
+++ b/src/lm_datasets/datasets/de/openlegaldata.py
@@ -1,4 +1,4 @@
-from lm_datasets.datasets.base import Availability, Genre
+from lm_datasets.datasets.base import Availability, Genre, License
 from lm_datasets.datasets.jsonl_dataset import JSONLDataset
 
 
@@ -11,6 +11,13 @@ class OpenLegalDataDataset(JSONLDataset):
     LANGUAGES = ["de"]
     GENRES = [Genre.LEGAL]
     AVAILIBILITY = Availability.ON_REQUEST
+    LICENSE = License(
+        "public domain",
+        url="https://www.gesetze-im-internet.de/urhg/__5.html",
+        commercial_use=True,
+        sharealike=False,
+        research_use=True,
+    )
 
     TOKENS = 9_700_000_000
 
diff --git a/src/lm_datasets/datasets/el/greek_legal_code.py b/src/lm_datasets/datasets/el/greek_legal_code.py
index 4d8f8a5..67537f9 100644
--- a/src/lm_datasets/datasets/el/greek_legal_code.py
+++ b/src/lm_datasets/datasets/el/greek_legal_code.py
@@ -1,22 +1,7 @@
-from lm_datasets.datasets.base import Availability, MB
+from lm_datasets.datasets.base import Availability, MB, License
 from lm_datasets.datasets.hf_dataset import HFDataset
 
 
-_CITATION = """
-@inproceedings{papaloukas-etal-2021-glc,
-    title = "Multi-granular Legal Topic Classification on Greek Legislation",
-    author = "Papaloukas, Christos and Chalkidis, Ilias and Athinaios, Konstantinos and Pantazi, Despina-Athanasia and Koubarakis, Manolis",
-    booktitle = "Proceedings of the Natural Legal Language Processing Workshop 2021",
-    year = "2021",
-    address = "Punta Cana, Dominican Republic",
-    publisher = "Association for Computational Linguistics",
-    url = "https://arxiv.org/abs/2109.15298",
-    doi = "10.48550/arXiv.2109.15298",
-    pages = "63--75"
-}
-"""  # noqa
-
-
 class GreekLegalCodeDataset(HFDataset):
     DATASET_ID = "greek_legal_code"
 
@@ -26,9 +11,20 @@ class GreekLegalCodeDataset(HFDataset):
         " origin of GLC is “Permanent Greek Legislation Code - Raptarchis”, a collection of Greek legislative "
         " documents classified into multi-level (from broader to more specialized) categories."
     )
-
+    CITATION = """@inproceedings{papaloukas-etal-2021-glc,
+        title = "Multi-granular Legal Topic Classification on Greek Legislation",
+        author = "Papaloukas, Christos and Chalkidis, Ilias and Athinaios, Konstantinos and Pantazi, Despina-Athanasia and Koubarakis, Manolis",
+        booktitle = "Proceedings of the Natural Legal Language Processing Workshop 2021",
+        year = "2021",
+        address = "Punta Cana, Dominican Republic",
+        publisher = "Association for Computational Linguistics",
+        url = "https://arxiv.org/abs/2109.15298",
+        doi = "10.48550/arXiv.2109.15298",
+        pages = "63--75"
+    }
+    """  # noqa
     HOMEPAGE = "https://huggingface.co/datasets/greek_legal_code"
-
+    LICENSE = License("unknown; likely publlic domain", url="https://www.secdigital.gov.gr/e-themis/")
     AVAILIBILITY = Availability.DIRECT_DOWNLOAD
 
     LANGUAGES = ["el"]
diff --git a/src/lm_datasets/datasets/el/greek_web_corpus.py b/src/lm_datasets/datasets/el/greek_web_corpus.py
index 1cae153..6b6f2a6 100644
--- a/src/lm_datasets/datasets/el/greek_web_corpus.py
+++ b/src/lm_datasets/datasets/el/greek_web_corpus.py
@@ -1,5 +1,5 @@
 from typing import Iterable
-from lm_datasets.datasets.base import GB, BaseDataset, Availability, QualityWarning
+from lm_datasets.datasets.base import GB, BaseDataset, Availability, QualityWarning, License
 from smart_open import open
 import logging
 
@@ -21,9 +21,9 @@ class GreekWebCorpus(BaseDataset):
     USED_BY = [
         "https://arxiv.org/abs/2304.00869",  # GreekBART
     ]
-
+    WEB_CRAWLED = True
     BYTES = 10 * GB
-
+    LICENSE = License("unknown; likely fair use")
     QUALITY_WARNINGS = [QualityWarning.SHORT_TEXT, QualityWarning.BAD_PUNCTUATION]
 
     def get_texts(self) -> Iterable[str]:
diff --git a/src/lm_datasets/datasets/en/dialogstudio.py b/src/lm_datasets/datasets/en/dialogstudio.py
index 26ae426..e7ae2a3 100644
--- a/src/lm_datasets/datasets/en/dialogstudio.py
+++ b/src/lm_datasets/datasets/en/dialogstudio.py
@@ -123,6 +123,8 @@ class DialogstudioDataset(HFDataset):
     GENRES = [Genre.DIALOGUE]
     LICENSE = "mixed"
 
+    DUMMY = True
+
     HF_DATASET_ID = "Salesforce/dialogstudio"
     HF_DATASET_SPLIT = "train"
 
diff --git a/src/lm_datasets/datasets/en/math_amps.py b/src/lm_datasets/datasets/en/math_amps.py
new file mode 100644
index 0000000..463ffb3
--- /dev/null
+++ b/src/lm_datasets/datasets/en/math_amps.py
@@ -0,0 +1,80 @@
+import json
+import logging
+import tarfile
+
+from lm_datasets.datasets.base import BaseDataset, License
+
+logger = logging.getLogger(__name__)
+
+
+class MathAMPSDataset(BaseDataset):
+    DATASET_ID = "math_amps"
+    TITLE = "Auxiliary Mathematics Problems and Solutions (AMPS) dataset"
+    HOMEPAGE = "https://github.com/hendrycks/math"
+    DESCRIPTION = (
+        "Our pretraining dataset, the Auxiliary Mathematics Problems and "
+        "Solutions (AMPS) dataset, has problems and step-by-step solutions typeset "
+        " in LATEX. AMPS contains over 100,000 problems pulled from Khan Academy and "
+        " approximately 5 million problems generated from manually designed Mathematica scripts. "
+        "Problems include various aspects of algebra, calculus, counting and statistics, geometry, "
+        "linear algebra, and number theory."
+    )
+    LANGUAGES = ["en"]
+    CITATION = """@article{hendrycksmath2021,
+  title={Measuring Mathematical Problem Solving With the MATH Dataset},
+  author={Dan Hendrycks and Collin Burns and Saurav Kadavath and Akul Arora and Steven Basart and Eric Tang and Dawn Song and Jacob Steinhardt},
+  journal={NeurIPS},
+  year={2021}
+}"""  # noqa
+    LICENSE = License("repository license is MIT; no specific data license")
+    DOWNLOAD_URLS = ["https://drive.google.com/file/d/1hQsua3TkpEmcJD_UWQx8dmNdEZPyxw23/view?usp=sharing"]
+
+    def download(self):
+        """
+        Manually download the archive from Google drive.
+        """
+        pass
+
+    def decompress(self):
+        """
+        gzip -d ...
+        """
+        pass
+
+    def get_texts(self):
+        # read from tar files
+        tar_file_paths = self.get_dataset_file_paths(needed_suffix=".tar")
+        for tar_fp in tar_file_paths:
+            logger.info(f"Extracting from {tar_fp}")
+
+            with tarfile.open(tar_fp) as tar_f:
+                while True:
+                    member = tar_f.next()
+
+                    if member is None:
+                        break
+
+                    if "README" in member.name:
+                        continue
+
+                    if member.name.endswith("json") or member.name.endswith(".txt"):
+                        with tar_f.extractfile(member) as member_f:
+                            if member.name.endswith(".json"):
+                                doc = json.load(member_f)
+
+                                # "problem"
+                                # "hints" -> list
+                                if "problem" in doc and "hints" in doc and len(doc) >= 2:
+                                    text = doc["title"] + self.title_delimiter if "title" in doc else ""
+
+                                    text += (
+                                        doc["problem"]
+                                        + self.paragraph_delimiter
+                                        + (self.paragraph_delimiter.join(doc["hints"]))
+                                    )
+                                else:
+                                    raise ValueError(f"Invalid JSON schema: {doc}")
+                            else:
+                                text = member_f.read().decode("utf-8")  # read directly text from *.txt file
+
+                            yield text
diff --git a/src/lm_datasets/datasets/en/pes2o.py b/src/lm_datasets/datasets/en/pes2o.py
index 6741764..ffe849c 100644
--- a/src/lm_datasets/datasets/en/pes2o.py
+++ b/src/lm_datasets/datasets/en/pes2o.py
@@ -1,4 +1,4 @@
-from lm_datasets.datasets.base import BILLION, Genre
+from lm_datasets.datasets.base import BILLION, Genre, License
 from lm_datasets.datasets.hf_dataset import HFDataset
 
 
@@ -21,7 +21,14 @@ class PeS2oDataset(HFDataset):
     )
     HOMEPAGE = "https://huggingface.co/datasets/allenai/peS2o"
     VERSION = "V2"
-
+    LICENSE = License(
+        name="Open Data Commons Attribution License (ODC-By) v1.0",
+        url="https://opendatacommons.org/licenses/by/1-0/",
+        attribution=True,
+        commercial_use=True,
+        research_use=True,
+        sharealike=False,
+    )
     LANGUAGES = ["en"]
     GENRES = [Genre.SCIENCE]
 
diff --git a/src/lm_datasets/datasets/en/pile_of_law.py b/src/lm_datasets/datasets/en/pile_of_law.py
new file mode 100644
index 0000000..57b956f
--- /dev/null
+++ b/src/lm_datasets/datasets/en/pile_of_law.py
@@ -0,0 +1,170 @@
+from lm_datasets.datasets.base import License, Availability, Genre
+from lm_datasets.datasets.hf_dataset import HFDataset
+
+# OK licenses
+# r_legaladvice,courtlistener_docket_entry_documents,atticus_contracts,courtlistener_opinions,tos,scotus_oral_arguments,exam_outlines,cfpb_creditcard_contracts,constitutions,congressional_hearings,un_debates
+# ['pile_of_law_r_legaladvice', 'pile_of_law_courtlistener_docket_entry_documents', 'pile_of_law_atticus_contracts', 'pile_of_law_courtlistener_opinions', 'pile_of_law_tos', 'pile_of_law_scotus_oral_arguments', 'pile_of_law_exam_outlines', 'pile_of_law_cfpb_creditcard_contracts', 'pile_of_law_constitutions', 'pile_of_law_congressional_hearings', 'pile_of_law_un_debates']
+# pile_of_law_r_legaladvice,pile_of_law_courtlistener_docket_entry_documents,pile_of_law_atticus_contracts,pile_of_law_courtlistener_opinions,pile_of_law_tos,pile_of_law_scotus_oral_arguments,pile_of_law_exam_outlines,pile_of_law_cfpb_creditcard_contracts,pile_of_law_constitutions,pile_of_law_congressional_hearings,pile_of_law_un_debates
+
+PILE_OF_LAW_SUBSETS_WITH_LICENSE = {
+    "r_legaladvice": License("Creative Commons Attribution 4.0 International"),
+    "courtlistener_docket_entry_documents": License("Underlying content is Public Domain."),
+    "atticus_contracts": License("CC BY 4.0"),
+    "courtlistener_opinions": License("Public domain"),
+    # "federal_register",
+    # "bva_opinions",
+    # "us_bills",
+    "cc_casebooks": License(
+        "Mixed; Most restrictive: CC BY-NC-SA 4.0", commercial_use=True
+    ),  # All CC, varying on exact restrictiveness. Most restrictive: CC BY-NC-SA 4.0. All licensing information preserved in individual documents.
+    "tos": License(
+        "Publicly available, unknown license. Assumed to be governed by fair use standards."
+    ),  # Publicly available, unknown license. Assumed to be governed by fair use standards.
+    # "euro_parl",
+    # "nlrb_decisions",
+    "scotus_oral_arguments": License("Public domain"),  # Public domain
+    # "cfr",
+    # "state_codes",
+    # "scotus_filings",
+    "exam_outlines": License(
+        "Publicly available, unknown license. Assumed to be governed by fair use standards."
+    ),  # Publicly available, unknown license. Assumed to be governed by fair use standards.
+    # "edgar",
+    "cfpb_creditcard_contracts": License(
+        "Publicly available, unknown license. Assumed to be governed by fair use standards."
+    ),  # # Publicly available, unknown license. Assumed to be governed by fair use standards.
+    "constitutions": License("CC BY-NC 3.0", commercial_use=False, attribution=True),  # CC BY-NC 3.09
+    "congressional_hearings": License("Public domain"),  # Public domain
+    # "oig",
+    # "olc_memos",
+    # "uscode",
+    # "founding_docs",
+    # "ftc_advisory_opinions",
+    "echr": License(
+        "Non-commercial, commercial use requires written permission",
+        commercial_use=False,
+        url="https://www.echr.coe.int/en/copyright-and-disclaimer",
+    ),  # Non-commercial, commercial use requires written permission
+    # "eurlex",
+    # "tax_rulings",
+    "un_debates": License("Public domain"),  # Public domain
+    # "fre",
+    # "frcp",
+    # "canadian_decisions",
+    # "eoir",
+    # "dol_ecab",
+    "icj-pcij": None,  # unkown
+    "uspto_office_actions": None,
+    "ed_policy_guidance": None,
+    "acus_reports": None,
+    # "hhs_alj_opinions",
+    # "sec_administrative_proceedings",
+    # "fmshrc_bluebooks",
+    "resource_contracts": None,
+    "medicaid_policy_guidance": None,
+    # "irs_legal_advice_memos",
+    "doj_guidance_documents": None,
+}
+
+
+class PileOfLawDataset(HFDataset):
+    """
+    NOTE: This dataset uses only selected subsets of the original "Pile of Law". See `HF_DATASET_CONFIGS`
+    """
+
+    DATASET_ID = "pile_of_law"
+    SOURCE_ID = "pile_of_law"
+    TITLE = "Pile of Law (selected subsets)"
+    HOMEPAGE = "https://huggingface.co/datasets/pile-of-law/pile-of-law"
+    DESCRIPTION = (
+        "We curate a large corpus of legal and administrative data. The utility of this data is twofold: "
+        "(1) to aggregate legal and administrative data sources that demonstrate different norms and legal "
+        "standards for data filtering; (2) to collect a dataset that can be used in the future for pretraining "
+        "legal-domain language models, a key direction in access-to-justice initiatives."
+    )
+    LANGUAGES = ["en"]  # Mainly English, but some other languages may appear in some portions of the data.
+    CITATION = """@misc{hendersonkrass2022pileoflaw,
+    url = {https://arxiv.org/abs/2207.00220},
+    author = {Henderson*, Peter and Krass*, Mark S. and Zheng, Lucia and Guha, Neel and Manning, Christopher D. and Jurafsky, Dan and Ho, Daniel E.},
+    title = {Pile of Law: Learning Responsible Data Filtering from the Law and a 256GB Open-Source Legal Dataset},
+    publisher = {arXiv},
+    year = {2022}
+    }"""  # noqa
+    GENRES = [Genre.LEGAL]
+    LICENSE = License(
+        "CreativeCommons Attribution-NonCommercial-ShareAlike 4.0 International. But individual sources may have other licenses. See paper for details.",
+        commercial_use=False,
+        sharealike=True,
+        attribution=True,
+    )
+    AVAILIBILITY = Availability.DIRECT_DOWNLOAD
+
+    HF_DATASET_ID = "pile-of-law/pile-of-law"
+    HF_DATASET_SPLIT = "train"
+    HF_DATASET_CONFIGS = [
+        "r_legaladvice",
+        "courtlistener_docket_entry_documents",
+        "atticus_contracts",
+        "courtlistener_opinions",
+        # "federal_register",
+        # "bva_opinions",
+        # "us_bills",
+        "cc_casebooks",
+        "tos",
+        # "euro_parl",
+        # "nlrb_decisions",
+        "scotus_oral_arguments",
+        # "cfr",
+        # "state_codes",
+        # "scotus_filings",
+        "exam_outlines",
+        # "edgar",
+        "cfpb_creditcard_contracts",
+        "constitutions",
+        "congressional_hearings",
+        # "oig",
+        # "olc_memos",
+        # "uscode",
+        # "founding_docs",
+        # "ftc_advisory_opinions",
+        "echr",
+        # "eurlex",
+        # "tax_rulings",
+        "un_debates",
+        # "fre",
+        # "frcp",
+        # "canadian_decisions",
+        # "eoir",
+        # "dol_ecab",
+        "icj-pcij",
+        "uspto_office_actions",
+        "ed_policy_guidance",
+        "acus_reports",
+        # "hhs_alj_opinions",
+        # "sec_administrative_proceedings",
+        # "fmshrc_bluebooks",
+        "resource_contracts",
+        "medicaid_policy_guidance",
+        # "irs_legal_advice_memos",
+        "doj_guidance_documents",
+    ]
+
+    streaming = True
+
+    def get_text_from_item(self, item) -> str:
+        return item[self.text_column_name]
+
+
+def get_class(subset, license_info):
+    class Dataset(PileOfLawDataset):
+        DATASET_ID = f"pile_of_law_{subset}"
+        TITLE = f"Pile of Law (subset: {subset})"
+        LICENSE = license_info
+        HF_DATASET_CONFIGS = [subset]
+        HAS_OVERLAP_WITH = ["pile_of_law"]
+
+    return Dataset
+
+
+def get_pile_of_law_auto_classes():
+    return [PileOfLawDataset] + [get_class(s, l) for s, l in PILE_OF_LAW_SUBSETS_WITH_LICENSE.items()]
diff --git a/src/lm_datasets/datasets/en/proof_pile.py b/src/lm_datasets/datasets/en/proof_pile.py
index 823e759..703eb83 100644
--- a/src/lm_datasets/datasets/en/proof_pile.py
+++ b/src/lm_datasets/datasets/en/proof_pile.py
@@ -1,4 +1,4 @@
-from lm_datasets.datasets.base import BILLION, Genre
+from lm_datasets.datasets.base import BILLION, Genre, License
 from lm_datasets.datasets.hf_dataset import HFDataset
 
 
@@ -15,11 +15,12 @@ class ProofPileDataset(HFDataset):
         " sources of both informal and formal mathematics, namely"
     )
     HOMEPAGE = "https://huggingface.co/datasets/hoskinson-center/proof-pile"
-
+    LICENSE = License("Apache 2.0 (probably code license instead of data license)")
     LANGUAGES = ["en"]
     GENRES = [Genre.SCIENCE, Genre.MATH]
 
     TOKENS = 8.3 * BILLION
+    HAS_OVERLAP_WITH = ["proof_pile_2", "math_amps", "pes2o"]
 
     HF_DATASET_ID = "hoskinson-center/proof-pile"
     HF_DATASET_SPLIT = "train"
diff --git a/src/lm_datasets/datasets/en/wikihow.py b/src/lm_datasets/datasets/en/wikihow.py
index 6e1e470..47b2266 100644
--- a/src/lm_datasets/datasets/en/wikihow.py
+++ b/src/lm_datasets/datasets/en/wikihow.py
@@ -1,4 +1,4 @@
-from lm_datasets.datasets.base import Availability, MB
+from lm_datasets.datasets.base import Availability, MB, License
 from lm_datasets.datasets.hf_dataset import HFDataset
 
 
@@ -11,7 +11,13 @@ class WikihowDataset(HFDataset):
     )
 
     HOMEPAGE = "https://github.com/mahnazkoupaee/WikiHow-Dataset"
-
+    LICENSE = License(
+        name="CC BY-NC-SA 3.0",
+        url="https://creativecommons.org/licenses/by-nc-sa/3.0/",
+        commercial_use=False,
+        sharealike=True,
+        research_use=True,
+    )
     AVAILIBILITY = Availability.DIRECT_DOWNLOAD
 
     LANGUAGES = ["en"]
diff --git a/src/lm_datasets/datasets/es/escorpius.py b/src/lm_datasets/datasets/es/escorpius.py
index af7af73..c498e92 100644
--- a/src/lm_datasets/datasets/es/escorpius.py
+++ b/src/lm_datasets/datasets/es/escorpius.py
@@ -21,6 +21,7 @@ class ESCorpiusDataset(HFDataset):
 
     LANGUAGES = ["es"]
     WEB_CRAWLED = True
+    DUMMY = True
 
     HF_DATASET_ID = "LHF/escorpius"
     HF_DATASET_SPLIT = "train"
diff --git a/src/lm_datasets/datasets/es/spanish_legal.py b/src/lm_datasets/datasets/es/spanish_legal.py
index 4347858..547ce07 100644
--- a/src/lm_datasets/datasets/es/spanish_legal.py
+++ b/src/lm_datasets/datasets/es/spanish_legal.py
@@ -1,6 +1,6 @@
 import logging
 import zipfile
-from lm_datasets.datasets.base import Availability, BaseDataset, Genre, GB, QualityWarning
+from lm_datasets.datasets.base import Availability, BaseDataset, Genre, GB, License, QualityWarning
 
 logger = logging.getLogger(__name__)
 
@@ -27,7 +27,14 @@ class SpanishLegalDataset(BaseDataset):
     AVAILIBITY = Availability.DIRECT_DOWNLOAD
     QUALITY_WARNINGS = [QualityWarning.BAD_LINEBREAKS, QualityWarning.SHORT_TEXT]
     BYTES = 8.9 * GB
-
+    LICENSE = License(
+        "Creative Commons Attribution 4.0 International",
+        url="https://zenodo.org/records/5495529",
+        attribution=True,
+        sharealike=False,
+        research_use=True,
+        commercial_use=True,
+    )
     USED_BY = ["https://huggingface.co/PlanTL-GOB-ES/RoBERTalex"]
 
     DOWNLOAD_URLS = [("https://zenodo.org/record/5495529/files/corpus.zip?download=1", "corpus.zip")]
diff --git a/src/lm_datasets/datasets/et/ekspress.py b/src/lm_datasets/datasets/et/ekspress.py
index 39d3e08..614dfd8 100644
--- a/src/lm_datasets/datasets/et/ekspress.py
+++ b/src/lm_datasets/datasets/et/ekspress.py
@@ -1,132 +1,92 @@
 import json
 import logging
-import argparse
-import os
-from typing import Iterable
-from smart_open import open
+import zipfile
 
-logging.basicConfig(
-    format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
-    datefmt="%Y-%m-%d %H:%M:%S",
-    level=logging.INFO,
-    handlers=[logging.StreamHandler()],
-)
-logger = logging.getLogger(__name__)
-
-# dataset metadata
-DATASET_ID = "et_ekspress"
-TITLE = "Ekspress news article archive (only Estonian) 1.0"
-HOMEPAGE = "https://www.clarin.si/repository/xmlui/handle/11356/1408"
-
-# language codes of the dataset (see ./resources/languages.json)
-LANGUAGES = ["et"]
-
-
-"""
-DOWNLOAD
------------
-
-Instruction
-
-- Run the command on the server:
-
-wget "https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1408/ee_articles_2009-2019.zip"
-
-- Extract the archive:
-
-unzip ee_articles_2009-2019.zip
-
-- Run python script:
-
-python ekspress.py /data/datasets/ele/et/ekspress/ee_articles_2009-2019/ \
-    /data/datasets/ele/et/ekspress/ekspress_processed.jsonl --limit 1000
-"""
-
-
-def file_iterator(data_path: str) -> Iterable[str]:
-    """
-    Iterates through all files in a given directory and outputs path names
-    """
+from lm_datasets.datasets.base import BaseDataset, License
 
-    for file in os.listdir(data_path):
-        if file.endswith(".json"):
-            filename = os.path.join(data_path, file)
-        yield filename
 
-
-def text_extractor(files_path: Iterable[str]) -> Iterable[str]:
-    """
-    Extracts the article texts from each json file.
-    """
-    for file in file_iterator(files_path):
-        with open(file, "r", encoding="utf8") as inp:
-            data = json.load(inp)
-            for article in data:
-                if article["channelLanguage"] == "nat":  # only include estonian
-                    if len(article["bodyText"]) > 250:  # at least over 250 characters
-                        yield article["bodyText"]
-                else:
-                    continue
-
-
-def save_texts_to_jsonl(
-    texts: Iterable[str],
-    output_file_path: str,
-    append: bool = False,
-    limit: int = 0,
-    min_length: int = 0,
-    print_write_progress: int = 1_000,
-):
-    """
-    Saves a list or generator of texts into a JSON-line file (each line is a valid JSON object)
-    """
-    mode = "a" if append else "w"
-    output_text_field = "text"
-
-    # Save as JSONL
-    logger.info(f"Writing output to {output_file_path} ({mode=})")
-
-    docs_count = 0
-
-    with open(output_file_path, mode) as f:
-        for i, text in enumerate(texts):
-            if min_length > 0 and len(text) < min_length:
-                # skip because of short text length
-                continue
-
-            f.write(json.dumps({output_text_field: text}, ensure_ascii=False) + "\n")
-            docs_count += 1
-
-            if i > 0 and (i % print_write_progress) == 0:
-                logger.info(f"Writen {i:,} lines ...")
-
-            if limit > 0 and docs_count >= limit:
-                logger.warning(f"Limit reached ({docs_count:,} docs)")
-                break
+logger = logging.getLogger(__name__)
 
 
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("input_path", help="Path to input file or directory")
-    parser.add_argument(
-        "output_path",
-        help="Output is saved in file path (JSONL file format)",
+class EkspressDataset(BaseDataset):
+    DATASET_ID = "ekspress"
+    TITLE = "Ekspress news article archive (only Estonian) 1.0"
+    HOMEPAGE = "https://www.clarin.si/repository/xmlui/handle/11356/1408"
+    DESCRIPTION = (
+        "The dataset is an archive of articles from the Ekspress Meedia news site from 2009-2019, "
+        "containing over 1.4M articles, mostly in Estonian language (1,115,120 articles) with some "
+        " in Russian (325,952 articles)."
+    )
+    LICENSE = License(
+        "Creative Commons - Attribution-NonCommercial-NoDerivatives 4.0 International (CC BY-NC-ND 4.0)",
+        commercial_use=False,
+        attribution=True,
+        derivates=False,
     )
-    parser.add_argument("--override", action="store_true", help="Override existing output files")
-    parser.add_argument("--download", action="store_true", help="Download dataset")
-    parser.add_argument("--limit", default=0, type=int, help="Limit dataset size (for debugging)")
-    parser.add_argument("--min_text_length", default=0, type=int, help="Min. text length (shorter texts are discarded)")
-    args = parser.parse_args()
-
-    if os.path.exists(args.output_path) and not args.override:
-        raise FileExistsError(f"Output exists already ({args.output_path}). Fix this error with --override")
-
-    # TODO this can may be optional
-    if not os.path.exists(args.input_path):
-        raise FileNotFoundError(f"Input does not exist: {args.input_path}")
-
-    # if args.download:
-    #     download()
 
-    texts = text_extractor(args.input_path)
-    save_texts_to_jsonl(texts, args.output_path, limit=args.limit, min_length=args.min_text_length)
+    LANGUAGES = ["et"]
+
+    CITATION = """@inproceedings{pollak-etal-2021-embeddia,
+    title = "{EMBEDDIA} Tools, Datasets and Challenges: Resources and Hackathon Contributions",
+    author = {Pollak, Senja  and
+      Robnik-{\v{S}}ikonja, Marko  and
+      Purver, Matthew  and
+      Boggia, Michele  and
+      Shekhar, Ravi  and
+      Pranji{\'c}, Marko  and
+      Salmela, Salla  and
+      Krustok, Ivar  and
+      Paju, Tarmo  and
+      Linden, Carl-Gustav  and
+      Lepp{\"a}nen, Leo  and
+      Zosa, Elaine  and
+      Ul{\v{c}}ar, Matej  and
+      Freienthal, Linda  and
+      Traat, Silver  and
+      Cabrera-Diego, Luis Adri{\'a}n  and
+      Martinc, Matej  and
+      Lavra{\v{c}}, Nada  and
+      {\v{S}}krlj, Bla{\v{z}}  and
+      {\v{Z}}nidar{\v{s}}i{\v{c}}, Martin  and
+      Pelicon, Andra{\v{z}}  and
+      Koloski, Boshko  and
+      Podpe{\v{c}}an, Vid  and
+      Kranjc, Janez  and
+      Sheehan, Shane  and
+      Boros, Emanuela  and
+      Moreno, Jose G.  and
+      Doucet, Antoine  and
+      Toivonen, Hannu},
+    booktitle = "Proceedings of the EACL Hackashop on News Media Content Analysis and Automated Report Generation",
+    month = apr,
+    year = "2021",
+    address = "Online",
+    publisher = "Association for Computational Linguistics",
+    url = "https://aclanthology.org/2021.hackashop-1.14",
+    pages = "99--109",
+    abstract = "This paper presents tools and data sources collected and released by the EMBEDDIA project, supported by the European Union{'}s Horizon 2020 research and innovation program. The collected resources were offered to participants of a hackathon organized as part of the EACL Hackashop on News Media Content Analysis and Automated Report Generation in February 2021. The hackathon had six participating teams who addressed different challenges, either from the list of proposed challenges or their own news-industry-related tasks. This paper goes beyond the scope of the hackathon, as it brings together in a coherent and compact form most of the resources developed, collected and released by the EMBEDDIA project. Moreover, it constitutes a handy source for news media industry and researchers in the fields of Natural Language Processing and Social Science.",
+}"""  # noqa
+
+    DOWNLOAD_URLS = ["https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1408/ee_articles_2009-2019.zip"]
+
+    def is_downloaded(self):
+        return len(self.get_dataset_file_paths(needed_suffix=".zip")) == len(self.DOWNLOAD_URLS)
+
+    def get_texts(self):
+        if not self.is_downloaded():
+            self.download()
+
+        # iterate over archives
+        for archive_fp in self.get_dataset_file_paths(needed_suffix=".zip"):
+            with zipfile.ZipFile(archive_fp, "r") as zf:
+                for zfn in zf.namelist():
+                    if zfn.endswith(".json"):
+                        logger.info(f"Reading from {zfn}")
+
+                        with zf.open(zfn) as file:
+                            data = json.load(file)
+                            for article in data:
+                                if article["channelLanguage"] == "nat":  # only include estonian
+                                    text = article["bodyText"]
+
+                                    yield text
diff --git a/src/lm_datasets/datasets/et/enc.py b/src/lm_datasets/datasets/et/enc.py
index 23264a3..da552a1 100644
--- a/src/lm_datasets/datasets/et/enc.py
+++ b/src/lm_datasets/datasets/et/enc.py
@@ -1,6 +1,6 @@
 import logging
 from typing import Iterable
-from lm_datasets.datasets.base import BaseDataset
+from lm_datasets.datasets.base import BaseDataset, License
 
 from smart_open import open
 
@@ -11,12 +11,27 @@
 class ENC2021Dataset(BaseDataset):
     DATASET_ID = "enc2021"
     TITLE = "Estonian National Corpus 2021"
+    DESCRIPTION = (
+        "Corpus is based on Estonian National Corpus 2013, which was renewed by Lexical Computing Ltd. "
+        "in 2017 and 2019 at the request of Estonian Language Institute."
+        "Subcorpora are: Estonian Reference Corpus 1990-2008, Estonian Web 2013, Estonian Web 2017, "
+        "Estonian Web 2019, Estonian Wikipedia 2017, Estonian Wikipedia 2019, Estonian Open Access Journals (DOAJ),"
+        " blogs, discussion, education, fiction, food, health, journals, news, religion, science, sex, society, sports.",
+    )
     HOMEPAGE = (
         "https://entu.keeleressursid.ee/shared/9939/EVKultjxSeFA2QhFkbE7fGGDGNT1zmJLOUGFK9hw53tq9Rx2YBTejI1IoKhy65zq"
     )
     LANGUAGES = ["et"]
     DOI = "10.5128/ERYa18.12"
-
+    LICENSE = License(
+        "Creative Commons Attribution-NonCommercial 4.0 International License",
+        url="http://arhiiv.rakenduslingvistika.ee/ajakirjad/index.php/aastaraamat/article/view/ERYa18.12",
+        attribution=True,
+        commercial_use=False,
+        research_use=True,
+        sharealike=False,
+        derivates=True,
+    )
     DOWNLOAD_URLS = [
         # (
         #     "https://entu.keeleressursid.ee/api2/file-24747?key=EVKultjxSeFA2QhFkbE7fGGDGNT1zmJLOUGFK9hw53tq9Rx2YBTejI1IoKhy65zq",
diff --git a/src/lm_datasets/datasets/et/estonian_reference_corpus.py b/src/lm_datasets/datasets/et/estonian_reference_corpus.py
index 0c4b169..17e8fdb 100644
--- a/src/lm_datasets/datasets/et/estonian_reference_corpus.py
+++ b/src/lm_datasets/datasets/et/estonian_reference_corpus.py
@@ -1,6 +1,6 @@
 from typing import Iterable
 import zipfile
-from lm_datasets.datasets.base import MILLION, BaseDataset, Availability, QualityWarning
+from lm_datasets.datasets.base import MILLION, BaseDataset, Availability, QualityWarning, License
 from lm_datasets.utils import remove_whitespaces_before_punctuation
 
 import logging
@@ -44,7 +44,13 @@ class EstonianReferenceCorpusDataset(BaseDataset):
     ]
     AVAILIBILITY = Availability.DIRECT_DOWNLOAD
 
-    LICENSE = "free for non-commercial use"
+    LICENSE = License(
+        "free for non-commercial use",
+        commercial_use=False,
+        research_use=True,
+        sharealike=False,
+        derivates=False,
+    )
 
     QUALITY_WARNINGS = [QualityWarning.BAD_WHITESPACES]
 
diff --git a/src/lm_datasets/datasets/eu/euscrawl.py b/src/lm_datasets/datasets/eu/euscrawl.py
index 64f559f..6fb17a1 100644
--- a/src/lm_datasets/datasets/eu/euscrawl.py
+++ b/src/lm_datasets/datasets/eu/euscrawl.py
@@ -1,9 +1,10 @@
-from lm_datasets.datasets.base import MILLION
+from lm_datasets.datasets.base import MILLION, License
 from lm_datasets.datasets.hf_dataset import HFDataset
 
 
 class EUSCrawlDataset(HFDataset):
     DATASET_ID = "euscrawl"
+    SOURCE_ID = "euscrawl"
     TITLE = "EusCrawl"
     LANGUAGES = ["eu"]
     DESCRIPTION = (
@@ -12,7 +13,9 @@ class EUSCrawlDataset(HFDataset):
         "scrapers to extract text from 33 Basque websites with high-quality content, resulting in cleaner text ",
         "compared to general purpose approaches.",
     )
-
+    LICENSE = License(
+        "mixed (see Tab. 2 in paper, e.g., CC-BY-NC-ND, CC-BY-NC-SA)", url="https://arxiv.org/pdf/2203.08111.pdf"
+    )
     TOKENS = 423 * MILLION
     WEB_CRAWLED = True
 
@@ -20,3 +23,24 @@ class EUSCrawlDataset(HFDataset):
     HF_DATASET_SPLIT = "train"
 
     text_column_name = "plain_text"
+
+    HAS_OVERLAP_WITH = ["wiki_eu"]
+
+
+class EUSCrawlFilteredDataset(EUSCrawlDataset):
+    DATASET_ID = "euscrawl_filtered"
+    TITLE = "EusCrawl (filtered: no Wikipedia, no NC-licenses)"
+    LICENSE = License(
+        "CC-BY-SA",
+        url="https://huggingface.co/datasets/HiTZ/euscrawl",
+        commercial_use=True,
+        attribution=True,
+        sharealike=False,
+        research_use=True,
+    )
+
+    def get_filter_func(self):
+        def filter_func(example):
+            return example["source"] != "wikipedia" and "nc" not in example["license"]
+
+        return filter_func
diff --git a/src/lm_datasets/datasets/fi/ylenews.py b/src/lm_datasets/datasets/fi/ylenews.py
index db8d27b..837ca06 100644
--- a/src/lm_datasets/datasets/fi/ylenews.py
+++ b/src/lm_datasets/datasets/fi/ylenews.py
@@ -1,5 +1,5 @@
 import json
-from lm_datasets.datasets.base import BaseDataset, Availability
+from lm_datasets.datasets.base import BaseDataset, Availability, License
 import zipfile
 import logging
 from tqdm.auto import tqdm
@@ -11,7 +11,13 @@ class YLENewsDataset(BaseDataset):
     DATASET_ID = "ylenews"
     TITLE = "Yle Finnish News Archive"
     HOMEPAGE = "http://urn.fi/urn:nbn:fi:lb-2021050401"
-
+    LICENSE = License(
+        "CLARIN ACA - NC (Academic - Non Commercial Use, Attribution, No Redistribution, Other)",
+        url="https://www.clarin.eu/content/licenses-and-clarin-categories",
+        commercial_use=False,
+        research_use=True,
+        distribution=False,
+    )
     AVAILIBILITY = Availability.SIGNIN_DOWNLOAD
 
     LANGUAGES = ["fi"]
diff --git a/src/lm_datasets/datasets/fr/cabernet.py b/src/lm_datasets/datasets/fr/cabernet.py
index 4f50597..a66593e 100644
--- a/src/lm_datasets/datasets/fr/cabernet.py
+++ b/src/lm_datasets/datasets/fr/cabernet.py
@@ -1,18 +1,48 @@
 import logging
 import tarfile
-from lm_datasets.datasets.base import Availability, BaseDataset
+from lm_datasets.datasets.base import Availability, BaseDataset, QualityWarning, License
 
 logger = logging.getLogger(__name__)
 
 
 class CaBeRnetDataset(BaseDataset):
+    """
+    Contents:
+    - Fiction
+    - Acad (Wikiepdia)
+    - Oral
+    - Pop (Ads)
+    - News
+    """
+
     DATASET_ID = "cabernet"
     TITLE = "CaBeRnet: a New French Balanced Reference Corpus"
     LANGUAGES = ["fr"]
     HOMEPAGE = "https://aclanthology.org/2020.cmlc-1.3/"
-
+    CITATION = """@inproceedings{popa-fabre-etal-2020-french,
+    title = "{F}rench Contextualized Word-Embeddings with a sip of {C}a{B}e{R}net: a New {F}rench Balanced Reference Corpus",
+    author = "Popa-Fabre, Murielle  and
+      Ortiz Su{\'a}rez, Pedro Javier  and
+      Sagot, Beno{\^\i}t  and
+      de la Clergerie, {\'E}ric",
+    booktitle = "Proceedings of the 8th Workshop on Challenges in the Management of Large Corpora",
+    month = may,
+    year = "2020",
+    address = "Marseille, France",
+    publisher = "European Language Ressources Association",
+    url = "https://aclanthology.org/2020.cmlc-1.3",
+    pages = "15--23",
+    abstract = "This paper investigates the impact of different types and size of training corpora on language models. By asking the fundamental question of quality versus quantity, we compare four French corpora by pre-training four different ELMos and evaluating them on dependency parsing, POS-tagging and Named Entities Recognition downstream tasks. We present and asses the relevance of a new balanced French corpus, CaBeRnet, that features a representative range of language usage, including a balanced variety of genres (oral transcriptions, newspapers, popular magazines, technical reports, fiction, academic texts), in oral and written styles. We hypothesize that a linguistically representative corpus will allow the language models to be more efficient, and therefore yield better evaluation scores on different evaluation sets and tasks. This paper offers three main contributions: (1) two newly built corpora: (a) CaBeRnet, a French Balanced Reference Corpus and (b) CBT-fr a domain-specific corpus having both oral and written style in youth literature, (2) five versions of ELMo pre-trained on differently built corpora, and (3) a whole array of computational results on downstream tasks that deepen our understanding of the effects of corpus balance and register in NLP evaluation.",
+    language = "English",
+    ISBN = "979-10-95546-61-0",
+}
+"""  # noqa
+    LICENSE = License(
+        "Creative Commons License", url="https://aclanthology.org/2020.cmlc-1.3/ (Section 1)", research_use=True
+    )
     AVAILIBITY = Availability.ON_REQUEST
     TOKENS = 711_792_861
+    QUALITY_WARNINGS = [QualityWarning.BAD_DOCUMENT_SPLITS]
 
     def decompress(self):
         # gzip -d CaBeRnet.tar.gz
@@ -30,15 +60,20 @@ def get_texts(self):
 
             # Read from txt files
             for member_fn in member_fns:
-                logger.info(f"Decompress: {member_fn}")
-                decompressed_member = tar_f.extractfile(member_fn).read()
-                txt = decompressed_member.decode()
+                if "_Acad.txt" in member_fn.name:
+                    # Skip "Acad" because it is Wikipedia
+                    continue
 
-                if txt:
-                    pass
+                if "_Pop.txt" in member_fn.name:
+                    # Skip "Pop" because it is ads
+                    continue
 
-                # for doc_text in txt.split("\n\n"):
-                #     yield doc_text
+                logger.info(f"Decompress: {member_fn}")
+                decompressed_member = tar_f.extractfile(member_fn).read()
+                txt_file_content = decompressed_member.decode()
 
-                break
-                print("x")
+                if txt_file_content:
+                    # File content does not contain proper document separators.
+                    # We use simply use triple line breaks (not perfect but should do the job).
+                    for doc_text in txt_file_content.split("\n\n\n"):
+                        yield doc_text
diff --git a/src/lm_datasets/datasets/ga/ga_bilingual_legistation.py b/src/lm_datasets/datasets/ga/ga_bilingual_legistation.py
index f7238dd..5f7020d 100644
--- a/src/lm_datasets/datasets/ga/ga_bilingual_legistation.py
+++ b/src/lm_datasets/datasets/ga/ga_bilingual_legistation.py
@@ -1,5 +1,5 @@
 import logging
-from lm_datasets.datasets.base import BaseDataset, MB, Availability
+from lm_datasets.datasets.base import BaseDataset, MB, Availability, License
 
 
 logger = logging.getLogger(__name__)
@@ -15,7 +15,7 @@ class GABilingualLegislationDataset(BaseDataset):
     DATASET_ID = "ga_bilingual_legistation"
     TITLE = "Irish legislation"
     HOMEPAGE = "https://portulanclarin.net/repository/browse/the-gaois-bilingual-corpus-of-english-irish-legislation-processed/daeac17c9e3511ea9b7f02420a000407b83de243dc0b469aab41084386c5b80f/"  # noqa
-
+    LICENSE = License("Open Under - PSI", url="https://elrc-share.eu/terms/openUnderPSI.html")
     LANGUAGES = ["ga"]
 
     AVAILIBILITY = Availability.DIRECT_DOWNLOAD
diff --git a/src/lm_datasets/datasets/ga/ga_universal_dependencies.py b/src/lm_datasets/datasets/ga/ga_universal_dependencies.py
index 7799915..5134911 100644
--- a/src/lm_datasets/datasets/ga/ga_universal_dependencies.py
+++ b/src/lm_datasets/datasets/ga/ga_universal_dependencies.py
@@ -1,7 +1,7 @@
 import logging
 import os
-from lm_datasets.datasets.base import BaseDataset, MB, Availability
-from lm_datasets.systems import get_path_by_system
+from lm_datasets.datasets.base import BaseDataset, MB, Availability, License
+from lm_datasets.utils.systems import get_path_by_system
 
 import wget
 
@@ -20,7 +20,13 @@ class GAUniversalDependenciesDataset(BaseDataset):
     TITLE = "Irish Universal Dependencies"
     HOMEPAGE = "https://universaldependencies.org/"
     AVAILIBILITY = Availability.DIRECT_DOWNLOAD
-
+    LICENSE = License(
+        "mixed (CC BY-SA 3.0 or CC BY-SA 4.0)",
+        attribution=True,
+        sharealike=True,
+        commercial_use=True,
+        research_use=True,
+    )
     LANGUAGES = ["ga"]
 
     DOWNLOAD_URLS = [
@@ -37,49 +43,17 @@ class GAUniversalDependenciesDataset(BaseDataset):
 
     BYTES = 9.7 * MB
 
-    def download(self):
-        output_dir = get_path_by_system(self.LOCAL_DIRS)
-
-        if not os.path.exists(output_dir):
-            os.makedirs(output_dir)
-
-        file_paths = []
-
-        for url in self.DOWNLOAD_URLS:
-            fn = url.split("/")[-1]
-            fp = os.path.join(output_dir, fn)
-            # try:
-            out_filename = wget.download(url, out=fp)
-            logger.info(f"Completed {out_filename}")
-
-            file_paths.append(fp)
-            # except HTTPError as e:
-            #     logger.error(f"Error {e}")
-
-        return file_paths
-
-    def get_dataset_file_paths(self):
-        # TODO handle datasets with multiple filles
-        dataset_dir = get_path_by_system(self.LOCAL_DIRS)
-
-        if not os.path.exists(dataset_dir):
-            return []
-
-        fps = [os.path.join(dataset_dir, f) for f in os.listdir(dataset_dir)]
-        fps = [fp for fp in fps if os.path.isfile(fp) and fp.endswith(".conllu")]
-
-        return fps
+    def is_downloaded(self):
+        return len(self.get_dataset_file_paths(needed_suffix=".conllu")) == len(self.DOWNLOAD_URLS)
 
     def get_texts(self):
         import conllu
 
-        file_paths = self.get_dataset_file_paths()
-
-        if len(file_paths) != len(self.DOWNLOAD_URLS):
-            file_paths = self.download()
+        if not self.is_downloaded():
+            self.download()
 
         # Parse CONLL files and extract sentences
-        for fp in file_paths:
+        for fp in self.get_dataset_file_paths(needed_suffix=".conllu"):
             logger.info(f"Reading {fp}")
 
             with open(fp) as f:
diff --git a/src/lm_datasets/datasets/hf_dataset.py b/src/lm_datasets/datasets/hf_dataset.py
index c741ff0..bd762d0 100644
--- a/src/lm_datasets/datasets/hf_dataset.py
+++ b/src/lm_datasets/datasets/hf_dataset.py
@@ -14,12 +14,15 @@ class HFDataset(BaseDataset):
     HF_DATASET_SPLIT: Optional[str] = None
     HF_DATASET_CONFIGS: Optional[List[str]] = None
     HF_DATA_DIR = None
+    HF_KWARGS = None
+    HF_REVISION: Optional[str] = None
 
     config_to_dataset: Optional[Dict] = None
     text_column_name = "text"
     title_column_name = None
     remove_columns = None
     streaming = False
+    keep_columns = False
 
     def __init__(self, **kwargs) -> None:
         super().__init__(**kwargs)
@@ -40,15 +43,30 @@ def download(self):
         for hf_config in self.get_hf_configs():
             logger.info(f"Downloading for {hf_config=}")
 
-            ds = load_dataset(
-                self.HF_DATASET_ID,
-                hf_config,
-                split=self.HF_DATASET_SPLIT,
-                data_dir=self.HF_DATA_DIR,
-                streaming=self.streaming,
-                use_auth_token=self.get_hf_auth_token(),
-                keep_in_memory=False,
-            )
+            if self.HF_KWARGS:
+                # use additional kwargs as defined by dataset class
+                ds = load_dataset(
+                    self.HF_DATASET_ID,
+                    hf_config,
+                    split=self.HF_DATASET_SPLIT,
+                    data_dir=self.HF_DATA_DIR,
+                    streaming=self.streaming,
+                    use_auth_token=self.get_hf_auth_token(),
+                    keep_in_memory=False,
+                    revision=self.HF_REVISION,
+                    **self.HF_KWARGS,
+                )
+            else:
+                ds = load_dataset(
+                    self.HF_DATASET_ID,
+                    hf_config,
+                    split=self.HF_DATASET_SPLIT,
+                    data_dir=self.HF_DATA_DIR,
+                    streaming=self.streaming,
+                    use_auth_token=self.get_hf_auth_token(),
+                    keep_in_memory=False,
+                    revision=self.HF_REVISION,
+                )
 
             # check dataset split
             if isinstance(ds, DatasetDict) and not self.HF_DATASET_SPLIT:
@@ -82,32 +100,33 @@ def get_filter_func(self):
     def get_text_from_item(self, item) -> str:
         return item[self.text_column_name]
 
+    def prepend_title(self, example):
+        example[self.text_column_name] = (
+            example[self.title_column_name] + self.title_delimiter + example[self.text_column_name]
+        )
+
+        return example
+
     def get_texts(self):
         self.download()
 
         # drop all non-text columns
         for ds_idx, config in enumerate(self.config_to_dataset):
             # remove non-text and non-title columns
-            columns_to_remove = set(self.config_to_dataset[config].column_names) - {self.text_column_name}
+            if not self.keep_columns:
+                columns_to_remove = set(self.config_to_dataset[config].column_names) - {self.text_column_name}
 
-            if self.title_column_name:
-                columns_to_remove = columns_to_remove - {self.title_column_name}
+                if self.title_column_name:
+                    columns_to_remove = columns_to_remove - {self.title_column_name}
 
-            logger.info(f"Removing columns (get texts): {columns_to_remove}")
+                logger.info("Removing columns (get texts): %s", columns_to_remove)
 
-            self.config_to_dataset[config] = self.config_to_dataset[config].remove_columns(columns_to_remove)
+                self.config_to_dataset[config] = self.config_to_dataset[config].remove_columns(columns_to_remove)
 
             if self.title_column_name:
                 logger.info(f"Prepending title to text column ({self.title_column_name=})")
 
-                def prepend_title(example):
-                    example[self.text_column_name] = (
-                        example[self.title_column_name] + self.title_delimiter + example[self.text_column_name]
-                    )
-
-                    return example
-
-                self.config_to_dataset[config] = self.config_to_dataset[config].map(prepend_title)
+                self.config_to_dataset[config] = self.config_to_dataset[config].map(self.prepend_title)
 
                 # remove title column
                 self.config_to_dataset[config] = self.config_to_dataset[config].remove_columns([self.title_column_name])
@@ -115,20 +134,8 @@ def prepend_title(example):
             ds_iterator = iter(self.config_to_dataset[config])
 
             for item in ds_iterator:
-                yield self.get_text_from_item(item)
-
-            # yield from self.config_to_dataset[config][self.text_column_name]
-
-            # if self.text_column_name != self.output_text_field:
-            #     # rename text column to output
-            #     self.config_to_dataset[config] = self.config_to_dataset[config].rename_column(
-            #         self.text_column_name, self.output_text_field
-            #     )
-
-            # write_mode = "a" if ds_idx > 0 else "w"
-
-            # write to JSON line files
-            # logger.info(f"Writing output to {self.get_output_file_path()}; {write_mode=}")
-            # self.config_to_dataset[config].to_json(
-            #     self.get_output_file_path(), lines=True, mode=write_mode, force_ascii=self.json_ensure_ascii
-            # )
+                if hasattr(self, "get_texts_from_item"):
+                    # multiple texts from a single item
+                    yield from self.get_texts_from_item(item)
+                else:
+                    yield self.get_text_from_item(item)
diff --git a/src/lm_datasets/datasets/hr/croatian_news_engri.py b/src/lm_datasets/datasets/hr/croatian_news_engri.py
index d38a01f..3be5ac4 100644
--- a/src/lm_datasets/datasets/hr/croatian_news_engri.py
+++ b/src/lm_datasets/datasets/hr/croatian_news_engri.py
@@ -7,7 +7,7 @@
 import logging
 
 
-from lm_datasets.datasets.base import BaseDataset, Availability
+from lm_datasets.datasets.base import BaseDataset, Availability, License
 
 
 from smart_open import open
@@ -19,7 +19,6 @@
 class CroatianNewsENGRIDataset(BaseDataset):
     """
 
-
     https://www.clarin.si/repository/xmlui/handle/11356/1416
 
     """
@@ -28,7 +27,12 @@ class CroatianNewsENGRIDataset(BaseDataset):
     TITLE = "Corpus of Croatian news portals ENGRI (2014-2018)"
     HOMEPAGE = "https://repository.pfri.uniri.hr/islandora/object/pfri%3A2156"
     AVAILIBILITY = Availability.DIRECT_DOWNLOAD
-
+    LICENSE = License(
+        "Creative Commons - Attribution-NonCommercial-ShareAlike 4.0 International (CC BY-NC-SA 4.0)",
+        url="https://creativecommons.org/licenses/by-nc-sa/4.0/",
+        commercial_use=False,
+        sharealike=True,
+    )
     DESCRIPTION = (
         "The corpus consists of texts collected from the most popular (based on the Reuters Institute Digital News ",
         "Report for 2018, retrieved from http://www.digitalnewsreport.org in April, 2019) news portals in Croatia ",
@@ -38,8 +42,6 @@ class CroatianNewsENGRIDataset(BaseDataset):
 
     LANGUAGES = ["hr"]
 
-    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/hr/croatian_news_engri"]
-
     TOKENS = 694799268
 
     @property
diff --git a/src/lm_datasets/datasets/hr/hrwac.py b/src/lm_datasets/datasets/hr/hrwac.py
index fe97b5e..a5d3015 100644
--- a/src/lm_datasets/datasets/hr/hrwac.py
+++ b/src/lm_datasets/datasets/hr/hrwac.py
@@ -1,5 +1,5 @@
 import logging
-from lm_datasets.datasets.base import BaseDataset, Availability
+from lm_datasets.datasets.base import BaseDataset, Availability, License
 
 import html
 
@@ -20,14 +20,20 @@ class HRWACDataset(BaseDataset):
     TITLE = "Croatian web corpus hrWaC 2.1"
     HOMEPAGE = "http://nlp.ffzg.hr/resources/corpora/hrwac/"
     AVAILIBILITY = Availability.DIRECT_DOWNLOAD
-
+    LICENSE = License(
+        "CC-BY-SA license",
+        url="https://creativecommons.org/licenses/by-sa/4.0/",
+        sharealike=True,
+        attribution=True,
+        commercial_use=True,
+        research_use=True,
+    )
     LANGUAGES = ["hr"]
 
     DOWNLOAD_URLS = [
         f"https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1064/hrWaC2.1.{i:02d}.xml.gz"
         for i in range(1, 14)
     ]
-    # LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/hr/hrWaC2.1"]
 
     TOKENS = 1_397_757_548
 
diff --git a/src/lm_datasets/datasets/hr/styria_news.py b/src/lm_datasets/datasets/hr/styria_news.py
index da002b9..f35e7da 100644
--- a/src/lm_datasets/datasets/hr/styria_news.py
+++ b/src/lm_datasets/datasets/hr/styria_news.py
@@ -2,7 +2,7 @@
 
 import pandas as pd
 
-from lm_datasets.datasets.base import BaseDataset, GB, Availability, Genre
+from lm_datasets.datasets.base import BaseDataset, GB, Availability, Genre, License
 
 
 logger = logging.getLogger(__name__)
@@ -19,7 +19,14 @@ class StyriaNewsDataset(BaseDataset):
     TITLE = "24sata news article archive 1.0"
     HOMEPAGE = "https://www.clarin.si/repository/xmlui/handle/11356/1410"
     AVAILIBILITY = Availability.DIRECT_DOWNLOAD
-
+    LICENSE = License(
+        "Creative Commons - Attribution-NonCommercial-NoDerivatives 4.0 International (CC BY-NC-ND 4.0)",
+        url="https://creativecommons.org/licenses/by-nc-nd/4.0/",
+        attribution=True,
+        commercial_use=False,
+        derivates=False,
+        sharealike=False,
+    )
     LANGUAGES = ["hr"]
     GENRES = [Genre.NEWS]
 
diff --git a/src/lm_datasets/datasets/it/itwac.py b/src/lm_datasets/datasets/it/itwac.py
index c05080e..3ba9964 100644
--- a/src/lm_datasets/datasets/it/itwac.py
+++ b/src/lm_datasets/datasets/it/itwac.py
@@ -1,6 +1,6 @@
 import logging
 from typing import Iterable
-from lm_datasets.datasets.base import BILLION, Availability, BaseDataset
+from lm_datasets.datasets.base import BILLION, Availability, BaseDataset, License
 from smart_open import open
 
 logger = logging.getLogger(__name__)
@@ -20,8 +20,9 @@ class ITWacDataset(BaseDataset):
     LANGUAGES = ["it"]
     AVAILIBILITY = Availability.ON_REQUEST
     HOMEPAGE = "https://docs.sslmit.unibo.it/doku.php?id=corpora:itwac"
-
+    LICENSE = License("unknown, likely research-only or fair use")
     TOKENS = 2 * BILLION
+    WEB_CRAWLED = True
 
     def decompress(self):
         #  7z x itwac_preproc.7z
diff --git a/src/lm_datasets/datasets/jsonl_dataset.py b/src/lm_datasets/datasets/jsonl_dataset.py
index 760c5e2..13b16ae 100644
--- a/src/lm_datasets/datasets/jsonl_dataset.py
+++ b/src/lm_datasets/datasets/jsonl_dataset.py
@@ -19,11 +19,15 @@ class JSONLDataset(BaseDataset):
     raw_jsonl_text_field = "text"
 
     def is_downloaded(self):
-        for fp in self.get_raw_jsonl_paths():
-            if not os.path.exists(fp):
-                return False
+        try:
+            for fp in self.get_raw_jsonl_paths():
+                if not os.path.exists(fp):
+                    return False
+
+            return True
 
-        return True
+        except FileNotFoundError:
+            return False
 
     def get_raw_jsonl_paths(self):
         if self.raw_jsonl_paths is None:
@@ -31,6 +35,9 @@ def get_raw_jsonl_paths(self):
 
         dataset_dir = self.get_local_dataset_dir()
 
+        if not os.path.exists(dataset_dir):
+            raise FileExistsError("Dataset directory does not exist")
+
         for fp in self.raw_jsonl_paths:
             if fp.startswith("*"):
                 # find based on suffix
@@ -58,6 +65,7 @@ def get_texts(self):
         """
         Iterate over all input files and read JSON from each line.
         """
+        processed_files = 0
         for fp in self.get_raw_jsonl_paths():
             logger.info(f"Reading from {fp}")
 
@@ -70,3 +78,8 @@ def get_texts(self):
             else:
                 with open(fp) as f:  # jsonl or jsonl.fz (via smart_open)
                     yield from self.get_texts_from_file_handler(f)
+
+            processed_files += 1
+
+        if processed_files == 0:
+            logger.warning("No file has been processed.")
diff --git a/src/lm_datasets/datasets/lt/seimas_lt_en.py b/src/lm_datasets/datasets/lt/seimas_lt_en.py
index ad18654..92a77f8 100644
--- a/src/lm_datasets/datasets/lt/seimas_lt_en.py
+++ b/src/lm_datasets/datasets/lt/seimas_lt_en.py
@@ -1,5 +1,5 @@
 import zipfile
-from lm_datasets.datasets.base import BaseDataset, Genre
+from lm_datasets.datasets.base import BaseDataset, Genre, License
 
 
 class SeismasLTENDataset(BaseDataset):
@@ -12,7 +12,7 @@ class SeismasLTENDataset(BaseDataset):
     ]
     LANGUAGES = ["lt"]
     GENRES = [Genre.GOVERNMENT]
-
+    LICENSE = License("Open under PSI", url="https://elrc-share.eu/terms/openUnderPSI.html")
     BYTES = 160 * 1024
 
     def get_texts(self):
diff --git a/src/lm_datasets/datasets/lv/state_related_latvian_web.py b/src/lm_datasets/datasets/lv/state_related_latvian_web.py
index 6402ee7..154eab2 100644
--- a/src/lm_datasets/datasets/lv/state_related_latvian_web.py
+++ b/src/lm_datasets/datasets/lv/state_related_latvian_web.py
@@ -1,7 +1,7 @@
 import logging
 import zipfile
 
-from lm_datasets.datasets.base import MB, Availability, BaseDataset
+from lm_datasets.datasets.base import MB, Availability, BaseDataset, License
 
 logger = logging.getLogger(__name__)
 
@@ -11,7 +11,7 @@ class StateRelatedLatvianWebDataset(BaseDataset):
     TITLE = "Corpus of State-related content from the Latvian Web (Processed)"
     HOMEPAGE = "http://catalog.elra.info/en-us/repository/browse/ELRA-W0169/"
     AVAILIBILITY = Availability.SIGNIN_DOWNLOAD
-
+    LICENSE = License("CC-BY-SA-4.0", attribution=True, sharealike=True, commercial_use=True, research_use=True)
     BYTES = 3.4 * MB
     LANGUAGES = ["lv"]
 
diff --git a/src/lm_datasets/datasets/mt/korpus_malti.py b/src/lm_datasets/datasets/mt/korpus_malti.py
index d4418f5..b88170f 100644
--- a/src/lm_datasets/datasets/mt/korpus_malti.py
+++ b/src/lm_datasets/datasets/mt/korpus_malti.py
@@ -1,8 +1,6 @@
 import json
 from pathlib import Path
-from lm_datasets.datasets.base import BaseDataset, Availability
-from lm_datasets.datasets.hf_dataset import HFDataset
-
+from lm_datasets.datasets.base import BaseDataset, Availability, License
 
 from tqdm.auto import tqdm
 import logging
@@ -10,76 +8,6 @@
 logger = logging.getLogger(__name__)
 
 
-class KorpusMaltiHFDataset(HFDataset):
-    DATASET_ID = "korpus_malti"
-
-    TITLE = "Korpus Malti"
-    HOMEPAGE = "https://huggingface.co/datasets/MLRS/korpus_malti"
-    AVAILIBILITY = Availability.DIRECT_DOWNLOAD
-
-    LANGUAGES = ["mt"]
-
-    DESCRIPTION = (
-        "General Corpora for the Maltese Language. This dataset is composed of texts ",
-        "from various genres/domains written in Maltese.",
-    )
-
-    PII = "I have not checked the data source for personally identifiable or sensitive information."
-
-    LICENSE = "cc-by-nc-sa-4.0"  # DFKI has a permission for LLM training with commercial license
-
-    # From paper https://aclanthology.org/2022.deeplo-1.10.pdf (Table 1)
-    TOKENS = 466_601_373 - 98_582_031 - 1_885_661  # full dataset - law_eu - wiki
-
-    HF_DATASET_ID = "MLRS/korpus_malti"
-    HF_DATASET_SPLIT = "train"
-    HF_DATASET_CONFIGS = [
-        "belles_lettres",
-        "blogs",
-        "comics",
-        "court",
-        "eu_docs",
-        "gov_docs",
-        "government_gazzette",
-        # "law_eu" # avoid duplication with EUR-LEx
-        "law_mt",
-        "legal",
-        "nonfiction",
-        "parliament",
-        "press_eu",
-        "press_mt",
-        "speeches",
-        "theses",
-        "umlib_oar",
-        "web_general",
-        # "wiki",  # avoid duplication with Wiki
-    ]
-    # text_column_name = "text"
-    # remove_columns = ["doc_id", "LICENSE", "uri", "date_built"]
-
-    def get_text_from_item(self, item) -> str:
-        sentences = item["text"]
-        text = self.sentence_delimiter.join(sentences)
-
-        return text
-
-    # def get_texts(self):
-    #     # load dataset for each config
-    #     for config in self.dataset_configs:
-    #         logger.info(f"Processing: {config}")
-    #         self.datasets[config] = load_dataset(
-    #             self.HF_DATASET_ID,
-    #             config,
-    #             streaming=False,
-    #             use_auth_token=self.get_hf_auth_token(),
-    #         )
-
-    #         for sentences in self.datasets[config]["train"]["text"]:
-    #             text = self.sentence_delimiter.join(sentences)
-
-    #             yield text
-
-
 class KorpusMaltiDataset(BaseDataset):
     DATASET_ID = "korpus_malti"
 
@@ -96,7 +24,14 @@ class KorpusMaltiDataset(BaseDataset):
 
     PII = "I have not checked the data source for personally identifiable or sensitive information."
 
-    LICENSE = "cc-by-nc-sa-4.0"  # DFKI has a permission for LLM training with commercial license
+    LICENSE = License(
+        "Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License (DFKI has a permission for LLM training with commercial license)",  # noqa
+        url="https://creativecommons.org/licenses/by-nc-sa/4.0/",
+        commercial_use=False,
+        attribution=True,
+        sharealike=True,
+        research_use=True,
+    )
 
     # From paper https://aclanthology.org/2022.deeplo-1.10.pdf (Table 1)
     TOKENS = 466_601_373 - 98_582_031 - 1_885_661  # full dataset - law_eu - wiki
diff --git a/src/lm_datasets/datasets/multilingual/colossal_oscar.py b/src/lm_datasets/datasets/multilingual/colossal_oscar.py
index 834c6c9..2d59f80 100644
--- a/src/lm_datasets/datasets/multilingual/colossal_oscar.py
+++ b/src/lm_datasets/datasets/multilingual/colossal_oscar.py
@@ -4,20 +4,41 @@
 
 import numpy as np
 
-from lm_datasets.datasets.base import Availability
+from lm_datasets.datasets.base import Availability, License
 from lm_datasets.datasets.jsonl_dataset import JSONLDataset
 
 
 logger = logging.getLogger(__name__)
 
 
-DEFAULT_OSCAR_MIN_HARMFUL_PP = 20.0
+# Magic numbers
+DEFAULT_OSCAR_MIN_HARMFUL_PP = 25.0
+DEFAULT_OSCAR_MAX_HARMFUL_PP = 100_000
 
 """
+Extract list:
+",".join([f"colossal_oscar_{dump}_ca" for dump in  OSCAR_DUMPS])
 'colossal_oscar_05-06-23_bg,colossal_oscar_05-06-23_cs,colossal_oscar_05-06-23_da,colossal_oscar_05-06-23_el,colossal_oscar_05-06-23_et,colossal_oscar_05-06-23_fi,colossal_oscar_05-06-23_fr,colossal_oscar_05-06-23_ga,colossal_oscar_05-06-23_hr,colossal_oscar_05-06-23_hu,colossal_oscar_05-06-23_lt,colossal_oscar_05-06-23_lv,colossal_oscar_05-06-23_mt,colossal_oscar_05-06-23_nl,colossal_oscar_05-06-23_pl,colossal_oscar_05-06-23_pt,colossal_oscar_05-06-23_ro,colossal_oscar_05-06-23_sk,colossal_oscar_05-06-23_sl,colossal_oscar_05-06-23_sv,colossal_oscar_05-06-23_uk,colossal_oscar_05-06-23_sr,colossal_oscar_05-06-23_sh,colossal_oscar_05-06-23_nn,colossal_oscar_05-06-23_no,colossal_oscar_05-06-23_eu,colossal_oscar_05-06-23_ca,colossal_oscar_05-06-23_gl,colossal_oscar_03-04-23_bg,colossal_oscar_03-04-23_cs,colossal_oscar_03-04-23_da,colossal_oscar_03-04-23_el,colossal_oscar_03-04-23_et,colossal_oscar_03-04-23_fi,colossal_oscar_03-04-23_fr,colossal_oscar_03-04-23_ga,colossal_oscar_03-04-23_hr,colossal_oscar_03-04-23_hu,colossal_oscar_03-04-23_lt,colossal_oscar_03-04-23_lv,colossal_oscar_03-04-23_mt,colossal_oscar_03-04-23_nl,colossal_oscar_03-04-23_pl,colossal_oscar_03-04-23_pt,colossal_oscar_03-04-23_ro,colossal_oscar_03-04-23_sk,colossal_oscar_03-04-23_sl,colossal_oscar_03-04-23_sv,colossal_oscar_03-04-23_uk,colossal_oscar_03-04-23_sr,colossal_oscar_03-04-23_sh,colossal_oscar_03-04-23_nn,colossal_oscar_03-04-23_no,colossal_oscar_03-04-23_eu,colossal_oscar_03-04-23_ca,colossal_oscar_03-04-23_gl'
 """  # noqa
-OSCAR_DUMPS = ["05-06-23", "03-04-23"]
+# all dumps: 2015-14  2016-40  2017-43  2018-47  2019-22  2020-24  2020-45  2021-49  2022-27  2022-49  2023-14  2023-23
+# first sample:
+# OSCAR_DUMPS = ["05-06-23", "03-04-23"]
+OSCAR_DUMPS = [
+    "2015-14",
+    "2016-40",
+    "2017-43",
+    "2018-47",
+    "2019-22",
+    "2020-24",
+    "2020-45",
+    "2021-49",
+    "2022-27",
+    "2022-49",
+    "2023-14",
+    "2023-23",
+]
 LANGUAGES = "bg cs da el et fi ga hr hu lt lv mt nl pl pt ro sk sl sv uk sr sh nn no eu ca gl".split(" ")  # removed fr
+LANGUAGES += ["en", "de", "fr", "es", "it"]  # EU top-5 languages
 
 EXCLUDE_CATEGORIES = {
     # See http://dsi.ut-capitole.fr/blacklists/index_en.php
@@ -27,6 +48,11 @@
     "dangerous_material",
     "phishing",
     "warez",
+    "ddos",
+    "hacking",
+    "malware",
+    "mixed_adult",
+    "sect",
 }
 
 
@@ -35,42 +61,46 @@ class ColossalOscarBaseDataset(JSONLDataset):
     Read OSCAR output from jsonl.zst files (as provided on HF)
     """
 
+    DATASET_ID = None
     SOURCE_ID = "colossal_oscar"
 
+    TITLE = "Colossal OSCAR 1"
     DESCRIPTION = "Colossal OSCAR 1"
     HOMEPAGE = "https://huggingface.co/datasets/oscar-corpus/colossal-oscar-1.0"
     AVAILIBILITY = Availability.SIGNIN_DOWNLOAD
 
-    LOCAL_DIRS = ["pegasus:/netscratch/mostendorff/experiments/eulm/data/colossal-oscar-1.0"]
-
-    DATASET_ID = "colossal_oscar_05-06-23_da"
     LANGUAGES = ["da"]
     DUMP_VERSION = "05-06-23"
+    WEB_CRAWLED = True
+
+    LICENSE = License(
+        "CommonCrawl terms of use; Only the annotations are distributed under a cc0-1.0 license",
+        url="https://commoncrawl.org/terms-of-use",
+        commercial_use=True,
+        research_use=True,
+    )
 
     min_harmful_pp = None
+    max_harmful_pp = None
 
     def __init__(self, **kwargs) -> None:
         super().__init__(**kwargs)
 
-        # set min_harmful_pp
-
-        # TODO hard coded
-        # resources_dir = os.path.join("/netscratch/mostendorff/experiments/eulm", "resources", "oscar")
+        self.min_harmful_pp = DEFAULT_OSCAR_MIN_HARMFUL_PP
+        self.max_harmful_pp = DEFAULT_OSCAR_MAX_HARMFUL_PP
 
-        # with open(os.path.join(resources_dir, "adult_pp_dict.pkl"), "rb") as f:
-        #     c_dict_means = pickle.load(f)
+        # Load language specific settings from config
+        if hasattr(self.config, "colossal_oscar_min_harmful_pp_by_language"):
+            colossal_oscar_min_harmful_pp_by_language = self.config.colossal_oscar_min_harmful_pp_by_language
 
-        # with open(os.path.join(resources_dir, "adult_pp_dict_full.pkl"), "rb") as file:
-        #     # Mean perplexity for documents with the Adult quality warning in the Oscar Dataset
-        #     dict_full = pickle.load(file)
+            if self.get_language_code() in colossal_oscar_min_harmful_pp_by_language:
+                self.min_harmful_pp = colossal_oscar_min_harmful_pp_by_language[self.get_language_code()]
 
-        dict_full = {}
+        if hasattr(self.config, "colossal_oscar_max_harmful_pp_by_language"):
+            colossal_oscar_max_harmful_pp_by_language = self.config.colossal_oscar_max_harmful_pp_by_language
 
-        self.min_harmful_pp = (
-            np.mean(dict_full[self.get_language_code()])
-            if self.get_language_code() in dict_full
-            else DEFAULT_OSCAR_MIN_HARMFUL_PP
-        )
+            if self.get_language_code() in colossal_oscar_max_harmful_pp_by_language:
+                self.max_harmful_pp = colossal_oscar_max_harmful_pp_by_language[self.get_language_code()]
 
         # logger.info(f"{self.min_harmful_pp=}")
 
@@ -83,7 +113,14 @@ def get_text_from_item(self, doc):
             and doc["metadata"]["harmful_pp"]
             and doc["metadata"]["harmful_pp"] < self.min_harmful_pp
         ):
-            self.counter.update({"filtered_harmful_pp": 1})
+            self.counter.update({"min_filtered_harmful_pp": 1})
+            return None
+        elif (
+            "harmful_pp" in doc["metadata"]
+            and doc["metadata"]["harmful_pp"]
+            and doc["metadata"]["harmful_pp"] > self.max_harmful_pp
+        ):
+            self.counter.update({"max_filtered_harmful_pp": 1})
             return None
         elif doc["metadata"]["categories"] and len(set(doc["metadata"]["categories"]) & EXCLUDE_CATEGORIES) > 0:
             self.counter.update({"filtered_categories": 1})
@@ -95,15 +132,22 @@ def get_raw_jsonl_paths(self):
         lang = self.get_language_code()
         dataset_path = Path(os.path.join(self.get_local_dataset_dir(), self.DUMP_VERSION, f"{lang}_meta"))
 
+        if not dataset_path.exists():
+            raise FileNotFoundError(f"Dataset path does not exist: {dataset_path}")
+
         return sorted([str(p) for p in dataset_path.glob("*.jsonl.zst")])
 
     def get_bytes(self):
-        return sum(os.stat(fp).st_size for fp in self.get_raw_jsonl_paths())
+        try:
+            return sum(os.stat(fp).st_size for fp in self.get_raw_jsonl_paths())
+        except FileNotFoundError:
+            return -1
 
 
 def get_colossal_oscar_class(lang, dump_version):
     class ColossalOscarDataset(ColossalOscarBaseDataset):
         DATASET_ID = f"colossal_oscar_{dump_version}_{lang}"
+        TITLE = f"Colossal OSCAR 1 ({lang}; {dump_version})"
         LANGUAGES = [lang]
         DUMP_VERSION = dump_version
 
diff --git a/src/lm_datasets/datasets/multilingual/curlicat.py b/src/lm_datasets/datasets/multilingual/curlicat.py
index 9e2f8d5..c2d8fdf 100644
--- a/src/lm_datasets/datasets/multilingual/curlicat.py
+++ b/src/lm_datasets/datasets/multilingual/curlicat.py
@@ -2,7 +2,7 @@
 from typing import Iterable
 from smart_open import open
 
-from lm_datasets.datasets.base import BaseDataset, Availability
+from lm_datasets.datasets.base import BaseDataset, Availability, License
 from tqdm.auto import tqdm
 
 from lm_datasets.io.conllu_file import get_texts_from_conllu_file
@@ -11,6 +11,10 @@
 
 
 class CurlicatBaseDataset(BaseDataset):
+    """
+    Licenses are mostly mixed. See https://aclanthology.org/2022.lrec-1.11.pdf
+    """
+
     SOURCE_ID = "curlicat"
     AVAILIBILITY = Availability.DIRECT_DOWNLOAD
 
@@ -68,16 +72,15 @@ class CurlicatBGDataset(CurlicatBaseDataset):
     DATASET_ID = "curlicat_bg"
     TITLE = "CURLICAT Bulgarian Corpus"
     DESCRIPTION = """The Bulgarian CURLICAT corpus consists of texts from different sources, provided with  appropriate licences for distribution. We used three general types of sources with regard to the metadata extraction: Bulgarian National Corpus (provided that they have redistributable licensing terms); some public repositories with open and copyright free data; blogs with redistributable licenses, open content websites, etc. The Bulgarian CURLICAT collection contains 113 087 documents, distributed in seven thematic domains: Culture, Education, European Union, Finance, Politics, Economics, and Science. For more information see the CURLICAT website (http:curlicat-project.eu/deliverables) """  # noqa
-
     HOMEPAGE = "https://elrc-share.eu/repository/browse/curlicat-bulgarian-corpus/fed6af2a590311ed9c1a00155d0267062ed273d01d2343f1b78d08d4d481679d/"  # noqa
-
     LANGUAGES = ["bg"]
-
-    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/bg/curlicat_bg/Bulgarian_Curlicat_corpus"]
-
     TOKENS = 35_319_695
-
-    # done
+    LICENSE = License(
+        name="CC-BY-SA-4.0",
+        url="https://elrc-share.eu/static/metashare/licences/CC-BY-SA-4.0.pdf",
+        attribution=True,
+        sharealike=True,
+    )
 
 
 class CurlicatHRDataset(CurlicatBaseDataset):
@@ -88,16 +91,11 @@ class CurlicatHRDataset(CurlicatBaseDataset):
 
     DATASET_ID = "curlicat_hr"
     TITLE = "CURLICAT Croatian Corpus"
-
     HOMEPAGE = "https://elrc-share.eu/repository/browse/curlicat-croatian-corpus/00815518592811ed9c1a00155d026706bc4c59740fce4f7986213e7eef133023/"  # noqa
-
     LANGUAGES = ["hr"]
-
     DOWNLOAD_URLS = ["https://zzl.ffzg.unizg.hr/files/5cd9d6151612489fafb2c1f116140519/hr_curlicat_2023-01-13.zip"]
-
     TOKENS = 49_007_508
-
-    # done
+    LICENSE = License("unknown")
 
 
 class CurlicatHUDataset(CurlicatBaseDataset):
@@ -106,13 +104,17 @@ class CurlicatHUDataset(CurlicatBaseDataset):
     HOMEPAGE = "https://elrc-share.eu/repository/browse/curlicat-hungarian-corpus/8b6c8dcb58ea11ed9c1a00155d02670679a453431c8147079e5a7d9b879a9729/"  # noqa
     LANGUAGES = ["hu"]
     TOKENS = 61_196_946
+    LICENSE = License(
+        "CC-BY-SA-4.0",
+        url="https://elrc-share.eu/static/metashare/licences/CC-BY-SA-4.0.pdf",
+        attribution=True,
+        sharealike=True,
+    )
 
     def get_curclicat_file_paths(self):
         # Hugarian files are stored in *.txt files subdirectories
         return self.get_dataset_file_paths(subdirectories=True, needed_suffix=".txt")
 
-    # extracted
-
 
 class CurlicatPLDataset(CurlicatBaseDataset):
     """
@@ -122,31 +124,36 @@ class CurlicatPLDataset(CurlicatBaseDataset):
 
     DATASET_ID = "curlicat_pl"
     TITLE = "CURLICAT Polish Corpus"
-
     HOMEPAGE = "https://elrc-share.eu/repository/browse/curlicat-polish-corpus/f63ae912553911ed9c1a00155d02670648c0a234e0314895b52169af2af57dd7/"  # noqa
-
     LANGUAGES = ["pl"]
     TOKENS = 59_301_782
 
-    # done
+    LICENSE = License(
+        name="CC-BY-SA-4.0",
+        url="https://elrc-share.eu/static/metashare/licences/CC-BY-SA-4.0.pdf",
+        attribution=True,
+        sharealike=True,
+    )
 
 
 class CurlicatRODataset(CurlicatBaseDataset):
     DATASET_ID = "curlicat_ro"
     TITLE = "CURLICAT Romanian Corpus"
     DESCRIPTION = """The corpus contains 26,477 files, which represent our contribution to the CURLICAT project. It contains texts from 7 domains: science, politics, culture, economy, health, education, nature. Each file has multiple levels of annotation: tokenized, lemmatized, morphologically annotated, dependency parsed, named entities, nominal phrases, IATE terms and automatic domain-specific terms were identified as well. All processing tools are available within the RELATE platform. The corpus was automatically anonymized. Alternate download location: https://relate.racai.ro/resources/curlicat/"""  # noqa
-
     HOMEPAGE = "https://elrc-share.eu/repository/browse/curlicat-romanian-corpus/8b6c8dca58ea11ed9c1a00155d026706fb03ef8b4c1847cfbe9cea869a82731e/"  # noqa
-
     LANGUAGES = ["ro"]  # TODO Romanian; Moldavian; Moldovan (ro)
-
     TOKENS = 94_925_454
-
     DOWNLOAD_URLS = ["https://relate.racai.ro/resources/curlicat/ro_curlicat_20221031.zip"]
 
     # See https://stackoverflow.com/questions/17912307/u-ufeff-in-python-string
     conllup_encoding = "utf-8-sig"
-    # done
+
+    LICENSE = License(
+        name="CC-BY-SA-4.0",
+        url="https://elrc-share.eu/static/metashare/licences/CC-BY-SA-4.0.pdf",
+        attribution=True,
+        sharealike=True,
+    )
 
 
 class CurlicatSKDataset(CurlicatBaseDataset):
@@ -155,21 +162,16 @@ class CurlicatSKDataset(CurlicatBaseDataset):
 
     DATASET_ID = "curlicat_sk"
     TITLE = "CURLICAT Slovak Corpus (3rd version)"
-
     HOMEPAGE = "https://elrc-share.eu/repository/browse/curlicat-slovak-corpus-v10/b419d7086ef411ed9c1a00155d0267066a930aa487824c08ba48f1183e993aca/"  # noqa
-
     LANGUAGES = ["sk"]
-
     DOWNLOAD_URLS = ["https://www.juls.savba.sk/data/curlicat/curlicat-sk-20221025-v1.0.conllup.xz"]
-
     TOKENS = 67_000_000
+    LICENSE = License("unknown")
 
     def decompress(self):
         # xz --decompress  curlicat-sk-20221025-v1.0.conllup.xz
         pass
 
-    # done
-
 
 class CurlicatSLDataset(CurlicatBaseDataset):
     DATASET_ID = "curlicat_sl"
@@ -177,5 +179,9 @@ class CurlicatSLDataset(CurlicatBaseDataset):
     HOMEPAGE = "https://elrc-share.eu/repository/browse/curlicat-slovenian-corpus/e549f298590711ed9c1a00155d026706db0d61d46f294d9a821307cf9c5df245/"  # noqa
     LANGUAGES = ["sl"]
     TOKENS = 43_481_563
-
-    # done
+    LICENSE = License(
+        name="CC-BY-SA-4.0",
+        url="https://elrc-share.eu/static/metashare/licences/CC-BY-SA-4.0.pdf",
+        attribution=True,
+        sharealike=True,
+    )
diff --git a/src/lm_datasets/datasets/multilingual/eurlex.py b/src/lm_datasets/datasets/multilingual/eurlex.py
index bda1754..dbd14f8 100644
--- a/src/lm_datasets/datasets/multilingual/eurlex.py
+++ b/src/lm_datasets/datasets/multilingual/eurlex.py
@@ -1,5 +1,5 @@
 import os
-from lm_datasets.datasets.base import Availability
+from lm_datasets.datasets.base import Availability, Genre, License
 from lm_datasets.datasets.jsonl_dataset import JSONLDataset
 
 # From https://docs.google.com/spreadsheets/d/1_rfLKa_Kq09YI0BPnfmoSL6-U3SHrW8tRwmk3-Qchzo/edit#gid=1698256287
@@ -38,8 +38,16 @@ class EURLexBaseDataset(JSONLDataset):
     DESCRIPTION = "EurlexResources: A Corpus Covering the Largest EURLEX Resources."
     HOMEPAGE = "https://huggingface.co/datasets/joelito/eurlex_resources"
     AVAILIBILITY = Availability.DIRECT_DOWNLOAD
-
-    LOCAL_DIRS = ["pegasus:/netscratch/mostendorff/experiments/eulm/data/docs_by_language"]
+    GENRES = [Genre.LEGAL]
+    LICENSE = License(
+        "Creative Commons Attribution 4.0 International licence",
+        url="https://eur-lex.europa.eu/content/legal-notice/legal-notice.html#2.%20droits",
+        attribution=True,
+        commercial_use=True,
+        sharealike=False,
+        research_use=True,
+        distribution=True,
+    )
 
     def get_raw_jsonl_paths(self):
         dataset_dir = self.get_local_dataset_dir()
diff --git a/src/lm_datasets/datasets/multilingual/legal_mc4.py b/src/lm_datasets/datasets/multilingual/legal_mc4.py
index e147efd..0caac5f 100644
--- a/src/lm_datasets/datasets/multilingual/legal_mc4.py
+++ b/src/lm_datasets/datasets/multilingual/legal_mc4.py
@@ -1,4 +1,4 @@
-from lm_datasets.datasets.base import Availability
+from lm_datasets.datasets.base import Availability, Genre, License
 from lm_datasets.datasets.hf_dataset import HFDataset
 
 
@@ -35,6 +35,13 @@ class LegalMC4BaseDataset(HFDataset):
     AVAILIBILITY = Availability.DIRECT_DOWNLOAD
     WEB_CRAWLED = True
     # DUMMY = True
+    LICENSE = License(
+        "AllenAI are releasing this dataset under the terms of ODC-BY. By using this, you are also bound by the Common Crawl terms of use in respect of the content contained in the dataset.",
+        commercial_use=True,
+        research_use=True,
+        attribution=True,
+    )
+    GENRES = [Genre.LEGAL]
 
     HF_DATASET_ID = "joelito/legal-mc4"
     HF_DATASET_SPLIT = "train"
diff --git a/src/lm_datasets/datasets/multilingual/macocu.py b/src/lm_datasets/datasets/multilingual/macocu.py
index b554c4e..264b095 100644
--- a/src/lm_datasets/datasets/multilingual/macocu.py
+++ b/src/lm_datasets/datasets/multilingual/macocu.py
@@ -2,7 +2,7 @@
 import logging
 from typing import Set
 import zipfile
-from lm_datasets.datasets.base import BaseDataset, Availability, QualityWarning
+from lm_datasets.datasets.base import BaseDataset, Availability, QualityWarning, License
 
 
 logger = logging.getLogger(__name__)
@@ -12,6 +12,14 @@ class MacocuBaseDataset(BaseDataset):
     SOURCE_ID = "macocu"
     AVAILIBILITY = Availability.DIRECT_DOWNLOAD
     WEB_CRAWLED = True
+    LICENSE = License(
+        "CC0-No Rights Reserved",
+        commercial_use=True,
+        research_use=True,
+        sharealike=False,
+        attribution=False,
+        distribution=True,
+    )
 
     languages_needed: Set[str] = None  # hbs_lat,en,hbs_cyr
 
diff --git a/src/lm_datasets/datasets/multilingual/oscar_dummy.py b/src/lm_datasets/datasets/multilingual/oscar_dummy.py
deleted file mode 100644
index 7d3c5d3..0000000
--- a/src/lm_datasets/datasets/multilingual/oscar_dummy.py
+++ /dev/null
@@ -1,229 +0,0 @@
-from lm_datasets.datasets.base import BaseDataset, Availability
-
-# From https://docs.google.com/spreadsheets/d/1_rfLKa_Kq09YI0BPnfmoSL6-U3SHrW8tRwmk3-Qchzo/edit#gid=430577929
-OSCAR_RAW_LANG_TO_TOKENS = """af	6,217,024
-am	40,262,809
-an	264
-ar	10,081,452,882
-arz	385,511
-as	24,395,215
-ast	10,917
-av	1,073
-az	316,850,330
-azb	2,029,729
-ba	26,036,637
-be	164,729,607
-bg	3,635,273,738
-bh	507
-bn	1,092,983,765
-bo	6,062,558
-bpy	346,947
-br	4,759,407
-bs	395
-bxr	698
-ca	2,240,460,836
-ce	1,051,752
-ceb	6,263,404
-ckb	61,334,746
-cs	9,717,378,559
-cv	3,039,925
-cy	52,250,043
-da	2,217,634,340
-de	73,848,586,648
-dsb	84
-dv	10,655,359
-el	7,691,622,692
-en	523,869,288,690
-eo	67,774,923
-es	63,388,237,965
-et	938,296,892
-eu	136,672,615
-fa	9,609,206,698
-fi	4,198,143,883
-fr	62,127,088,294
-fy	9,885,788
-ga	9,054,923
-gd	18,458
-gl	38,345,625
-gn	260
-gom	121,035
-gsw	34,328
-gu	417,001,705
-he	1,697,158,891
-hi	2,475,605,444
-hr	3,542,961
-hsb	15,827
-ht	20,671
-hu	16,013,364,289
-hy	336,045,041
-ia	9,384
-id	3,228,020,221
-ie	0
-ilo	4,411
-io	2,598
-is	294,471,539
-it	36,327,274,203
-ja	4,401,059,165
-jbo	312,315
-jv	3,286
-ka	373,935,158
-kk	214,679,857
-km	59,880,231
-kn	124,924,350
-ko	3,435,866,935
-krc	8,385
-ku	25,921,607
-kv	5,909
-kw	80
-ky	32,062,783
-la	307,865
-lb	2,514,838
-lez	60,634
-li	169
-lmo	3,269
-lo	10,659,203
-lt	1,674,362,574
-lv	845,459,899
-mai	467
-mg	1,416,430
-mhr	1,615,215
-min	305,589
-mk	389,344,425
-ml	236,597,838
-mn	454,350,415
-mr	252,706,331
-mrj	60,180
-ms	238,477
-mt	149,886
-multi	1,251,676,406
-mwl	54
-my	82,433,836
-mzn	16,115
-nah	279
-nds	1,612,175
-ne	278,901,036
-new	229,703
-nl	19,564,553,306
-nn	575,518
-no	373,160,033
-oc	34,701
-or	31,963,340
-os	3,935,964
-pa	104,235,418
-pl	18,073,705,588
-pms	510,087
-pnb	11,844,695
-ps	30,196,179
-pt	15,172,557,311
-qu	13
-ro	6,302,600,833
-ru	78,032,029,344
-sa	2,479,345
-sah	4,288,051
-sd	14,667,207
-sh	166,517
-si	172,755,385
-sk	2,704,716,280
-sl	192,816,743
-so	51
-sq	462,694,599
-sr	632,781,822
-su	258
-sv	6,993,719,601
-sw	164,459
-ta	738,824,392
-te	201,575,815
-tg	76,987,285
-th	2,224,483,018
-tk	325,786
-tl	110,560,444
-tr	8,290,890,087
-tt	59,253,765
-ug	14,659,554
-uk	3,183,842,018
-ur	434,023,273
-uz	1,665,960
-vi	22,424,984,210
-vo	49,968
-wa	6,347
-war	19,665
-wuu	1,199
-x-eml	329
-xal	27
-xmf	283,316
-yi	14,287,370
-yo	2,396
-zh	44,378,380,161"""
-
-RAW_FILTER_LANGUAGES = """bg
-cs
-da
-de
-el
-en
-es
-et
-fi
-fr
-ga
-hr
-hu
-it
-lt
-lv
-mt
-nl
-pl
-pt
-ro
-sk
-sl
-sv
-code
-uk
-sr
-sh
-nn
-no"""
-
-
-class OscarBaseDataset(BaseDataset):
-    SOURCE_ID = "oscar"
-    DESCRIPTION = (
-        "The OSCAR project (Open Super-large Crawled Aggregated coRpus) is an Open Source project aiming to provide"
-        " web-based multilingual resources and datasets for Machine Learning (ML) and Artificial Intelligence (AI)"
-        " applications."
-    )
-    HOMEPAGE = "https://oscar-project.org/"
-    AVAILIBILITY = Availability.ON_REQUEST
-    WEB_CRAWLED = True
-    DUMMY = True
-
-
-def get_oscar_dummy_cls_by_language(lang, tokens):
-    class OscarLanguageDataset(OscarBaseDataset):
-        TOKENS = tokens
-        DATASET_ID = "oscar_" + lang
-        TITLE = "oscar_" + lang
-        LANGUAGES = [lang]
-
-    return OscarLanguageDataset
-
-
-def get_oscar_dummy_classes():
-    """
-    Generate dummy dataset classes with token count based on OSCAR 23.01
-    """
-    pass
-
-    lang_to_tokens = {
-        row.split("\t")[0]: int(row.split("\t")[1].replace(",", "")) for row in OSCAR_RAW_LANG_TO_TOKENS.splitlines()
-    }
-
-    filter_langs = set(RAW_FILTER_LANGUAGES.splitlines())
-
-    return [
-        get_oscar_dummy_cls_by_language(lang, tokens * 10)  # TODO token count like OSCAR based on ten CC dumps
-        for lang, tokens in lang_to_tokens.items()
-        if lang in filter_langs
-    ]
diff --git a/src/lm_datasets/datasets/multilingual/redpajama.py b/src/lm_datasets/datasets/multilingual/redpajama.py
index abc591e..210f0b6 100644
--- a/src/lm_datasets/datasets/multilingual/redpajama.py
+++ b/src/lm_datasets/datasets/multilingual/redpajama.py
@@ -1,5 +1,5 @@
 import os
-from lm_datasets.datasets.base import Availability, BILLION
+from lm_datasets.datasets.base import Availability, BILLION, License
 from lm_datasets.datasets.hf_dataset import HFDataset
 
 
@@ -42,25 +42,23 @@ class RedPajamaBaseDataset(HFDataset):
     HOMEPAGE = "https://huggingface.co/datasets/togethercomputer/RedPajama-Data-1T"
     AVAILIBILITY = Availability.DIRECT_DOWNLOAD
 
-    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/multilingual/redpajama"]
-
     HF_DATASET_ID = "togethercomputer/RedPajama-Data-1T"
     HF_DATASET_SPLIT = "train"
 
     text_column_name = "text"
 
-    def __init__(self, **kwargs) -> None:
-        super().__init__(**kwargs)
-
+    def download(self):
         # set download diretory
         os.environ["RED_PAJAMA_DATA_DIR"] = self.get_local_dataset_dir()
 
+        super().download()
+
 
 class RedPajamaBookDataset(RedPajamaBaseDataset):
     DATASET_ID = "redpajama_book"
     HF_DATASET_CONFIGS = ["book"]
     TOKENS = 26 * BILLION
-
+    LICENSE = License("partially copyrighted/pirated", commercial_use=False, research_use=False, distribution=False)
     LANGUAGES = ["en"]
 
 
@@ -69,14 +67,24 @@ class RedPajamaArxivDataset(RedPajamaBaseDataset):
     HF_DATASET_CONFIGS = ["arxiv"]
     TOKENS = 28 * BILLION
 
+    LICENSE = License("mixed arxiv licenses", research_use=True)
     LANGUAGES = ["en"]
 
+    HAS_OVERLAP_WITH = ["pes2o"]
+
 
 class RedPajamaStackexchangeDataset(RedPajamaBaseDataset):
     DATASET_ID = "redpajama_stackexchange"
     HF_DATASET_CONFIGS = ["stackexchange"]
     TOKENS = 20 * BILLION
-
+    LICENSE = License(
+        "cc-by-sa 4.0",
+        url="https://creativecommons.org/licenses/by-sa/4.0/",
+        attribution=True,
+        commercial_use=True,
+        research_use=True,
+        sharealike=False,
+    )
     LANGUAGES = ["en"]
 
 
diff --git a/src/lm_datasets/datasets/multilingual/wikimedia.py b/src/lm_datasets/datasets/multilingual/wikimedia.py
index 76ff63a..a037964 100644
--- a/src/lm_datasets/datasets/multilingual/wikimedia.py
+++ b/src/lm_datasets/datasets/multilingual/wikimedia.py
@@ -1,6 +1,6 @@
 import os
 import logging
-from lm_datasets.datasets.base import Availability, BaseDataset, TOKENS_PER_BYTE
+from lm_datasets.datasets.base import Availability, BaseDataset, TOKENS_PER_BYTE, License
 from lm_datasets.utils.wikimedia import parse_and_clean_wikicode
 
 
@@ -220,7 +220,13 @@ class WikimediaBaseDataset(BaseDataset):
     AVAILIBILITY = Availability.DIRECT_DOWNLOAD
     # DUMMY = True
     SOURCE_ID = None  # wiki,wikibooks,wikiquote,wikinews,wikisource,wikivoyage
-
+    LICENSE = License(
+        "Creative Commons Attribution-ShareAlike 4.0 International License",
+        url="https://en.wikipedia.org/wiki/Wikipedia:Text_of_the_Creative_Commons_Attribution-ShareAlike_4.0_International_License",
+        commercial_use=True,
+        research_use=True,
+        attribution=True,
+    )
     download_base_url = "https://dumps.wikimedia.org/"
     dump_date = "20230801"
 
diff --git a/src/lm_datasets/datasets/nl/sonar.py b/src/lm_datasets/datasets/nl/sonar.py
index e8125d7..5453421 100644
--- a/src/lm_datasets/datasets/nl/sonar.py
+++ b/src/lm_datasets/datasets/nl/sonar.py
@@ -3,7 +3,7 @@
 
 from multiprocessing.pool import Pool
 
-from lm_datasets.datasets.base import BaseDataset, Availability, QualityWarning
+from lm_datasets.datasets.base import BaseDataset, Availability, QualityWarning, License
 from lm_datasets.utils import remove_whitespaces_before_punctuation
 
 
@@ -17,7 +17,6 @@
 
 class SonarBaseDataset(BaseDataset):
     SOURCE_ID = "sonar"
-
     TITLE = "SoNaR Corpus NC 1.2"
     HOMEPAGE = "https://taalmaterialen.ivdnt.org/download/tstc-sonar-corpus/"
     DESCRIPTION = (
@@ -28,17 +27,13 @@ class SonarBaseDataset(BaseDataset):
     )
 
     AVAILIBILITY = Availability.SIGNIN_DOWNLOAD
-
     LANGUAGES = ["nl"]
-
     USED_BY = ["GroNLP/bert-base-dutch-cased"]
-
     TOKENS = 0
     # at least 1,030,000 docs
-
+    LICENSE = License("unknown, likely research only or fair use")
     QUALITY_WARNINGS = [QualityWarning.BAD_WHITESPACES]
 
-    # LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/nl/sonar/SONAR500/FoLiA/WR-P-E-A_discussion_lists"]
     # subsets = ["WR-P-E-J_wikipedia"]
 
     # # Excluded
diff --git a/src/lm_datasets/datasets/nl/sonar_new_media.py b/src/lm_datasets/datasets/nl/sonar_new_media.py
index 876b46f..608b65f 100644
--- a/src/lm_datasets/datasets/nl/sonar_new_media.py
+++ b/src/lm_datasets/datasets/nl/sonar_new_media.py
@@ -9,9 +9,8 @@
 logger = logging.getLogger(__name__)
 
 
+# deprecated/dummy -> extraction invalid
 class SonarNewMediaDataset(BaseDataset):
-    """ """
-
     DATASET_ID = "sonar_new_media"
     TITLE = "SoNaR Nieuwe Media Corpus (Version 1.0)"
     HOMEPAGE = "https://taalmaterialen.ivdnt.org/download/tstc-sonar-nieuwe-media-corpus-1/"
@@ -19,14 +18,11 @@ class SonarNewMediaDataset(BaseDataset):
         "The SoNaR New Media Corpus 1.0 contains texts from new media (sms, tweets and ",
         "chat messages) that were collected within the STEVIN-project SoNaR.",
     )
-
     AVAILIBILITY = Availability.SIGNIN_DOWNLOAD
-
     QUALITY_WARNINGS = [QualityWarning.SHORT_TEXT]
-
     LANGUAGES = ["nl"]
-
     TOKENS = 36_000_000
+    DUMMY = True
 
     def get_texts(self):
         from bs4 import BeautifulSoup
diff --git a/src/lm_datasets/datasets/no/maalfrid_2021.py b/src/lm_datasets/datasets/no/maalfrid_2021.py
index 7a0b5ba..61e7d2e 100644
--- a/src/lm_datasets/datasets/no/maalfrid_2021.py
+++ b/src/lm_datasets/datasets/no/maalfrid_2021.py
@@ -1,14 +1,15 @@
-from lm_datasets.datasets.base import BaseDataset, BILLION
+from lm_datasets.datasets.base import BaseDataset, BILLION, Availability, License
 
 
 class Maalfrid2021Dataset(BaseDataset):
     DATASET_ID = "maalfrid_2021"
     TITLE = "Målfrid 2021 - Freely Available Documents from Norwegian State Institutions"
     HOMEPAGE = "https://hdl.handle.net/21.11146/69"
-    AVAILIBILITY = "Yes - it has a direct download link or links"
-
+    AVAILIBILITY = Availability.DIRECT_DOWNLOAD
+    LICENSE = License(
+        "Norwegian Licence for Open Government Data (NLOD)", url="https://data.norge.no/nlod/en/2.0", attribution=True
+    )
     LANGUAGES = ["nb", "nn", "en", "se", "other"]
-
     DESCRIPTION = (
         "This corpus consists of documents from 339 internet domains of Norwegian state institutions and comprises"
         " approximately 4.1 billion tokens in total, which makes it one of the largest freely available resources for"
@@ -22,4 +23,6 @@ class Maalfrid2021Dataset(BaseDataset):
     TOKENS = 4.1 * BILLION
 
     DOWNLOAD_URLS = ["https://www.nb.no/sbfil/tekst/maalfrid_2021/maalfrid_2021.tar.gz"]
-    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/no/maalfrid_2021"]
+
+    DUMMY = True
+    HAS_OVERLAP_WITH = ["norwegian_cc"]
diff --git a/src/lm_datasets/datasets/no/nak.py b/src/lm_datasets/datasets/no/nak.py
index efddb24..af4c3c7 100644
--- a/src/lm_datasets/datasets/no/nak.py
+++ b/src/lm_datasets/datasets/no/nak.py
@@ -1,17 +1,23 @@
 import os
 import tarfile
-from lm_datasets.datasets.base import BaseDataset, BILLION, MILLION
-from lm_datasets.systems import get_path_by_system
+from lm_datasets.datasets.base import BaseDataset, BILLION, MILLION, Availability, License, Genre
+from lm_datasets.utils.systems import get_path_by_system
 
 
 class NAKDataset(BaseDataset):
+    # TODO overlap with norwegian_cc?
     DATASET_ID = "nak"
     TITLE = "Norwegian Newspaper Corpus"
     HOMEPAGE = "https://hdl.handle.net/21.11146/4"
-    AVAILIBILITY = "Yes - it has a direct download link or links"
-
+    AVAILIBILITY = Availability.DIRECT_DOWNLOAD
+    LICENSE = License(
+        "Creative_Commons-BY-NC (CC-BY-NC)",
+        url="https://creativecommons.org/licenses/by-nc/4.0/",
+        commercial_use=False,
+        attribution=True,
+    )
     LANGUAGES = ["nb", "nn"]
-
+    GENRES = [Genre.NEWS]
     DESCRIPTION = (
         "This version of Norwegian Newspaper Corpus contains text from 1998 to 2019. The corpus contains approximately"
         " 1,68 billion words for Norwegian Bokmål, and about 68 million words for Norwegian Nynorsk."
@@ -32,7 +38,6 @@ class NAKDataset(BaseDataset):
         "https://www.nb.no/sbfil/tekst/nak_2013.tar",
         "https://www.nb.no/sbfil/tekst/nak_2012.tar",
     ]
-    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/no/nak"]
 
     def get_dataset_file_paths(self):
         # TODO handle datasets with multiple filles
diff --git a/src/lm_datasets/datasets/no/nbdigital.py b/src/lm_datasets/datasets/no/nbdigital.py
index 9f5edc7..4475f11 100644
--- a/src/lm_datasets/datasets/no/nbdigital.py
+++ b/src/lm_datasets/datasets/no/nbdigital.py
@@ -1,11 +1,12 @@
-from lm_datasets.datasets.base import BaseDataset, GB
+from lm_datasets.datasets.base import BaseDataset, GB, Availability, License
 
 
 class NBDigitalDataset(BaseDataset):
+    # TODO overlap with norwegian_cc?
     DATASET_ID = "nbdigital"
     TITLE = "Public Domain Texts from NBdigital"
     HOMEPAGE = "https://hdl.handle.net/21.11146/34"
-    AVAILIBILITY = "Yes - it has a direct download link or links"
+    AVAILIBILITY = Availability.DIRECT_DOWNLOAD
 
     LANGUAGES = ["nb", "nn"]
 
@@ -15,10 +16,9 @@ class NBDigitalDataset(BaseDataset):
         " public institutions for publically available material)."
     )
     PII = "I have not checked the data source for personally identifiable or sensitive information."
-    LICENSE = "public domain"
+    LICENSE = License("Creative_Commons-ZERO (CC-ZERO)", url="https://creativecommons.org/publicdomain/zero/1.0/")
 
     # Size: 26344 books
     BYTES = 2.7 * GB
 
     DOWNLOAD_URLS = ["https://www.nb.no/sbfil/tekst/20150526_nbdig_txt01.tar.gz"]
-    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/multilingual"]
diff --git a/src/lm_datasets/datasets/no/norwegian_cc.py b/src/lm_datasets/datasets/no/norwegian_cc.py
index b8d22e0..d4b87b1 100644
--- a/src/lm_datasets/datasets/no/norwegian_cc.py
+++ b/src/lm_datasets/datasets/no/norwegian_cc.py
@@ -1,8 +1,19 @@
-from lm_datasets.datasets.base import Availability
+from lm_datasets.datasets.base import Availability, License
 from lm_datasets.datasets.hf_dataset import HFDataset
 
 
 class NorwegianCCBaseDataset(HFDataset):
+    """
+    Licenses:
+
+    government_nb, government_nn, parliament, publicreports, lovdata_cd_*, maalfrid_* 	NLOD 2.0
+    newspapers_ocr, newspapers_pdf, books 	CC0 1.0
+    newspapers_online_nb, newspapers_online_nn 	CC BY-NC 2.0
+    opensubtitles, wikipedia 	CC BY-SA 3.0
+
+    Source: https://huggingface.co/datasets/NbAiLab/NCC#license
+    """
+
     DATASET_ID = "norwegian_cc"
     SOURCE_ID = "norwegian_cc"
 
@@ -18,7 +29,12 @@ class NorwegianCCBaseDataset(HFDataset):
     )
 
     PII = "I have not checked the data source for personally identifiable or sensitive information."
-    LICENSE = "mixed"
+    LICENSE = License(
+        "mixed (NLOD 2.0, CC0 1.0, CC BY-NC 2.0, CC BY-SA 3.0)",
+        url="https://huggingface.co/datasets/NbAiLab/NCC#license",
+        attribution=True,
+        commercial_use=False,
+    )
 
     HF_DATASET_ID = "NbAiLab/NCC"
     HF_DATASET_SPLIT = "train"
@@ -51,7 +67,7 @@ class NorwegianCCNODataset(NorwegianCCBaseDataset):
 
 
 class NorwegianCCNNDataset(NorwegianCCBaseDataset):
-    DATASET_ID = "norwegian_cc_no"
+    DATASET_ID = "norwegian_cc_nn"
     LANGUAGES = ["nn"]
 
     # total: BYTES = 45 * GB
diff --git a/src/lm_datasets/datasets/no/parlamint.py b/src/lm_datasets/datasets/no/parlamint.py
index cbc9f55..214a2d7 100644
--- a/src/lm_datasets/datasets/no/parlamint.py
+++ b/src/lm_datasets/datasets/no/parlamint.py
@@ -20,6 +20,5 @@ class ParlaMintDataset(BaseDataset):
     # Size: 975 million tokens
     TOKENS = 975 * MILLION
 
-    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/no/parlamint"]
-
     DUMMY = True
+    HAS_OVERLAP_WITH = ["norwegian_cc"]
diff --git a/src/lm_datasets/datasets/no/parliamentary_proceedings.py b/src/lm_datasets/datasets/no/parliamentary_proceedings.py
index 8158241..a4d9a95 100644
--- a/src/lm_datasets/datasets/no/parliamentary_proceedings.py
+++ b/src/lm_datasets/datasets/no/parliamentary_proceedings.py
@@ -20,6 +20,5 @@ class ParliamentaryProceedingsDataset(BaseDataset):
     # Size: 1.5 billion tokens
     TOKENS = 1.5 * BILLION
 
-    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/no/parliamentary_proceedings"]
-
     DUMMY = True
+    HAS_OVERLAP_WITH = ["norwegian_cc"]
diff --git a/src/lm_datasets/datasets/no/sakspapir_nno.py b/src/lm_datasets/datasets/no/sakspapir_nno.py
index 64960b1..4dd38d6 100644
--- a/src/lm_datasets/datasets/no/sakspapir_nno.py
+++ b/src/lm_datasets/datasets/no/sakspapir_nno.py
@@ -1,11 +1,11 @@
-from lm_datasets.datasets.base import BaseDataset, MILLION
+from lm_datasets.datasets.base import BaseDataset, MILLION, Availability, License
 
 
 class SakspapirNNODataset(BaseDataset):
     DATASET_ID = "sakspapir_nno"
     TITLE = "Legal Documents from Norwegian Nynorsk Municipialities"
     HOMEPAGE = "https://hdl.handle.net/21.11146/60"
-    AVAILIBILITY = "Yes - it has a direct download link or links"
+    AVAILIBILITY = Availability.DIRECT_DOWNLOAD
 
     LANGUAGES = ["nn", "nb"]
 
@@ -15,12 +15,11 @@ class SakspapirNNODataset(BaseDataset):
         " Municipalities (see lnk.no)."
     )
     PII = "I have not checked the data source for personally identifiable or sensitive information."
-    LICENSE = "open license"
+    LICENSE = License("Creative_Commons-ZERO (CC-ZERO)", url="https://creativecommons.org/publicdomain/zero/1.0/")
 
     # Size: 127 million tokens
     TOKENS = 127 * MILLION
 
     DOWNLOAD_URLS = ["https://www.nb.no/sbfil/tekst/sakspapir_nno/sakspapir_nno_01.tar.gz"]
-    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/no/sakspapir_nno"]
 
     DUMMY = True
diff --git a/src/lm_datasets/datasets/parquet_dataset.py b/src/lm_datasets/datasets/parquet_dataset.py
new file mode 100644
index 0000000..218c534
--- /dev/null
+++ b/src/lm_datasets/datasets/parquet_dataset.py
@@ -0,0 +1,17 @@
+from lm_datasets.datasets.base import BaseDataset
+
+
+class ParquetDataset(BaseDataset):
+    SINGLE_OUTPUT_FILE = False
+
+    def get_texts(self):
+        if self.output_format == "parquet":
+            yield from self.generate_texts_from_output(shuffled=False)
+        else:
+            raise ValueError("Dataset is already in parquet format; cannot generate texts for a different format!")
+
+    def extract_plaintext(self):
+        if self.output_format == "parquet":
+            raise ValueError("Dataset is already in parquet format; no text extraction needed!")
+        else:
+            super().extract_plaintext()
diff --git a/src/lm_datasets/datasets/pl/pl_nkjp.py b/src/lm_datasets/datasets/pl/pl_nkjp.py
index 6a421d8..1529879 100644
--- a/src/lm_datasets/datasets/pl/pl_nkjp.py
+++ b/src/lm_datasets/datasets/pl/pl_nkjp.py
@@ -4,7 +4,7 @@
 import xml.etree.cElementTree as ET
 
 from tqdm.auto import tqdm
-from lm_datasets.datasets.base import MILLION, BaseDataset, QualityWarning
+from lm_datasets.datasets.base import MILLION, BaseDataset, QualityWarning, Availability, License
 
 logger = logging.getLogger(__name__)
 
@@ -40,7 +40,7 @@ class NKJPPodkorpusMilionowyDataset(BaseDataset):
     DATASET_ID = "pl_nkjp"
     TITLE = "NKJP-PodkorpusMilionowy-1.2 (National Corpus of Polish)"
     HOMEPAGE = "http://clip.ipipan.waw.pl/NationalCorpusOfPolish"
-
+    LICENSE = License("CC-BY", attribution=True)
     LANGUAGES = ["pl"]
     DOWNLOAD_URLS = [
         (
@@ -49,7 +49,7 @@ class NKJPPodkorpusMilionowyDataset(BaseDataset):
         )
     ]
     TOKENS = 1 * MILLION
-    LICENSE = "CC-BY"
+    AVAILIBILITY = Availability.DIRECT_DOWNLOAD
     QUALITY_WARNINGS = [QualityWarning.BAD_LINEBREAKS]
 
     def decompress(self):
diff --git a/src/lm_datasets/datasets/pl/pl_parliamentary_corpus.py b/src/lm_datasets/datasets/pl/pl_parliamentary_corpus.py
index 07ca4d2..457da55 100644
--- a/src/lm_datasets/datasets/pl/pl_parliamentary_corpus.py
+++ b/src/lm_datasets/datasets/pl/pl_parliamentary_corpus.py
@@ -4,7 +4,7 @@
 
 import xml.etree.cElementTree as ET
 
-from lm_datasets.datasets.base import BaseDataset, Genre
+from lm_datasets.datasets.base import BaseDataset, Genre, License, Availability
 
 
 logger = logging.getLogger(__name__)
@@ -42,9 +42,9 @@ class PLParliamentaryCorpusDataset(BaseDataset):
         " Texts in the PPC corpus cover the period of a hundred years from 1919 to 2019."
     )
     HOMEPAGE = "http://clip.ipipan.waw.pl/PPC"
-
+    AVAILIBILITY = Availability.DIRECT_DOWNLOAD
     LANGUAGES = ["pl"]
-    LICENSE = "CC-BY"
+    LICENSE = License("CC-BY")
 
     GENRES = [Genre.GOVERNMENT]
 
diff --git a/src/lm_datasets/datasets/pt/brwac.py b/src/lm_datasets/datasets/pt/brwac.py
index 9a80ae5..1370da7 100644
--- a/src/lm_datasets/datasets/pt/brwac.py
+++ b/src/lm_datasets/datasets/pt/brwac.py
@@ -1,4 +1,4 @@
-from lm_datasets.datasets.base import Availability, BILLION
+from lm_datasets.datasets.base import Availability, BILLION, License
 from lm_datasets.datasets.hf_dataset import HFDataset
 
 
@@ -29,17 +29,28 @@ class BrWacDataset(HFDataset):
         "2.68 billion tokens and 5.79 million types. Please note that this resource is available"
         "solely for academic research purposes, and you agreed not to use it for any commercial applications."
     )
+    CITATION = """@inproceedings{wagner2018brwac,
+  title={The brwac corpus: A new open resource for brazilian portuguese},
+  author={Wagner Filho, Jorge A and Wilkens, Rodrigo and Idiart, Marco and Villavicencio, Aline},
+  booktitle={Proceedings of the Eleventh International Conference on Language Resources and Evaluation (LREC 2018)},
+  year={2018}
+}
+"""  # noqa
+    LICENSE = License("research-only", commercial_use=False, research_use=True)
 
     TOKENS = 2.6 * BILLION
 
     HF_DATASET_ID = "brwac"
     HF_DATASET_SPLIT = "train"
-    HF_DATA_DIR = "/netscratch/ortiz/corpora/ELE/pt/"
+    HF_DATA_DIR = "/netscratch/ortiz/corpora/ELE/pt/brwac/"
 
     text_column_name = "text"
-    remove_columns = ["doc_id", "uri"]
-
-    title_column_name = "title"
+    keep_columns = True
+    streaming = True
+    #  title_column_name = "title" ## do not use auto-concatenate
 
     def get_text_from_item(self, item) -> str:
-        return item["title"] + self.title_delimiter + (self.paragraph_delimiter.join(item["text"]["paragraphs"]))
+        paragraphs = [self.sentence_delimiter.join(sentences) for sentences in item["text"]["paragraphs"]]
+        text = item["title"] + self.title_delimiter + (self.paragraph_delimiter.join(paragraphs))
+
+        return text
diff --git a/src/lm_datasets/datasets/pt/parlamento_pt.py b/src/lm_datasets/datasets/pt/parlamento_pt.py
index 795d453..82a7e1e 100644
--- a/src/lm_datasets/datasets/pt/parlamento_pt.py
+++ b/src/lm_datasets/datasets/pt/parlamento_pt.py
@@ -1,4 +1,4 @@
-from lm_datasets.datasets.base import Availability, GB
+from lm_datasets.datasets.base import Availability, GB, License
 from lm_datasets.datasets.hf_dataset import HFDataset
 
 import logging
@@ -11,7 +11,10 @@ class ParlamentoPtDataset(HFDataset):
     TITLE = "ParlamentoPT"
     HOMEPAGE = "https://huggingface.co/datasets/PORTULAN/parlamento-pt"
     AVAILIBILITY = Availability.DIRECT_DOWNLOAD
-
+    LICENSE = License(
+        "open data (Portuguese Parliament portal policy)",
+        url="https://www.parlamento.pt/Cidadania/Paginas/DadosAbertos.aspx",
+    )
     LANGUAGES = ["pt"]
 
     DESCRIPTION = """
diff --git a/src/lm_datasets/datasets/ro/marcell_legislative_subcorpus_v2.py b/src/lm_datasets/datasets/ro/marcell_legislative_subcorpus_v2.py
index f55164e..ac19b7c 100644
--- a/src/lm_datasets/datasets/ro/marcell_legislative_subcorpus_v2.py
+++ b/src/lm_datasets/datasets/ro/marcell_legislative_subcorpus_v2.py
@@ -1,7 +1,7 @@
 import logging
 import zipfile
 from tqdm.auto import tqdm
-from lm_datasets.datasets.base import BaseDataset, MILLION, Genre
+from lm_datasets.datasets.base import BaseDataset, MILLION, Genre, Availability, License
 
 logger = logging.getLogger(__name__)
 
@@ -10,8 +10,15 @@ class MarcellLegislativeSubcorpusV2Dataset(BaseDataset):
     DATASET_ID = "marcell_legislative_subcorpus_v2"
     TITLE = "MARCELL Romanian legislative subcorpus v2"
     HOMEPAGE = "https://elrc-share.eu/repository/browse/marcell-romanian-legislative-subcorpus-v2/2da548428b9d11eb9c1a00155d026706ce94a6b59ffc4b0e9fb5cd9cebe6889e/"  # noqa
-    AVAILIBILITY = "Yes - it has a direct download link or links"
-
+    AVAILIBILITY = Availability.DIRECT_DOWNLOAD
+    LICENSE = License(
+        "CC0-1.0",
+        url="https://elrc-share.eu/static/metashare/licences/CC0-1.0.pdf",
+        commercial_use=True,
+        attribution=False,
+        research_use=False,
+        sharealike=False,
+    )
     LANGUAGES = ["ro"]
     GENRES = [Genre.LEGAL]
 
diff --git a/src/lm_datasets/datasets/sk/sk_court_decisions.py b/src/lm_datasets/datasets/sk/sk_court_decisions.py
index 2c4194d..f94a856 100644
--- a/src/lm_datasets/datasets/sk/sk_court_decisions.py
+++ b/src/lm_datasets/datasets/sk/sk_court_decisions.py
@@ -1,19 +1,21 @@
 from typing import List
-from lm_datasets.datasets.base import BaseDataset, Genre
+from lm_datasets.datasets.base import BaseDataset, Genre, License, Availability
 
 
 class SKCourtDecisionsDataset(BaseDataset):
     DATASET_ID = "sk_court_decisions"
     TITLE = "od-justice 2.0"
-    HOMEPAGE = "https://www.juls.savba.sk/data/od-justice/od-justice-2.0.ver.xz"
-    AVAILIBILITY = "Yes - it has a direct download link or links"
-
+    HOMEPAGE = "https://www.juls.savba.sk/justicecorp.html"
+    DOWNLOAD_URLS = ["https://www.juls.savba.sk/data/od-justice/od-justice-2.0.ver.xz"]
+    AVAILIBILITY = Availability.DIRECT_DOWNLOAD
+    LICENSE = License("open data", url="https://obcan.justice.sk/opendata")
     LANGUAGES = ["sk"]
     GENRES = [Genre.LEGAL]
-
-    DESCRIPTION = "Slovak court decisions"
+    DESCRIPTION = (
+        "Slovak court decisions. The corpus is based on data made available by the "
+        "Ministry of Justice of the Slovak Republic."
+    )
     PII = "No"
-    LICENSE = "public domain"
 
     TOKENS = 10_618_105_036
 
diff --git a/src/lm_datasets/datasets/sk/sk_laws.py b/src/lm_datasets/datasets/sk/sk_laws.py
index a785d43..8fb6c69 100644
--- a/src/lm_datasets/datasets/sk/sk_laws.py
+++ b/src/lm_datasets/datasets/sk/sk_laws.py
@@ -1,11 +1,11 @@
 from typing import List
-from lm_datasets.datasets.base import BaseDataset, Genre
+from lm_datasets.datasets.base import BaseDataset, Genre, License
 
 
 class SKLawsDataset(BaseDataset):
     DATASET_ID = "sk_laws"
     TITLE = "Korpus slovenských právnych predpisov v1.9"
-    HOMEPAGE = "https://www.juls.savba.sk/data/marcell/legal-sk-20220322-1.9.ver.xz"
+    DOWNLOAD_URLS = ["https://www.juls.savba.sk/data/marcell/legal-sk-20220322-1.9.ver.xz"]
     AVAILIBILITY = "Yes - it has a direct download link or links"
 
     LANGUAGES = ["sk"]
@@ -13,7 +13,7 @@ class SKLawsDataset(BaseDataset):
 
     DESCRIPTION = "Slovak body of laws (1955-2022)"
     PII = "No"
-    LICENSE = "public domain"
+    LICENSE = License("public domain")
 
     TOKENS = 45_000_000
 
diff --git a/src/lm_datasets/datasets/sl/academic_slovene_kas.py b/src/lm_datasets/datasets/sl/academic_slovene_kas.py
index 9531e3a..5923700 100644
--- a/src/lm_datasets/datasets/sl/academic_slovene_kas.py
+++ b/src/lm_datasets/datasets/sl/academic_slovene_kas.py
@@ -1,14 +1,7 @@
-"""
-https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1448/kas.json.tar.gz
-
-academic_slovene_kas
-
-"""
-
 import json
 import logging
 
-from lm_datasets.datasets.base import BaseDataset, Availability
+from lm_datasets.datasets.base import BaseDataset, Availability, License
 
 from tqdm.auto import tqdm
 from smart_open import open
@@ -28,11 +21,16 @@ class AcademicSloveneKASDataset(BaseDataset):
     TITLE = "Corpus of academic Slovene KAS 2.0"
     HOMEPAGE = "https://www.clarin.si/repository/xmlui/handle/11356/1448"
     AVAILIBILITY = Availability.DIRECT_DOWNLOAD
-
+    LICENSE = License(
+        "CLARIN.SI Licence ACA ID-BY-NC-INF-NORED 1.0",
+        url="https://clarin.si/repository/xmlui/page/licence-aca-id-by-nc-inf-nored-1.0",
+        attribution=True,
+        commercial_use=False,
+        distribution=False,
+        informing=True,
+    )
     LANGUAGES = ["sl"]
 
-    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/sl/academic_slovene_kas"]
-
     DOWNLOAD_URLS = ["https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1448/kas.json.tar.gz"]
 
     TOKENS = 1_496_079_001
diff --git a/src/lm_datasets/datasets/sl/cc_gigafida.py b/src/lm_datasets/datasets/sl/cc_gigafida.py
index 2889ef9..465b9d0 100644
--- a/src/lm_datasets/datasets/sl/cc_gigafida.py
+++ b/src/lm_datasets/datasets/sl/cc_gigafida.py
@@ -1,14 +1,18 @@
-from lm_datasets.datasets.base import BaseDataset, Availability
+from lm_datasets.datasets.base import BaseDataset, Availability, License
 
 
 class CCGigaFidaDataset(BaseDataset):
-    """ """
-
     DATASET_ID = "cc_gigafida"
     TITLE = "Written corpus ccGigafida 1.0"
     HOMEPAGE = "https://www.clarin.si/repository/xmlui/handle/11356/1035"
     AVAILIBILITY = Availability.DIRECT_DOWNLOAD
-
+    LICENSE = License(
+        "Creative Commons - Attribution-NonCommercial-ShareAlike 4.0 International (CC BY-NC-SA 4.0)",
+        url="https://creativecommons.org/licenses/by-nc-sa/4.0/",
+        commercial_use=False,
+        sharealike=True,
+        attribution=True,
+    )
     LANGUAGES = ["sl"]
 
     DOWNLOAD_URLS = ["https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1035/ccGigafida-text.zip"]
diff --git a/src/lm_datasets/datasets/sl/slwac_web.py b/src/lm_datasets/datasets/sl/slwac_web.py
index 0e51e9e..2ff1837 100644
--- a/src/lm_datasets/datasets/sl/slwac_web.py
+++ b/src/lm_datasets/datasets/sl/slwac_web.py
@@ -1,5 +1,5 @@
 from typing import List
-from lm_datasets.datasets.base import BILLION, BaseDataset, Availability
+from lm_datasets.datasets.base import BILLION, BaseDataset, Availability, License
 from lm_datasets.utils import generate_texts_from_tab_columns_in_xml
 from smart_open import open
 
@@ -9,7 +9,9 @@ class SLWaCWebDataset(BaseDataset):
     HOMEPAGE = "http://nlp.ffzg.hr/resources/corpora/slwac/"
     DOWNLOAD_URLS = ["http://nlp.ffzg.hr/data/corpora/slwac2.0.gz"]
     LANGUAGES = ["sl"]
-
+    LICENSE = License(
+        "CC-BY-SA license", attribution=True, sharealike=True, url="http://creativecommons.org/licenses/by-sa/4.0/"
+    )
     TITLE = "slWaC web corpus"
     DESCRIPTION = (
         "slWaC is a web corpus collected from the .si top-level domain in 2011 and 2014. The corpus is"
diff --git a/src/lm_datasets/datasets/sr/srpkor.py b/src/lm_datasets/datasets/sr/srpkor.py
index 08c9254..419e08a 100644
--- a/src/lm_datasets/datasets/sr/srpkor.py
+++ b/src/lm_datasets/datasets/sr/srpkor.py
@@ -1,6 +1,6 @@
 import json
 from typing import Iterable
-from lm_datasets.datasets.base import BaseDataset, Availability, QualityWarning
+from lm_datasets.datasets.base import BaseDataset, Availability, QualityWarning, License
 import zipfile
 import logging
 
@@ -13,18 +13,21 @@ class SrpKorDataset(BaseDataset):
 
     AVAILIBILITY = Availability.ON_REQUEST
 
+    LICENSE = License("Do not redistribute, DFKI has permission to use it for pre-training LLMs")
+
     LANGUAGES = ["sr"]
     QUALITY_WARNINGS = [QualityWarning.SHORT_TEXT]
 
     def get_texts(self) -> Iterable[str]:
         # read from OpenGptx-JeRTeh.zip
         zip_fp = self.get_dataset_file_paths(needed_suffix=".zip", single_file=True)
-        logger.info(f"Extracting from {zip_fp}")
+        logger.info("Extracting from %s", zip_fp)
+
         with zipfile.ZipFile(zip_fp) as zf:
             member_fns = zf.namelist()
             for fn in member_fns:
                 if fn.endswith(".json"):
-                    logger.info(f"Reading from {fn}")
+                    logger.info("Reading from %s", fn)
 
                     with zf.open(fn) as member_f:
                         obj = json.load(member_f)
diff --git a/src/lm_datasets/datasets/sv/sv_gigaword.py b/src/lm_datasets/datasets/sv/sv_gigaword.py
index 29747ec..c09f7fd 100644
--- a/src/lm_datasets/datasets/sv/sv_gigaword.py
+++ b/src/lm_datasets/datasets/sv/sv_gigaword.py
@@ -5,7 +5,7 @@
 
 import tarfile
 
-from lm_datasets.datasets.base import BaseDataset, Availability, BILLION
+from lm_datasets.datasets.base import BaseDataset, Availability, BILLION, QualityWarning, License
 
 
 logger = logging.getLogger(__name__)
@@ -25,6 +25,17 @@ class SVGigawordDataset(BaseDataset):
 
     AVAILIBILITY = Availability.DIRECT_DOWNLOAD
 
+    LICENSE = License(
+        "BY-SA 4.0",
+        attribution=True,
+        sharealike=False,
+        commercial_use=True,
+        research_use=True,
+        url="http://creativecommons.org/licenses/by-sa/4.0/",
+    )
+
+    QUALITY_WARNINGS = [QualityWarning.BAD_WHITESPACES]
+
     DOWNLOAD_URLS = [
         "https://spraakbanken.gu.se/lb/resurser/meningsmangder/gigaword-1950-59.tar",
         "https://spraakbanken.gu.se/lb/resurser/meningsmangder/gigaword-1960-69.tar",
@@ -47,7 +58,7 @@ def get_texts(self):
         # read from tar files
         tar_file_paths = self.get_dataset_file_paths(needed_suffix=".tar")
         for tar_fp in tar_file_paths:
-            logger.info(f"Extracting from {tar_fp}")
+            logger.info("Extracting from %s", tar_fp)
 
             with tarfile.open(tar_fp) as tar_f:
                 members = [m for m in tar_f.getmembers() if m.name.endswith(".xml.bz2")]
@@ -68,6 +79,15 @@ def get_texts(self):
                             text += el.replace("\n", " ")
                         text = text.replace("  ", " ").strip()
 
-                        # print(text)
-
                         yield text
+
+                        # Free up memory as describe in:
+                        # https://stackoverflow.com/questions/12160418/why-is-lxml-etree-iterparse-eating-up-all-my-memory  # noqa
+                        element.clear()
+
+                        # Also eliminate now-empty references from the root node to elem
+                        for ancestor in element.xpath("ancestor-or-self::*"):
+                            while ancestor.getprevious() is not None:
+                                del ancestor.getparent()[0]
+
+                        # print(text)
diff --git a/src/lm_datasets/datasets/uk/uk_laws.py b/src/lm_datasets/datasets/uk/uk_laws.py
index d0b1abc..8d9e7c8 100644
--- a/src/lm_datasets/datasets/uk/uk_laws.py
+++ b/src/lm_datasets/datasets/uk/uk_laws.py
@@ -1,7 +1,7 @@
 import logging
 from smart_open import open
 
-from lm_datasets.datasets.base import Availability, BaseDataset
+from lm_datasets.datasets.base import Availability, BaseDataset, License
 
 
 logger = logging.getLogger(__name__)
@@ -17,6 +17,8 @@ class UKLawsDataset(BaseDataset):
 
     AVAILIBILITY = Availability.DIRECT_DOWNLOAD
 
+    LICENSE = License("Unknown, likely public domain")
+
     DOWNLOAD_URLS = ["https://lang.org.ua/static/downloads/corpora/laws.txt.tokenized.bz2"]
 
     def is_downloaded(self):
@@ -28,7 +30,7 @@ def get_texts(self):
 
         file_path = self.get_dataset_file_paths(needed_suffix=".tokenized.bz2", single_file=True)
 
-        logger.info(f"Reading from {file_path}")
+        logger.info("Reading from %s", file_path)
 
         with open(file_path) as f:
             text = ""
diff --git a/src/lm_datasets/extract_plaintext.py b/src/lm_datasets/extract_plaintext.py
deleted file mode 100644
index 6760868..0000000
--- a/src/lm_datasets/extract_plaintext.py
+++ /dev/null
@@ -1,143 +0,0 @@
-import argparse
-
-import logging
-
-from lm_datasets.utils import get_auto_workers, get_bytes_from_int_or_string
-
-from .datasets.dataset_registry import get_registered_dataset_classes
-from .datasets.base import BaseDataset
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("datasets", help="Name of datasets to process (comma separated)")
-    parser.add_argument(
-        "output_dir",
-        help="Output is saved in this directory (<language code>/<source name>.<jsonl/parquet>)",
-    )
-    parser.add_argument("--override", action="store_true", help="Override existing output files")
-    parser.add_argument(
-        "--ignore_errors",
-        action="store_true",
-        help="Ignore dataset-level errors (use when processing multiple datasets)",
-    )
-    parser.add_argument("--json_ensure_ascii", action="store_true", help="Escape non-ASCII characters in JSON output")
-    parser.add_argument("--output_format", default="jsonl", type=str, help="Output format (jsonl,parquet)")
-    parser.add_argument(
-        "--output_compression",
-        default=None,
-        type=str,
-        help="""Compression of output (jsonl: "gzip"; parquet: "NONE", "SNAPPY", "GZIP", "BROTLI", "LZ4", "ZSTD")""",
-    )
-    parser.add_argument(
-        "--output_batch_size",
-        default=1000,
-        type=int,
-        help="""Write batch size; smaller batch size = more accurate splitts but slower""",
-    )
-    parser.add_argument(
-        "--max_output_chunk_uncompressed_bytes",
-        default="10GB",
-        type=str,
-        help="Chunks are splitted if they exceed this byte count (<n>, <n>KB, <n>MB, or <n>GB)",
-    )
-    parser.add_argument(
-        "--workers",
-        default=1,
-        type=int,
-        help="Number of workers for parallel processing",
-    )
-    parser.add_argument(
-        "--raw_datasets_dir",
-        default=None,
-        type=str,
-        help="Dataset files are read from this directory",
-    )
-    parser.add_argument(
-        "--extra_dataset_registries",
-        default=None,
-        type=str,
-        help="List of Python packages to load dataset registries",
-    )
-    parser.add_argument(
-        "--log_file",
-        default=None,
-        type=str,
-        help="Log file is saved at this path",
-    )
-    parser.add_argument("--verbose", action="store_true", help="Enable verbose logging (log level = debug)")
-    parser.add_argument("--limit", default=0, type=int, help="Limit dataset size (for debugging)")
-    parser.add_argument(
-        "--skip_items",
-        default=0,
-        type=int,
-        help="Skip N items (depending on dataset: directories, subsets, files, documents) (for debugging)",
-    )
-    parser.add_argument("--hf_auth_token", default=None, type=str, help="HuggingFace auth token")
-    args = parser.parse_args()
-
-    log_handlers = [logging.StreamHandler()]
-
-    if args.log_file:
-        log_handlers.append(logging.FileHandler(args.log_file))
-
-    logging.basicConfig(
-        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
-        datefmt="%Y-%m-%d %H:%M:%S",
-        level=logging.DEBUG if args.verbose else logging.INFO,
-        handlers=log_handlers,
-    )
-    logger = logging.getLogger(__name__)
-
-    id_to_dataset_class = {cls.DATASET_ID: cls for cls in get_registered_dataset_classes(args.extra_dataset_registries)}
-
-    datasets_list = args.datasets.split(",")
-
-    if len(datasets_list) == 1 and datasets_list[0] == "all":
-        # Get list of all non-dummy datasets
-        datasets_list = id_to_dataset_class.keys()
-
-    # Iterate over datasets
-    for i, dataset_id in enumerate(datasets_list, 1):
-        if dataset_id in id_to_dataset_class:
-            dataset_cls = id_to_dataset_class[dataset_id]
-        else:
-            raise ValueError(f"Unknown dataset ID: {dataset_id} (available: {id_to_dataset_class.keys()})")
-
-        logger.info(f"Dataset ID: {dataset_id} ({i} / {len(datasets_list)})")
-
-        try:
-            dataset: BaseDataset = dataset_cls(
-                raw_datasets_dir=args.raw_datasets_dir,
-                output_dir=args.output_dir,
-                workers=get_auto_workers(args.workers),
-                limit=args.limit,
-                override_output=args.override,
-                output_format=args.output_format,
-                output_compression=args.output_compression,
-                output_batch_size=args.output_batch_size,
-                json_ensure_ascii=args.json_ensure_ascii,
-                skip_items=args.skip_items,
-                max_output_chunk_uncompressed_bytes=get_bytes_from_int_or_string(
-                    args.max_output_chunk_uncompressed_bytes
-                ),
-            )
-
-            if dataset.is_dummy():
-                logger.warning(f"Skipping {dataset_id} (cannot extract from dummy datasets)")
-                continue
-
-            try:
-                dataset.extract_plaintext()
-            except FileExistsError as e:
-                logger.error(f"Cannot extract text from {dataset_id}: {e}")
-        except KeyboardInterrupt as e:
-            logger.error(f"Stopping... {e}")
-            break
-        except BaseException as e:
-            if args.ignore_errors:
-                logger.error(f"Unexpected error occured with {dataset_id}: {e}")
-            else:
-                raise e
-
-    logger.info("Done")
diff --git a/src/lm_datasets/extract_text.py b/src/lm_datasets/extract_text.py
new file mode 100644
index 0000000..e898c00
--- /dev/null
+++ b/src/lm_datasets/extract_text.py
@@ -0,0 +1,77 @@
+import argparse
+
+import logging
+
+from lm_datasets.utils import get_auto_workers, get_bytes_from_int_or_string
+from lm_datasets.utils.settings import DEFAULT_MIN_TEXT_LENGTH
+
+from .datasets.dataset_registry import (
+    get_dataset_class_by_id,
+    get_registered_dataset_classes,
+    get_registered_dataset_ids,
+)
+from .datasets.base import BaseDataset
+from .utils.config import Config, get_common_argparser, parse_args_and_get_config
+
+
+def extract_text(config: Config):
+    logger = config.init_logger(__name__)
+
+    datasets_list = config.datasets.split(",")
+
+    if len(datasets_list) == 1:
+        if datasets_list[0] == "all":
+            # Get list of all regsitered datasets
+            datasets_list = get_registered_dataset_ids(config.extra_dataset_registries)
+
+        elif datasets_list[0] == "all_from_source":
+            # Get registered datasets based on source
+            if config.source_id is None:
+                raise ValueError("The argument or config `source_id` must be set.")
+
+            datasets_list = get_registered_dataset_ids(
+                config.extra_dataset_registries, needed_source_id=config.source_id
+            )
+
+    # Iterate over datasets
+    for i, dataset_id in enumerate(datasets_list, 1):
+        logger.info(f"Dataset ID: {dataset_id} ({i} / {len(datasets_list)})")
+
+        try:
+            dataset_cls = get_dataset_class_by_id(dataset_id, config.extra_dataset_registries)
+            dataset: BaseDataset = dataset_cls(
+                raw_datasets_dir=config.raw_datasets_dir,
+                output_dir=config.output_dir,
+                workers=get_auto_workers(config.workers),
+                limit=config.limit,
+                override_output=config.override,
+                output_format=config.output_format,
+                output_compression=config.output_compression,
+                output_batch_size=config.output_batch_size,
+                json_ensure_ascii=config.json_ensure_ascii,
+                skip_items=config.skip_items,
+                max_output_chunk_uncompressed_bytes=get_bytes_from_int_or_string(
+                    config.max_output_chunk_uncompressed_bytes
+                ),
+                min_length=config.min_text_length,
+                config=config,
+            )
+
+            if dataset.is_dummy():
+                logger.warning(f"Skipping {dataset_id} (cannot extract from dummy datasets)")
+                continue
+
+            try:
+                dataset.extract_plaintext()
+            except FileExistsError as e:
+                logger.error(f"Cannot extract text from {dataset_id}: {e}")
+        except KeyboardInterrupt as e:
+            logger.error(f"Stopping... {e}")
+            break
+        except BaseException as e:
+            if config.ignore_errors:
+                logger.error(f"Unexpected error occured with {dataset_id}: {e}")
+            else:
+                raise e
+
+    logger.info("Done")
diff --git a/src/lm_datasets/hf_tokenize_parquet_dataset.py b/src/lm_datasets/hf_tokenize_parquet_dataset.py
new file mode 100644
index 0000000..26d0a5d
--- /dev/null
+++ b/src/lm_datasets/hf_tokenize_parquet_dataset.py
@@ -0,0 +1,367 @@
+from dataclasses import dataclass, field
+from itertools import chain
+import logging
+import os
+from pathlib import Path
+import sys
+from typing import Optional
+from transformers import AutoTokenizer
+from transformers import (
+    AutoTokenizer,
+    HfArgumentParser,
+)
+from tqdm.auto import tqdm
+import signal
+
+import multiprocessing
+from pathlib import Path
+from typing import Iterable, List
+
+from itertools import chain
+
+import pyarrow as pa
+from pyarrow.dataset import dataset as pa_dataset
+
+from pyarrow.dataset import write_dataset, ParquetFileFormat
+from pyarrow import RecordBatch
+
+from transformers import AutoTokenizer
+from tqdm.auto import tqdm
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ScriptArguments:
+    """
+    Arguments for this script
+    """
+
+    tokenizer_name_or_path: str = field(
+        metadata={"help": ("Tokenizer name or path")},
+    )
+    dataset_dir: str = field(
+        metadata={"help": ("Path to input dataset (where *.parquet files are located, see dataset_glob)")},
+    )
+    output_path_prefix: str = field(
+        metadata={
+            "help": (
+                "Tokenized dataset is saved to this file path prefix on the disk (to each file name `.part-*-of-*.parquet` is appended)"
+            )
+        },
+    )
+    dataset_glob: str = field(
+        default="*.parquet",
+        metadata={"help": ("Glob for selecting files from datset dir")},
+    )
+    # output_format: Literal["jsonl", "hf"] = field(
+    #     default="jsonl",
+    #     metadata={"help": ("Output format")},
+    # )
+    # composed_dataset_split: Literal["train", "validation"] = field(
+    #     default="train",
+    #     metadata={"help": ("Split of composed dataset")},
+    # )
+    text_column_name: str = field(
+        default="text",
+        metadata={"help": ("Name of text column in input dataaset")},
+    )
+    max_seq_length: int = field(
+        default=None,
+        metadata={"help": ("Max. input sequence length (default = from tokenizer/model settings)")},
+    )
+    limit: int = field(
+        default=0,
+        metadata={"help": ("Limits number of input examples (only for debugg; 0 = no limit)")},
+    )
+    batch_size: int = field(
+        default=10_000,
+        metadata={"help": ("Batch size of reading/processing/writting data")},
+    )
+    output_max_rows_per_file: int = field(
+        default=1024 * 1024,
+        metadata={"help": ("Output max_rows_per_file")},
+    )
+    output_max_rows_per_group: int = field(
+        default=10 * 1024,
+        metadata={"help": ("Output max_rows_per_group:")},
+    )
+    num_proc: int = field(
+        default=None,
+        metadata={"help": ("Number of parallel processes (default all cpus)")},
+    )
+    # print_progress: int = field(
+    #     default=10_000,
+    #     metadata={"help": ("Print progress after every N examples")},
+    # )
+    use_fast_tokenizer: bool = field(
+        default=False,
+        metadata={"help": ("Enable fast tokenizer (produces sometimes different results)")},
+    )
+    do_group: bool = field(
+        default=False,
+        metadata={"help": ("Group tokenized samples into same-length samples based on `max_seq_length`")},
+    )
+    # skip_tokenization: bool = field(
+    #     default=False,
+    #     metadata={
+    #         "help": (
+    #             "Skip tokenization step (apply only grouping)"
+    #         )
+    #     },
+    # )
+    override: bool = field(
+        default=False,
+        metadata={"help": ("Override existing output")},
+    )
+    return_special_tokens_mask: bool = field(
+        default=False,
+        metadata={"help": ("return_special_tokens_mask (for MLM data)")},
+    )
+    return_token_type_ids: bool = field(
+        default=False,
+        metadata={"help": ("return_token_type_ids (for MLM data)")},
+    )
+    verbose: bool = field(
+        default=False,
+        metadata={"help": ("Verbose logging (debug)")},
+    )
+
+
+def generate_texts(
+    pyarrow_dataset,
+    reader_batch_size: int,
+    print_progress: bool = True,
+    row_limit: Optional[int] = None,
+    text_column_name: str = "text",
+) -> Iterable[List[str]]:
+    """
+    Reads PyArrow dataset in batches and generates texts
+    """
+
+    batch_iter = pyarrow_dataset.to_batches(columns=[text_column_name], batch_size=reader_batch_size)
+    max_batches = round(row_limit / reader_batch_size) if row_limit is not None and row_limit > 0 else None
+
+    if print_progress:
+        # get total number of batches
+        total_rows = pyarrow_dataset.count_rows()
+        logger.info("Total input dataset rows: %i", total_rows)
+
+        total_batches = round(total_rows / reader_batch_size)
+
+        if max_batches is not None and max_batches > 0:
+            total_batches = min(total_batches, max_batches)
+
+        batch_iter = tqdm(batch_iter, total=total_batches, desc="Iterating over input")
+
+    for batch_i, batch in enumerate(batch_iter):
+        texts = batch[0].tolist()
+
+        logger.debug("texts loaded %i", len(texts))
+
+        yield texts
+
+        if max_batches is not None and batch_i >= max_batches:
+            break
+
+    logger.info("Dataset completed.")
+
+
+if __name__ == "__main__":
+    parser = HfArgumentParser((ScriptArguments))
+    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
+        # If we pass only one argument to the script and it's the path to a json file,
+        # let's parse it to get our arguments.
+        args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))[0]
+    else:
+        args = parser.parse_args_into_dataclasses()[0]
+
+    # Setup logging
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        handlers=[logging.StreamHandler(sys.stdout)],
+        level=logging.DEBUG if args.verbose else logging.INFO,
+    )
+
+    if args.do_group and args.max_seq_length is None:
+        raise ValueError("Cannot `do_group` while `max_seq_length` is none.")
+
+    # if os.path.exists(args.output_path) and not args.override:
+    #     raise FileExistsError("Output path exists already: %s", args.output_path
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+    # Settings
+    output_path_prefix = Path(args.output_path_prefix)
+    do_group = args.do_group
+    text_column_name = args.text_column_name
+    max_seq_length = args.max_seq_length
+    return_special_tokens_mask = args.return_special_tokens_mask  # for MLM
+    return_token_type_ids = args.return_token_type_ids
+    print_progress = True
+    batch_size = args.batch_size
+    num_proc = args.num_proc
+    override = args.override
+    limit = args.limit
+
+    # Output data schema
+    pa_columns = [
+        ("input_ids", pa.list_(pa.uint32())),
+        ("attention_mask", pa.list_(pa.uint8())),
+    ]
+    if return_token_type_ids:
+        pa_columns.append(("token_type_ids", pa.list_(pa.uint8())))
+
+    if return_special_tokens_mask:
+        pa_columns.append(("special_tokens_mask", pa.list_(pa.uint8())))
+
+    tokenized_data_schema = pa.schema(pa_columns)
+
+    logger.info(
+        "Initializing dataset from %s (glob: %s)",
+        args.dataset_dir,
+        args.dataset_glob,
+    )
+    dataset_file_paths = list(sorted(Path(args.dataset_dir).glob(args.dataset_glob)))
+
+    if not dataset_file_paths:
+        raise FileNotFoundError(f"No files found in dataset dir: {args.dataset_dir} (glob: {args.dataset_glob})")
+
+    logger.info("Files found in dataset dir: %i", len(dataset_file_paths))
+    ds = pa_dataset(source=dataset_file_paths, format="parquet")
+
+    logger.info("Initialize tokenizer from %s", args.tokenizer_name_or_path)
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.tokenizer_name_or_path,
+        fast=args.use_fast_tokenizer,
+    )
+
+    def group_texts(examples):
+        """
+        Main data processing function that will concatenate all texts from our dataset and generate chunks of max_seq_length.
+        """
+        # Concatenate all texts.
+        concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}
+        total_length = len(concatenated_examples[list(examples.keys())[0]])
+        # We drop the small remainder, and if the total_length < max_seq_length  we exclude this batch and return an empty dict.
+        # We could add padding if the model supported it instead of this drop, you can customize this part to your needs.
+        total_length = (total_length // max_seq_length) * max_seq_length
+        # Split by chunks of max_len.
+        result = {
+            k: [t[i : i + max_seq_length] for i in range(0, total_length, max_seq_length)]
+            for k, t in concatenated_examples.items()
+        }
+        logger.debug("grouped: %i", len(result["input_ids"]))
+
+        # This is done on the model level
+        # if causal_lm:
+        #     result["labels"] = result["input_ids"].copy()
+
+        return result
+
+    def tokenize_texts(list_of_text: List[str]) -> dict:
+        """
+        Perform the actual tokenization (is called by worker threads)
+
+        Tokenized data is grouped if `do_group` is enabled.
+        """
+        tokenizer_out = tokenizer(list_of_text, return_special_tokens_mask=return_special_tokens_mask)
+        logger.debug("tokenized: %i", len(list_of_text))
+
+        if do_group:
+            tokenizer_out = group_texts(tokenizer_out)
+
+        return tokenizer_out
+
+    def init_worker():
+        signal.signal(signal.SIGINT, signal.SIG_IGN)
+
+    def tokenize_in_parallel(pyarrow_text_dataset, pyarrow_schema=None, processes=None) -> Iterable[RecordBatch]:
+        with multiprocessing.Pool(processes=processes, initializer=init_worker) as pool:
+            try:
+                for grouped_tokenized_batch in pool.imap(
+                    tokenize_texts,
+                    generate_texts(
+                        pyarrow_text_dataset,
+                        reader_batch_size=batch_size,
+                        print_progress=print_progress,
+                        row_limit=limit,
+                        text_column_name=text_column_name,
+                    ),
+                    chunksize=1,
+                ):
+                    record_batch = RecordBatch.from_pydict(grouped_tokenized_batch, schema=pyarrow_schema)
+
+                    logger.debug("record_batch: %i", len(record_batch))
+
+                    yield record_batch
+
+                logger.info("All texts tokenized")
+
+                # close pool
+                # pool.terminate()
+                # pool.join()
+                pool.close()
+                pool.join()
+
+            except KeyboardInterrupt:
+                logger.error("Stoppping ...")
+
+                pool.terminate()
+                pool.join()
+
+    # Output settings
+    output_base_dir = output_path_prefix.parent
+    output_file_name_prefix = output_path_prefix.name
+
+    file_options = ParquetFileFormat().make_write_options(compression="zstd")
+    output_file_names = []  # written file names are saved here via `file_visitor`
+
+    def file_visitor(written_file):
+        """
+        Keep track of written files (for later renaming)
+        """
+        file_name = Path(written_file.path).name
+
+        logger.debug("Writing to %s", file_name)
+
+        output_file_names.append(file_name)
+
+    # Write from iterator
+    write_dataset(
+        data=tokenize_in_parallel(
+            pyarrow_text_dataset=ds,
+            pyarrow_schema=tokenized_data_schema,
+            processes=num_proc,
+        ),
+        base_dir=output_base_dir,
+        basename_template=output_file_name_prefix + ".part-{i}.parquet",
+        format="parquet",
+        file_options=file_options,
+        # partitioning="",
+        max_rows_per_group=args.output_max_rows_per_group,
+        max_rows_per_file=args.output_max_rows_per_file,
+        create_dir=True,
+        schema=tokenized_data_schema,
+        file_visitor=file_visitor,
+        # create_dir=True,
+        existing_data_behavior="overwrite_or_ignore"
+        if override
+        else "error",  # error, overwrite_or_ignore, delete_matching
+    )
+
+    # rename files with total number of parts
+    for part, output_file_name in enumerate(output_file_names):
+        # replace: .part-{i}. with part-{i:04d}-of-{len(output_file_names):04d}.
+        new_file_name = output_file_name.replace(
+            f".part-{part}.", f".part-{1+part:04d}-of-{len(output_file_names):04d}."
+        )
+        logger.info(f"Renaming {output_file_name} to {new_file_name}")
+
+        os.rename(
+            os.path.join(output_base_dir, output_file_name),
+            os.path.join(output_base_dir, new_file_name),
+        )
+
+    logger.info("done")
diff --git a/src/lm_datasets/io/parquet.py b/src/lm_datasets/io/parquet.py
index ed4b41e..3bb8509 100644
--- a/src/lm_datasets/io/parquet.py
+++ b/src/lm_datasets/io/parquet.py
@@ -1,11 +1,17 @@
+import asyncio
 import os
+import types
 import pyarrow.parquet as pq
 import logging
 
 import itertools
-from typing import Any, Generator, Tuple
+from typing import Any, Generator, Iterable, Iterator, List, Optional, Tuple, Union
 
 import pyarrow as pa
+import polars as pl
+
+from itertools import islice
+
 
 logger = logging.getLogger(__name__)
 
@@ -27,69 +33,174 @@ def open_parquet_file_with_retries(file_path, retries: int = 2):
     return f
 
 
-def get_parquet_batches(texts: Generator[str, Any, None], schema, batch_size):
-    while True:
-        arr = pa.array(itertools.islice(texts, batch_size))
-        batch = pa.RecordBatch.from_arrays([arr], schema=schema)
+def chunked(generator, size):
+    """Read parts of the generator, pause each time after a chunk"""
+    # islice returns results until 'size',
+    # make_chunk gets repeatedly called by iter(callable).
+    if isinstance(generator, types.GeneratorType):
+        gen = iter(generator)
+    else:
+        # async genator
+        gen = aiter(generator)
 
-        if not batch:
-            break
-        yield batch
+    make_chunk = lambda: list(islice(gen, size))
+    return iter(make_chunk, [])
+
+
+# def get_parquet_batches(rows_iterator: Iterator[Union[str, dict]], schema, batch_size):
+#     async_iterator = get_async_parquet_batches(rows_iterator, schema, batch_size)
+#     while True:
+#         try:
+#             yield asyncio.run(async_iterator.__anext__())
+#         except StopAsyncIteration:
+#             break
+
+# def sync_iter(asyn_iter):
+#     while True:
+#         try:
+#             yield asyncio.run(asyn_iter.__anext__())
+#         except StopAsyncIteration:
+#             break
+
+
+def get_parquet_batches(rows_iterator: Iterator[Union[str, dict]], schema, batch_size):
+    batched_columns = [[] for _ in range(len(schema.names))]
+
+    # rows_iterator = sync_iter(rows_iterator)
+
+    for py_row in rows_iterator:
+        # while True:
+        #     py_row = await anext(rows_iterator)
+
+        if isinstance(py_row, tuple):
+            for col_idx in range(len(schema.names)):
+                batched_columns[col_idx].append(py_row[col_idx])
+        else:
+            # the only column is the text column
+            batched_columns[0].append(py_row)
+
+        if len(batched_columns[0]) == batch_size:
+            yield pa.RecordBatch.from_arrays(
+                [pa.array(batched_column) for batched_column in batched_columns], schema=schema
+            )
+
+            # reset
+            batched_columns = [[] for _ in range(len(schema.names))]
+
+    # last batch
+    if len(batched_columns[0]) > 0:
+        yield pa.RecordBatch.from_arrays(
+            [pa.array(batched_column) for batched_column in batched_columns], schema=schema
+        )
+
+    # else:
+    #     # text and non-text columns
+    #     arrays = [pa.array([row[col] for row in py_batch]) for col in range(len(schema.names))]
+
+    # pa_batch = pa.RecordBatch.from_arrays(arrays, schema=schema)
+    # pa_batch = pa.RecordBatch.from_pylist(py_batch, schema=schema)
+
+    # yield pa_batch
+
+    # while True:
+    #     arr = pa.array(itertools.islice(texts, batch_size))
+    #     batch = pa.RecordBatch.from_arrays([arr], schema=schema)
+
+    #     if not batch:
+    #         break
+    #     yield batch
 
 
 def save_texts_to_parquet_chunks(
     texts: Generator[str, Any, None],
     schema,
-    max_chunk_uncompressed_bytes: int,
     output_path_func: callable,
+    max_chunk_uncompressed_bytes: Optional[int] = None,
+    max_chunk_rows: Optional[int] = None,
     compression: str = "ZSTD",
     batch_size: int = 1024,
     print_write_progress: int = 10_000,
     limit: int = 0,
 ) -> Tuple[int, int]:
+    return write_to_parquet_chunks(
+        data=get_parquet_batches(texts, schema=schema, batch_size=batch_size),
+        schema=schema,
+        output_path_func=output_path_func,
+        max_chunk_uncompressed_bytes=max_chunk_uncompressed_bytes,
+        max_chunk_rows=max_chunk_rows,
+        compression=compression,
+        print_write_progress=print_write_progress,
+        limit=limit,
+    )
+
+
+def write_to_parquet_chunks(
+    data: Iterable[pa.RecordBatch],
+    schema,
+    output_path_func: callable,
+    max_chunk_uncompressed_bytes: Optional[int] = None,
+    max_chunk_rows: Optional[int] = None,
+    compression: str = "ZSTD",
+    print_write_progress: int = 10_000,
+    limit: int = 0,
+) -> Tuple[int, int]:
+    """
+    This could be replaced by `pa.write_dataset` -- however their implementation
+    does not support crtl+c (?) thus we do our own implementation.
+    """
     max_chunks = 9999
+    chunk_rows = None
+    chunk_fp = None
     total_rows = 0
     total_bytes = 0
-    batch_iter = get_parquet_batches(texts, schema=schema, batch_size=batch_size)
+    # batch_iter = get_parquet_batches(texts, schema=schema, batch_size=batch_size)
     limit_reached = False
 
+    if max_chunk_uncompressed_bytes is not None and max_chunk_rows is not None:
+        raise ValueError("Cannot set both `max_chunk_uncompressed_bytes` and `max_chunk_rows`")
+    elif max_chunk_uncompressed_bytes or max_chunk_rows:
+        do_chunks = True
+    else:
+        do_chunks = False
+
     for chunk_i in range(1, max_chunks + 1):
         chunk_rows = 0
         chunk_nbytes = 0
-        chunk_buffer_size = 0
-        chunk_fp = output_path_func(chunk_i) if max_chunk_uncompressed_bytes > 0 else output_path_func()
+        # chunk_buffer_size = 0
+        chunk_fp = output_path_func(chunk_i) if do_chunks else output_path_func()
 
         logger.info(f"Writing to {chunk_fp}")
 
         with pq.ParquetWriter(chunk_fp, schema=schema, compression=compression) as writer:
             try:
                 while True:
-                    batch = next(batch_iter)
+                    batch = next(data)
+
                     writer.write_batch(batch)
                     total_rows += len(batch)
                     total_bytes = batch.nbytes
                     chunk_rows += len(batch)
                     chunk_nbytes += batch.nbytes
-                    chunk_buffer_size += batch.get_total_buffer_size()
+                    # chunk_buffer_size += batch.get_total_buffer_size()
 
-                    if total_rows > 0 and (total_rows % print_write_progress) == 0:
-                        logger.info(f"Written {total_rows:,} rows ...")
+                    if total_rows > 0 and print_write_progress > 0 and (total_rows % print_write_progress) == 0:
+                        logger.info(f"Written %s rows ...", total_rows)
 
                     if limit > 0 and total_rows >= limit:
-                        logger.warning(f"Limit reached ({total_rows:,} rows)")
+                        logger.warning(f"Limit reached (%s rows)", total_rows)
                         limit_reached = True
                         break
 
-                    if max_chunk_uncompressed_bytes > 0 and chunk_nbytes >= max_chunk_uncompressed_bytes:
+                    if (max_chunk_uncompressed_bytes is not None and chunk_nbytes >= max_chunk_uncompressed_bytes) or (
+                        max_chunk_rows is not None and chunk_rows >= max_chunk_rows
+                    ):
                         # if chunk_buffer_size >= max_chunk_bytes_with_safety:
-                        logger.info(
-                            f"Chunk {chunk_i} completed (rows: {chunk_rows:,}; nbytes: {chunk_nbytes:,}; buffer size:"
-                            f" {chunk_buffer_size:,})"
-                        )
+                        logger.info(f"Chunk {chunk_i} completed (rows: {chunk_rows:,}; nbytes: {chunk_nbytes:,})")
+                        # buffer size: {chunk_buffer_size:,})"
                         # logger.info(f"Chunk size on disk: {os.stat(chunk_fp).st_size:,} bytes")
                         break
 
-            except StopIteration:
+            except (StopIteration, StopAsyncIteration):
                 logger.info(f"All rows written ({total_rows=}; {total_bytes=})")
                 break
 
@@ -99,10 +210,55 @@ def save_texts_to_parquet_chunks(
 
     total_chunks = chunk_i
 
-    if max_chunk_uncompressed_bytes > 0:
+    if chunk_rows == 0:
+        logger.warning("Last chunk is empty. Removing file: %s", chunk_fp)
+        os.remove(chunk_fp)
+        total_chunks -= 1
+
+    if do_chunks:
         # Rename files with total number of chunks
         for chunk_i in range(1, total_chunks + 1):
-            logger.info(f"Renaming to {output_path_func(chunk_i, total_chunks)}")
-            os.rename(output_path_func(chunk_i), output_path_func(chunk_i, total_chunks))
+            new_chunk_file_path = output_path_func(chunk_i, total_chunks)
+            logger.info("Renaming to %s", new_chunk_file_path)
+            os.rename(output_path_func(chunk_i), new_chunk_file_path)
 
     return total_rows, total_chunks
+
+
+def get_selected_row_groups(
+    parquet_file: pq.ParquetFile, file_offset: int, file_limit: int
+) -> Tuple[List[int], Union[dict, None]]:
+    """
+    Find the row groups that should be read for `file_offset` and `file_limit`
+    """
+    group_idx_to_offset_last_row: dict[int, Tuple[int, int]] = {}
+    row_groups = []
+    group_offset = 0
+
+    logger.debug("Selecting the row groups within offset=%s and limit=%s", file_offset, file_limit)
+
+    if file_limit == 0 and file_offset == 0:
+        return list(range(parquet_file.num_row_groups)), None
+
+    # Iterate over all row groups
+    for group_idx in range(parquet_file.num_row_groups):
+        # Fetch row group metadata
+        group_metadata = parquet_file.metadata.row_group(group_idx)
+        group_num_rows = group_metadata.num_rows
+        group_last_row = group_offset + group_num_rows
+        group_idx_to_offset_last_row[group_idx] = (group_offset, group_last_row)
+
+        logger.debug("Row group #%s: %s - %s", group_idx, group_offset, group_last_row)
+
+        # Row indicies of group based be part of selected rows
+        if (file_offset >= group_offset and file_offset < group_last_row) or (
+            file_offset < group_offset and file_limit > group_offset
+        ):
+            row_groups.append(group_idx)
+
+        # increase offset for next group
+        group_offset = group_last_row
+
+    logger.debug("Selected row groups: %s", row_groups)
+
+    return row_groups, group_idx_to_offset_last_row
diff --git a/src/lm_datasets/print_stats.py b/src/lm_datasets/print_stats.py
index 942e5c0..0f407fb 100644
--- a/src/lm_datasets/print_stats.py
+++ b/src/lm_datasets/print_stats.py
@@ -1,106 +1,32 @@
-import argparse
+from .utils.config import Config
+from .utils.dataframe import get_datasets_as_dataframe
 
-import logging
 
-import pandas as pd
+def print_stats(config: Config) -> str:
+    logger = config.init_logger(__name__)
 
-from .datasets.dataset_registry import get_registered_dataset_classes
-from .datasets.base import BaseDataset
-
-logging.basicConfig(
-    format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
-    datefmt="%Y-%m-%d %H:%M:%S",
-    level=logging.INFO,
-    handlers=[logging.StreamHandler()],
-)
-logger = logging.getLogger(__name__)
-
-
-def get_datasets_as_dataframe(
-        output_dir=None,
-        output_format="jsonl",
-        raw_datasets_dir=None,
-        extra_dataset_registries=None
-    ):
-
-    # convert datasets to table rows
-    rows = []
-
-    for ds_cls in get_registered_dataset_classes(extra_dataset_registries=extra_dataset_registries):
-        ds: BaseDataset = ds_cls(
-            output_dir=output_dir,
-            output_format=output_format,
-            raw_datasets_dir=raw_datasets_dir,
-        )
-        row = dict(
-            dataset_id=ds.DATASET_ID,
-            source_id=ds.get_source_id(),
-            title=ds.TITLE,
-            homepage=ds.HOMEPAGE,
-            language=ds.get_language_code(),
-            tokens=ds.get_tokens(),
-            bytes=ds.BYTES,
-            is_downloaded=None,
-            web_crawled=1 if ds.WEB_CRAWLED else 0,
-            availibility=ds.AVAILIBILITY,
-            dummy=1 if ds.DUMMY else 0,
-            has_output_file=None,
-        )
-
-        if raw_datasets_dir:
-            row["is_downloaded"] = 1 if ds.is_downloaded() else 0
-
-        if output_dir:
-            row["has_output_file"] = 1 if ds.has_output_files(min_file_size=1024) else 0
-
-        rows.append(row)
-
-    df = pd.DataFrame(rows)
-
-    return df
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    # parser.add_argument("dataset", help="Name of dataset to process")
-    parser.add_argument("--print_format", default="csv", type=str, help="Print format (tsv,csv,md)")
-    parser.add_argument(
-        "--raw_datasets_dir",
-        default=None,
-        type=str,
-        help="Dataset files are read from this directory (needed for `is_downloaded` field)",
-    )
-    parser.add_argument(
-        "--output_dir",
-        default=None,
-        type=str,
-        help="Processed dataset are saved in this directory (need for `has_output_file` field)",
-    )
-    parser.add_argument(
-        "--output_format",
-        default="jsonl",
-        type=str,
-        help="Format of processed dataset (need for `has_output_file` field)",
-    )
-    parser.add_argument(
-        "--extra_dataset_registries",
-        default=None,
-        type=str,
-        help="List of Python packages to load dataset registries",
-    )
-    args = parser.parse_args()
-    print_format = args.print_format
+    print_format = config.print_format
 
     df = get_datasets_as_dataframe(
-        output_dir=args.output_dir if args.output_dir else "/dev/null",
-        output_format=args.output_format,
-        raw_datasets_dir=args.raw_datasets_dir,
-        extra_dataset_registries=args.extra_dataset_registries,
+        output_dir=config.output_dir if config.output_dir else "/dev/null",
+        output_format=config.output_format,
+        shuffled_output_dir=config.shuffled_output_dir,
+        raw_datasets_dir=config.raw_datasets_dir,
+        extra_dataset_registries=config.extra_dataset_registries,
+        rows_count=config.rows_count,
+        shuffled_rows_count=config.shuffled_rows_count,
+        output_compression=config.output_compression,
+        limit=config.limit,
+        exclude_dummy_datasets=config.exclude_dummy_datasets,
+        show_progress=True,
+        metrics_dir=config.metrics_dir,
+        token_estimation_path=config.token_estimation_path,
+        config=config,
+        extra_columns=config.extra_columns.split(",") if config.extra_columns else None,
     )
 
-    # to_markdown
-    # to_csv
     to_kwargs = dict(index=False)
+
     if print_format == "tsv":
         out = df.to_csv(sep="\t", **to_kwargs)
     elif print_format == "csv":
@@ -110,4 +36,19 @@ def get_datasets_as_dataframe(
     else:
         raise ValueError("Unsupported output format: %s" % print_format)
 
+    # Print to stdout
     print(out)
+
+    if config.save_to:
+        to_kwargs["path_or_buf"] = config.save_to
+
+        if print_format == "tsv":
+            df.to_csv(sep="\t", **to_kwargs)
+        elif print_format == "csv":
+            df.to_csv(**to_kwargs)
+        elif print_format == "md":
+            df.to_markdown(**to_kwargs)
+
+        logger.info("Saved to %s", config.save_to)
+
+    return out
diff --git a/src/lm_datasets/shuffle.py b/src/lm_datasets/shuffle.py
deleted file mode 100644
index edfcf08..0000000
--- a/src/lm_datasets/shuffle.py
+++ /dev/null
@@ -1,206 +0,0 @@
-"""
-
-input: parquet files, configs
-
-shuffle + train/test split
-
-"""
-import argparse
-
-import os
-import logging
-
-
-from pathlib import Path
-
-from .datasets.dataset_registry import get_registered_dataset_classes
-from .datasets.base import BaseDataset, GB
-
-from datasets import load_dataset
-
-from tqdm.auto import tqdm
-
-import pyarrow.parquet as pq
-
-import polars as pl
-
-
-DEFAULT_MIN_FILE_SIZE_FOR_BUFFERED_SHUFFLING = 5 * GB
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("datasets", help="Name of datasets to shuffle (comma separated)")
-    parser.add_argument(
-        "--output_dir",
-        default=None,
-        type=str,
-        help="Unshuffled datasets are loaded from this directory",
-    )
-    parser.add_argument(
-        "--shuffled_output_dir",
-        default=None,
-        type=str,
-        help="Shuffled dataset are saved in this directory",
-    )
-    parser.add_argument(
-        "--output_format",
-        default="parquet",
-        type=str,
-        help="Format of processed dataset",
-    )
-    parser.add_argument(
-        "--output_compression",
-        default=None,
-        type=str,
-        help="""Compression of output (jsonl: "gzip"; parquet: "NONE", "SNAPPY", "GZIP", "BROTLI", "LZ4", "ZSTD")""",
-    )
-    parser.add_argument(
-        "--seed",
-        default=0,
-        type=int,
-        help="Random seed",
-    )
-    parser.add_argument(
-        "--shuffle_buffer_size",
-        default=1_000_000,
-        type=int,
-        help="Number of items in buffer to be shuffled at once (larger buffer = more memory but better shuffing)",
-    )
-    parser.add_argument(
-        "--min_file_size_for_buffered_shuffling",
-        default=DEFAULT_MIN_FILE_SIZE_FOR_BUFFERED_SHUFFLING,
-        type=int,
-        help="Min. file size bytes for buffered shuffling (default: 5GB; set to 0 to disable)",
-    )
-    parser.add_argument(
-        "--log_file",
-        default=None,
-        type=str,
-        help="Log file is saved at this path",
-    )
-    parser.add_argument("--verbose", action="store_true", help="Enable verbose logging (log level = debug)")
-    parser.add_argument("--override", action="store_true", help="Override existing output files")
-    parser.add_argument(
-        "--skip_large_datasets",
-        action="store_true",
-        help="Skip datasets with bytes > --min_file_size_for_buffered_shuffling",
-    )
-    args = parser.parse_args()
-
-    log_handlers = [logging.StreamHandler()]
-
-    if args.log_file:
-        log_handlers.append(logging.FileHandler(args.log_file))
-
-    logging.basicConfig(
-        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
-        datefmt="%Y-%m-%d %H:%M:%S",
-        level=logging.DEBUG if args.verbose else logging.INFO,
-        handlers=log_handlers,
-    )
-    logger = logging.getLogger(__name__)
-
-    if args.output_format != "parquet":
-        raise ValueError(f"Output format is not supported (currently only parquet is supported; {args.output_format=})")
-
-    seed = args.seed
-    logger.info(f"Seed: {seed}")
-
-    shuffle_buffer_size = args.shuffle_buffer_size
-
-    id_to_dataset_class = {cls.DATASET_ID: cls for cls in get_registered_dataset_classes()}
-
-    datasets_list = args.datasets.split(",")
-
-    if len(datasets_list) == 1 and datasets_list[0] == "all":
-        # Get list of all non-dummy datasets
-        datasets_list = id_to_dataset_class.keys()
-
-    min_file_size_for_buffered_shuffling = args.min_file_size_for_buffered_shuffling
-
-    # Iterate over datasets
-    for i, dataset_id in enumerate(datasets_list, 1):
-        if dataset_id in id_to_dataset_class:
-            dataset_cls = id_to_dataset_class[dataset_id]
-        else:
-            raise ValueError(f"Unknown dataset ID: {dataset_id} (available: {id_to_dataset_class.keys()})")
-
-        logger.info(f"Dataset ID: {dataset_id} ({i} / {len(datasets_list)})")
-
-        dataset: BaseDataset = dataset_cls(
-            output_dir=args.output_dir,
-            output_format=args.output_format,
-            output_batch_size=shuffle_buffer_size,
-            output_compression=args.output_compression,
-        )
-
-        if not dataset.has_output_files():
-            logger.warning(f"Skipping {dataset_id}: cannot shuffle dataset without processed output files")
-            continue
-
-        output_fps = dataset.get_output_file_paths()
-
-        for i, output_fp in enumerate(sorted(output_fps), i):
-            logger.info(f"Shuffling {output_fp} ({i} / {len(output_fps)} files of {dataset_id})")
-
-            output_file_name = Path(output_fp).name
-
-            shuffled_output_file_path = os.path.join(
-                args.shuffled_output_dir, output_file_name.replace(".parquet", ".shuffled.parquet")
-            )
-
-            assert shuffled_output_file_path != output_fp
-
-            if os.path.exists(shuffled_output_file_path):
-                if args.override:
-                    logger.warning("Overriding existing shuffled output file")
-                else:
-                    logger.warning(
-                        f"Skipping {dataset_id}: Shuffled output file exist already ({shuffled_output_file_path})"
-                    )
-                    continue
-
-            # File size
-            file_stats = os.stat(output_fp)
-
-            if min_file_size_for_buffered_shuffling > 0 and file_stats.st_size > min_file_size_for_buffered_shuffling:
-                # File is too large to be shuffled all at once => shuffle in chunks
-                if args.skip_large_datasets:
-                    logger.info(f"Skip because too large dataset ({file_stats.st_size=})")
-                    continue
-
-                # Reading meta from parquet (for progress bar)
-                logger.info("Reading metadata ...")
-                metadata = pq.read_metadata(output_fp)
-                docs_count = metadata.num_rows
-
-                logger.info("Initializing HF streaming dataset ...")
-                hf_dataset = load_dataset("parquet", data_files={"train": output_fp}, split="train", streaming=True)
-
-                logger.info("Shuffling and writing to new file ...")
-
-                def generate_text():
-                    for item in tqdm(hf_dataset.shuffle(buffer_size=shuffle_buffer_size, seed=seed), total=docs_count):
-                        yield str(item[dataset.output_text_field])  # force to str
-
-                # Writer
-                shuffled_docs_count = dataset.save_texts_to_parquet(
-                    generate_text(), file_path=shuffled_output_file_path, apply_filter=False
-                )
-
-                if docs_count != shuffled_docs_count:
-                    logger.error(
-                        f"Original and shuffled docs count do not match: {docs_count=} vs {shuffled_docs_count=}"
-                    )
-
-            else:
-                # Shuffle in memory
-                logger.info("Initializing PL in-memory dataframe ...")
-                df = pl.read_parquet(output_fp)
-                logger.info("Shuffling and writing to new file ...")
-                df = df.sample(fraction=1, shuffle=True, seed=seed).write_parquet(
-                    shuffled_output_file_path, compression="zstd"
-                )
-
-    logger.info("Done")
diff --git a/src/lm_datasets/shuffle_datasets.py b/src/lm_datasets/shuffle_datasets.py
new file mode 100644
index 0000000..298452a
--- /dev/null
+++ b/src/lm_datasets/shuffle_datasets.py
@@ -0,0 +1,140 @@
+import os
+
+from .datasets.dataset_registry import (
+    get_dataset_class_by_id,
+    get_registered_dataset_ids,
+)
+from .datasets.base import BaseDataset
+from .utils.config import Config
+
+from datasets import load_dataset
+
+from tqdm.auto import tqdm
+
+import pyarrow.parquet as pq
+
+import polars as pl
+
+
+def shuffle_datasets(config: Config):
+    """
+    input: parquet files, configs
+
+    output: shuffled files
+    """
+
+    logger = config.init_logger(__name__)
+
+    if config.output_format != "parquet":
+        raise ValueError(
+            f"Output format is not supported (currently only parquet is supported; {config.output_format=})"
+        )
+
+    seed = config.seed
+    logger.info(f"Seed: {seed}")
+
+    shuffle_buffer_size = config.shuffle_buffer_size
+
+    datasets_list = config.datasets.split(",")
+
+    if len(datasets_list) == 1:
+        if datasets_list[0] == "all":
+            # Get list of all regsitered datasets
+            datasets_list = get_registered_dataset_ids(config.extra_dataset_registries)
+
+        elif datasets_list[0] == "all_from_source":
+            # Get registered datasets based on source
+            if config.source_id is None:
+                raise ValueError("The argument --source_id must be set.")
+
+            datasets_list = get_registered_dataset_ids(
+                config.extra_dataset_registries, needed_source_id=config.source_id
+            )
+
+    min_file_size_for_buffered_shuffling = config.min_file_size_for_buffered_shuffling
+
+    # Iterate over datasets
+    for i, dataset_id in enumerate(datasets_list, 1):
+        logger.info(f"Dataset ID: {dataset_id} ({i} / {len(datasets_list)})")
+
+        dataset_cls = get_dataset_class_by_id(dataset_id, config.extra_dataset_registries)
+        dataset: BaseDataset = dataset_cls(
+            output_dir=config.output_dir,
+            output_format=config.output_format,
+            output_batch_size=shuffle_buffer_size,
+            output_compression=config.output_compression,
+            shuffled_output_dir=config.shuffled_output_dir,
+            config=config,
+        )
+
+        if not dataset.has_output_files():
+            logger.warning(f"Skipping {dataset_id}: cannot shuffle dataset without processed output files")
+            continue
+
+        unshuffled_output_file_paths = dataset.get_output_file_paths()
+
+        for i, unshuffled_file_path in enumerate(sorted(unshuffled_output_file_paths), i):
+            logger.info(
+                f"Shuffling {unshuffled_file_path} ({i} / {len(unshuffled_output_file_paths)} files of {dataset_id})"
+            )
+
+            shuffled_output_file_path = dataset.get_shuffled_output_file_path(unshuffled_file_path)
+
+            assert str(shuffled_output_file_path) != str(unshuffled_file_path)
+
+            if os.path.exists(shuffled_output_file_path):
+                if config.override:
+                    logger.warning("Overriding existing shuffled output file")
+                else:
+                    logger.warning(
+                        f"Skipping {dataset_id}: Shuffled output file exist already ({shuffled_output_file_path})"
+                    )
+                    continue
+
+            # File size
+            file_stats = os.stat(unshuffled_file_path)
+
+            if min_file_size_for_buffered_shuffling > 0 and file_stats.st_size > min_file_size_for_buffered_shuffling:
+                # File is too large to be shuffled all at once => shuffle in chunks
+                if config.skip_large_datasets:
+                    logger.info(f"Skip because too large dataset ({file_stats.st_size=})")
+                    continue
+
+                # Reading meta from parquet (for progress bar)
+                logger.info("Reading metadata ...")
+                metadata = pq.read_metadata(unshuffled_file_path)
+                docs_count = metadata.num_rows
+
+                logger.info("Initializing HF streaming dataset ...")
+                hf_dataset = load_dataset(
+                    "parquet", data_files={"train": unshuffled_file_path}, split="train", streaming=True
+                )
+
+                logger.info("Shuffling and writing to new file ...")
+
+                def generate_text():
+                    for item in tqdm(hf_dataset.shuffle(buffer_size=shuffle_buffer_size, seed=seed), total=docs_count):
+                        yield str(item[dataset.get_output_text_field()])  # force to str
+
+                # Writer
+                shuffled_docs_count, saved_chunks = dataset.save_texts_to_parquet(
+                    generate_text(), file_path=shuffled_output_file_path, apply_filter=False
+                )
+
+                if docs_count != shuffled_docs_count:
+                    logger.error(
+                        f"Original and shuffled docs count do not match: {docs_count=} vs {shuffled_docs_count=}"
+                    )
+
+            else:
+                # Shuffle in memory
+                selected_columns = [dataset.get_output_text_field()]
+                logger.info("Initializing PL in-memory dataframe (%s) ...", selected_columns)
+                df = pl.read_parquet(unshuffled_file_path, columns=selected_columns)
+
+                logger.info("Shuffling and writing to new file ...")
+                df = df.sample(fraction=1, shuffle=True, seed=seed).write_parquet(
+                    shuffled_output_file_path, compression="zstd"
+                )
+
+    logger.info("Done")
diff --git a/src/lm_datasets/train_sp_tokenizer.py b/src/lm_datasets/train_sp_tokenizer.py
new file mode 100644
index 0000000..204e808
--- /dev/null
+++ b/src/lm_datasets/train_sp_tokenizer.py
@@ -0,0 +1,252 @@
+import io
+import json
+import math
+import os
+from pathlib import Path
+import time
+from typing import Iterable
+from pyarrow import parquet as pq
+from lm_datasets.utils import get_auto_workers
+from lm_datasets.utils.config import Config
+import sentencepiece as spm
+from transformers.convert_slow_tokenizer import import_protobuf
+
+
+def train_sp_tokenizer(config: Config):
+    logger = config.init_logger(__name__)
+
+    # from_pretrained_model_path = ""  # existing SP model
+    tokenizer_train_ratio = config.tokenizer_ratio
+    text_field_name = (
+        config.text_field_name
+    )  # "dataset_id"  # bugfix <--- dataset was created with bad column names "dataset_id" and "text" are mixed up
+    batch_size = config.input_batch_size
+    composed_dataset_path = Path(config.composed_dataset_dir)
+    train_stats_path = composed_dataset_path / "train_stats.json"
+    output_tokenizer_path = Path(config.output_tokenizer_path)
+    train_pq_file_paths = list(sorted(composed_dataset_path.glob("train.part-*.parquet")))
+
+    if not composed_dataset_path.exists():
+        raise FileNotFoundError("Composed dataset path does not exists: %s" % composed_dataset_path)
+
+    if output_tokenizer_path.exists() and not config.override:
+        raise FileExistsError("Output exists already at: %s (fix with --overide)" % output_tokenizer_path)
+
+    if not output_tokenizer_path.parent.exists():
+        os.makedirs(output_tokenizer_path.parent)
+
+    if not train_stats_path.exists():
+        raise FileNotFoundError("Training data stats file does not exists: %s" % train_stats_path)
+
+    with open(train_stats_path) as f:
+        train_stats = json.load(f)
+
+    dataset_id_to_train_row_count = {
+        ds_id: stats["split_to_offset_and_limit"]["train"][1] - stats["split_to_offset_and_limit"]["train"][0]
+        for ds_id, stats in train_stats["dataset_id_to_stats"].items()
+    }
+    total_train_row_count = sum(dataset_id_to_train_row_count.values())
+    tokenizer_train_row_count = math.floor(tokenizer_train_ratio * total_train_row_count)
+
+    logger.info("Dataset has %i total rows", total_train_row_count)
+
+    def generate_texts(pq_file_paths, row_limit=0, text_column="text", sentence_splitting: bool = False) -> Iterable:
+        """
+        Stream text data from parquet files
+        """
+        pq_dataset = pq.ParquetDataset(path_or_paths=pq_file_paths)
+
+        logger.info("Dataset initialized with %i fragments", len(pq_dataset.fragments))
+
+        rows = 0
+        for frament in pq_dataset.fragments:
+            for batch in frament.to_batches(batch_size=batch_size, columns=[text_column]):
+                for text in batch.columns[0]:
+                    text = text.as_py()  # cast to py
+
+                    if sentence_splitting:
+                        for sentence in text.splitlines():  # very basic sentence splitting. TODO is this OK for code?
+                            yield sentence
+
+                            # rows += 1
+
+                            # if row_limit > 0 and rows >= row_limit:
+                            #     # sentence loop
+                            #     break
+                    else:
+                        yield text
+
+                    rows += 1
+
+                    if row_limit > 0 and rows >= row_limit:
+                        # text loop
+                        break
+
+                if row_limit > 0 and rows >= row_limit:
+                    # batch loop
+                    break
+
+            if row_limit > 0 and rows >= row_limit:
+                # frament loop
+                break
+
+        logger.info("Texts generated: %i", rows)
+
+    # Source tokenizer
+    if config.source_tokenizer_path is not None:
+        model_pb2 = import_protobuf()
+
+        source_model_proto = model_pb2.ModelProto()
+        with open(config.source_tokenizer_path, "rb") as f:
+            source_model_proto.ParseFromString(f.read())
+        # source_proto = source_model_proto
+
+        # print(source_model_proto.trainer_spec)
+        source_trainer_spec = source_model_proto.trainer_spec
+
+        source_trainer_spec_keys = [f.name for f in source_trainer_spec.DESCRIPTOR.fields]
+        source_trainer_spec_dict = {k: getattr(source_trainer_spec, k) for k in source_trainer_spec_keys}
+
+        logger.info("Source tokenizer loaded from: %s", config.source_tokenizer_path)
+        logger.info("Source train specs: %s", source_trainer_spec_dict)
+
+    else:
+        logger.info("Source tokenizer was not provided")
+
+        source_trainer_spec_dict = {}
+
+    # Text generator
+    logger.info("Tokenizer will be trained on %i input rows", tokenizer_train_row_count)
+
+    texts = generate_texts(
+        train_pq_file_paths,
+        row_limit=tokenizer_train_row_count,
+        text_column=text_field_name,
+        sentence_splitting=config.sentence_splitting,
+    )
+
+    # Write model to memory
+    model_writer = io.BytesIO()
+
+    # Trainer settings
+
+    # These settings are loaded from the source model
+    allowed_source_kwargs = [
+        # user settings
+        # "vocab_size",
+        # "num_threads",
+        # "model_type",
+        # options
+        # "enable_differential_privacy",
+        "self_test_sample_size",
+        # "differential_privacy_noise_level",
+        # "differential_privacy_clipping_threshold",
+        "character_coverage",
+        "input_sentence_size",
+        "shuffle_input_sentence",
+        "seed_sentencepiece_size",
+        "shrinking_factor",
+        "max_sentence_length",
+        "num_sub_iterations",
+        "max_sentencepiece_length",
+        "split_by_unicode_script",
+        "split_by_number",
+        "split_by_whitespace",
+        "treat_whitespace_as_suffix",
+        # "allow_whitespace_only_pieces",
+        "split_digits",
+        "byte_fallback",
+        "vocabulary_output_piece_score",
+        "hard_vocab_limit",
+        "use_all_vocab",
+        "unk_id",
+        "bos_id",
+        "eos_id",
+        "pad_id",
+        "unk_piece",
+        "bos_piece",
+        "eos_piece",
+        "pad_piece",
+        "unk_surface",
+        "train_extremely_large_corpus",
+    ]
+
+    # Default settings (based on Mistral)
+    trainer_kwargs = dict(
+        # User settings (via config)
+        sentence_iterator=texts,
+        model_writer=model_writer,
+        num_threads=get_auto_workers(
+            config.workers
+        ),  # BPE doesn't support multi-threaded training. num_threads are not used.
+        model_type=config.tokenizer_model_type,  # model algorithm: unigram, bpe, word or char
+        vocab_size=config.tokenizer_vocab_size,
+        # Mistral settings
+        enable_differential_privacy=False,
+        self_test_sample_size=0,
+        differential_privacy_noise_level=0.0,
+        differential_privacy_clipping_threshold=0,
+        character_coverage=0.9999499917030334,
+        input_sentence_size=200_000_000,  # (maximum size of sentences the trainer loads), default: 0
+        shuffle_input_sentence=True,  # (Randomly sample input sentences in advance. Valid when --input_sentence_size > 0)
+        # mining_sentence_size=0,  # unkown
+        # training_sentence_size=0,
+        seed_sentencepiece_size=1000000,  # sp default
+        shrinking_factor=0.75,  # sp default
+        max_sentence_length=4192,  # sp default
+        # num_threads=80,
+        num_sub_iterations=2,  # sp default
+        max_sentencepiece_length=16,  # sp default
+        split_by_unicode_script=True,  # sp default
+        split_by_number=True,  # sp default
+        split_by_whitespace=True,  # sp default
+        treat_whitespace_as_suffix=False,  # sp default
+        allow_whitespace_only_pieces=True,
+        split_digits=True,
+        # pretokenization_delimiter='',  # unknown
+        # 'control_symbols': [], 'user_defined_symbols': [], 'required_chars': '',
+        byte_fallback=True,
+        vocabulary_output_piece_score=True,
+        hard_vocab_limit=True,  # sp default
+        use_all_vocab=False,  # sp default
+        unk_id=0,
+        bos_id=1,
+        eos_id=2,
+        pad_id=-1,
+        unk_piece="<unk>",
+        bos_piece="
",
+        eos_piece="
",
+        pad_piece="<pad>",
+        unk_surface=" ⁇ ",  # sp default
+        train_extremely_large_corpus=False,  # sp default
+        # model_prefix=model_prefix,  # `model_prefix`and `model_writer` cannot be defined at the same time!
+    )
+
+    # set settings from source tokenizer
+    if source_trainer_spec_dict:
+        for k in allowed_source_kwargs:
+            trainer_kwargs[k] = source_trainer_spec_dict[k]
+
+    start_time = time.perf_counter()
+
+    logger.info("SP train settings: %s", trainer_kwargs)
+
+    # train with combination of user settings and source model settings
+    train_res = spm.SentencePieceTrainer.train(
+        **trainer_kwargs,
+    )
+
+    # Serialize the model as file
+    with open(output_tokenizer_path, "wb") as f:
+        f.write(model_writer.getvalue())
+
+    logger.info("Saved to %s", config.output_tokenizer_path)
+
+    # Directly load the model from serialized model.
+    sp = spm.SentencePieceProcessor(model_proto=model_writer.getvalue())
+    logger.info("Example output: %s", sp.encode("this is test"))
+
+    end_time = time.perf_counter()
+    elapsed_time = end_time - start_time
+
+    logger.info(f"done (training took {elapsed_time=:0.2f} seconds)")
diff --git a/src/lm_datasets/utils/__init__.py b/src/lm_datasets/utils/__init__.py
index 2ac6108..29fc0c2 100644
--- a/src/lm_datasets/utils/__init__.py
+++ b/src/lm_datasets/utils/__init__.py
@@ -126,7 +126,9 @@ def get_parquet_compression(request_compression: str) -> str:
 def get_bytes_from_int_or_string(bytes_int_or_str: Union[int, str]) -> int:
     from lm_datasets.datasets.base import GB, KB, MB
 
-    if isinstance(bytes_int_or_str, int):
+    if bytes_int_or_str is None:
+        return None
+    elif isinstance(bytes_int_or_str, int):
         return bytes_int_or_str
     elif isinstance(bytes_int_or_str, str):
         if bytes_int_or_str.endswith("KB"):
@@ -135,6 +137,8 @@ def get_bytes_from_int_or_string(bytes_int_or_str: Union[int, str]) -> int:
             return int(bytes_int_or_str.rstrip("MB")) * MB
         elif bytes_int_or_str.endswith("GB"):
             return int(bytes_int_or_str.rstrip("GB")) * GB
+        elif not bytes_int_or_str:
+            return None
         else:
             return int(bytes_int_or_str)
     else:
diff --git a/src/lm_datasets/utils/config.py b/src/lm_datasets/utils/config.py
new file mode 100644
index 0000000..a731b7a
--- /dev/null
+++ b/src/lm_datasets/utils/config.py
@@ -0,0 +1,135 @@
+import argparse
+from typing import List, Iterable, Literal, Union
+import yaml
+import logging
+
+
+logger = logging.getLogger(__name__)
+
+
+def get_common_argparser(required_configs: bool = False):
+    common_parser = argparse.ArgumentParser()
+
+    common_parser.add_argument(
+        "--raw_datasets_dir",
+        default=None,
+        type=str,
+        help="Dataset files are read from this directory (needed for `is_downloaded` field)",
+    )
+    common_parser.add_argument(
+        "--output_dir",
+        default=None,
+        type=str,
+        help="Processed dataset are saved in this directory (need for `has_output_file` field)",
+    )
+    common_parser.add_argument(
+        "--output_format",
+        default="jsonl",
+        type=str,
+        help="Format of processed dataset (jsonl,parquet; need for `has_output_file` field)",
+    )
+    common_parser.add_argument(
+        "--extra_dataset_registries",
+        default=None,
+        type=str,
+        help="List of Python packages to load dataset registries",
+    )
+    common_parser.add_argument(
+        "-c",
+        "--configs",
+        nargs="+",
+        help=(
+            "Paths to one or more YAML-config files (duplicated settings are override based on config order; "
+            "cmd args override configs)"
+        ),
+        default=None,
+        dest="config_paths",
+        required=required_configs,
+    )
+    common_parser.add_argument(
+        "--log_file",
+        default=None,
+        type=str,
+        help="Log file is saved at this path",
+    )
+    common_parser.add_argument("--verbose", action="store_true", help="Enable verbose logging (log level = debug)")
+    return common_parser
+
+
+class Config:
+    composed_dataset_dir = None  # composed dataset (train/val split) is saved into this directory
+    local_dirs_by_dataset_id = {}
+    local_dirs_by_source_id = {}
+    sampling_factor_by_dataset_id = {}
+    sampling_factor_by_source_id = {}
+    sampling_factor_by_language = {}
+
+    only_selected_datasets: bool = False
+    selected_dataset_ids: List[str] = []
+    selected_source_ids: List[str] = []
+
+    validation_ratio = 0.005  # number of documents in the split: len(dataset) * ratio
+    validation_min_total_docs = 1_000  # to be used as validation set, the dataset must have at least n docs
+    validation_max_split_docs = 1_000  # number of documents in validation split are capped at this numbers
+    validation_min_split_docs = 10  # split must have at least this number of documents, otherwise it will be discarded
+    tokenizer_train_ratio = 0.1  # % of train data used for tokenizer training
+
+    # Vocab size should divisble by 8
+    # - Jan's recommendation: 250680
+    # - NVIDIA recommendation for multilingual models: 256000
+    tokenizer_vocab_size: int = 256000
+    tokenizer_model_type: Literal["bpe", "unigram", "word", "char"] = "bpe"  # SP model types
+
+    seed: int = 0
+
+    extra_dataset_registries: Union[None, str, List[str]] = None
+    extra_dataset_classes: Union[None, List] = None
+    use_default_dataset_registry: bool = True
+
+    verbose = False
+    log_file = None
+
+    def __init__(self, **entries):
+        self.__dict__.update(entries)
+
+    def init_logger(self, logger_name):
+        log_handlers = [logging.StreamHandler()]
+
+        if self.log_file:
+            log_handlers.append(logging.FileHandler(self.log_file))
+
+        logging.basicConfig(
+            format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
+            datefmt="%Y-%m-%d %H:%M:%S",
+            level=logging.DEBUG if self.verbose else logging.INFO,
+            handlers=log_handlers,
+        )
+        logger = logging.getLogger(logger_name)
+
+        return logger
+
+
+def get_config_from_paths(config_paths: Iterable, override: dict = None) -> Config:
+    config = {}
+
+    if config_paths:
+        for config_path in config_paths:
+            logger.info("Loading config: %s", config_path)
+            with open(config_path) as f:
+                _config = yaml.safe_load(f)
+
+                config.update(_config)
+    if override:
+        # Override with args
+        config.update(override)
+
+    config = Config(**config)
+
+    return config
+
+
+def parse_args_and_get_config(parser):
+    args = parser.parse_args()
+    config = get_config_from_paths(args.config_paths, override=args.__dict__)
+
+    return config
diff --git a/src/lm_datasets/utils/dataframe.py b/src/lm_datasets/utils/dataframe.py
new file mode 100644
index 0000000..11ef76e
--- /dev/null
+++ b/src/lm_datasets/utils/dataframe.py
@@ -0,0 +1,203 @@
+import logging
+import json
+from pathlib import Path
+from typing import Iterable, List, Optional
+import pandas as pd
+from tqdm.auto import tqdm
+from lm_datasets.datasets.dataset_registry import get_registered_dataset_classes
+from lm_datasets.datasets.base import TOKENS_PER_BYTE, BaseDataset, License
+from lm_datasets.utils.config import Config
+
+logger = logging.getLogger(__name__)
+
+
+def stringify_list(list: Iterable, sep: str = ",") -> str:
+    return sep.join([str(e) for e in list])
+
+
+DEFAULT_TOKENS_PER_WHITESPACE = 2.33  # based on EURO dataset
+
+AVAILABLE_DATAFRAME_COLUMNS = {
+    "license": lambda ds: ds.LICENSE,
+    "license_commercial_use": lambda ds: ds.LICENSE.commercial_use
+    if ds.LICENSE is not None and isinstance(ds.LICENSE, License)
+    else None,
+    "license_sharealike": lambda ds: ds.LICENSE.sharealike
+    if ds.LICENSE is not None and isinstance(ds.LICENSE, License)
+    else None,
+    "overlap": lambda ds: ",".join(ds.HAS_OVERLAP_WITH),
+    "quality_warnings": lambda ds: stringify_list(ds.QUALITY_WARNINGS),
+    "genres": lambda ds: stringify_list(ds.GENRES),
+    "languages": lambda ds: stringify_list(ds.LANGUAGES),
+}
+
+
+def get_dataframe_row_from_dataset(
+    ds: BaseDataset,
+    rows_count=False,
+    shuffled_rows_count=False,
+    output_compression=False,
+    metrics: Optional[dict] = None,
+    estimated_tokens_per_whitespace: Optional[float] = None,
+    extra_columns: Optional[Iterable[str]] = None,
+) -> dict:
+    row = dict(
+        dataset_id=ds.DATASET_ID,
+        source_id=ds.get_source_id(),
+        title=ds.TITLE,
+        homepage=ds.HOMEPAGE,
+        language=ds.get_language_code(),
+        tokens=ds.get_tokens(),
+        bytes=ds.BYTES,
+        is_downloaded=None,
+        web_crawled=1 if ds.WEB_CRAWLED else 0,
+        availibility=ds.AVAILIBILITY,
+        dummy=1 if ds.DUMMY else 0,
+        has_output_files=None,
+        sampling_factor=ds.get_sampling_factor(),
+    )
+
+    if extra_columns:
+        for col_name in extra_columns:
+            if col_name in AVAILABLE_DATAFRAME_COLUMNS:
+                row[col_name] = AVAILABLE_DATAFRAME_COLUMNS[col_name](ds)
+            else:
+                raise ValueError(f"Unsupported column: {col_name}")
+
+    if ds.raw_datasets_dir:
+        row["is_downloaded"] = 1 if ds.is_downloaded() else 0
+
+    if ds.output_dir:
+        row["has_output_files"] = 1 if ds.has_output_files(min_file_size=1024) else 0
+
+        if rows_count:
+            row["rows_count"] = ds.get_output_rows_count()
+
+        if output_compression:
+            row["output_compression"] = ds.get_compression_from_output_files()
+
+    if ds.shuffled_output_dir is not None:
+        row["has_shuffled_output_files"] = 1 if ds.has_output_files(min_file_size=1024, shuffled=True) else 0
+
+        if shuffled_rows_count:
+            row["shuffled_rows_count"] = ds.get_output_rows_count(shuffled=True)
+
+        if output_compression:
+            row["shuffled_output_compression"] = ds.get_compression_from_output_files(shuffled=True)
+
+        # Estimations are very bad -> do not use
+        # row["estimated_bytes"] = ds.get_estimated_bytes_from_output(shuffled=True, read_first_n_rows=10_000)
+        # row["estimated_tokens"] = int(row["estimated_bytes"] * TOKENS_PER_BYTE)
+
+        # TODOs
+        # has_shuffled_output_files
+        # number of (shuffled) output_files
+        # number of output_rows
+        # compression (gzip/parquet) ---> make everything in the same compression to enable sampling based on byte size
+
+    if metrics:
+        row.update(metrics)
+
+        if estimated_tokens_per_whitespace:
+            row.update(
+                {
+                    "estimated_tokens_per_whitespace": estimated_tokens_per_whitespace,
+                    "total_estimated_tokens": metrics["whitespace_count"] * estimated_tokens_per_whitespace,
+                }
+            )
+
+    return row
+
+
+def get_datasets_as_dataframe(
+    output_dir=None,
+    output_format="jsonl",
+    raw_datasets_dir=None,
+    extra_dataset_registries=None,
+    shuffled_output_dir=None,
+    rows_count=False,
+    shuffled_rows_count=False,
+    output_compression=False,
+    limit: int = 0,
+    exclude_dummy_datasets: bool = False,
+    show_progress: bool = False,
+    workers: int = 1,
+    metrics_dir: Optional[str] = None,
+    token_estimation_path=None,
+    config: Optional[Config] = None,
+    extra_columns: Optional[Iterable[str]] = None,
+):
+    # convert datasets to table rows
+    rows = []
+    ds_clss = get_registered_dataset_classes(extra_dataset_registries=extra_dataset_registries)
+
+    if limit > 0:
+        ds_clss = ds_clss[:limit]
+
+    if workers > 1:
+        raise NotImplementedError
+
+    ds_clss_iterator = tqdm(ds_clss, desc="Loading dataset details") if show_progress else ds_clss
+
+    # Metrics (whitespaces, byte count)
+    dataset_id_to_metrics = {}
+
+    if metrics_dir:
+        for json_path in Path(metrics_dir).glob("*.json"):
+            with open(json_path) as f:
+                dataset_id_to_metrics.update(json.load(f))
+
+    # Load token estimations (tokens per whitespace)
+    dataset_id_to_tokens_per_whitespace = {}
+
+    if token_estimation_path:
+        dataset_id_to_tokens_per_whitespace = {}
+
+        with open(token_estimation_path) as f:
+            dataset_id_to_metrics_with_tokens = json.load(f)
+
+            for dataset_id, metrics in dataset_id_to_metrics_with_tokens.items():
+                # use whitespace_count or bytee?
+                if (
+                    "whitespace_count" in metrics
+                    and metrics["whitespace_count"] > 0
+                    and "tokens_count" in metrics
+                    and metrics["tokens_count"] > 0
+                ):
+                    dataset_id_to_tokens_per_whitespace[dataset_id] = (
+                        metrics["tokens_count"] / metrics["whitespace_count"]
+                    )
+
+    for ds_cls in ds_clss_iterator:
+        dataset_id = ds_cls.DATASET_ID
+        ds: BaseDataset = ds_cls(
+            output_dir=output_dir,
+            output_format=output_format,
+            raw_datasets_dir=raw_datasets_dir,
+            shuffled_output_dir=shuffled_output_dir,
+            config=config,
+        )
+
+        if config.only_selected_datasets and not ds.is_selected():
+            logger.info("Skip %s (not part of selected datasets or sources)", dataset_id)
+            continue
+
+        row = get_dataframe_row_from_dataset(
+            ds,
+            rows_count=rows_count,
+            shuffled_rows_count=shuffled_rows_count,
+            output_compression=output_compression,
+            metrics=dataset_id_to_metrics[dataset_id] if dataset_id in dataset_id_to_metrics else None,
+            estimated_tokens_per_whitespace=dataset_id_to_tokens_per_whitespace[dataset_id]
+            if dataset_id in dataset_id_to_tokens_per_whitespace
+            else None,
+            extra_columns=extra_columns,
+        )
+
+        if exclude_dummy_datasets and ds.is_dummy():
+            logger.info("Skip %s (dummy dataset)", dataset_id)
+            continue
+
+        rows.append(row)
+
+    return pd.DataFrame(rows)
diff --git a/src/lm_datasets/utils/dataset_generator.py b/src/lm_datasets/utils/dataset_generator.py
new file mode 100644
index 0000000..716e047
--- /dev/null
+++ b/src/lm_datasets/utils/dataset_generator.py
@@ -0,0 +1,463 @@
+import asyncio
+import math
+from pathlib import Path
+import random
+from typing import Iterable, Iterator, List, Optional, Tuple, Union
+from lm_datasets.datasets.base import BaseDataset
+from lm_datasets.datasets.dataset_registry import get_registered_dataset_classes
+from lm_datasets.utils.config import Config
+
+from tqdm.auto import tqdm
+import numpy as np
+import pyarrow as pa
+from enum import Enum
+import logging
+
+
+logger = logging.getLogger(__name__)
+
+
+class DatasetSplit(str, Enum):
+    """
+    Full dataset:
+
+    [-------------]
+
+    first n rows as validation:
+
+    [xxxx][-------]
+
+    last 1-n rows as train set:
+
+    [----][xxxxxxx]
+
+    first n rows of train set are the tokenizer train set:
+
+    [----][xx][---]
+
+    Sampling: Up/down sampling is only applied to train and tokenizer_train set.
+    """
+
+    FULL = "full"
+    TRAIN = "train"
+    TOKENIZER_TRAIN = "tokenizer_train"
+    VALIDATION = "validation"
+
+    # def __str__(self):
+    #     return str(self.value)
+
+
+DatasetSplitOffsetLimit = dict[DatasetSplit, Tuple[int, int]]
+
+
+def iter_random_indices(
+    rng: np.random.Generator,
+    num_sources: int,
+    probabilities: List[float],
+    random_batch_size: int = 100,
+) -> Iterator[int]:
+    """
+    Get an infinite iterator that randomly samples the index of the source to pick examples from.
+    """
+    assert len(probabilities) == num_sources
+
+    while True:
+        yield from (int(i) for i in rng.choice(num_sources, size=random_batch_size, p=probabilities))
+
+
+def interleave_dataset_iterators(
+    datasets_iterators: List[Iterable],
+    probabilities: List[float],
+    seed: int = 0,
+    random_indicies_batch_size: int = 100,
+    datasets_prefixes: Optional[List] = None,
+) -> Iterator:
+    """
+    Interleave dataset iterators.
+
+    Implementation based on Huggingface's datasets library
+
+    https://github.com/huggingface/datasets/blob/main/src/datasets/iterable_dataset.py#L31
+    """
+    datasets_count = len(datasets_iterators)
+    random_generator = np.random.default_rng(seed)
+
+    indices_iterator = iter_random_indices(
+        random_generator,
+        num_sources=datasets_count,
+        probabilities=probabilities,
+        random_batch_size=random_indicies_batch_size,
+    )
+
+    is_exhausted = np.full(datasets_count, False)
+
+    # Iterators with zero probability may occurc (e.g., dataset size = 0)
+    # Mark these iterators as already exhausted to avoid infinite loops.
+    zero_probabilities = np.array(probabilities) == 0
+    is_exhausted[zero_probabilities] = np.full_like(is_exhausted[zero_probabilities], True)
+
+    for dataset_idx in indices_iterator:
+        try:  # let's pick one sample from the iterator at index i
+            dataset_iterator = datasets_iterators[dataset_idx]
+            dataset_sample = next(dataset_iterator)
+
+            # dataset_sample = await anext(dataset_iterator)
+            # dataset_sample = anext(dataset_iterator)
+            # dataset_sample = asyncio.run(dataset_iterator.__anext__())
+
+            #         while True:
+            # try:
+            #     yield asyncio.run(interleaved_iterator.__anext__())
+            # except StopAsyncIteration:
+            #     break
+
+            if datasets_prefixes:
+                # prepend returned sample with dataset specific value, e.g., dataset IDs.
+                yield datasets_prefixes[dataset_idx], dataset_sample
+                # yield {"text": dataset_sample, "dataset_id": datasets_prefixes[dataset_idx]}
+            else:
+                yield dataset_sample
+
+        except (StopIteration, StopAsyncIteration):
+            # here it means that the i-th iterabledataset is empty, i.e we never have the occasion to yield an element of the i-th dataset.
+            # we still check if the stopping criteria is met and if we break out of the loop in case of an oversampling strategy
+            is_exhausted[dataset_idx] = True
+
+            if np.all(is_exhausted):
+                # if the stopping criteria is met, break the main for loop
+                break
+
+
+def sampled_dataset_iterator(
+    iterator_func, true_dataset_rows, target_dataset_rows, offset, limit, **iterator_func_kwargs
+) -> Iterator:
+    if target_dataset_rows > true_dataset_rows:
+        # upsampling
+        full_reads = math.floor(target_dataset_rows / true_dataset_rows)
+        last_partial_read = (target_dataset_rows / true_dataset_rows) - full_reads
+        last_partial_read_n = math.floor(last_partial_read * true_dataset_rows)
+
+        logger.debug("Upsample dataset. Full reads: %s times; Last read: %s rows", full_reads, last_partial_read_n)
+
+        # full read (no limit): repeated reading from the dataset
+        for _ in range(full_reads):
+            for sample in iterator_func(offset=offset, limit=limit, **iterator_func_kwargs):
+                yield sample
+
+        # last partial read (with limit)
+        if last_partial_read_n > 0:
+            adjusted_limit = offset + last_partial_read_n
+            for sample in iterator_func(offset=offset, limit=adjusted_limit, **iterator_func_kwargs):
+                yield sample
+
+    else:
+        # no upsampling, simply read with given offset + limit
+        adjusted_limit = limit - (true_dataset_rows - target_dataset_rows)
+
+        for sample in iterator_func(offset=offset, limit=adjusted_limit, **iterator_func_kwargs):
+            yield sample
+
+
+class DatasetGenerator(object):
+    """
+    Dataset generator class
+
+    Generate function needs to be implemented as a class to make the
+    collected statistics accessible for an iterator.
+    """
+
+    dataset_id_to_stats: dict[str, dict] = {}
+    dataset_id_to_args: dict[str, dict] = {}
+
+    list_of_dataset_iterators = []
+    list_of_dataset_rows = []
+    list_of_dataset_ids = []
+
+    def __init__(
+        self,
+        config: Config,
+        shuffled_output_dir,
+        output_format,
+        save_to_dir: Union[str, Path],
+        split: DatasetSplit = DatasetSplit.TRAIN,
+    ) -> None:
+        self.config = config
+        self.shuffled_output_dir = shuffled_output_dir
+        self.output_format = output_format
+        self.split = split
+
+        if isinstance(save_to_dir, str):
+            save_to_dir = Path(save_to_dir)
+
+        self.save_to_dir = save_to_dir
+
+    def prepare_datasets(self, use_sampling: bool = False, print_progress: bool = False):
+        available_dataset_classes = get_registered_dataset_classes(
+            extra_dataset_registries=self.config.extra_dataset_registries,
+            extra_dataset_classes=self.config.extra_dataset_classes,
+            use_default_registry=self.config.use_default_dataset_registry,
+        )
+        selected_dataset_ids = self.config.selected_dataset_ids
+        selected_source_ids = self.config.selected_source_ids
+
+        self.list_of_dataset_rows = []
+        self.list_of_dataset_ids = []
+        self.dataset_id_to_stats = {}
+        self.dataset_id_to_args = {}
+
+        # Construct iterators and get dataset sizes
+        if print_progress:
+            available_dataset_classes = tqdm(available_dataset_classes, desc="Prepare datasets")
+
+        for ds_cls in available_dataset_classes:
+            ds: BaseDataset = ds_cls(
+                output_dir=None,
+                output_format=self.output_format,
+                raw_datasets_dir=None,
+                shuffled_output_dir=self.shuffled_output_dir,
+                config=self.config,
+            )
+            dataset_id = ds.DATASET_ID
+            source_id = ds.get_source_id()
+
+            if dataset_id in selected_dataset_ids or source_id in selected_source_ids:
+                split_offset_limit = get_splits_as_offsets_and_limits(ds, use_shuffled_output=True)
+                offset, limit = split_offset_limit[self.split]
+                dataset_rows = limit - offset
+
+                if use_sampling:
+                    # Expected dataset size depends on sampling factor -> influences probabilty
+                    expected_dataset_rows = math.floor(dataset_rows * ds.get_sampling_factor())
+                else:
+                    expected_dataset_rows = dataset_rows
+
+                if expected_dataset_rows > 0:
+                    if dataset_id in self.dataset_id_to_args:
+                        raise ValueError(f"Duplicated dataset ID: {dataset_id}")
+
+                    self.dataset_id_to_args[dataset_id] = dict(
+                        iterator_func=ds.generate_texts_from_output,
+                        true_dataset_rows=dataset_rows,
+                        target_dataset_rows=expected_dataset_rows,
+                        offset=offset,
+                        limit=limit,
+                    )
+                    self.list_of_dataset_rows.append(expected_dataset_rows)
+                    self.list_of_dataset_ids.append(dataset_id)
+
+                    # save stats
+                    self.dataset_id_to_stats[dataset_id] = {
+                        "dataset_rows": dataset_rows,
+                        "expected_dataset_rows": expected_dataset_rows,
+                        "split_to_offset_and_limit": split_offset_limit,
+                    }
+                else:
+                    logger.warning("Dataset is zero expected rows: %s", dataset_id)
+
+    def generate_texts_from_single_dataset(
+        self,
+        dataset_id,
+        reader_batch_size: int = 100,
+        print_progress: bool = False,
+    ) -> Iterator:
+        if not self.dataset_id_to_args:
+            raise ValueError("Dataset args are not set. Did you run `prepare_datasets`?")
+
+        args = self.dataset_id_to_args[dataset_id]
+        dataset_iterator = sampled_dataset_iterator(
+            shuffled=True,
+            shuffle_output_file_paths=True,
+            reader_implementation="pyarrow",
+            batch_size=reader_batch_size,
+            **args,
+        )
+
+        if print_progress:
+            dataset_iterator = tqdm(
+                dataset_iterator, desc=f"Generating from {dataset_id}", total=args["limit"] - args["offset"]
+            )
+
+        yield from dataset_iterator
+
+    def generate_texts_from_interleaved_datasets(
+        self,
+        random_batch_size: int = 100,
+        reader_batch_size: int = 100,
+        generate_dataset_ids: bool = False,
+        print_progress: bool = False,
+    ) -> Iterator:
+        """
+        Generate texts based on split or full datasets and the given config.
+        """
+        if not self.dataset_id_to_args:
+            raise ValueError("Dataset args are not set. Did you run `prepare_datasets`?")
+
+        self.list_of_dataset_iterators = []
+
+        # Iterate over prepared datasets
+        for args in self.dataset_id_to_args.values():
+            dataset_iterator = sampled_dataset_iterator(
+                shuffled=True,
+                shuffle_output_file_paths=True,
+                reader_implementation="pyarrow",
+                # reader_implementation="polars_read_parquet",
+                batch_size=reader_batch_size,
+                **args,
+            )
+
+            self.list_of_dataset_iterators.append(dataset_iterator)
+
+        assert len(self.list_of_dataset_iterators) == len(self.dataset_id_to_args)
+        assert len(self.list_of_dataset_iterators) == len(self.list_of_dataset_ids)
+        assert len(self.list_of_dataset_iterators) == len(self.list_of_dataset_rows)
+
+        if len(self.list_of_dataset_iterators) == 0:
+            raise ValueError(
+                "No dataset has been selected. Did you set the config for `selected_dataset_ids` and `selected_source_ids`?"
+            )
+
+        # Random order of datasets
+        dataset_indices = list(range(len(self.list_of_dataset_iterators)))
+        random.seed(self.config.seed)
+        random.shuffle(dataset_indices)
+
+        def select_by_indicies(list, indices):
+            return [list[idx] for idx in indices]
+
+        # Compute probabilites
+        total_rows = sum(self.list_of_dataset_rows)
+        probabilities = [dataset_rows / total_rows for dataset_rows in self.list_of_dataset_rows]
+
+        interleaved_iterator = interleave_dataset_iterators(
+            select_by_indicies(self.list_of_dataset_iterators, dataset_indices),
+            select_by_indicies(probabilities, dataset_indices),
+            datasets_prefixes=pa.array(select_by_indicies(self.list_of_dataset_ids, dataset_indices), type=pa.string())
+            if generate_dataset_ids
+            else None,
+            seed=self.config.seed,
+            random_indicies_batch_size=random_batch_size,
+        )
+
+        if print_progress:
+            interleaved_iterator = tqdm(interleaved_iterator, total=total_rows, desc="Interleave datasets")
+
+        for sample in interleaved_iterator:
+            yield sample
+
+    def save_to_path(self, part: Optional[int] = None, total_parts: Optional[int] = None, suffix=""):
+        split_str = self.split + suffix
+        output_extension = "." + self.output_format
+
+        if total_parts is None:
+            fn = f"{split_str}.part-{part:04d}{output_extension}"
+        else:
+            fn = f"{split_str}.part-{part:04d}-of-{total_parts:04d}{output_extension}"
+
+        return self.save_to_dir / fn
+
+
+def generate_texts_from_dataset(
+    dataset: BaseDataset,
+    split: DatasetSplit = None,
+    split_offset_limit: Optional[DatasetSplitOffsetLimit] = None,
+    use_shuffled_output: bool = True,
+):
+    """
+
+    Full dataset:
+
+    [-------------]
+
+    first n rows as validation:
+
+    [xxxx][-------]
+
+    last 1-n rows as train set:
+
+    [----][xxxxxxx]
+
+    first n rows of train set are the tokenizer train set:
+
+    [----][xx][---]
+
+    """
+
+    if split_offset_limit is None:
+        split_offset_limit = get_splits_as_offsets_and_limits(dataset, use_shuffled_output=use_shuffled_output)
+
+    dataset_offset, dataset_limit = split_offset_limit[split]
+
+    logger.info("Generating texts for split=%s with offset=%s and limit=%s", split, dataset_offset, dataset_limit)
+
+    yield from dataset.generate_texts_from_output(
+        shuffled=use_shuffled_output,
+        shuffle_output_file_paths=use_shuffled_output,
+        offset=dataset_offset,
+        limit=dataset_limit,
+    )
+
+
+def get_splits_as_offsets_and_limits(dataset: BaseDataset, use_shuffled_output: bool = True) -> DatasetSplitOffsetLimit:
+    """
+    Generate splits as offsets and limit depending on total rows and defined split settings (ratio, min. doc count etc.)
+    """
+    if dataset.config is None:
+        raise ValueError("dataset config is not set")
+
+    n_docs = dataset.get_output_rows_count(shuffled=use_shuffled_output)  # total dataset size
+    validation_offset = 0  # validation rows are always the first n rows
+    train_limit = n_docs  # train rows are always the last n rows
+
+    # Compute offsets + limits for split selection
+    if dataset.HAS_PREDEFINED_VALIDATION_SET:
+        logger.warning("No validation split - because a predefined validation set exists: %s", dataset.DATASET_ID)
+        validation_n_docs = 0
+    elif dataset.config.validation_min_total_docs is not None and n_docs < dataset.config.validation_min_total_docs:
+        logger.warning(
+            "No validation split for %s - because dataset has too few docs (given: %s; expected: %s).",
+            dataset.DATASET_ID,
+            n_docs,
+            dataset.config.validation_min_total_docs,
+        )
+        validation_n_docs = 0
+    else:
+        validation_n_docs = min(
+            math.floor(dataset.config.validation_ratio * n_docs),
+            dataset.config.validation_max_split_docs if dataset.config.validation_max_split_docs else math.inf,
+        )
+
+        if dataset.config.validation_min_split_docs and validation_n_docs < dataset.config.validation_min_split_docs:
+            logger.warning(
+                "No validation split for %s - because split would be too small (given: %s; expected: %s).",
+                dataset.DATASET_ID,
+                validation_n_docs,
+                dataset.config.validation_min_split_docs,
+            )
+            validation_n_docs = 0
+
+    if validation_n_docs > 0:
+        train_offset = validation_n_docs  # + 1
+    else:
+        train_offset = 0
+
+    train_n_docs = n_docs - validation_n_docs
+    tokenizer_train_n_docs = math.floor(train_n_docs * dataset.config.tokenizer_train_ratio)
+
+    # Training dataset for tokenizer (by default: 10% of training set)
+    tokenizer_train_offset = train_offset
+    tokenizer_train_limit = train_offset + tokenizer_train_n_docs  # + 1
+
+    logger.debug("validation_n_docs: %s", validation_n_docs)
+    logger.debug("train_n_docs: %s", train_n_docs)
+    logger.debug("tokenizer_train_n_docs: %s", tokenizer_train_n_docs)
+
+    splits: DatasetSplitOffsetLimit = {
+        DatasetSplit.FULL: (0, n_docs),
+        DatasetSplit.VALIDATION: (validation_offset, validation_n_docs),
+        DatasetSplit.TRAIN: (train_offset, train_limit),
+        DatasetSplit.TOKENIZER_TRAIN: (tokenizer_train_offset, tokenizer_train_limit),
+    }
+
+    logger.debug("splits: %s", splits)
+
+    return splits
diff --git a/src/lm_datasets/utils/languages.py b/src/lm_datasets/utils/languages.py
new file mode 100644
index 0000000..38300f7
--- /dev/null
+++ b/src/lm_datasets/utils/languages.py
@@ -0,0 +1,187 @@
+# Unified language names (e.g., "Catalan; Valencian" => "Catalan"):
+# ca,ny,dv,ff,el,ht,he,kl,ki,ky,kj,li,mr,nv,oj,cu,os,pa,ps,ro,to,za
+LANGUAGE_CODE_TO_NAME = {
+    "ab": "Abkhaz",
+    "aa": "Afar",
+    "af": "Afrikaans",
+    "ak": "Akan",
+    "sq": "Albanian",
+    "am": "Amharic",
+    "ar": "Arabic",
+    "an": "Aragonese",
+    "hy": "Armenian",
+    "as": "Assamese",
+    "av": "Avaric",
+    "ae": "Avestan",
+    "ay": "Aymara",
+    "az": "Azerbaijani",
+    "bm": "Bambara",
+    "ba": "Bashkir",
+    "eu": "Basque",
+    "be": "Belarusian",
+    "bn": "Bengali",
+    "bh": "Bihari",
+    "bi": "Bislama",
+    "bs": "Bosnian",
+    "br": "Breton",
+    "bg": "Bulgarian",
+    "my": "Burmese",
+    "ca": "Catalan",  # Catalan; Valencian
+    "ch": "Chamorro",
+    "ce": "Chechen",
+    "ny": "Chichewa",  # Chichewa; Chewa; Nyanja
+    "zh": "Chinese",
+    "cv": "Chuvash",
+    "kw": "Cornish",
+    "co": "Corsican",
+    "cr": "Cree",
+    "hr": "Croatian",
+    "cs": "Czech",
+    "da": "Danish",
+    "dv": "Dhivehi",  # Divehi; Dhivehi; Maldivian;
+    "nl": "Dutch",
+    "en": "English",
+    "eo": "Esperanto",
+    "et": "Estonian",
+    "ee": "Ewe",
+    "fo": "Faroese",
+    "fj": "Fijian",
+    "fi": "Finnish",
+    "fr": "French",
+    "ff": "Fula",  # Fula; Fulah; Pulaar; Pular
+    "gl": "Galician",
+    "ka": "Georgian",
+    "de": "German",
+    "el": "Greek",  # Greek, Modern
+    "gn": "Guaraní",
+    "gu": "Gujarati",
+    "ht": "Haitian",  # Haitian; Haitian Creole
+    "ha": "Hausa",
+    "he": "Hebrew",  # Hebrew (modern)
+    "hz": "Herero",
+    "hi": "Hindi",
+    "ho": "Hiri Motu",
+    "hu": "Hungarian",
+    "ia": "Interlingua",
+    "id": "Indonesian",
+    "ie": "Interlingue",
+    "ga": "Irish",
+    "ig": "Igbo",
+    "ik": "Inupiaq",
+    "io": "Ido",
+    "is": "Icelandic",
+    "it": "Italian",
+    "iu": "Inuktitut",
+    "ja": "Japanese",
+    "jv": "Javanese",
+    "kl": "Kalaallisut",  # Kalaallisut, Greenlandic
+    "kn": "Kannada",
+    "kr": "Kanuri",
+    "ks": "Kashmiri",
+    "kk": "Kazakh",
+    "km": "Khmer",
+    "ki": "Kikuyu",  # Kikuyu, Gikuyu
+    "rw": "Kinyarwanda",
+    "ky": "Kirghiz",  # Kirghiz, Kyrgyz
+    "kv": "Komi",
+    "kg": "Kongo",
+    "ko": "Korean",
+    "ku": "Kurdish",
+    "kj": "Kwanyama",  # Kwanyama, Kuanyama
+    "la": "Latin",
+    "lb": "Luxembourgish",  # Luxembourgish, Letzeburgesch
+    "lg": "Luganda",
+    "li": "Limburgish",  # Limburgish, Limburgan, Limburger
+    "ln": "Lingala",
+    "lo": "Lao",
+    "lt": "Lithuanian",
+    "lu": "Luba-Katanga",
+    "lv": "Latvian",
+    "gv": "Manx",
+    "mk": "Macedonian",
+    "mg": "Malagasy",
+    "ms": "Malay",
+    "ml": "Malayalam",
+    "mt": "Maltese",
+    "mi": "Māori",
+    "mr": "Marathi",  # Marathi (Marāṭhī)
+    "mh": "Marshallese",
+    "mn": "Mongolian",
+    "na": "Nauru",
+    "nv": "Navajo",  # Navajo, Navaho
+    "nb": "Norwegian Bokmål",
+    "nd": "North Ndebele",
+    "ne": "Nepali",
+    "ng": "Ndonga",
+    "nn": "Norwegian Nynorsk",
+    "no": "Norwegian",
+    "ii": "Nuosu",
+    "nr": "South Ndebele",
+    "oc": "Occitan",
+    "oj": "Ojibwe",  # Ojibwe, Ojibwa
+    "cu": "Church Slavonic",  # Old Church Slavonic, Church Slavic, Church Slavonic, Old Bulgarian, Old Slavonic
+    "om": "Oromo",
+    "or": "Oriya",
+    "os": "Ossetian",  # Ossetian, Ossetic
+    "pa": "Panjabi",  # Panjabi, Punjabi
+    "pi": "Pāli",
+    "fa": "Persian",
+    "pl": "Polish",
+    "ps": "Pashto",  # Pashto, Pushto
+    "pt": "Portuguese",
+    "qu": "Quechua",
+    "rm": "Romansh",
+    "rn": "Kirundi",
+    "ro": "Romanian",  # Romanian, Moldavian, Moldovan
+    "ru": "Russian",
+    "sa": "Sanskrit",
+    "sc": "Sardinian",
+    "sd": "Sindhi",
+    "se": "Northern Sami",
+    "sm": "Samoan",
+    "sg": "Sango",
+    "sr": "Serbian",
+    "sh": "Serbo-Croatian",
+    "gd": "Gaelic",
+    "sn": "Shona",
+    "si": "Sinhalese",
+    "sk": "Slovak",
+    "sl": "Slovene",
+    "so": "Somali",
+    "st": "Southern Sotho",
+    "es": "Spanish",
+    "su": "Sundanese",
+    "sw": "Swahili",
+    "ss": "Swati",
+    "sv": "Swedish",
+    "ta": "Tamil",
+    "te": "Telugu",
+    "tg": "Tajik",
+    "th": "Thai",
+    "ti": "Tigrinya",
+    "bo": "Tibetan",
+    "tk": "Turkmen",
+    "tl": "Tagalog",
+    "tn": "Tswana",
+    "to": "Tonga",  # Tonga (Tonga Islands)
+    "tr": "Turkish",
+    "ts": "Tsonga",
+    "tt": "Tatar",
+    "tw": "Twi",
+    "ty": "Tahitian",
+    "ug": "Uighur",
+    "uk": "Ukrainian",
+    "ur": "Urdu",
+    "uz": "Uzbek",
+    "ve": "Venda",
+    "vi": "Vietnamese",
+    "vo": "Volapük",
+    "wa": "Walloon",
+    "cy": "Welsh",
+    "wo": "Wolof",
+    "fy": "Western Frisian",
+    "xh": "Xhosa",
+    "yi": "Yiddish",
+    "yo": "Yoruba",
+    "za": "Zhuang",  # Zhuang, Chuang
+}
diff --git a/src/lm_datasets/utils/settings.py b/src/lm_datasets/utils/settings.py
index f2b0fb2..3d2d095 100644
--- a/src/lm_datasets/utils/settings.py
+++ b/src/lm_datasets/utils/settings.py
@@ -1,6 +1,14 @@
 # This file contains the settings for the project.
 import logging
 
+MILLION = 1_000_000
+BILLION = 1_000_000_000
+
+KB = 1024
+MB = 1024 * 1024
+GB = 1024 * 1024 * 1024
+TB = 1024 * 1024 * 1024
+
 
 LANGUAGES = {
     # 24 languages
@@ -60,6 +68,8 @@
 }
 WIKI_MIN_TEXT_LENGTH = 5_000
 
+DEFAULT_MIN_TEXT_LENGTH = 250
+
 OSCAR_MIN_HARMFUL_PP_BY_LANGUAGE = {
     "en": 10.0,
     "fr": 10.0,
@@ -84,3 +94,41 @@
 DEFAULT_VOCAB_SIZE = 159900
 DEFAULT_TOKENIZER_RATIO = 0.10
 DEFAULT_TEST_RATIO = 0.01
+
+EURO_LANGUAGES = [
+    "en",
+    "de",
+    "es",
+    "fr",
+    "it",
+    "pl",
+    "nl",
+    "pt",
+    "cs",
+    "el",
+    "hu",
+    "sk",
+    "sv",
+    "bg",
+    "ro",
+    "uk",
+    "fi",
+    "da",
+    "sl",
+    "hr",
+    "lt",
+    "et",
+    "no",
+    "lv",
+    "mt",
+    "sr",
+    "ca",
+    "ga",
+    "eu",
+    "nn",
+    "gl",
+    "sh",
+]
+
+
+DEFAULT_MIN_FILE_SIZE_FOR_BUFFERED_SHUFFLING = 5 * GB
diff --git a/src/lm_datasets/systems.py b/src/lm_datasets/utils/systems.py
similarity index 84%
rename from src/lm_datasets/systems.py
rename to src/lm_datasets/utils/systems.py
index b0a6ffa..c2778d8 100644
--- a/src/lm_datasets/systems.py
+++ b/src/lm_datasets/utils/systems.py
@@ -3,6 +3,8 @@
 from enum import Enum
 from typing import List
 
+from lm_datasets.utils.config import Config
+
 
 logger = logging.getLogger(__name__)
 
@@ -24,7 +26,9 @@ def get_current_system(allow_default: bool = True) -> System:
         raise ValueError("Cannot determine system")
 
 
-def get_path_by_system(possible_paths: List[str], allow_default: bool = True, default_path: str = "/dev/null") -> str:
+def get_path_by_system(
+    possible_paths: List[str], allow_default: bool = True, default_path: str = "/dev/null", config: Config = None
+) -> str:
     current_system = get_current_system(allow_default)
 
     for raw_path in possible_paths:
diff --git a/src/lm_datasets/viewer/app.py b/src/lm_datasets/viewer/app.py
new file mode 100644
index 0000000..90044e5
--- /dev/null
+++ b/src/lm_datasets/viewer/app.py
@@ -0,0 +1,453 @@
+"""
+Streamlit app
+"""
+import math
+import os
+from pathlib import Path
+import streamlit as st
+import numpy as np
+import pandas as pd
+import textwrap
+import sys
+import pyarrow.parquet as pq
+import random
+import logging
+from textwrap import TextWrapper
+import argparse
+
+from viewer_utils import millify, sizeof_fmt
+
+from lm_datasets.datasets.dataset_registry import get_registered_dataset_classes
+from lm_datasets.datasets.base import BaseDataset
+from lm_datasets.utils.dataframe import get_datasets_as_dataframe
+from lm_datasets.utils.config import get_common_argparser, parse_args_and_get_config
+from lm_datasets.utils.dataframe import get_datasets_as_dataframe
+
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+    level=logging.INFO,
+    handlers=[logging.StreamHandler()],
+)
+logger = logging.getLogger(__name__)
+
+logger.info("Starting streamlit app")
+
+# parser = argparse.ArgumentParser()
+parser = argparse.ArgumentParser(parents=[get_common_argparser()], add_help=False)
+
+parser.add_argument(
+    "--shuffled_output_dir",
+    default=None,
+    type=str,
+    help="Shuffled dataset are saved in this directory",
+)
+parser.add_argument(
+    "--dataframe_cache_path",
+    default=None,
+    type=str,
+    help="Save and load dataframe cache file to this path (no cache if not set)",
+)
+parser.add_argument(
+    "--only_selected_datasets",
+    action="store_true",
+    help="Include only datasets there were explicitly selected (via config)",
+)
+parser.add_argument("--rows_count", action="store_true", help="Extract number of rows from output files")
+# args = parser.parse_args()
+config = parse_args_and_get_config(parser)
+
+raw_datasets_dir = config.raw_datasets_dir
+output_dir = config.output_dir
+
+styl = """
+<style>
+    textarea:disabled {
+    background-color: #fff;
+    color: #000;
+    opacity: 1;
+    -webkit-text-fill-color: #000;
+    }
+    div [data-testid=stImage]{
+            text-align: center;
+            display: block;
+            margin-left: auto;
+            margin-right: auto;
+            width: 100%;
+    }
+</style>
+"""
+
+app_state = st.experimental_get_query_params()
+# print(app_state)
+start = True
+loaded = True
+
+
+st.set_page_config(layout="wide")
+
+id_to_dataset_class = {
+    cls.DATASET_ID: cls
+    for cls in get_registered_dataset_classes(extra_dataset_registries=config.extra_dataset_registries)
+}
+
+logger.info("Loading df ...")
+
+
+# @st.cache_data
+def load_data():
+    if config.dataframe_cache_path is None or (
+        config.dataframe_cache_path is not None and not os.path.exists(config.dataframe_cache_path)
+    ):
+        _df = get_datasets_as_dataframe(
+            output_dir=output_dir,
+            shuffled_output_dir=config.shuffled_output_dir,
+            raw_datasets_dir=raw_datasets_dir,
+            output_format="parquet",
+            extra_dataset_registries=config.extra_dataset_registries,
+            rows_count=config.rows_count,
+            config=config,
+        )
+    else:
+        # load from disk
+        logger.info("Reading df from disk: %s", config.dataframe_cache_path)
+        _df = pd.read_csv(config.dataframe_cache_path, index_col=None)
+
+    # save to disk?
+    if config.dataframe_cache_path and not os.path.exists(config.dataframe_cache_path):
+        logger.info("Writing df to disk: %s", config.dataframe_cache_path)
+        _df.to_csv(config.dataframe_cache_path, index=False)
+
+    return _df
+
+
+df = load_data()
+
+logger.info("Filtering df ...")
+# selector_has_output_file = df.has_output_file == 1
+# df = df[selector_has_output_file]
+
+# Exclude dummies
+selector_is_non_dummy = df.dummy != 1
+df = df[selector_is_non_dummy]
+
+# Exclude oscar
+selector_is_non_oscar = ~df.source_id.str.contains("oscar")
+df = df[selector_is_non_oscar]
+
+# df = df.head()
+
+# Exclude opengptx
+# selector_is_non_opengptx = ~df.source_id.str.contains("opengptx")
+# df = df[selector_is_non_opengptx]
+
+# Exclude code
+# selector_is_non_code = df.language != "code"
+# df = df[selector_is_non_code]
+
+logger.info(f"All columns loaded: {df.columns}")
+
+selected_columns = [
+    "dataset_id",
+    # "title",
+    # "source_id",
+    "language",
+    "tokens",
+    # "web_crawled",
+]
+
+if config.output_dir:
+    selected_columns.append("has_output_files")
+
+if config.shuffled_output_dir:
+    selected_columns.append("has_shuffled_output_files")
+
+if config.rows_count:
+    selected_columns.append("rows_count")
+
+    if config.shuffled_output_dir:
+        selected_columns.append("shuffled_rows_count")
+
+
+# Sidebar
+# st.sidebar.image(
+#     str(
+#         Path(__file__).parent.parent.parent.parent
+#         / "images/A_colorful_parrot_sitting_on_a_pile_of_books__whit-removebg-preview.png"
+#     ),
+#     width=75,
+# )
+st.sidebar.markdown(
+    """<center>
+<h2><a href="https://github.com/malteos/lm-datasets">lm-datasets</a><br />Dataset Viewer</h2>
+</center>""",
+    unsafe_allow_html=True,
+)
+st.sidebar.subheader("")
+
+
+mode_options = ["overview", "stats", "dataset_details"]
+mode_index = 0
+
+if "mode" in app_state:
+    # print("app_state mode = ", app_state["mode"])
+    mode_index = mode_options.index(app_state["mode"][0])
+    # print("index ", mode_index)
+
+mode = st.sidebar.selectbox(
+    "Mode",
+    mode_options,
+    index=mode_index,
+    # format_func=lambda a: a
+)
+
+app_state["mode"] = mode
+st.experimental_set_query_params(**app_state)
+
+if mode == "overview":
+    st.header("Overview")
+
+    if "dataset_id" in app_state:
+        del app_state["dataset_id"]
+
+    if "selected_row_group" in app_state:
+        del app_state["selected_row_group"]
+
+    if "offset" in app_state:
+        del app_state["offset"]
+
+    st.experimental_set_query_params(**app_state)
+
+    # st.table(
+    #     data=pd.DataFrame([dict(a=1), dict(a=4)]),
+    #     # use_container_width=True,
+    #     # height=(len(df) + 1) * 35 + 3
+    # )
+
+    # View selected columbs of dataframe
+    st.dataframe(data=df[selected_columns], use_container_width=True, height=(len(df) + 1) * 35 + 3)
+
+
+elif mode == "stats":
+    st.header("Statistics")
+
+    # Summary
+    st.markdown(f"*Total datasets*: {len(df)}")
+    # st.markdown(f"*Total bytes*: {sizeof_fmt(df.bytes.sum())}")
+    st.markdown(f"*Total tokens*: {millify(df.tokens.sum())}")
+    st.markdown(f"*Unique languages*: {len(df.language.unique())}")
+    st.markdown(f"*Unique sources*: {len(df.source_id.unique())}")
+
+    # Plots
+    import plotly.express as px
+
+    tokens_by_language = df.groupby("language")["tokens"].sum().sort_values(ascending=False).reset_index()
+    tokens_by_web_crawled = df.groupby("web_crawled")["tokens"].sum().sort_values(ascending=False).reset_index()
+
+    fig = px.bar(tokens_by_language, x="language", y="tokens", title="Tokens by language")
+    st.plotly_chart(fig, use_container_width=True)
+
+    fig = px.pie(
+        tokens_by_web_crawled, names="web_crawled", values="tokens", title="Tokens by Web-crawled (1) or not (0)"
+    )
+    st.plotly_chart(fig, use_container_width=True)
+
+elif mode == "dataset_details":
+    dataset_id_options = list(df.dataset_id.unique())
+    dataset_id_index = 0
+
+    if "dataset_id" in app_state:
+        logger.info("app_state dataset_id = %s", app_state["dataset_id"])
+        dataset_id_index = dataset_id_options.index(app_state["dataset_id"][0])
+        logger.info("dataset_id_index = %s", dataset_id_index)
+
+    # selection = "opengptx_costep_de"  #  dataset_ids[0]
+
+    selected_dataset_id = st.sidebar.selectbox(
+        "Dataset", dataset_id_options, index=dataset_id_index, format_func=lambda a: a
+    )
+
+    app_state["dataset_id"] = selected_dataset_id
+    st.experimental_set_query_params(**app_state)
+
+    # Main
+    ds: BaseDataset = id_to_dataset_class[selected_dataset_id](
+        output_dir=output_dir,
+        shuffled_output_dir=config.shuffled_output_dir,
+        raw_datasets_dir=raw_datasets_dir,
+        output_format="parquet",
+        config=config,
+    )
+
+    st.title("Dataset: " + selected_dataset_id)
+
+    st.markdown(f"*Description*: {ds.DESCRIPTION}")
+    st.markdown(f"*Homepage*: {ds.HOMEPAGE}")
+    st.markdown(f"*Estimated tokens*: {ds.get_tokens()}")
+    st.markdown(f"*Language*: {ds.get_language_code()}")
+    st.markdown(f"*License*: {ds.LICENSE}")
+    st.markdown(f"*Known quality issues*: {ds.QUALITY_WARNINGS}")
+    st.markdown(f"*Genres*: {ds.GENRES}")
+
+    # st.markdown(f"*Is downloaded?*: {ds.is_downloaded()}")
+
+    st.markdown(styl, unsafe_allow_html=True)
+
+    if config.shuffled_output_dir:
+        use_shuffled_output_files = st.sidebar.checkbox("Shuffled output files")
+    else:
+        use_shuffled_output_files = False
+
+    if ds.has_output_files(min_file_size=271, shuffled=use_shuffled_output_files):  # empty parquet files have 270 bytes
+        st.header("Text preview:")
+
+        TEXT_COLUMN_NAME = "text"
+        MAX_TEXT_LENGTH = 50_000
+
+        output_file_paths = list(sorted(ds.get_output_file_paths(shuffled=use_shuffled_output_files)))
+
+        output_file_index = st.sidebar.number_input(
+            f"Output file index (max: {len(output_file_paths) - 1})",
+            value=0,
+            min_value=0,
+            max_value=len(output_file_paths) - 1,
+            step=1,
+        )
+
+        pq_fp = str(output_file_paths[output_file_index])  # using only part 1 if multiple chunks
+
+        st.markdown(f"*Output file*: `{pq_fp}`")
+
+        with open(pq_fp, "rb") as file_handler:
+            logger.info(f"Opening {pq_fp}")
+
+            parquet_file = pq.ParquetFile(file_handler)
+
+            logger.info("Reading metadata from parquet file ...")
+
+            metadata = parquet_file.metadata
+            num_row_groups = metadata.num_row_groups
+            num_rows = metadata.num_rows  # pq_file.metadata.num_rows  # TODO
+
+            rows_per_group = math.ceil(num_rows / num_row_groups)
+
+            # print(metadata)
+
+            # parquet_file = pq.ParquetFile(pq_fp)
+
+            max_row_group = num_row_groups - 1
+            # row_group_initial = 0
+
+            if "row_group" not in st.session_state:
+                if "selected_row_group" in app_state:
+                    st.session_state.row_group = int(app_state["selected_row_group"][0])
+                else:
+                    st.session_state.row_group = 0
+
+            if st.sidebar.button("Random offset"):
+                # random_row_group =
+                # row_group_initial = random_row_group
+                st.session_state.row_group = random.randint(0, max_row_group)
+                # st.sidebar.write(f"rand = {random_row_group}")
+                use_random_offset = True
+            else:
+                use_random_offset = False
+
+            selected_row_group = st.sidebar.number_input(
+                "Row group (Max: %d)" % max_row_group,
+                min_value=0,
+                max_value=max_row_group,
+                # value=row_group_initial,
+                step=1,
+                key="row_group",
+            )
+            app_state["selected_row_group"] = selected_row_group
+            st.experimental_set_query_params(**app_state)
+
+            st.markdown(f"*Row group*: `{selected_row_group}`")
+
+            logger.info("Reading group")
+            # try:
+            row_group = parquet_file.read_row_group(selected_row_group)
+            # except OSError:
+            #     # Dirty bugfix: Some times reading at the first try does not work => simply retry
+            #     parquet_file = pq.ParquetFile(pq_fp)
+            #     row_group = parquet_file.read_row_group(selected_row_group)
+
+            texts = row_group[TEXT_COLUMN_NAME]
+
+            step = 50
+            max_offset = max(1, len(texts) - step)
+
+            if "offset" not in st.session_state:
+                if "selected_offset" in app_state:
+                    st.session_state.offset = int(app_state["selected_offset"][0])
+                else:
+                    st.session_state.offset = 0
+
+            if use_random_offset:
+                st.session_state.offset = random.randint(0, max_offset)
+
+            # else:
+            #     # if "offset" in app_state:
+            #     #     st.session_state.offset = int(app_state["offset"][0])
+            #     # else:
+            #     st.session_state.offset = 0
+
+            offset = st.sidebar.number_input(
+                "Offset (Max: %d)" % max_offset,
+                min_value=0,
+                max_value=max_offset,
+                # value=0,
+                step=step,
+                key="offset",
+            )
+            app_state["selected_offset"] = offset
+            st.experimental_set_query_params(**app_state)
+
+            st.markdown(f"*Offset*: `{offset}`")
+
+            # Plain text viewer
+            for idx in range(offset, min(offset + step, len(texts))):
+                # st.text("        ")
+                st.subheader(f"#{selected_row_group}.{idx}")
+                # st.text("        ")
+
+                text = str(texts[idx])
+
+                # wrapper = TextWrapper(width=200, max_lines=50, replace_whitespace=False, drop_whitespace=False)
+
+                # # text = "".join(["Foobar..."] * 1000)
+
+                # wrapped_lines = wrapper.wrap(text)
+                # wrapper_text = """\n<mark>&lt;wrap&gt;</mark> """.join(wrapped_lines)
+
+                if len(text) > MAX_TEXT_LENGTH:
+                    text = text[:MAX_TEXT_LENGTH] + " ... [truncated]"
+
+                # # st.text(textwrap.fill(, width=120))
+                # # st.text(f"<pre>{wrapper_text}</pre>")
+                # # wrapper_text = text
+
+                # # st.write(f"<pre>{wrapper_text}</pre>", unsafe_allow_html=True)
+                # st.text(wrapper_text)
+
+                # st.text(text)
+                st.text_area(
+                    label=f"{selected_row_group}.{idx}",
+                    value=text,
+                    height=450,
+                    disabled=True,
+                    label_visibility="hidden",
+                )
+
+                st.divider()
+
+    else:
+        st.error("Preview not possible!")
+        st.markdown(
+            f"The dataset's output file does not exist (or is empty): `{ds.get_output_file_paths(shuffled=use_shuffled_output_files)}`"
+        )
+
+else:
+    st.error("Invalid mode selected.")
diff --git a/src/lm_datasets/viewer/ngrok-app.py b/src/lm_datasets/viewer/ngrok-app.py
new file mode 100644
index 0000000..bba876a
--- /dev/null
+++ b/src/lm_datasets/viewer/ngrok-app.py
@@ -0,0 +1,47 @@
+#!/usr/bin/env python
+#
+# For basic launch, run:
+# python ngrok.py
+import os
+from pathlib import Path
+import sys
+import asyncio, ngrok
+import click
+
+from streamlit.web.bootstrap import run
+
+NGROK_PORT = int(os.environ.get("NGROK_PORT", "8501"))
+NGROK_DOMAIN = os.environ.get("NGROK_DOMAIN", None)
+
+
+async def setup_tunnel():
+    listen = f"localhost:{NGROK_PORT}"
+    session = await ngrok.NgrokSessionBuilder().authtoken_from_env().connect()
+    if NGROK_DOMAIN:
+        tunnel = await (
+            session.http_endpoint()
+            .domain(NGROK_DOMAIN)
+            .listen()  # .domain('<name>.ngrok.app') # if on a paid plan, set a custom static domain
+        )
+    else:
+        tunnel = await (
+            session.http_endpoint().listen()  # .domain('<name>.ngrok.app') # if on a paid plan, set a custom static domain
+        )
+    click.secho(
+        f"Forwarding to {listen} from ingress url: {tunnel.url()}",
+        fg="green",
+        bold=True,
+    )
+    tunnel.forward(listen)  # forward_tcp
+
+
+try:
+    # Check if asyncio loop is already running. If so, piggyback on it to run the ngrok tunnel.
+    running_loop = asyncio.get_running_loop()
+    running_loop.create_task(setup_tunnel())
+except RuntimeError:
+    # No existing loop is running, so we can run the ngrok tunnel on a new loop.
+    asyncio.run(setup_tunnel())
+
+# forward sys.argv
+run(str(Path(__file__).parent / "app.py"), command_line=None, args=sys.argv[2:], flag_options={})
diff --git a/src/lm_datasets/viewer/viewer_utils.py b/src/lm_datasets/viewer/viewer_utils.py
new file mode 100644
index 0000000..6acc88b726
--- /dev/null
+++ b/src/lm_datasets/viewer/viewer_utils.py
@@ -0,0 +1,18 @@
+import math
+
+millnames = ["", " k", " M", " B", " T"]
+
+
+def millify(n):
+    n = float(n)
+    millidx = max(0, min(len(millnames) - 1, int(math.floor(0 if n == 0 else math.log10(abs(n)) / 3))))
+
+    return "{:.0f}{}".format(n / 10 ** (3 * millidx), millnames[millidx])
+
+
+def sizeof_fmt(num, suffix="B"):
+    for unit in ("", "Ki", "Mi", "Gi", "Ti", "Pi", "Ei", "Zi"):
+        if abs(num) < 1024.0:
+            return f"{num:3.1f}{unit}{suffix}"
+        num /= 1024.0
+    return f"{num:.1f}Yi{suffix}"


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 0000000..38bcdb1
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,26 @@
+from pathlib import Path
+import pytest
+import random
+import os
+
+TESTS_DIR = Path(__file__).parent
+FIXTURES_DIR = TESTS_DIR / "fixtures"
+RANDOM_SEED = 0
+
+random.seed(RANDOM_SEED)
+
+
+if os.getenv("_PYTEST_RAISE", "0") != "0":
+
+    @pytest.hookimpl(tryfirst=True)
+    def pytest_exception_interact(call):
+        raise call.excinfo.value
+
+    @pytest.hookimpl(tryfirst=True)
+    def pytest_internalerror(excinfo):
+        raise excinfo.value
+
+
+@pytest.fixture(scope="module")
+def raw_datasets_dir():
+    return FIXTURES_DIR / "raw_datasets"
diff --git a/tests/dummy_datasets.py b/tests/dummy_datasets.py
new file mode 100644
index 0000000..ca1eb70
--- /dev/null
+++ b/tests/dummy_datasets.py
@@ -0,0 +1,76 @@
+from lm_datasets.datasets.base import BaseDataset
+from lm_datasets.datasets.dataset_registry import get_registered_dataset_classes
+from lm_datasets.utils.config import Config
+
+
+class DummyBaseDataset(BaseDataset):
+    """
+    A dummy dataset for debugging and unit tests.
+    """
+
+    DATASET_ID = None
+    SOURCE_ID = "dummy"
+    SIZE = None
+    PREFIX = ""
+
+    def get_output_rows_count(self, shuffled: bool = False) -> int:
+        return self.SIZE
+
+    def get_texts(self):
+        """
+        Each dataset sample is string reflected the sample index, optinally with a datatset prefix.
+        """
+        for i in range(self.SIZE):
+            text = self.PREFIX + str(i)
+
+            yield text
+
+    def get_output_extension(self, with_dot: bool = True, shuffled: bool = False) -> str:
+        # use the same file names for shuffled and unshuffled
+        return super().get_output_extension(with_dot=with_dot, shuffled=False)
+
+
+def get_dummy_dataset_cls(ds_size: int, text_prefix=""):
+    class DummyDataset(DummyBaseDataset):
+        DATASET_ID = "dummy_%s%s" % (text_prefix, ds_size)
+        SIZE = ds_size
+        PREFIX = text_prefix
+
+    return DummyDataset
+
+
+def save_texts_for_temp_datasets(
+    config: Config,
+    temp_dir,
+    max_output_chunk_rows=None,
+    max_output_chunk_uncompressed_bytes=None,
+    output_batch_size=5,
+    output_format="parquet",
+):
+    list_of_texts = []
+    list_of_dataset_ids = []
+
+    for ds_cls in get_registered_dataset_classes(
+        extra_dataset_registries=config.extra_dataset_registries,
+        extra_dataset_classes=config.extra_dataset_classes,
+        use_default_registry=config.use_default_dataset_registry,
+    ):
+        ds: BaseDataset = ds_cls(
+            output_dir=temp_dir,
+            shuffled_output_dir=temp_dir,
+            max_output_chunk_rows=max_output_chunk_rows,
+            max_output_chunk_uncompressed_bytes=max_output_chunk_uncompressed_bytes,
+            output_batch_size=output_batch_size,
+            output_format=output_format,
+            config=config,
+        )
+
+        # Generate data
+        texts = list(ds.get_texts())
+
+        saved_docs_count, saved_chunks = ds.save_texts_to_parquet(texts, apply_filter=False)
+
+        list_of_texts.append(texts)
+        list_of_dataset_ids.append(ds.DATASET_ID)
+
+    return list_of_texts, list_of_dataset_ids
diff --git a/tests/test_compose_dataset_benchmark.py b/tests/test_compose_dataset_benchmark.py
new file mode 100644
index 0000000..1188dff
--- /dev/null
+++ b/tests/test_compose_dataset_benchmark.py
@@ -0,0 +1,109 @@
+import asyncio
+import logging
+from pathlib import Path
+import tempfile
+import time
+import pyarrow as pa
+from lm_datasets.io.parquet import save_texts_to_parquet_chunks
+
+from lm_datasets.utils.config import Config, get_config_from_paths
+from lm_datasets.utils.dataset_generator import DatasetGenerator, DatasetSplit
+
+logger = logging.getLogger(__name__)
+
+
+def test_compose_dataset():
+    with tempfile.TemporaryDirectory() as temp_dir:
+        config: Config = get_config_from_paths(
+            ["../eulm/lm_datasets_configs/euro_dataset.yml", "../eulm/lm_datasets_configs/pegasus.yml"],
+            dict(
+                interleave_random_batch_size=1_000,
+                input_batch_size=100,
+                # output_batch_size=100,  # 10_000
+                output_batch_size=10_000,
+                use_sampling=True,
+                save_dataset_ids=True,
+                limit=1_000_000,
+                shuffled_output_dir="/data/datasets/lm-datasets_data/euro_dataset_v1_shuffled",
+                split=DatasetSplit.TRAIN,
+                extra_dataset_registries="internal_lm_datasets.dataset_registry",
+                selected_source_ids=[
+                    # "colossal_oscar"
+                    "norwegian_cc",
+                ],
+                selected_dataset_ids=[],
+            ),
+        )
+
+        output_text_field = "text"
+        output_dataset_id_field = "dataset_id"
+        output_compression = "ZSTD"
+        max_output_chunk_uncompressed_bytes = 10 * 1024 * 1024 * 1024  # 10 GB
+
+        dataset_generator = DatasetGenerator(
+            config,
+            shuffled_output_dir=config.shuffled_output_dir,
+            output_format="parquet",
+            save_to_dir=Path(temp_dir),
+            split=config.split,
+        )
+        dataset_generator.prepare_datasets(
+            use_sampling=config.use_sampling,
+            print_progress=True,
+        )
+
+        # By default, the only column is the text column
+        parquet_columns = [
+            (output_text_field, pa.string()),
+        ]
+
+        if config.save_dataset_ids:
+            # Add `dataset_id` column
+            parquet_columns.append((output_dataset_id_field, pa.string()))
+
+        parquet_schema = pa.schema(parquet_columns)
+
+        start_time = time.perf_counter()
+
+        # Save texts from iterator to disk
+        save_texts_to_parquet_chunks(
+            dataset_generator.generate_texts_from_interleaved_datasets(
+                generate_dataset_ids=config.save_dataset_ids,
+                print_progress=True,
+                random_batch_size=config.interleave_random_batch_size,
+                reader_batch_size=config.input_batch_size,
+            ),
+            parquet_schema,
+            output_path_func=dataset_generator.save_to_path,
+            max_chunk_uncompressed_bytes=max_output_chunk_uncompressed_bytes,
+            compression=output_compression,
+            batch_size=config.output_batch_size,
+            print_write_progress=0,
+            limit=config.limit,
+        )
+
+        end_time = time.perf_counter()
+        elapsed_time = end_time - start_time
+
+        logger.info("done")
+
+        logger.info(f"{start_time=} (after preparation)")
+        logger.info(f"{end_time=}")
+        logger.info(f"{elapsed_time=:0.2f} seconds")
+        logger.info(f"{config.limit/(elapsed_time)=}")
+
+        pass
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+        # level=logging.DEBUG,
+        level=logging.INFO,
+        handlers=[logging.StreamHandler()],
+    )
+
+    # asyncio.run(
+    test_compose_dataset()
+    # )
diff --git a/tests/test_config.py b/tests/test_config.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/test_generate_texts_from_output.py b/tests/test_generate_texts_from_output.py
new file mode 100644
index 0000000..364b423
--- /dev/null
+++ b/tests/test_generate_texts_from_output.py
@@ -0,0 +1,124 @@
+import tempfile
+from .dummy_datasets import get_dummy_dataset_cls
+from lm_datasets.datasets.base import BaseDataset
+from lm_datasets.utils.config import Config
+
+
+def _test_generate_texts_from_output(
+    dataset_size=100,
+    offset=0,
+    limit=0,
+    seed=0,
+    shuffle_output_file_paths=False,
+    compare_text_indicies=None,
+    max_output_chunk_rows=12,
+    max_output_chunk_uncompressed_bytes=None,
+):
+    config = Config()
+    config.seed = seed
+
+    ds_cls = get_dummy_dataset_cls(dataset_size)
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        ds: BaseDataset = ds_cls(
+            output_dir=temp_dir,
+            shuffled_output_dir=temp_dir,
+            max_output_chunk_rows=max_output_chunk_rows,
+            max_output_chunk_uncompressed_bytes=max_output_chunk_uncompressed_bytes,
+            output_batch_size=5,
+            output_format="parquet",
+            config=config,
+        )
+
+        # Generate data
+        texts = list(ds.get_texts())
+
+        saved_docs_count, saved_chunks = ds.save_texts_to_parquet(texts, apply_filter=False)
+
+        assert saved_docs_count == dataset_size
+
+        # read again
+        output_texts = list(
+            ds.generate_texts_from_output(
+                shuffled=False,
+                offset=offset,
+                limit=limit,
+                shuffle_output_file_paths=shuffle_output_file_paths,
+                reader_implementation="pyarrow",
+            )
+        )
+
+        if limit > 0:
+            expected_texts_count = limit - offset
+        elif offset > 0:
+            expected_texts_count = len(texts) - offset
+        else:
+            expected_texts_count = len(texts)
+
+        assert expected_texts_count == len(output_texts)
+
+        if compare_text_indicies:
+            for original_text_idx, output_text_idx in compare_text_indicies:
+                assert texts[original_text_idx] == output_texts[output_text_idx]
+
+
+def test_1():
+    _test_generate_texts_from_output(dataset_size=100, offset=0, limit=0, compare_text_indicies=[(0, 0), (-1, -1)])
+
+
+def test_2():
+    _test_generate_texts_from_output(dataset_size=100, offset=0, limit=19, compare_text_indicies=[(0, 0), (19 - 1, -1)])
+
+
+def test_3():
+    _test_generate_texts_from_output(
+        dataset_size=100,
+        offset=0,
+        limit=0,
+        shuffle_output_file_paths=True,
+        seed=0,
+        compare_text_indicies=[(60, 0), (-1, -1)],
+    )
+
+
+def test_4():
+    _test_generate_texts_from_output(
+        dataset_size=100,
+        offset=10,
+        limit=0,
+        shuffle_output_file_paths=True,
+        seed=1,
+        compare_text_indicies=[(55, 0), (29, -1)],
+    )
+
+
+def test_5():
+    _test_generate_texts_from_output(
+        dataset_size=1000,
+        offset=10,
+        limit=200,
+        shuffle_output_file_paths=True,
+        compare_text_indicies=[(95, 0), (734, -1)],
+        max_output_chunk_uncompressed_bytes=512,
+        max_output_chunk_rows=None,
+    )
+
+
+if __name__ == "__main__":
+    import logging
+
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+        level=logging.DEBUG,
+        handlers=[logging.StreamHandler()],
+    )
+    logger = logging.getLogger(__name__)
+
+    test_1()
+    test_2()
+    test_3()
+    test_4()
+    test_5()  # TODO
+
+    print("done")
diff --git a/tests/test_interleave_datasets.py b/tests/test_interleave_datasets.py
new file mode 100644
index 0000000..e13e9cc
--- /dev/null
+++ b/tests/test_interleave_datasets.py
@@ -0,0 +1,218 @@
+from collections import Counter
+import tempfile
+
+import pytest
+from lm_datasets.utils.config import Config
+from lm_datasets.utils.dataset_generator import DatasetGenerator, DatasetSplit
+from tests.dummy_datasets import get_dummy_dataset_cls, save_texts_for_temp_datasets
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def _test_interval_datasets(
+    dataset_sizes,
+    dataset_prefixes,
+    sampling_factors,
+    expected_output_rows,
+    split=DatasetSplit.FULL,
+    seed=0,
+    shuffle_output_file_paths=False,
+    max_output_chunk_rows=12,
+    max_output_chunk_uncompressed_bytes=None,
+    output_batch_size=5,
+    output_format="parquet",
+    validation_ratio=0.25,
+    tokenizer_train_ratio=0.1,
+    validation_min_total_docs=None,
+    validation_max_split_docs=None,
+    validation_min_split_docs=None,
+):
+    ds_clss = [get_dummy_dataset_cls(size, prefix) for size, prefix in zip(dataset_sizes, dataset_prefixes)]
+
+    config = Config()
+    config.seed = seed
+    config.extra_dataset_classes = ds_clss
+    config.use_default_dataset_registry = False
+    config.selected_source_ids = ["dummy"]
+    config.validation_ratio = validation_ratio  # number of documents in the split: len(dataset) * ratio
+    config.validation_min_total_docs = (
+        validation_min_total_docs  # to be used as validation set, the dataset must have at least n docs
+    )
+    config.validation_max_split_docs = (
+        validation_max_split_docs  # number of documents in validation split are capped at this numbers
+    )
+    config.validation_min_split_docs = (
+        validation_min_split_docs  # split must have at least this number of documents, otherwise it will be discarded
+    )
+    config.tokenizer_train_ratio = tokenizer_train_ratio  # % of train data used for tokenizer training
+
+    # expected_rows_per_dataset_id = {"dummy_b_100": 200, "dummy_a_100": 150, "dummy_c_100": 10}
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        list_of_saved_texts, list_of_dataset_ids = save_texts_for_temp_datasets(
+            config,
+            temp_dir,
+            max_output_chunk_rows=max_output_chunk_rows,
+            max_output_chunk_uncompressed_bytes=max_output_chunk_uncompressed_bytes,
+            output_format=output_format,
+            output_batch_size=output_batch_size,
+        )
+
+        config.sampling_factor_by_dataset_id = {dsid: f for dsid, f in zip(list_of_dataset_ids, sampling_factors)}
+
+        dataset_generator = DatasetGenerator(
+            config,
+            split=split,
+            shuffled_output_dir=temp_dir,
+            output_format=output_format,
+            save_to_dir=temp_dir,
+        )
+        dataset_generator.prepare_datasets(
+            use_sampling=True,
+            print_progress=True,
+        )
+        output_rows = list(
+            dataset_generator.generate_texts_from_interleaved_datasets(
+                generate_dataset_ids=True,
+                print_progress=True,
+            )
+        )
+        rows_per_dataset_id = Counter([ds_id.as_py() for ds_id, sample in output_rows])
+
+        logger.debug(rows_per_dataset_id)
+
+        # compare expected and true row count
+        expected_rows_per_dataset_id = {
+            dsid: row_count for dsid, row_count in zip(list_of_dataset_ids, expected_output_rows)
+        }
+        for dataset_id, row_count in expected_rows_per_dataset_id.items():
+            assert (
+                rows_per_dataset_id[dataset_id] == row_count
+            ), f"compare expected and true row count for {dataset_id}. given: {rows_per_dataset_id[dataset_id]}; expected: {row_count}"
+
+    logger.debug("test completed")
+
+
+def test_simple():
+    _test_interval_datasets(
+        dataset_sizes=[100, 100, 100],
+        dataset_prefixes=["a", "b", "c"],
+        sampling_factors=[1.0, 1.5, 0.1],
+        expected_output_rows=[100, 150, 10],
+        split=DatasetSplit.FULL,
+    )
+
+
+def test_with_sampling():
+    _test_interval_datasets(
+        dataset_sizes=[100, 51, 93],
+        dataset_prefixes=["a", "b", "c"],
+        sampling_factors=[3.5, 1.5, 0.1],
+        expected_output_rows=[350, 76, 9],
+        split=DatasetSplit.FULL,
+    )
+
+
+def test_train_split():
+    _test_interval_datasets(
+        dataset_sizes=[100, 100, 100],
+        dataset_prefixes=["a", "b", "c"],
+        sampling_factors=[1.0, 1.0, 1.0],
+        expected_output_rows=[75, 75, 75],
+        split=DatasetSplit.TRAIN,
+        validation_ratio=0.25,
+        tokenizer_train_ratio=0.1,
+    )
+
+
+def test_validation_split():
+    _test_interval_datasets(
+        dataset_sizes=[200, 200, 200],
+        dataset_prefixes=["a", "b", "c"],
+        sampling_factors=[0.5, 2.0, 3.0],
+        expected_output_rows=[25, 100, 150],
+        split=DatasetSplit.VALIDATION,
+        validation_ratio=0.25,
+        tokenizer_train_ratio=0.1,
+    )
+
+
+def test_tokenizer_train_split():
+    _test_interval_datasets(
+        dataset_sizes=[200, 200, 200],
+        dataset_prefixes=["a", "b", "c"],
+        sampling_factors=[0.5, 2.0, 3.0],
+        expected_output_rows=[7, 30, 45],
+        split=DatasetSplit.TOKENIZER_TRAIN,
+        validation_ratio=0.25,
+        tokenizer_train_ratio=0.1,
+    )
+
+
+def test_validation_with_max_split_docs():
+    _test_interval_datasets(
+        dataset_sizes=[1000, 1000, 100],
+        dataset_prefixes=["a", "b", "c"],
+        sampling_factors=[1.0, 1.0, 1.0],
+        expected_output_rows=[100, 100, 30],
+        split=DatasetSplit.VALIDATION,
+        validation_ratio=0.3,
+        tokenizer_train_ratio=0.1,
+        validation_max_split_docs=100,
+    )
+
+
+def test_validation_min_split_docs():
+    _test_interval_datasets(
+        dataset_sizes=[1000, 1000, 100],
+        dataset_prefixes=["a", "b", "c"],
+        sampling_factors=[1.0, 1.0, 1.0],
+        expected_output_rows=[300, 300, 0],
+        split=DatasetSplit.VALIDATION,
+        validation_ratio=0.3,
+        tokenizer_train_ratio=0.1,
+        validation_min_split_docs=50,
+    )
+
+
+def test_validation_min_total_docs():
+    _test_interval_datasets(
+        dataset_sizes=[1000, 1000, 100],
+        dataset_prefixes=["a", "b", "c"],
+        sampling_factors=[1.0, 1.0, 1.0],
+        expected_output_rows=[300, 300, 0],
+        split=DatasetSplit.VALIDATION,
+        validation_ratio=0.3,
+        tokenizer_train_ratio=0.1,
+        validation_min_total_docs=100 + 1,
+    )
+
+
+@pytest.mark.skip()
+def test_large():
+    factor = 1000
+    _test_interval_datasets(
+        dataset_sizes=[100] * factor,
+        dataset_prefixes=[f"{i}_" for i in range(factor)],
+        sampling_factors=[1.0] * factor,
+        expected_output_rows=[7_500] * factor,
+        split=DatasetSplit.TRAIN,
+    )
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+        level=logging.DEBUG,
+        # level=logging.INFO,
+        handlers=[logging.StreamHandler()],
+    )
+
+    # test_simple()
+    test_train_split()
+    test_validation_min_total_docs()
+    test_tokenizer_train_split()
+    # test_large()
+    print("done")
diff --git a/tests/test_iterate_over_shuffled_datasets.py b/tests/test_iterate_over_shuffled_datasets.py
index 082ac0e..f06f66f 100644
--- a/tests/test_iterate_over_shuffled_datasets.py
+++ b/tests/test_iterate_over_shuffled_datasets.py
@@ -1,8 +1,9 @@
 import os
 import pytest
 from dataclasses import dataclass
-from lm_datasets.datasets.dataset_registry import get_registered_dataset_classes
+from lm_datasets.datasets.dataset_registry import get_dataset_class_by_id, get_registered_dataset_classes
 from lm_datasets.datasets.base import BaseDataset
+from lm_datasets.utils.config import Config
 
 import pyarrow.parquet as pq
 
@@ -11,8 +12,11 @@
 
 @pytest.mark.skipif(not os.path.exists(SHUFFLED_OUTPUT_DIR), reason="test file not exists")
 def test_iterate_over_datasets():
-    id_to_dataset_class = {cls.DATASET_ID: cls for cls in get_registered_dataset_classes()}
-    dataset_cls = id_to_dataset_class["openlegaldata"]
+    # id_to_dataset_class = {cls.DATASET_ID: cls for cls in get_registered_dataset_classes()}
+    dataset_id = "openlegaldata"
+    dataset_id = "starcoder"
+
+    dataset_cls = get_dataset_class_by_id(dataset_id)
 
     @dataclass
     class Arguments(object):
@@ -27,44 +31,52 @@ class Arguments(object):
         train_ratio = 0.05  # number of rows to use for tokenizer training
 
     args = Arguments()
+    config = Config()
+    config.local_dirs_by_dataset_id = {"starcoder": "/netscratch/ortiz/corpora/starcoder/starcoderdata"}
 
     dataset: BaseDataset = dataset_cls(
         output_dir=args.output_dir,
         output_format=args.output_format,
         output_compression=args.output_compression,
         shuffled_output_dir=args.shuffled_output_dir,
+        config=config,
     )
 
     batch_size = 1000
     # text_column_name = "text"
 
-    def generate_texts():
-        for fp in sorted(dataset.get_chunked_output_file_paths(shuffled=True)):
-            pq_fp = fp
+    def generate_texts_from_output():
+        output_paths = list(dataset.get_output_file_paths(shuffled=True))
+
+        if output_paths:
+            for fp in sorted(output_paths):
+                pq_fp = fp
 
-            with open(pq_fp, "rb") as file_handler:
-                # pq_file = open_parquet_file_with_retries(dataset.get_output_file_path(), retries=2)
-                pq_file = pq.ParquetFile(file_handler)
+                with open(pq_fp, "rb") as file_handler:
+                    # pq_file = open_parquet_file_with_retries(dataset.get_output_file_path(), retries=2)
+                    pq_file = pq.ParquetFile(file_handler)
 
-                print(f"rows {pq_file.metadata.num_rows}")
+                    print(f"rows {pq_file.metadata.num_rows}")
 
-                rows = 0
-                max_rows = int(pq_file.metadata.num_rows * args.train_ratio)
-                skip_rows = min(int(pq_file.metadata.num_rows * args.skip_ratio), args.skip_n)
+                    rows = 0
+                    max_rows = int(pq_file.metadata.num_rows * args.train_ratio)
+                    skip_rows = min(int(pq_file.metadata.num_rows * args.skip_ratio), args.skip_n)
 
-                print(f"reading {max_rows} rows")
-                print(f"skipping {skip_rows} rows")
+                    print(f"reading {max_rows} rows")
+                    print(f"skipping {skip_rows} rows")
 
-                for batch in pq_file.iter_batches(batch_size):
-                    for text in batch[dataset.output_text_field]:
-                        yield text
+                    for batch in pq_file.iter_batches(batch_size):
+                        for text in batch[dataset.get_output_text_field()]:
+                            yield text
 
-                        rows += 1
+                            rows += 1
 
-                print()
-            break
+                    print()
+                break
+        else:
+            print("missing output_paths")
 
-    for text in generate_texts():
+    for text in generate_texts_from_output():
         print(text)
 
     print("x")
diff --git a/tests/test_openlegaldata.py b/tests/test_openlegaldata.py
deleted file mode 100644
index b0f407a..0000000
--- a/tests/test_openlegaldata.py
+++ /dev/null
@@ -1,23 +0,0 @@
-from smart_open import open
-import json
-import pytest
-import os
-
-DATASET_FILE_PATH = "/netscratch/mostendorff/experiments/eulm/data/ele/de/openlegaldata/merged_v1_v2.jsonl.gz"
-
-
-@pytest.mark.skipif(not os.path.exists(DATASET_FILE_PATH), reason="test file not exists")
-def test_openlegaldata():
-    docs_count = 1_658_611
-    docs = []
-
-    # after 250/1659 row groups
-    # rows per group = 1000
-
-    with open(DATASET_FILE_PATH) as f:
-        for idx, line in enumerate(f):
-            if idx > docs_count - 1_000:
-                doc = json.loads(line)
-                docs.append(doc)
-
-    print("done")
diff --git a/tests/test_read_parquet.py b/tests/test_read_parquet.py
index a55c4c2..9b52f0b 100644
--- a/tests/test_read_parquet.py
+++ b/tests/test_read_parquet.py
@@ -1,152 +1,192 @@
+import asyncio
+import itertools
+import logging
 from pathlib import Path
+import tempfile
+import time
+import pyarrow as pa
 import pyarrow.parquet as pq
-import os
-import pytest
+from lm_datasets.io.parquet import save_texts_to_parquet_chunks
 
-from tqdm.auto import tqdm
+from lm_datasets.utils.config import Config, get_config_from_paths
+from lm_datasets.utils.dataset_generator import DatasetGenerator, DatasetSplit
 
+logger = logging.getLogger(__name__)
 
-@pytest.mark.skip(reason="requires dataset files")
-def test_read_pq():
-    fp = (  # macocu_el.parquet
-        "/netscratch/mostendorff/experiments/eulm/data/lm_datasets_texts/opengptx_costep_de.parquet"  # noqa
-    )
-    fp = "/netscratch/mostendorff/experiments/eulm/data/lm_datasets_texts/opengptx_dissertations_de.parquet"  # noqa
-    fp = "/netscratch/mostendorff/experiments/eulm/data/lm_datasets_texts/opengptx_dissertations_de.parquet"  # noqa
-    fp = "/netscratch/mostendorff/experiments/eulm/data/lm_datasets_texts_shuffled/opengptx_bt_drucksachen.shuffled.parquet"  # noqa
 
-    assert os.path.exists(fp)
+def sync_dataset_iterator(file_path, batch_size, output_text_field):
+    with open(file_path, "rb") as file_handler:
+        pq_file = pq.ParquetFile(file_handler)
+        logger.info("Generate from %s with bs=%i", file_path, batch_size)
+
+        for pq_batch in pq_file.iter_batches(columns=[output_text_field], batch_size=batch_size, use_threads=False):
+            # logger.info(f"{file_path} , {len(pq_batch[0])=}")
+            for v in pq_batch[0]:
+                yield v
 
-    metadata = pq.read_metadata(fp)
-    print(metadata)
 
-    parquet_file = pq.ParquetFile(fp)
+async def async_dataset_iterator(file_path, batch_size, output_text_field):
+    with open(file_path, "rb") as file_handler:
+        pq_file = pq.ParquetFile(file_handler)
+        logger.info("Generate from %s with bs=%i", file_path, batch_size)
 
-    rg = parquet_file.read_row_group(0)
-    texts = rg["text"]
+        for pq_batch in pq_file.iter_batches(columns=[output_text_field], batch_size=batch_size, use_threads=False):
+            # logger.info(f"{file_path} , {len(pq_batch[0])=}")
+            for v in pq_batch[0]:
+                yield v
 
-    print(str(texts[0])[:100])
 
-    print(metadata)
+def sync_convert_to_batches(iterable, batch_size):
+    py_batch = []
 
+    for sample in iterable:
+        py_batch.append(sample)
 
-@pytest.mark.skip(reason="requires dataset files")
-def test_polars_read():
-    fps = [
-        "/netscratch/mostendorff/experiments/eulm/data/lm_datasets_texts_shuffled/opengptx_bt_drucksachen.shuffled.parquet"  # noqa
-    ]  # 6445
-    fps = Path("/netscratch/mostendorff/experiments/eulm/data/lm_datasets_texts/").glob("*.parquet")
-    fps = list(sorted(fps))
+        if len(py_batch) >= batch_size:
+            yield pa.array(py_batch)
+            py_batch = []  # reset
 
-    for i, fp in enumerate(tqdm(fps)):
-        fp = str(fp)
+    if py_batch:
+        yield pa.array(py_batch)
 
-        print(i, fp)
 
-        if "__bak" in fp or "part-" in fp or i < 137:
-            print("-- skip")
-            continue
+async def async_convert_to_batches(iterable, batch_size):
+    py_batch = []
 
-        fp = "/netscratch/mostendorff/experiments/eulm/data/lm_datasets_texts/pes2o.parquet"
+    async for sample in iterable:
+        py_batch.append(sample)
 
-    print("done")
+        if len(py_batch) >= batch_size:
+            yield pa.array(py_batch)
+            py_batch = []  # reset
 
+    if py_batch:
+        yield pa.array(py_batch)
 
-@pytest.mark.skip(reason="requires dataset files")
-def test_parquet_os_error():
-    fps = ["/netscratch/mostendorff/experiments/eulm/data/lm_datasets_texts/danewsroom.parquet"]
 
-    fps = Path("/netscratch/mostendorff/experiments/eulm/data/lm_datasets_texts/").glob("*.parquet")
-    fps = list(sorted(fps))
+# async def test_read_parquet():
+#     output = []
 
-    for i, fp in enumerate(tqdm(fps)):
-        fp = str(fp)
+#     async for batch in convert_to_batches(generate_from_parquet(), batch_size=5):
+#         output.append(batch)
 
-        print(i, fp)
+#     logger.info(f"{len(output)=}")
 
-        if "__bak" in fp or "part-" in fp:  # or i < 137:
-            print("-- skip")
-            continue
 
-        # if ""
-        print(fp)
+async def async_generate_from_parquet(file_limit):
+    fps = list(sorted(Path("/data/datasets/lm-datasets_data/euro_dataset_v1_shuffled").glob("*.parquet")))[
+        :file_limit
+    ]  # colossal_oscar_2023-
+    print(f"{len(fps)=}")
+    output_text_field = "text"
+    batch_size = 10
 
-        # import pyarrow as pa
-        # t = pq.read_table(fp, memory_map=True)
+    dataset_iterators = [async_dataset_iterator(file_path, batch_size, output_text_field) for file_path in fps]
 
-        ok = 0
-        errors = 0
+    for dataset_idx in itertools.cycle(range(len(dataset_iterators))):
+        try:
+            iterator = dataset_iterators[dataset_idx]
+            # sample = next(iterator)
+            # await sample = iterator.__anext__()
+            sample = await anext(iterator)
 
-        for i in range(4):
-            try:
-                # with open(fp, "rb") as f:
-                #     pq.read_schema(f)
+            yield sample
+        except (StopIteration, StopAsyncIteration):
+            break
 
-                with open(fp, "rb") as f:
-                    parquet_file = pq.ParquetFile(f)
-                    for _, record_batch in enumerate(parquet_file.iter_batches(batch_size=100)):
-                        # pa_table = pa.Table.from_batches([record_batch])
-                        for text in record_batch["text"]:
-                            text = str(text)
-                            break
-                        break
 
-                # tab = pq.read_table(fp, memory_map=True)
-                # for batch in tab.to_batches():
-                #     for text in batch["text"]:
-                #         text = str(text)
-                #         break
-                #     break
+def sync_generate_from_parquet(file_limit):
+    fps = list(sorted(Path("/data/datasets/lm-datasets_data/euro_dataset_v1_shuffled").glob("*.parquet")))[
+        :file_limit
+    ]  # colossal_oscar_2023-
+    print(f"{len(fps)=}")
+    output_text_field = "text"
+    batch_size = 10
 
-                # from datasets import load_dataset
-                # ds = load_dataset("parquet", data_files={"train": fp+"1"}, split="train", streaming=True)
-                # next(iter(ds))
+    dataset_iterators = [sync_dataset_iterator(file_path, batch_size, output_text_field) for file_path in fps]
 
-                # with  pq.ParquetFile(fp) as parquet_file:
-                #     rg = parquet_file.read_row_group(0)  # columns=["text"]
-                #     texts = rg["text"]
-                #     text = str(texts[0])
+    for dataset_idx in itertools.cycle(range(len(dataset_iterators))):
+        try:
+            iterator = dataset_iterators[dataset_idx]
+            sample = next(iterator)
+            # await sample = iterator.__anext__()
+            # sample = await anext(iterator)
 
-                # parquet_file = pq.ParquetFile(fp)
+            yield sample
+        except (StopIteration, StopAsyncIteration):
+            break
 
-                # texts = rg["text"]
-                # text = str(texts[0])
 
-                # df = pl.scan_parquet(fp)
-                # text = df.row(0)
+async def async_save_texts(text_iterator, limit):
+    output = []
 
-                # print(parquet_file.closed)
+    async for batch in async_convert_to_batches(text_iterator, batch_size=5):
+        output.append(batch)
 
-                # for batch in parquet_file.iter_batches(row_groups=[0]):
-                #     text = str(batch["text"][0])
-                #     break
+        if len(output) > limit:
+            print("limit reached")
+            break
 
-                # if "x" in text:
-                #     pass
+    logger.info(f"{len(output)=}")
 
-                ok += 1
-                # print(i, "OK")
-            except OSError as e:
-                errors += 1
-                print(i, "ERROR", fp)
-                raise e
 
-        print("stats = ", ok, errors, fp)
+def sync_save_texts(text_iterator, limit):
+    output = []
 
-        if errors:
-            print("#### error")
+    for batch in sync_convert_to_batches(text_iterator, batch_size=5):
+        output.append(batch)
 
-        print("")
+        if len(output) > limit:
+            print("limit reached")
+            break
 
-    print("done")
+    logger.info(f"{len(output)=}")
 
 
-def test_read_groups():
-    # TODO loop over files and read random groups
-    pass
+async def async_generate_and_save(limit, file_limit):
+    await async_save_texts(
+        async_generate_from_parquet(file_limit),
+        limit,
+    )
+
+
+def sync_generate_and_save(limit, file_limit):
+    sync_save_texts(
+        sync_generate_from_parquet(file_limit),
+        limit,
+    )
 
 
 if __name__ == "__main__":
-    # test_read_pq()
-    test_parquet_os_error()
-    # test_polars_read()
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+        # level=logging.DEBUG,
+        level=logging.INFO,
+        handlers=[logging.StreamHandler()],
+    )
+
+    start_time = time.perf_counter()
+    file_limit = 100
+    limit = 100_000
+    # asyncio.run(async_generate_and_save(limit, file_limit))
+    sync_generate_and_save(limit, file_limit)
+
+    end_time = time.perf_counter()
+    elapsed_time = end_time - start_time
+
+    logger.info("done")
+
+    logger.info(f"{start_time=}")
+    logger.info(f"{end_time=}")
+    logger.info(f"{elapsed_time=:0.2f} seconds")  # approx. 20 seconds
+
+    # limit=10000
+    # file_limit=20
+    # elapsed_time=8.72 seconds --- sync
+    # elapsed_time=9.03 seconds -- async
+
+    # limit=100_000
+    # file_limit=100
+    # async:  elapsed_time=18.46 seconds
+    # sync: elapsed_time=18.17 seconds
diff --git a/tests/test_shuffle_parquet.py b/tests/test_shuffle_parquet.py
deleted file mode 100644
index a8a4ad0..0000000
--- a/tests/test_shuffle_parquet.py
+++ /dev/null
@@ -1,33 +0,0 @@
-import os
-import pytest
-
-DATASET_FILE_PATH = (
-    "/netscratch/mostendorff/experiments/eulm/data/lm_datasets_texts_shuffled/opengptx_bt_drucksachen.shuffled.parquet"
-)
-
-
-@pytest.mark.skipif(not os.path.exists(DATASET_FILE_PATH), reason="test file not exists")
-def test_shuffle_parquet():
-    import polars as pl
-
-    fp = DATASET_FILE_PATH
-
-    df = pl.scan_parquet(fp)
-
-    df = df.with_columns(pl.all().shuffle(seed=1)).with_row_count()
-    train_fraction: float = 0.75
-
-    df_train = df.filter(pl.col("row_nr") < pl.col("row_nr").max() * train_fraction)
-    df_test = df.filter(pl.col("row_nr") >= pl.col("row_nr").max() * train_fraction)
-
-    df_train.collect().write_parquet("./shuffled_train.parquet")
-
-    df_train.sink_parquet("./shuffled_train.parquet")
-    df_test.write_parquet("./shuffled_train.parquet")
-
-    df = df.sample(fraction=1, shuffle=True)
-
-    print("done")
-
-
-# df.write_parquet(shuffled_output_file_path)
diff --git a/tests/test_split_dataset.py b/tests/test_split_dataset.py
new file mode 100644
index 0000000..a1f54c5
--- /dev/null
+++ b/tests/test_split_dataset.py
@@ -0,0 +1,148 @@
+import math
+import random
+import tempfile
+from lm_datasets.datasets.base import BaseDataset
+from lm_datasets.utils.config import Config
+
+from lm_datasets.utils.dataset_generator import (
+    generate_texts_from_dataset,
+    DatasetSplit,
+    get_splits_as_offsets_and_limits,
+)
+
+from .dummy_datasets import get_dummy_dataset_cls
+
+
+def _test_train_validation_split(
+    dataset_size,
+    validation_ratio,
+    tokenizer_train_ratio,
+    max_output_chunk_rows=None,
+    max_output_chunk_uncompressed_bytes=None,
+):
+    random.seed(0)
+
+    config = Config()
+    config.validation_ratio = validation_ratio
+    config.validation_min_total_docs = 0
+    config.validation_min_split_docs = 0
+    config.tokenizer_train_ratio = tokenizer_train_ratio
+
+    expected_validation_size = math.floor(dataset_size * validation_ratio)
+    expected_train_size = dataset_size - expected_validation_size
+    expected_tokenizer_train_size = math.floor(expected_train_size * tokenizer_train_ratio)
+
+    ds_cls = get_dummy_dataset_cls(dataset_size)
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        ds: BaseDataset = ds_cls(
+            output_dir=temp_dir,
+            shuffled_output_dir=temp_dir,
+            max_output_chunk_rows=max_output_chunk_rows,
+            max_output_chunk_uncompressed_bytes=max_output_chunk_uncompressed_bytes,
+            output_batch_size=5,
+            output_format="parquet",
+            config=config,
+        )
+
+        # Generate data
+        texts = list(ds.get_texts())
+
+        shuffled_docs_count, saved_chunks = ds.save_texts_to_parquet(texts, apply_filter=False)
+
+        assert shuffled_docs_count == dataset_size
+
+        use_shuffled_output = False
+        split_offset_limit = get_splits_as_offsets_and_limits(ds, use_shuffled_output=use_shuffled_output)
+
+        full_texts = list(
+            generate_texts_from_dataset(
+                ds,
+                split=DatasetSplit.FULL,
+                split_offset_limit=split_offset_limit,
+                use_shuffled_output=use_shuffled_output,
+            )
+        )
+        print("full_texts", len(full_texts))
+
+        train_texts = list(
+            generate_texts_from_dataset(
+                ds,
+                split=DatasetSplit.TRAIN,
+                split_offset_limit=split_offset_limit,
+                use_shuffled_output=use_shuffled_output,
+            )
+        )
+        print("train_texts", len(train_texts), expected_train_size)
+
+        tokenizer_train_texts = list(
+            generate_texts_from_dataset(
+                ds,
+                split=DatasetSplit.TOKENIZER_TRAIN,
+                split_offset_limit=split_offset_limit,
+                use_shuffled_output=use_shuffled_output,
+            )
+        )
+        print("tokenizer_train_texts", len(tokenizer_train_texts), expected_tokenizer_train_size)
+
+        assert len(tokenizer_train_texts) == expected_tokenizer_train_size
+
+        val_texts = list(
+            generate_texts_from_dataset(
+                ds,
+                split=DatasetSplit.VALIDATION,
+                split_offset_limit=split_offset_limit,
+                use_shuffled_output=use_shuffled_output,
+            )
+        )
+        print("val_texts", len(val_texts), expected_validation_size)
+
+        assert full_texts == texts
+
+        assert len(full_texts) == dataset_size
+        assert len(train_texts) + len(val_texts) == dataset_size
+
+        # no overlap train-val
+        assert len(set(val_texts) & set(train_texts)) == 0
+
+        assert len(train_texts) == expected_train_size
+        assert len(val_texts) == expected_validation_size
+
+        assert len(set(tokenizer_train_texts) & set(train_texts)) == expected_tokenizer_train_size
+
+
+def test_train_validation_split_1000_02_01():
+    _test_train_validation_split(
+        dataset_size=1000, validation_ratio=0.2, tokenizer_train_ratio=0.1, max_output_chunk_rows=21
+    )
+
+
+def test_train_validation_split_512_033_033():
+    _test_train_validation_split(
+        dataset_size=512,
+        validation_ratio=0.33,
+        tokenizer_train_ratio=0.33,
+        max_output_chunk_uncompressed_bytes=1024 * 2.5,
+    )
+
+
+def test_train_validation_split_100_025_033():
+    _test_train_validation_split(
+        dataset_size=100, validation_ratio=0.25, tokenizer_train_ratio=0.33, max_output_chunk_rows=10
+    )
+
+
+if __name__ == "__main__":
+    import logging
+
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+        level=logging.DEBUG,
+        handlers=[logging.StreamHandler()],
+    )
+    logger = logging.getLogger(__name__)
+
+    test_train_validation_split_1000_02_01()
+    test_train_validation_split_512_033_033()
+    test_train_validation_split_100_025_033()
diff --git a/tests/test_textwrap.py b/tests/test_textwrap.py
deleted file mode 100644
index 263f41f..0000000
--- a/tests/test_textwrap.py
+++ /dev/null
@@ -1,13 +0,0 @@
-from textwrap import TextWrapper
-
-wrapper = TextWrapper(width=100, max_lines=5)
-
-text = "".join(["Foobar..."] * 1000)
-
-wrapped_lines = wrapper.wrap(text)
-wrapper_text = """\n<mark>&gt;wrap&lt;</mark> """.join(wrapped_lines)
-print(wrapper_text)
-
-print(text)
-
-print("done")
diff --git a/tests/test_write_parquet_chunks.py b/tests/test_write_parquet_chunks.py
index 14a0815..bc61b73 100644
--- a/tests/test_write_parquet_chunks.py
+++ b/tests/test_write_parquet_chunks.py
@@ -19,12 +19,15 @@
 # TMP_DIR = Path(__file__).parent.parent / "data/tmp"
 
 
-def get_texts(n: int, min_len: int, max_len: int):
+def get_texts(n: int, min_len: int, max_len: int, extra_column=False):
     for i in range(n):
         random_len = randint(min_len, max_len)
         random_str = "".join(random.choices(string.ascii_lowercase, k=random_len))
 
-        yield random_str
+        if extra_column:
+            yield {"text": random_str, "extra_column": "EXTRA_COLUMN"}
+        else:
+            yield random_str
 
 
 def test_write_parquet_chunks():
@@ -52,15 +55,39 @@ def get_output_path(part=None, total_parts=None) -> str:
             print(f"Removed {fp}")
 
         output_text_field = "text"
-        schema = pa.schema(
+        text_only_schema = pa.schema(
             [
                 (output_text_field, pa.string()),
             ]
         )
+        text_and_extra_col_schema = pa.schema(
+            [
+                (output_text_field, pa.string()),
+                ("extra_column", pa.string()),
+            ]
+        )
+
+        # text + extra column
+
+        saved_docs, saved_chunks = save_texts_to_parquet_chunks(
+            texts=get_texts(
+                n=10_000, min_len=250, max_len=10_000, extra_column=True
+            ),  # approx 50 MB  (none compression)
+            schema=text_and_extra_col_schema,
+            max_chunk_uncompressed_bytes=10 * 1024 * 1024,  # 10 MB
+            output_path_func=get_output_path,
+            compression="ZSTD",
+            batch_size=24,
+        )
+
+        assert saved_docs == 10_000
+        assert saved_chunks == 5
+
+        # text only
 
         saved_docs, saved_chunks = save_texts_to_parquet_chunks(
             texts=get_texts(n=10_000, min_len=250, max_len=10_000),  # approx 50 MB  (none compression)
-            schema=schema,
+            schema=text_only_schema,
             max_chunk_uncompressed_bytes=10 * 1024 * 1024,  # 10 MB
             output_path_func=get_output_path,
             compression="ZSTD",
@@ -72,7 +99,7 @@ def get_output_path(part=None, total_parts=None) -> str:
 
         limited_saved_docs, limited_saved_chunks = save_texts_to_parquet_chunks(
             texts=get_texts(n=10_000, min_len=250, max_len=10_000),  # approx 50 MB  (none compression)
-            schema=schema,
+            schema=text_only_schema,
             max_chunk_uncompressed_bytes=10 * 1024 * 1024,  # 10 MB
             output_path_func=get_output_path,
             compression="ZSTD",
@@ -85,7 +112,7 @@ def get_output_path(part=None, total_parts=None) -> str:
 
         single_file_saved_docs, single_file_saved_chunks = save_texts_to_parquet_chunks(
             texts=get_texts(n=1_000, min_len=250, max_len=10_000),  # approx 5 MB (none compression)
-            schema=schema,
+            schema=text_only_schema,
             max_chunk_uncompressed_bytes=0,
             output_path_func=get_output_path,
             compression="ZSTD",